Функциональность виртуального фона в виде кроссплатформенного SDK: исследования, разработка, запуск на разных платформах

Введение

В связи с пандемией и масштабным переходом многих сотрудников на удаленную работу резко возросла популярность мессенджеров, платформ для видеозвонков, стриминговых сервисов которые помогают организовать эффективную удаленную работу.  

Сотрудникам пришлось осваивать технологии и подстраиваться под работу из дома, где, зачастую нет хорошего освещения (для качественного видео стриминга), красивого одноцветного фона, дети и родственники регулярно появляющиеся в кадре и отвлекающие внимание. Появилась функциональность виртуального фона (также известная как замена или размытие фона в видеозвонках, или перерождение хромакея в умный хромакей, который не требует однотонного фона), это когда автоматически в режиме реального времени в видео заменяется фон на заранее выбранную картинку или размывается чтобы не было видно что происходит за человеком. Одни из первых такую функциональность ввели компании Zoom, Skype, Google.

Популярность росла, пользователи уже не просто интересовались, а требовали наличия виртуального фона во всех программах связанных с видеозвонкам. Собственная разработка - долго  и дорого, а также требует определенных компетенций команды, а готовых решений на рынке просто не было. В этой статье мы расскажем про путь создания собственного кроссплатформенного решения позволяющего заменять/вырезать фон в видео с веб камер в режиме реального времени.

Существующие подходы

Традиционно для удаления фона применяли технику «зеленого экрана», когда пользователь ставит позади себя однотонное полотно, цвет которого считается фоновым. В таком случае достаточно вырезать пиксели, соответствующие основному цвету, и некоторый объем вокруг этой точки, чтобы исключить «протекание» фона из-за неравномерности освещения экрана или шумов камеры. Такое решение позволяет получить хорошую картинку, но требует от пользователя определенных усилий для подготовки места съемки. Это не будет большой проблемой для людей, регулярно ведущих трансляции, но для повседневного или рабочего общения в чатах хотелось бы обойтись вообще без подготовки.

При удалении фона нам необходимо классифицировать каждый пиксель как “человек” или “фон”, что полностью совпадает с постановкой задачи сегментации изображения. И тут на сцену выходят свёрточные нейронные сети. С момента их дебюта на соревновании по классификации изображений ImageNet Large Scale Visual Recognition Challenge прошло много времени и им покорились многие задачи компьютерного зрения. Именно методы на основе нейронных сетей уверенно держат верхние позиции во всех бенчмарках по сегментации изображений, что делает их очевидным кандидатом для решения задачи определения фона.

Перед нами стояло две основных задачи. Во-первых, подготовить набор данных для решения именно нашей проблемы. Во-вторых, разработать архитектуру, пригодную для запуска на достаточно слабом железе, включая ограниченные по мощности процессоры ноутбуков и мобильные устройства на Android и iOS.

Существующие наборы данных

Для того, чтобы лучше понять требования к обучающим данным, мы изучили существующие датасеты для проблем, похожих на нашу.

MS COCO

Стандартный датасет, использующийся как бенчмарк в работах по детекции и семантической сегментации изображений. Они содержит достаточно большое количество примеров с людьми, но эти изображения далеки от тех, что обычно встречаются при использовании веб-камер. Качество масок тоже оставляет желать лучшего, особенно на границах контура человека.

Supervisely Person

Набор данных, созданный специально для определения людей на изображениях. Источником изображений послужили фотографии с фотостоков. Хотя визуальное качество масок в этом датасете достаточно высокое, исходные изображения слишком далеки от того, что обычно можно увидеть на кадрах с веб-камер: фон размыт и фокус сделан на человеке, сцена хорошо освещена, разрешение снимка высокое, шума практически нет. Обученные на таком датасете модели не подходят для решения нашей задачи.

AISegment.com - Matting Human Datasets

