Что такое Mediapipe, или Как мы делали инференс в реальном времени на Android и iOS

Привет! Я Андрей Татаринов, директор 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 и реализовать четыре метода:

1. ​​GetContract(). Проверяем типы входящих и выходящих данных на соответствие заранее заданным. Практика показала, что если этот метод не реализовывать, ничего не сломается. 

2.  Open(). Инициализируем калькулятор при запуске.

3. Process(). Запускаем вычисление ноды, используя понятие контекст — некой сущности, в которой записаны текущие переменные в текущем моменте запуска пайплайна.

4. 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, так что мы решили не сдаваться.

В итоге мы реализовали следующий граф:

Какую еще полезную информацию мы можем сообщить миру?

1. В самом Mediapipe реализованы калькуляторы, выполняющие инференс только на Tensorflow. Мы лично пользовались только tflite, но якобы есть поддержка еще saved_model и frozen_graph. Если вы, например, любитель PyTorch, будьте готовы либо конвертировать вашу модель в tflite, либо писать собственный калькулятор для инференса.

2. Разработчики без стеснения пушат в мастер, так что, если хотите быть уверенными в успешности сборки, фиксируйтесь на каком-нибудь теге.

3. Будьте готовы к сражениям с разными делегатами, OpenGl и OpenCl, если ваша модель отличается от некоего «стандартного набора». Без шуток, не так просто найти и сконвертить модель детекции, запускающуюся в Mediapipe с поддержкой ГПУ.

4. Писать кастомные калькуляторы достаточно просто; благо, в репозитории есть много примеров, которые можно взять за основу.

5. Не игнорируйте setup_{}.sh-скрипты — они корректно настраивают OpenCV и Android SDK/NDK для Mediapipe.

6. Ну, и в общем и целом, описанный выше граф вполне реально реализовать в 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. Безусловно, это не панацея, время на работу с ним вы все равно потратите. Однако за счет большого количества уже готовых и оптимизированных модулей для работы с данными или моделями, это будет явно проще и быстрее, чем писать все с нуля. А учитывая существование этой статьи — это пройдет почти безболезненно 😉