Привет! Я Андрей Татаринов, директор AGIMA.AI. Мы занимаемся проектами в области машинного обучения и анализа данных. В этой статье расскажу, как мы использовали фреймворк Mediapipe для iOS и Android, запускали его на десктопе, писали кастомные калькуляторы и в поддержку сообщества.
Корпорация Google анонсировала Mediapipe на CVPR в 2019 году. Оригинальная статья предлагает простую идею: давайте рассматривать процесс работы с данными как граф, вершины которого — это модули обработки данных. Помимо удобного способа описания последовательности работы с данными, Mediapipe обещает сбилдить ваше решение в качестве библиотек для iOS и Android. И это очень актуально!
Запускать ML-модели и процессить видео на конечных устройствах сложно. Если самостоятельно реализовывать обработку видео в реальном времени и «в лоб» запускать нейросети, то легко попасть в ситуацию неконтролируемого роста объема работы: приложение «утечет» по памяти, появятся глюки или скопится очередь из кадров.
Поэтому Mediapipe пользуются компании-разработчики мобильных приложений. У них нет ресурсов на самостоятельный процессинг видео. Да и зачем их искать, если можно взять готовую компоненту из Mediapipe и «приклеить» к ней собственную модель или иную необходимую фишку? Фреймворк реализует графическое преобразование видео, позволяет комбинировать ML-ные компоненты с механической обработкой и на выходе получать кросс-платформенное решение для мобильных и десктопных устройств.
Mediapipe компенсирует ограниченную функциональность мобильных устройств и на выходе дает видео приемлемого качества. Идеально, не так ли?
Задачи, которые решает Mediapipe из коробки
Как проще всего познакомиться с Mediapipe? Запустить один из уже разработанных ими примеров, а потом попробовать немного его поменять.
Что умеет делать Mediapipe из коробки:
распознавать лица и делать Face Mesh;
определять радужные оболочки глаза: зрачки и контуры глаза;
находить руки, ноги, определять позы;
сегментировать волосы и селфи;
запускать модель детекции и трекать ее предсказания;
мгновенно идентифицировать движения;
Objectron: определять 3D-объекты по 2D-изображениям;
KNIFT: сопоставлять признаки на основе шаблонов;
AutoFlip: конвейер автоматической обрезки видео.
Сразу замечу, что стандартного двухстадийного пайплайна «детекция + классификация» в нем не реализовано. Но его несложно собрать из готовых модулей обработки данных. В общем и целом, можно даже написать и свой собственный модуль :)
Из чего состоит Mediapipe
Mediapipe состоит из трех основных структурных компонентов: калькуляторов (вершин графа), входных/выходных пакетов и графов вычислений.
Калькуляторы
Калькулятор — вершина графа, код, выполняющий трансформацию над входными данными и отдающий выходные.
Чтобы создать калькулятор, нужно отнаследоваться от CalculatorBase и реализовать четыре метода:
GetContract(). Проверяем типы входящих и выходящих данных на соответствие заранее заданным. Практика показала, что если этот метод не реализовывать, ничего не сломается.
Open(). Инициализируем калькулятор при запуске.
Process(). Запускаем вычисление ноды, используя понятие контекст — некой сущности, в которой записаны текущие переменные в текущем моменте запуска пайплайна.
Close(). Освобождаем ресурсы калькулятора.
Например, если стоит задача прогнать некую модельку классификации на входном изображении, то нужно:
Проинициализировать эту модель в Open().
В Process() из контекста вытащить входное изображение и использовать на неё модель, получить предсказание и отдать его в выход.
Освободить ресурсы в Close().
(Калькулятор: пример кода для модели классификации на входном изображении).
class ClassificationCalculator : public Node {
public:
// ожидаем на входе либо тип Image, либо тип GpuBuffer
// на выходе отдаем список из объектов Classification
static constexpr Input<OneOf<mediapipe::Image,
mediapipe::ImageFrame>> kInImage{"IMAGE"};
static constexpr Input<GpuBuffer>::Optional
kInImageGpu{"IMAGE_GPU"};
static constexpr Output<ClassificationList>
kOutClassificationList{"CLASSIFICATIONS"};
MEDIAPIPE_NODE_CONTRACT(kInImage, kInImageGpu,
kOutClassificationList);
absl::Status Open(CalculatorContext* cc) override;
absl::Status Process(CalculatorContext* cc) override;
absl::Status Close(CalculatorContext* cc) override;
absl::Status SomeCalculator::Open(CalculatorContext* cc) {
// вот здесь загружаем модель в private переменную
return absl::OkStatus();
}
// Copied some code from image_to_tensor_calculator.cc
absl::Status SomeCalculator::Process(CalculatorContext* cc) {
// проверяем, что в данном контексте имеется вход
if ((kInImage(cc).IsConnected() && kInImage(cc).IsEmpty()) ||
(kInImageGpu(cc).IsConnected() && kInImageGpu(cc).IsEmpty())) {
// Timestamp bound update happens automatically.
return absl::OkStatus();
}
ASSIGN_OR_RETURN(auto image_ptr, GetInputImage(cc));
ASSIGN_OR_RETURN(auto classification_list_ptr, GetClassificationList(cc));
// прогоняем модель на картинке image_ptr
// записываем результаты в выходные данные
auto out_classification_list = absl::make_unique<ClassificationList>();
Classification* classification = out_classification_list->add_classification();
classification->set_index(0);
classification->set_score(0.8);
classification->set_label("class0")
classification->set_display_name("Dog");
____kOutClassificationList(cc).Send(std::move(out_classification_list));
return absl::OkStatus();
}
absl::Status SomeCalculator::Close(CalculatorContext* cc) {
// освобождаем ресурсы (удаляем модель из памяти)
return absl::OkStatus();
}Входные/выходные пакеты и графы вычислений
Граф — совокупность входных данных и вершин (калькуляторов), реализующий какой-то пайплайн вычисления. Граф определяется в формате TensorFlow Graph Text и может быть описан в виде файла .pbtxt.
В калькулятор можно подключать опции, описываемые протоколом Protobuf, которые можно изменять. Так можно указывать размер изображения для ресайза.
Например, картинку сверху можно было бы описать следующим графом:
input_stream: "input1"
input_stream: "input2"
input_stream: "input3"
output_stream: "output1"
output_stream: "output2"
node {
calculator: "CalculatorCalculator1"
input_stream: "INPUT_TAG1:input1"
input_stream: "INPUT_TAG2:input2"
input_stream: "INPUT_TAG3:input3"
output_stream: "OUTPUT_TAG:output"
}
node {
calculator: "CalculatorCalculator2"
input_stream: "INPUT_TAG:output"
output_stream: "OUTPUT_TAG1:output1"
output_stream: "OUTPUT_TAG2:output2"
}В Mediapipe принят следующий синтаксис записи данных: TAG:variable_name. Теги позволяют калькулятором детектить нужные ему входные переменные, получаемые из контекста. Например, в коде выше с калькулятором мы записываем тег или как «IMAGE», который хранится на CPU, или как «IMAGE_GPU», который хранится на GPU. В зависимости от того, какие переменные были поданы калькулятору, будет делаться Process.
Mediapipe даже имеет собственный сервис, в который можно загрузить граф и посмотреть красивую визуализацию: https://viz.mediapipe.dev/.
Не так подробно и более схематично умеет визуализировать графы и Netron (netron.app).
Стандартно написанные калькуляторы
В Mediapipe определено много стандартных типов и калькуляторов, реализующих разную логику работы. В качестве примеров типов можно указать видео, изображения, аудио, тексты, временные ряды. Далее идут тензоры — принятые типы в машинном обучении, с помощью которых можно получать объекты результатов предсказаний модели: Detections, Classifications, Landmarks (кейпоинты) — всё это тоже типы. Примеры можно посмотреть здесь.
Калькуляторы позволяют совершать над этими типами разного рода трансформации, преобразовывающие данные из одного типа в другой.
Все калькуляторы (и подграфы) определены здесь. Большую часть из них можно найти в папке calculators/ или gpu/. Какой-либо внешней документации по ним нет, но сам код почти везде задокументирован: рядом с определением калькулятора в виде комментариев описывается сам калькулятор. Помимо этого, рядом с калькулятором (записанные в файлах .cc и .h) может лежать и .proto-файл — они определяют так называемые опции калькуляторов и записаны в формате Protobuf.
В качестве примера можно рассмотреть задачу детекции объектов на входном изображении некоторой ML-моделью. Что для этого нужно:
Получить входную картинку и преобразовать её размеры в 640х640 (пусть наша модель ожидает входной тензор такого размера — эти числа нужно указывать в опцию соответствующего калькулятора).
Превратить сжатую картинку в тензор.
Прогнать этот тензор через модель и получить сырые предсказания.
Сырые предсказания сконвертировать в понятный для Mediapipe тип Detections.
Получить этот тип на выходе графа.
В Mediapipe все эти действия превращаются в элегантный граф — последовательность вершин с определенными типами входных/выходных данных и параметрами вершины:
# Задаем входящий и выходящий потоки (выходящий, например, мы можем ловить коллбеком в приложении)
input_stream: "image"
output_stream: "detections"
# Входное изображение превращается в изображение размера 640х640 с
# сохранением aspect ratio. Остальная часть картинки заливается нулевыми
# пикселями (преобразование, имеющее название леттербоксинг).
# Затем изображение трансформируется в тензор, требуемый моделью TensorFlow Lite.
node {
calculator: "ImageToTensorCalculator"
input_stream: "IMAGE:image"
output_stream: "TENSORS:input_tensor"
output_stream: "LETTERBOX_PADDING:letterbox_padding"
options: {
[mediapipe.ImageToTensorCalculatorOptions.ext] {
output_tensor_width: 640
output_tensor_height: 640
keep_aspect_ratio: true
output_tensor_float_range {
min: 0.0
max: 255.0
}
border_mode: BORDER_ZERO
}
}
}
# Здесь просто делаем инференс tflite модельки
# Есть поддержка делегатов gpu, tflite, nnapi и тп
node {
calculator: "InferenceCalculator"
input_stream: "TENSORS:input_tensor"
output_stream: "TENSORS:output_tensor"
options: {
[mediapipe.InferenceCalculatorOptions.ext] {
model_path: "model.tflite"
delegate { gpu {} }
}
}
}
# Декодирует тензоры, полученные моделью TensorFlow Lite в выход
# типа Detections. Каждый Detection -- это минимум Label, # Label_id, Score и Bounding Box,описывающий обнаруженный объект # в формате (xmin, ymin, xmax, ymax), где координаты нормализованы # в диапазоне [0,1]
node {
calculator: "TensorsToDetectionsCalculator"
input_stream: "TENSORS:output_tensor"
output_stream: "DETECTIONS:raw_detections"
options: {
[mediapipe.TensorsToDetectionsCalculatorOptions.ext] {
num_classes: 1
num_boxes: 2
num_coords: 4
min_score_thresh: 0.5
}
}
}
# Превращает полученные моделью координаты детекции относительно ресайзнутой
# леттербоксингом картинки в координаты детекции относительно # оригинального изображения
node {
calculator: "DetectionLetterboxRemovalCalculator"
input_stream: "DETECTIONS:raw_detections"
input_stream: "LETTERBOX_PADDING:letterbox_padding"
output_stream: "DETECTIONS:detections"
}
Полученные Detections можно отдавать обратно в контекст и работать с ними дальше (например, сохранить их в файлы с разметкой или отрисовать на полученной вначале картинке).
Какой граф делали мы
Теперь давайте перейдем от общего к частному.
Перед нами стояла задача перенести двухстадийный пайплайн (детекция + классификация) в риалтайм для мобильных телефонов. Причем после модели классификации в некоторых случаях была необходимость запускать C++ код, обрабатывающий изображение определенным алгоритмом.
Как модель детекции мы использовали Yolov5S, как модель классификации — MobileNetV2.
Очень быстро выяснились следующие вещи:
кастомные графы почему-то в Python не работают;
General community discussion around MediaPipe скорее мертва, чем жива;
для поддержки в Github Issue в целом нормально копипастить фразу «не могли бы вы предоставить логи об ошибке» (после того, как ты их предоставил, естественно) и в целом игнорировать твой вопрос месяцами.
С другой стороны, их примеры вполне работали на CPP, а библиотеки успешно конвертировались для iOS и Android, так что мы решили не сдаваться.
В итоге мы реализовали следующий граф:
Какую еще полезную информацию мы можем сообщить миру?
В самом Mediapipe реализованы калькуляторы, выполняющие инференс только на Tensorflow. Мы лично пользовались только tflite, но якобы есть поддержка еще saved_model и frozen_graph. Если вы, например, любитель PyTorch, будьте готовы либо конвертировать вашу модель в tflite, либо писать собственный калькулятор для инференса.
Разработчики без стеснения пушат в мастер, так что, если хотите быть уверенными в успешности сборки, фиксируйтесь на каком-нибудь теге.
Будьте готовы к сражениям с разными делегатами, OpenGl и OpenCl, если ваша модель отличается от некоего «стандартного набора». Без шуток, не так просто найти и сконвертить модель детекции, запускающуюся в Mediapipe с поддержкой ГПУ.
Писать кастомные калькуляторы достаточно просто; благо, в репозитории есть много примеров, которые можно взять за основу.
Не игнорируйте setup_{}.sh-скрипты — они корректно настраивают OpenCV и Android SDK/NDK для Mediapipe.
Ну, и в общем и целом, описанный выше граф вполне реально реализовать в Mediapipe и сбилдить из него библиотеки для Android и iOS, работающие с приличным FPS.
Создание библиотек с графом Mediapipe для iOS и Android
Для сборки решения Mediapipe использует Bazel. Как и с построеним графа, с процессом билда мы тоже решили пойти от простого к сложному. Итак, что нужно сделать, чтобы сбилдить решение с кастомным графом для Linux:
Написать этот граф и, например, положить в его новую папочку mediapipe/graphs/custom_graph/cutom_graph.pbtxt (по аналогии с этим решением). К слову, начиная с версии 0.8.11 не обязательно всё складывать внутрь //mediapipe.
В папке с этим графом в BUILD прописываем вручную все нужные зависимости для него (в основном — калькуляторы) в cc_library (для примера из предыдущего пункта).
Положить в папочку mediapipe/examples/desktop/custom_example/ BUILD файл, в котором определить cc_binary, зависящий от калькуляторов из предыдущего пункта и CPP кода, запускающего этот граф (пример).
Сбилдить командой:
bazel build -c opt --define MEDIAPIPE_DISABLE_GPU=1 mediapipe/examples/desktop/custon_example:{cc_binary name}После этого сбилдится cc_binary (в логах вы увидите путь к нему), и останется только запустить этот граф командой:
GLOG_logtostderr=1 {cc_binary_path} --calculator_graph_config_file=mediapipe/graphs/custom_graph/custom_graph.pbtxt \
--input_side_packets=input_video_path=<input video path>,output_video_path=<output video path>Вроде бы ничего сложного, и, в принципе, оригинал этой инструкции есть и в документации.
Важно еще понимать, что граф и его запуск — это разные вещи. В примере выше мы запускаем граф, который на вход ожидает путь к видео (и указываем его в аргументе --input_side_packets=input_video_path), а сам же граф запускается внутри кода C++.
Последнее, на что хотелось бы обратить внимание при Linux-сборке — сторонние зависимости. В одном из наших калькуляторов использовалась внешняя библиотека, написанная на C++. Ее мы добавляли по аналогии с OpenCV: заранее сбилдили, создали в WORKSPACE-файле new_local_repository (у OpenCV — http_archive ) и записали в папочку third_party/ соответствующий BUILD файл. И указали эту зависимость в BUILD-файле калькулятора.
Предположим, что у нас не возникло ошибок на этом этапе и можно двигаться дальше!
Во-первых, следует определиться с кодом приложения, в котором вы этот граф собираетесь запускать. Для Android уже есть два примера, которые рекомендует официальная документация (раз и два). Для iOS есть инструкция и шаблон, так что на этом этапе iOS-девелопер все-таки понадобится.
Не забудьте разобраться, какое именно изображение поступает в граф — Image или Image_GPU (это зависит от кода приложения). Проставьте делегаты в зависимости от того, какие ускорители вы можете/хотите использовать на устройстве. И самое главное, проверьте на соответствие input_stream и output_stream в графе и приложении.
Далее уже идут особенности билдов, и для Android и iOS они, естественно, различаются.
Android
Для Android, как минимум не требуется XCode и Mac.
В целом, существующая инструкция описывает максимально подробно процесс билда AAR под Android. После AAR необходимо еще сбилдить граф в .binarypb. Для этого нужно в BUILD-файл добавить граф и выполнить команду bazel build -c opt mediapipe/graphs/custom_graph:custom_graph.
Потом поместить все модели, txt-файлы с классами и сам граф в app/src/main/assets в приложении (мы брали первый пример), а в app/src/main/java/com/example/myfacedetectionapp/MainActivity.java правильно расставить названия для input_stream и output_stream, а также сбилженного графа.
Если дополнительных библиотек к Mediapipe подключать не требуется, то вас можно поздравить! Однако, как я рассказывал выше, сторонняя зависимость в нашей сборке была. Как и в Linux, добавляли мы ее по аналогии с OpenCV. Сбилдили СPP-код заранее для arm-v8a и armeabi-v7 (здесь нам помогла эта статья). Заполнили WORKSPACE, добавили в third_party/ и… ничего не произошло!
Несмотря на то, что зависимость была проставлена в BUILD-файле калькулятора, приложение в Android AAR ее не видело. Все оказалось чуточку сложнее: AAR билдится с помощью написанного разработчиками специального Bazel-правила, которое вы можете найти здесь.
Когда мы добавили native.cc_library с нашей зависимостью туда (помним, что всё по аналогии с OpenCV), по ощущениям стало лучше, но Bazel почему-то отказывался находить сбилженные под Android исходники. Добавление флагов --linkopt=-L{path} с путем к исходникам эту проблему решило — и приложение стало работать корректно.
iOS
Вот и мы подобрались к iOS. Если вы нашли какой-нибудь Mac и успешно установили на него Xcode — дело за малым! Конкретно в нашем кейсе была необходимость создать не приложение, а библиотеку, которая бы использовалась в приложении заказчика. Снова создаем BUILD-файл под наш фреймворк. В качестве примера можно взять вот этот. Мы заменяли правило ios_application на ios_framework, для него необходимо сделать Header-файл с методами библиотеки и подсунуть в аргумент hdrs. Настроить bundle_id поможет эта инструкция. Далее останется только правильно расставить ссылки на граф, модели и txt-файлы с классами для моделей и запустить билд-командой bazel build --config=ios_arm64 :FrameworkName --verbose_failures из папки с BUILD-файлом.
Билдить отдельно граф или что-либо еще дополнительно не нужно — в логах Bazel покажет путь к ZIP-архиву, в котором уже лежат все необходимые файлы.
Те, кто внимательно следили за нашей историей, наверное, задаются вопросом: а как мы здесь добавили стороннюю зависимость? Так вот, здесь впервые на нашем пути нам крупно повезло: у нашей стороней завимости существовал Wrapper для iOS, который мы просто корректно прилинковали в аргумент deps правила ios_framework (и положили в папочку с приложением).
Заключение
Надеюсь, вам было интересно! Я не пытался рассказать о Mediapipe всё и не претендовал на полное исчерпывающее русскоязычное руководство по Mediapipe. Посвятив вас в контекст в первой половине статьи, во второй я пытался рассказать о нашем пути к инференсу в риалтайме, скрасив его полезными советами, ссылками и трюками. По крайней мере вам не придется набивать те же шишки, что уже набили мы.
Если у вас нет большого желания или ресурсов ввязываться в сложный мир риалтаймового инференса ML-моделей на iOS и Android, я рекомендую вам рассмотреть Mediapipe. Безусловно, это не панацея, время на работу с ним вы все равно потратите. Однако за счет большого количества уже готовых и оптимизированных модулей для работы с данными или моделями, это будет явно проще и быстрее, чем писать все с нуля. А учитывая существование этой статьи — это пройдет почти безболезненно ;)