Датасет для решения очень похожей на нашу проблемы - portrait matting. Основных отличий от нашей постановки два. Во-первых, вместо жесткой классификации тут требуется определить значение прозрачности в диапазоне [0; 1], что позволяет более качественно разделять фон и человека в сложных ситуациях. Например, четко определить границы волос довольно трудно и использование мягкого несколько улучшает визуальные характеристики результата. Во-вторых, авторы датасета ограничились портретными фотографиями с характерным ограничением разнообразия поз в них. В частности, обученные на этом датасете модели не могут корректно определять руки.

Сбор собственного набора данных

После изучения датасетов и нескольких безуспешных попыток обучить на них модель с приемлемым качеством мы решили, что лучшим способом добиться результата будет обучить модель на картинках напрямую из целевого домена. Хотелось получить датасет напрямую с веб-камер, чтобы визуальные характеристики изображений, которые модель увидит во время обучения, совпадали с кадрами, на которых будет производиться вывод. Первыми “жертвами” стали сотрудники нашей компании, которых мы попросили на добровольной основе записать для нас небольшие ролики с демонстрацией различных положений.

Количества и разнообразия собственных данных не хватило для того, чтобы значительно улучшить результаты модели. Погрустив над тем, что расширить датасет путем сбора данных с реальных людей не представляется экономически возможным, мы решили обратиться к самой большой коллекции видео на Земле – YouTube.

Мы собрали около 300 роликов из видео-блогов и записей стримов, достаточно похожих по качеству съемки на целевой домен. При сборе мы старались больше внимания уделять роликам, снятым в плохих условиях на плохие камеры, и следили, чтобы профессионально настроенные сетапы были в меньшинстве. Сбор данных и их разметка шла итеративно по мере улучшения модели, и каждый новый раунд разметки улучшал результат даже без изменения модели. На данный момент датасет состоит из примерно 12000 картинок с качественными масками.

Кроме самого набора данных, важную роль в улучшении результатов сыграл подбор существующих и разработка собственных аугментаций. Для построения пайплайна мы использовали библиотеку Albumentations. В ней реализован большой набор аугментаций и предлагаются абстракции для построения цепочек аугментаций (с возможностью ветвления), каждая из которых активируется с определенной вероятностью. Были использованы цветовые (смещение цветов в RGB, гауссов аддитивный и мультипликативный шумы, перестановка каналов, симуляция шума сенсора камеры, изменения яркости и насыщенности, блюр и т.д.) и пространственные (поворот, отражение) преобразования. Мы самостоятельно реализовали аугментации для замены фона и наложения синусоидального шума. Последний помогал при наличии мерцания на изображении из-за особенностей поведения сенсора камеры при включенном искусственном освещении.

Аугментации не отразились существенно на метриках, зато значительным образом улучшили стабильность поведения модели «в естественной среде».

Модель

Хотя данные в нашем случае были ключом к получению качественной сегментации, модель все еще играет важную роль. Особенно когда важна максимальная производительность и просто взять готовую сеть целиком из статьи нельзя. В основу архитектуры мы положили схему“encoder-decoder”. Первая часть сети (encoder) поэтапно извлекает признаки из исходного изображения, снижая пространственное разрешение карты признаков в 2 раза на каждом этапе обработки. Итоговое разрешение карты признаков может быть в 32 раза меньше исходного и предсказать достаточно мелкие детали на таком маленьком разрешении затруднительно. Поэтому вторая часть сети использует признаки из исходного изображения и поэтапно восстанавливает разрешение до исходного.

Первая версия архитектуры представляла из себя энкодер-декодер с использованием миниатюрной DenseNet в качестве энкодера и последовательности слоёв билинейной интерполяции. Хотя производительность сети нас устраивала, модель могла предсказывать только общий контур сидящей смирно фигуры. Сложный фон или поднятые руки сбивали модель с толку.

Эксперименты с сетью UNet показали высокую точность, однако скорость вывода даже на GPU была слишком низкой. Мы попробовали разные варианты наивного масштабирования: уменьшение разрешения входа, уменьшение глубины карт признаков, удаление стадий декодирования, но качество падало быстрее, чем увеличивалась скорость модели. В итоге мы заменили обычные свертки на depthwise-separable варианты. Это увеличило скорость до приемлемых показателей без существенной потери в качестве. Такая модель первой преодолела рубеж качества, достаточного для деплоя и демонстрации.

