Как разрабатывали модели замены фона для видеосозвонов в SaluteJazz

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

На нашей конференции I'ML спикеры из SberDevices Дмитрий Балиев и Давид Нурдинов рассказали о том, как занимались этой задачей для сервиса SaluteJazz (ранее известном как SberJazz). А теперь мы для Хабра сделали текстовую версию их доклада. Интересно может быть и тем, кто занимается сегментацией в ML, и тем, кто в целом работает с Computer Vision, и просто тем, кому любопытно узнать «что стоит за кнопкой замены фона».

Вступление

Дмитрий Балиев: Всем привет! Меня зовут Дима, я лид команды, которая разрабатывает системы компьютерного зрения для платформы SaluteJazz. Сегодня со мной Давид, он разрабатывал непосредственно модели замены фона. 

О чем вкратце доклад: как получить нужные модели, чтобы в примере на открывающей картинке перейти от варианта слева к варианту справа.

Надеюсь, у меня получилось вас заинтересовать. Поехали!

Что такое замена фона?

Давид Нурдинов: Это действительно очень хороший вопрос. Все довольно просто. Задача в том, чтобы оставить человека, но заменить один фон на другой.

Как достичь желаемого эффекта? Тут нам поможет формула для замены фона.

Все подписано на изображении. Если говорить формально, то вся задача — поиск α, которая подскажет нам, где человек, а где — фон. 

Выглядит просто, но мы живем в реальном мире, где у каждого пользователя свое устройство, и это накладывает свои ограничения:

• Скорость выполнения. Каждый кадр модель должна отрабатывать за определенное время. В данном случае это время меньше 25 мс. 

• Вес модели. Здесь не было строгих ограничений, но вообще чем меньше, тем лучше. 

• Достойное качество. Не хотелось бы, чтобы фон проглядывал через нашу картинку.

Дмитрий: Для решения задачи нам нужно определиться со следующим:

1. Подход к решению. Какую вообще задачу машинного обучения мы решаем?

2. Набор данных.

3. Архитектура модели.

4. Метрики для оценки результата.

Подходы к решению задачи

Маттинг

Давид: α в формуле выше рассматривается как маска прозрачности, где каждый пиксель — шкала от 0 до 1, которая показывает прозрачность пикселя. Рассмотрим пример:

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

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

Дмитрий: Что мы хотим в идеале получить от маттинга? Картинку, когда мы берем изображение, убираем все, что относилось к фону, и останется передний план, который можно будет перенести на новое изображение. Если делать это наивно, то просто возьмем альфа-канал, умножим его на изображение и получим что-то похожее на передний план. Но на границе, где альфа-канал не равен 1 или 0, у нас появятся артефакты.

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

Семантическая сегментация

Давид: Он действительно есть. Суть семантической сегментации в следующем: если раньше мы рассматривали эту маску как некий показатель прозрачности пикселя, то теперь мы решаем задачу классификации. Говорим маске: если пиксель принадлежит объекту на переднем плане, то присваиваем 1, если фону, то 0. Тогда нет никаких проблем с кейсом где-то между. Собирать такие данные проще простого. Человеку нужен только телефон, чтобы сфотографировать себя, и неплохое освещение.

Но не всё так просто. За более простой подход мы расплачиваемся потерей детализации. На картинке слева — результат маски с сегментацией, а справа — с маттингом. Видно, что на месте волос с маттингом виден каждый локон, а сегментационная маска потеряла детали.

Сопоставив плюсы и минусы подходов, мы в итоге остановились на семантической сегментации. 

Данные для обучения

Дмитрий: С постановкой задачи разобрались. Теперь мы знаем, что будем обучать сегментационные модели. Давайте попробуем собрать датасет, с которым можно будет решить эту задачу. Какие могут быть вводные на старте? 

