Привет, Хабр! В прошлой статье мы обсудили, что такое faceswap и довольно подробно разобрали существующие подходы. В этой статье мы хотим остановиться на том, как именно мы в Sber AI решаем эту задачу, а также погрузить вас в детали реализации нашего решения. Любое обучение модели зависит от трех составляющих — данных, самой модели и процесса обучения. В статье мы бы хотели поговорить про все эти компоненты, а также про дополнительные задачи и их решения, которые позволили нашему итоговому алгоритму переноса лица выглядеть качественно как для изображений, так и для видео. В таких сложных задачах, как перенос лица с одного изображения на другое, выбор подхода всегда играет ключевую роль. Мы решили, что не будем придумывать что-то кардинально новое, а в качестве baseline решения воспользуемся одним из уже существующих подходов, но с дополнительными улучшениями. Поскольку перед нами стояла задача построить модель, которая будет осуществлять перенос только по одному изображению человека, мы сразу отмели многие существующие решения. Из тех, что нам показались интересными были модели FaceShifter, SimSwap и One Shot Face Swapping on Megapixels.  Первые два подхода идейно довольно похожи, но отличаются деталями архитектуры и подходами к обучению. В основе третьего подхода лежит StyleGAN2 — авторы пытаются подобрать на основе двух изображений такой вектор в латентном пространстве модели, чтобы по нему можно было сгенерировать FaceSwap с первого изображения на второе. Подход интересен тем, что позволяет генерировать изображения высокого разрешения, однако в силу устройства StyleGAN2 — генерация лиц с поворотами содержит значительное количество артефактов, а значит и перенос на видео не сможет получиться достаточно качественным, так что этот подход мы не стали рассматривать. В итоге, за основу мы взяли здесь подход из статьи FaceShifter, а именно его часть — AEI-Net, но также взяли несколько интересных идей из других статей. На всякий случай напомним, что работает AEI-Net так: 
Рисунок 1 — Схема AEI-Net — модели переноса лиц на изображениях На верхнем уровне данная архитектура состоит из трех частей:  Мы попробовали модифицировать каждую из данных частей, сохранив общую структуру модели. Рассмотрим эксперименты, которые мы провели. Эксперименты с Identity Encoder Изучив все возможные подходы к получению identity вектора по фотографии человека мы решили, что все же не хотим отходить от модели ArcFace, поскольку она на данный момент является неоспоримой SoTA в данной задаче. Однако существует несколько различных реализаций этой модели, поэтому мы использовали следующие подходы: В итоге, после многих экспериментов, мы остановились на использовании одной версии модели ArcFace, в том числе потому, что использование сразу нескольких энкодеров для получения векторов, которые надо потом усреднить — делало процесс переноса более долгим.  Эксперименты с энкодером для   Чтобы понять, как можно поменять энкодер, давайте подумаем — а какая вообще перед нами стоит задача? Мы хотим доставать признаки из , причем они должны содержать данные с разных уровней детализации, чтобы у генератора была подробная информация про  — иначе он просто не сможет сгенерировать изображение, достаточно на  похожее.  Мы попробовали следующие подходы: 
Рисунок 2 — Результаты переноса при использовании различных энкодеров для  Ниже в таблице приведены метрики для моделей с разными энкодерами с учетом применения алгоритмов постобработки (блендинга). Все метрики рассчитываются так, как описано в статьях FaceShifter, SimSwap, Hifiswap, Smooth-Swap. Не углубляясь в математические формулировки, используемые метрики можно описать следующим образом: Эксперименты с итоговым генератором В случае с самим генератором, нам необходимо смешивать признаки  разного уровня, вектор  и то, что получилось на этапе генерации на предыдущем шаге. В оригинальной архитектуре в основе всей работы генератора лежит AAD ResBlk, состоящий из нескольких AAD блоков, внутри каждого из которых мы забираем информацию из id-вектора с помощью AdaIN, а из  с помощью Spade блока (рисунок 3). 
Рисунок 3 — AAD генератор В ходе экспериментов мы также проверили следующие гипотезы: Заменить смешивание id-вектора AdaIN на Adaconv, Attention в AAD блоке. Использование AdaConv и Attention вместо AdaIN не дало нам прироста по качеству, так что мы остановились на тех же AdaIN и Spade блоках, которые были изначально. Изменить число AAD блоков внутри AAD ResBlk Использование меньшего количества блоков позволило значительно ускорить модель, поскольку самой долгой составляющей является работа именно генератора, однако это привело к небольшому снижению качества работы модели. Чтобы облегчить обучение более легкой версии модели мы прибегли к использованию knowledge distillation — добавили loss, который бы требовал, чтобы выход более легкой версии модели был близок к тому, что выдает обычная модель — здесь нам помог обычный l2-loss, а также perceptual loss.  Также мы решили построить более тяжелую версию модели, с тремя AAD блоками, что действительно позволило улучшить качество генерации, однако замедлило работу самой модели. На рисунке показаны примеры работы все трёх версий модели. 
Рисунок 4 — Перенос с разным числом AAD блоков В данной таблице можно посмотреть на получающиеся метрики с учетом блендинга, для моделей трех размеров. Видно, что более тяжелая модель лучше переносит “identity” человека, однако приводит к снижению других метрик. Мы в качестве основной модели используем модель с 2 блоками. Выбрать правильный лосс для модели — это очень важно, ведь он указывает нам, чего именно мы хотим достичь. В данной задаче, помимо лоссов, которые были в оригинальной статье, мы также использовали ряд дополнительных функций для улучшения работы нашей модели. Однако, для начала давайте вспомним, какие функции мы уже упоминали: Во-первых, мы модифицировали лосс реконструкции, использовав идею из статьи SimSwap. В оригинале, идея этого лосса в том, что если мы даем модели две одинаковые фотографии человека — мы не хотим, чтобы модель с фотографией что-то делала. Однако тут можно пойти дальше — не обязательно, чтобы , достаточно, чтобы  и  были фотографиями одного человека — тогда мы все еще будем хотеть, чтобы  никак не менялся в результате переноса. Поскольку мы использовали датасеты, где для каждого человека у нас было по несколько фотографий, реализовать такую модификацию лосса стало возможным.  Также мы решили, что очень важной составляющей в визуальном восприятии результатов переноса являются глаза — тем более, когда речь идет о переносе на видео, и важно, чтобы человек на переносе смотрел туда же, куда и до него. Поэтому мы решили добавить специальный лосс на глаза, попробовав несколько подходов для корректировки направления взгляда: В ходе экспериментов мы остановились на первом подходе, так как именно он давал значительные улучшения в плане реалистичности и сохранения направления взгляда, как у . Помимо этого, мы пробовали довольно много лоссов для улучшения переноса identity: К сожалению, все эти подходы практически не дали улучшения в качестве. В качестве данных мы использовали несколько датасетов:  VGG Face 2 Наш основной датасет, состоящий из 3.31 млн изображений 9131 личности. Содержит большую вариацию в позах, возрасте, этнической принадлежности людей и качестве самих изображений.  CelebA-HQ  Датасет, содержащий 30.000 изображений лиц в высоком качестве.  Для повышения качества данных мы предварительно убираем все изображения с размерностью ниже 250×250 пикселей. Затем мы обрезаем лицо на каждом изображении с помощью нашего детектора — важно использовать одну и ту же версию кропа как на стадии обучения, так и на инференсе. Также важно, чтобы детектор был хороший — если детекция лиц на соседних кадрах видео не будет близка — перенос будет содержать артефакты. После того, как модель обучена, она умеет переносить личность человека с одного изображения на другое. Однако, это только половина дела — наша модель умеет осуществлять перенос только для кропнутых кадров головы. Как же осуществить перенос личности человека с одного изображения на видео и что для этого нужно? В первую очередь необходимо понять, какие вообще есть составляющие у переноса с изображения на видео: Тут возникает сразу несколько вопросов, например — как правильно вырезать лицо с видео, если там больше одного человека? Как вставить перенос обратно? А можно ли улучшать качество переноса уже после того, как отработала наша основная модель? На эти вопросы мы и постараемся ответить. На каждом кадре видео мы не только ищем лица людей и вырезаем интересующую нас область лица, но и запоминаем матрицу преобразования — она поможет нам вставить уже измененное лицо в изначальное место на изображении. Однако, если мы будем вставлять картинку, полученную моделью, целиком, могут появиться визуальные дефекты — будет четко заметна граница наложения на изначальный кадр. Этот эффект возникает как из-за неполного соответствия яркости картинок, так и из-за возможной размытости картинки, сгенерированной моделью. Поэтому необходимо обеспечить плавный переход от изначального кадра к полученному переносу, и в этом нам помогают сегментационные маски.  Маску лица можно представить в виде матрицы, определяющей, какие пиксели принадлежат лицу, а какие нет. Таким образом, мы можем определить точное расположение лица и переносить только его. При использовании бинарной маски место наложения лица всё ещё может оставаться заметным, поэтому мы дополнительно размываем маску у границ. Процесс наложения можно представить в виде смешивания пикселей изначального и измененного кадров.   
Рисунок 5 — Как устроен перенос на основе масок лица Также при использовании указанного выше подхода мы столкнулись со следующей проблемой — иногда у  и  не совпадают пропорции лица, так как модель старается сохранить форму . Если итоговое лицо на  сильно шире, чем то, которое было на , то перенос будет лишь частичный, и мы не сохраним форму лица как на .  Поэтому мы решили отслеживать ключевые точки для сгенерированного лица и лица из видео, и, если координаты меток будут значительно отличаться, дополнительно модифицировать маску переноса. В случае если лицо, полученное моделью при наложении, полностью покрывает лицо на видео, мы увеличиваем маску, тем самым создавая эффект переноса не только лица, но и формы головы. В обратном случае мы наоборот уменьшаем маску и усиливаем степень размытия, перенося только центральную часть лица.  Помимо хорошего переноса identity нам необходимо учитывать и качество самой картинки. Поэтому было решено добавить постобработку в виде этапа суперразрешения (SuperResolution). Так как модель должна обрабатывать каждый кадр видео, скорость инференса в данном случае была крайне важна. Решено было остановиться на подходе, схожим с face-renovation у HiFaceGAN — на стадии обучения мы предварительно ухудшаем изображения за счет понижения размерности, небольшого размытия и деформации и подаем на вход модели. Затем мы используем модель для восстановления изначального качества изображения. Использование SuperResolution модели действительно позволило улучшить качество итогового переноса.  
Рисунок 6 — Пример работы super resolution модуля Скорость работы всего пайплайна зависит, как мы уже обсудили, от нескольких вещей — как быстро мы находимо лица на видео, скорость работы модели, скорость получения масок и вставки результата переноса в оригинальные кадры, скорость добавления звука. Добавив обработку кадров батчами, переписав некоторые функции через библиотеку kornia, поместив большую часть операций на CUDA и проведя оптимизацию добавления звука, нам удалось ускорить общий пайплайн всего переноса в два раза — до менее чем 20 секунд на 15 секундное видео на Tesla V100 (что является близким к квазиоптимальному решению в смысле вычислительной эффективности). Использование SuperResolution модели также может добавить несколько дополнительных секунд на весь пайплайн.  Последнее, о чем мы еще не поговорили — это детали переноса на видео. Например, что будет, если на видео показаны несколько человек? Как сделать так, чтобы на каждом кадре мы переносили лицо только на одного человека? И как этого человека выбрать?
Допустим, на видео присутствует несколько человек: Если не указывать никакое лицо в качестве референса, то перенос будет осуществляться на случайного человека из видео, а если в кадре только один человек — перенос будет осуществляться на него (вектор ArcFace не вычисляется). Здесь можно заметить, что такой подход может существенно снизить производительность работы пайплайна, поскольку для каждого лица в кадре нам придется вычислять identity вектор с помощью ArcFace. В связи с этим, если в видео присутствует несколько человек, выполнение пайплайна может занять больше времени. На gif показан результат работы нашей модели.   
Рисунок 7 — Результат переноса нашей модели Все проведенные нами эксперименты действительно позволили получить хорошо работающую модель, с помощью которой можно осуществлять качественный перенос как на фото, так и на видео. Чтобы не быть голословными, мы посчитали различные метрики, и сравнили нашу модель с другими современными методами, реализующими технологию FaceSwap.  В данной таблице мы можем наблюдать сравнение нашего метода до блендинга (то есть выход самой модели) с другими методами — все значения метрик для других методов взяты из статей. Здесь в качестве сравнения мы используем нашу обычную модель с U-Net энкодером и двумя AAD блоками. А также мы посчитали метрики самостоятельно для каждого метода, однако уже после блендинга. Для этого мы использовали предоставленные видео в датасете FaceForensics++.  По метрикам видно, что наша модель показывает наилучшее качество сохранения identity среди всех моделей, а также очень хорошо сохраняет направление взгляда и в целом глаза. Мы провели достаточно большую работу и провели множество различных экспериментов, как с архитектурой модели, так и с другими составляющими обучения — в данной статье мы рассказали лишь о наиболее основных из них. В итоге, мы получили действительно хорошо работающую модель, которую можно применять для осуществления FaceSwap как на изображениях, так и на видео. Если вам понравилась данная статья и вам стало интересно попробовать работу наших моделей самостоятельно, то вы можете изучить код подробнее в репозитории GitHub или попробовать осуществить перенос с помощью Google Colab. Команда: Даниил Чесаков, Анастасия Мальцева, Александр Грошев, Андрей Кузнецов, Денис Димитров Ссылки: DataScientist