Параллельно мы оценивали легковесные модели на основе “dilated”свертки, однако получить быстрые модели с достаточно высоким качеством в короткие сроки не получалось, и мы забросили это направление исследований полностью.

Вдохновением для последующих версий модели послужила работа по оценке глубины с помощью UNet-подобной сети, в которой авторы предложили ассиметричную архитектуру с легким энкодером и тяжелым декодером. Еще одной важной идеей было получение предсказания в меньшем разрешении по сравнению с входной картинкой. Мы взяли эти две идеи и реализовали свою собственную модель, использовав легковесные боттлнеки и несколько упростив декодер. От расчета ошибки на нескольких разных масштабах отказались, так как не наблюдали его влияния на качество. Вместо него мы добавили в сеть дополнительную ветку, которая использовала последнюю карту признаков для расчета ошибки только на границах маски. Она позволила немного улучшить стабильность предсказания на границах.

Деплой модели

После обучения модели её необходимо подготовить к разворачиванию в продакшене. На первых этапах предполагалось развертывать модель в составе нативных приложений на платформе Windows с возможность запуска нейронки на CPU и на GPU. К сожалению, деплоить модель напрямую с помощью PyTorch не представляется возможным из-за маленькой скорости инференса на CPU. Кроме того, поддержка GPU Windows ограничена карточками nvidia, бинарь библиотеки весит около 3 Гб и на стороне пользователя требуется установка CuDNN, так как её нельзя распространять в составе собственных решений.

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

WinML

В первую очередь наш взгляд упал на WinML – библиотека для вывода нейронок от Microsoft, интегрированная в Windows 10 начиная с версии 2004. Оно позволяет запускать модели на CPU и GPU, причем любых вендоров. Библиотека понимает только сериализованные в ONNX модельки. Благо, что в PyTorch конвертация модели делается в пару строчек.

С помощью WinML мы смогли быстро получить первые рабочие прототипы и продолжить исследование доступных средств для запуска нейронных сетей.

OpenVINO

Фреймворк от Intel для оптимизации и вывода нейронок на процессорах x86 и интегрированных в процессоры GPU. OpenVINO реализует большое количество операторов, особенно необходимых для компьютерного зрения, так что напороться на несовместимость шансы не так велики. Все это великолепие сопровождается отличной документацией, тулами для тестирования и бенчмаркинга моделей, большим зоопарком уже готовых для использования моделей. Напрямую из PyTorch модель экспортировать нельзя, сначала нужно сконвертировать её в ONNX, Разумеется, поддерживается интеграция в C++ проекты.

После демонстрации десктопного SDK мы получили запрос на добавление возможности запуска модели на мобильных устройствах и в веб-браузерах.

MediaPipe

Для начала мы попробовали библиотеку MediaPipe, созданную в недрах Google. Основные сценарий её использования это реализация пайплайнов потоковой обработки видео и аудио в мультимедийных приложениях. Вы описываете источники данных (камера, микрофон), последовательность манипуляций над данными в виде графа, а фреймворк берет на себя заботы по обслуживанию процесса обработки. За вывод нейронок отвечает TensorFlow Lite. Библиотека поддерживает десктопные и мобильные платформы, но с оговорками, о которых мы поговорим ниже.

Так как этот фреймворк не поддерживает импорт из ONNX и PyTorch, сначала нужно решить вопрос перегонки модели в понятный ему формат. После изучения опыта хождения по граблям других людей, мы пошли по не самому простому пути, но с наиболее предсказуемым результатом – повторной реализации модели на TensorFlow и импорта весов из PyTorch. Процесс этот несложный и достаточно механический. Требуется только помнить о том, что TF по умолчанию использует NHWC формат тензоров, а PyTorch – NCHW. Эта несовместимость лечится простым транспонированием весов сверточных слоёв при импорте. Портировав таким образом модель на TensorFlow и убедившись в правильности процесса путем сравнения ответов сети, мы преобразовали сеть в формат TensorFlow Lite и принялись за запуск MediaPipe на мобильных девайсах.