У нас есть какое-то количество статей, по которым мы можем посмотреть стандартные датасеты, на которых люди обучают модели и делают для себя бенчмарки. Мы можем посмотреть в сторону датасетов из смежных областей: MS COCO, CIHP, Pascal VOC, Multi-Human Parsing v2, Supervise.ly.

Раньше мы занимались задачей keypoint detection. Некоторые датасеты там тоже содержат сегментационные маски, можем взять их в обучение. Даже с таким простым подходом получится собрать уже неплохой набор данных. Будет много картинок, например, один датасет MS COCO — это больше ста тысяч изображений, то есть достаточное количество примеров для обучения модели. Пример разметки можно увидеть на слайде.

Есть люди в кадре, они размечены масками. Кажется, это подходит для обучения моделей. 

Еще один важный момент — маттинговые датасеты тоже можно использовать для обучения сегментационной задаче. Вот пример датасета P3M:

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

Как мы будем использовать маттинговые датасеты? Мы просто берем альфа-канал, бинаризуем его по некоторому порогу (например, 0,5) и получаем на выходе сегментационную маску, состоящую из нулей и единиц, которые мы можем использовать для обучения сегментационных моделей.

Архитектура модели

Давид: Мы определились с постановкой задачи, собрали датасет. Пора выбрать архитектуру модели. Какие мы пробовали? На слайде представлена только часть из них.

В итоге победил старый добрый надежный Unet. Хоть этой архитектуре уже почти 10 лет (первая статья вышла в 2015 году), она до сих пор не утратила актуальности. Она и сейчас используется во многих доменах. При сегментации Unet бьет некоторые сетки помладше. Используется в генеративных моделях и других пайплайнах вроде денойзинга видео или изображений.

Метрики для обучения

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

Мы можем использовать стандартные классификационные метрики, чтобы измерять качество. Например, precision и recall. Можем построить кривую ROC AUC. Из более «картиночных» метрик можем взять IoU.

Мы же остановились на метрике Dice, и вот почему.  Если посмотреть на эту метрику применительно к задачам классификации, то окажется, что под Dice скрывается старый добрый F1 Score, который зарекомендовал себя как метрика, которая хорошо и балансно показывает качество моделей классификации. Поэтому будем использовать эту метрику как нашу основную для оценки модели.

Пайплайн

Что мы получаем? Есть некоторый пайплайн, как получить первые модели.

Берем данные → Обучаем модель Unet → Метрики → Выбираем лучшую → Смотрим глазами и оцениваем

Profit?

Получилось ли хорошо? На самом деле — спорно. С одной стороны, мы практически сразу получили модель, которая по метрикам на доменном датасете, который хорошо отражал именно видеозвонки, работала немного лучше, чем открытое решение, на которое мы ориентировались. На тот момент это был MediaPipe от Google. При этом в датасете, который мы собирали с разными интересными сложными кейсами, были проблемы. 

^{Запись доклада с таймкодом на соответствующем месте}

Например, кого-то сетка просто не видела, и там была пустая маска. У кого-то внезапно начинал появляться фон в углу изображения. И была довольно большая проблема с поднятыми руками — человек поднимал руки и исчезал.

Новый план

Давид: Что первое приходит в голову?

1. Почистить данные.

2. Собрать данные под доменные кейсы.

3. Пересмотреть аугментации.

4. Улучшить модель.

5. Улучшить метрики.

Хорошие данные для обучения

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

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

Другой пример — из датасета Supervise.ly. Здесь наоборот: качество масок отличное, прекрасно детализированное. Но почему-то было решено не размечать маской случай, когда человек что-то держит в руке. Во многих кейсах получается, что есть маска человека, а посередине — дыра. Мы не хотим использовать такие кейсы для обучения наших моделей. 

Посмотрев на отобранные датасеты, мы поняли, что большинство оказались либо плохо размеченными, либо плохо соответствующими нашему домену. Единственный хороший датасет из оставшихся — сет P3M10K. Здесь все хорошо и с разметкой, и с попаданием в домен. 

Повторим, что можно улучшить в датасетах:

1. Хорошая база.

2. Покрытие кейсов.

3. Negative Samples (то есть когда человека нет в кадре, сетка тоже должна хорошо работать).

Давид: Я могу тебе помочь с первым и третьим пунктами. Наши коллеги из SberDevices собрали EasyPortrait — очень хороший датасет для парсинга лиц. Он есть в открытом доступе. Мы переделали его под наш случай. 

Negative Samples тоже довольно хорошо ищутся, есть куча разных датасетов. Например, Houses Dataset с интерьерами, которые примерно напоминают помещения людей во время созвонов. 

Дмитрий: Посмотрев более критично на сбор данных, мы можем прийти к новому сетапу, на котором хотим обучать модели. Для базового обучения возьмем EasyPortrait, который даст нам 40 тысяч хорошо размеченных изображений. Добавим P3M10K, Negative Samples, а также некоторое количество внутренних данных, собранных под наш домен. Чем лучше и разнообразнее датасет, тем он полезнее для обучения нейросети. Мы собираем кейсы, пользуясь краудсорсингом. 

Аугментации данных

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

Также наши эксперименты показали, что важен порядок аугментации. Особенно в отношении того, в какие моменты мы делаем операции вроде resize и crop. Некоторые аугментации работают лучше до них, некоторые — после. 

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

Вот небольшой джентльменский набор аугментаций:

• Цветовые. RGB shift, contrast transformation, channel shuffle, sharpening.

• Геометрические. Crop, affine transform, perspective.

Это самый минимум, а вот в полноценном пайплайне, который мы используем, больше 20 разных видов аугментаций. Это тоже позволяет немного поднимать метрики для моделей. 

Улучшаем архитектуру

Давид: Итак, мы взяли обычный Unet. Он тем и хорош, что это очень базовая идея — энкодер и декодер, соединенные некоторыми skip connections. Это позволяет нам менять абсолютно все. Начали с энкодера и взяли вместо стандартного MobileNet v3. У него тоже скоро юбилей, но это до сих пор очень неплохой энкодер, особенно для реального времени. 

В качестве декодера, то есть блоков вместе с subsample, мы взяли блоки от MobileOne. Это блоки из статьи Apple. Они очень хорошо зашли. После репараметризации они превращаются в обычные Conv 2d, что довольно быстро работает. А вот ботлнек практически не меняли.

После всех этих манипуляций получилась другая сетка. Это вроде бы все еще Unet, но напоминает парадокс «корабля Тесея» — после замены деталей это все еще тот корабль или нет? 

Улучшаем метрику

Я уже говорил, что у нас нет метрики, которая бы четко сказала: смотрите, на границе качество у вас проседает. Поэтому мы начали использовать ее сами. В этом нам помогли всего лишь две математические операции: dilation и erosion.

Этот алгоритм неким ядром проходится по картинке. Dilation разжимает ее, как видно на маске слева. А erosion сужает. Вычли одно из другого и получили маску интереса справа. Теперь мы можем использовать эту маску с другими формулами. Теперь у нас есть не просто Dice, а Dice_border. Мы кидаем в Dice не просто маску, а маску, помноженную на шаблон границы. Делаем то же самое с predict и получаем неплохое понимание того, что происходит на границе.

После всех манипуляций пайплайн стал выглядеть так:

Новые данные → Фильтруем → Обучаем «Unet» → Метрики → Выбираем лучшую

→ Проверяем

После всего этого мы как будто одержали победу: руки человека остаются видны при поднятии. Метрика выросла. Фон заменен, и старый не проглядывает сквозь новый.  

Технические цели тоже достигнуты:

• Время выполнения занимает 23 мс. 

• Вес модели составляет 17,5 Мб.

• Качество достойное.

Дмитрий: Мы сделали интересный вывод из такой итерации. Он касается того, как отличается работа с моделями сегментации в науке и при разработке продукта.

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

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

Давид: Если бы мы жили в прекрасной сказке, то да. Но нет. После первой же победы обычно следует проигрыш. Давайте вспомним требования, в которые нам нужно было уложиться:

• Скорость выполнения. Меньше 25 мс. 

• Вес модели. Чем меньше, тем лучше. 

• Достойное качество. 

На бумаге 25 мс казались чем-то адекватным. А вот в конечном пайплайне это медленно работало. В итоге пришли к тому, что эта метрика должна быть меньше 12 мс.

Вроде бы 17,5 Мб — это не 17,5 Гб. Вполне резонно, но когда ты сталкиваешься с готовым продуктом вроде SaluteJazz, в котором уже помимо сегментации что-то есть, то решаешь ужаться. Требование стало меньше 5 Мб.

К качеству вопросов ни у кого не было, поэтому попросили оставить его как есть. Но желательно улучшить 🙂

Дмитрий: Интересные вводные. У нас появилось несколько планов.

План А:

• порезать модель руками

План Б:

• прунинг + fine-tune

План В:

• взять удачную модель из плана А;

• вернуть метрику с помощью дистилляции

План Г:

• квантизация весов

Облегчение модели

Давид: Начнем с самого простого. Что мы можем изменить:

• блоки энкодера;

• блоки декодера;

• блоки ботлнеков;

• глубину модели.

По соотношению трейдофа лучше всего оказались эксперименты с глубиной Unet. 

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

Прунинг модели

Дмитрий: Метрики у модели все-таки просели после такого подрезания. Давайте попробуем это изменить. Первый подход — прунинг модели. В чем идея? 

Обучаем модель → Убираем веса → Fine-tune → Легкая модель 

Шаги 2 и 3 повторяем несколько раз.  

Вообще про прунинг можно говорить много. 

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

Кроме того, важно выбрать правильный метод отбора весов, которые мы будем удалять. 

Оказалось, что самый простой метод — magnitude pruning — оказался самым качественным и по метрикам, и по удобству работы, с ним удобно итерироваться. Поиграв с разными параметрами прунинга, мы пришли к довольно интересным результатам. Взяли большую модель, которая работала не очень быстро (18,5 мс), модель поменьше (10 мс), но она существенно проседала по метрикам. Попробовали получить примерно такую же по скорости модель при помощи прунинга. 

Убрали 50% весов от изначальной модели, получили время инференса 9,75 мс. Это полностью нас устроило. Метрика при этом просела не так сильно. Кажется, что метод работает, и это победа. 

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

Дистилляция модели

Давид: Дистилляция — это техника, которая позволяет нам взять маленькую модель с неплохим временем выполнения, но с более низким качеством, и модель с хорошим качеством, но плохим временем исполнения. Попробуем взять лучшее из двух миров.

Назовем маленькую модель Student. Попытаемся перелить в нее знания из более крупной модели, которую назовем Teacher. Существует несколько подходов, которые позволяют это сделать. Первый называется дистилляцией ответов (response distillation).

В чем суть? Есть базовый пайплайн, где мы учим таргетную модель Student. У нее есть loss, ground truth. В этот пайплайн вмешивается модель Teacher, которая смотрит на данные, которые получает Student, и выдает свой ответ. Все это добавляется в конечный loss. Теперь Student не только предсказывает ground truth с помощью кросс-энтропии. Он еще и пытается мимикрировать под ответы Teacher, что выливается в новый loss.

Второй способ — feature distillation — немного хитрее. Если в первом мы добивались того, чтобы Student как можно точнее отдавал ответы, то есть мимикрировал под ответы Teacher, то во втором подходе мы хотим, чтобы ученик думал как учитель. 

Что это значит? Мы берем выходы со скрытых слоев учителя, сопоставляем их со скрытыми слоями ученика и хотим, чтобы их ответы были как можно ближе. Основные losses с данными ground truth вообще не меняются. 

Можно объединять оба метода, что мы и делали. В итоге не только получилось отыграть качество до результатов оригинальной модели, но и превзойти его! Казалось бы, это победа. С другой стороны, не все так просто. Мы поняли, что дистилляция не работает так, как это было заявлено. Дистилляция больше себя показала как техника регуляризации, которая позволяет учить модель дольше до того, как она начнет переобучаться. 

Второй инсайт: архитектура учителя должна не сильно отличаться от архитектуры ученика. Тогда будет лучше переход внутренних фич. 

Наконец, учитель и ученик должны видеть одни и те же данные. Есть разные подходы. Например, когда ученик видит один круг изображений, а учитель — другой. Так делать не надо. 

Теперь наш пайплайн немного изменился:

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

Квантизация модели

Дмитрий: В этом подходе все довольно просто. Есть два способа квантизации, которые мы можем применить. Первый вариант более безопасен — попробуем сжать веса из 32-битных float в 16-битные. В результате получим модели примерно в два раза легче. Скорее всего не будет и просадки по метрике. 

Есть более сложный вариант, который подразумевает квантизацию уже в 8-битные int. Он позволяет сделать модель еще меньше, в 4 раза. Но это происходит с потерей качества и требует дополнительных вещей во время обучения. Например, quantization aware training. 

При этом такой подход потенциально быстрее на устройствах, которые поддерживают быстрое исполнение в работе с 8-битными матрицами. Но поскольку наше таргетное устройство — это встроенный GPU на вашем ноутбуке, у которого нет поддержки таких операций, то мы выбрали квантизацию 16 бит и тем самым получили почти бесплатное облегчение модели. 

Результаты

• Время выполнения. 10 мс.

• Вес модели. 1,8 Мб.

• Достойное качество.

Постобработка результатов

Давид: На этом мы могли закончить, но все это время мы забывали об одном неоспоримом преимуществе: мы работаем с видеорядом. Обработав его, можно сделать еще и постобработку. 

Уменьшаем дрожание маски от кадра к кадру

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

^{Запись доклада с таймкодом на соответствующем месте}

Это поможет нам с дрожанием камеры. Модель отлично справилась 🙂

Дмитрий: На левом кадре есть дрожание границы маски, а на правом дрожания гораздо меньше. 

Размытие фона: не все так просто

Хотим сегодня также рассказать про размытие фона. Если делать это наивно, то у нас получится результат, как на картинке, когда вокруг человека возникает неприятный эффект гало.

Отчего он появляется? Мы считаем размытый фон так: берем кусок картинки и делаем свертку с некоторым ядром, например, гауссовским. Это проще представить так. Берем для каждого пикселя некоторую окрестность и считаем взвешенную сумму по окрестным пикселям. 

Картинка получается размытой. Из-за того, что человек размывается вместе с фоном, вблизи границы человека происходит проникновение красных пикселей на фон. Эти пиксели и попадают в наше неприятное гало. 

Как это убрать? Сделаем чуть более сложный алгоритм. У нас есть маска, которую посчитала наша сетка. Давайте из взвешенного суммирования уберем все пиксели, которые относились к маске человека. Или перемножим наши веса с этой маской. Тогда на выходе получим странную картинку, но это и есть то, что нам нужно.

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

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

Такой маленький элемент постобработки делает итоговую замену фона лучше.

Итоговый пайплайн

Новые данные → Фильтруем → Обучаем «Unet» → Метрики → Выбираем лучшую → Оптимизация → Постобработка → Проверяем → ??? → Profit!

Давид: Вот что мы в итоге обсудили:

• Общий пайплайн замены фона.

• Оптимизация моделей: как мы смогли в 10 раз облегчить модель, а качество даже повысили.

• Постобработка.

Будем рады ответить на ваши вопросы!