На Android нам пришлось самостоятельно реализовать несколько операций (вроде склейки предыдущей маски с текущим изображением) для обработки изображений на OpenGL, но в целом запуск модели прошел гладко и занял около двух дней. На не самых продвинутых устройствах мы смогли получить уверенные 25-30 FPS, что нас несомненно обрадовало. К тому же, к тому времени у нас были пути дальнейшей оптимизации модели.

К сожалению, с iOS все пошло не так гладко. Главной проблемой стала политика Apple, из-за которой мобильные устройства так и не получили поддержку вычислительных шейдеров OpenGL и пришлось экспериментировать с конвертацией текстур из OpenGL в Metal и обратно. Быстро реализовать это не получилось, поэтому мы отказались от использования MediaPipe на iOS.

CoreML

CoreML является родным для экосистемы Apple фреймворком для запуска нейронных сетей. Он поддерживает все железки этой компании, включая встроенные в SoC NPU (Neural Processing Unit). Процесс портирования оказался гладким, так как Apple позаботилась о путях прямой конвертации модели из PyTorch в CoreML. В итоге портированная модель показала около 30 FPS.

Web

Многие приложения сейчас с самого начала делаются как веб-ориентированные (тем более есть замечательный WebRTC), поэтому упускать возможность предложить им нашу модель не хочется.  Мы попробовали три опции: ONNX.js, TensorFlow.js и TensorFlow Lite. Опыт, к сожалению, был довольно болезненный и удовлетворительно на данный момент работает только одно решение.

Самый простой путь к запуску модели, в теории, представляет ONNX.js, который позволяет подгружать ONNX-модели и исполнять их на GPU или CPU. Картину портит только то, что библиотека не так активно поддерживается и поддерживает импорт ограниченного набора операторов и только из старых версий ONNX IR. Из-за этого могут быть проблемы при экспорте моделей из новых версий PyTorch. Например, в нашем случае проблемой стали слои транспонированной свертки и апсемплинга. Быстро побороть их не получилось мы пошли исследовать другие варианты.

TensorFlow.js представляет собой полноценный фреймворк для обучения и запуска нейронных сетей в JavaScript. В частности, можно импортировать графы из обычного Tensorflow. Мы взяли модель, подготовленную для мобильной версии, и без проблем смогли запустить модель на GPU через WebGL. Проблемы возникли с выводом на CPU при использовании XNNPACK – операция транспонированной свертки давала неверные результаты.

Последним исследованным вариантом был TensorFlow Lite, скомпилированный в WebAssembly с помощью emscripten. WebAssembly – это стандарт, определяющий независимый от платформы байткод, который может исполняться в песочнице. Он был создан в первую очередь для ускорения программ, которым недостаточно производительности JavaScript в браузерном окружении. Тулчейн emscripten позволяет компилировать программы на C/C++ в WebAssembly. К нашему счастью, TensorFlow Lite как раз написан на C++, что позволяет запустить его напрямую в браузере. Скомпилированная версия TensorFlow Lite быстро завелась и показала достаточную для деплоя производительность. Кроме того, такой способ деплоя открывает возможность использовать другие библиотеки, например, OpenCV для пре- и постпроцессинга изображений, чем мы с радостью воспользовались.

Главным препятствием на пути интеграции стало отсутствие встроенных средств шифрования моделей во всех рассмотренных нами фреймворках. Именно по этой причине мы остановились на варианте с компиляцией TensorFlow Lite, так как он позволяет реализовать шифрование модели самостоятельно.

Что получилось

Сейчас у нас подготовлены четыре версии SDK для всех основных платформ с быстрой, стабильной и качественной сегментацией фона. Есть интеграции с коммерческими проектами. Мы продолжаем развивать это направление. Вот сайт продукта: https://tomsksoft.com/solutions/cross-platform-virtual-backgroung-sdk

А вот пример работы самой сегментации на десктопе: