Как работает FaceSwap. Часть 2. Разработка от Sber AI25.01.2022 17:47

Привет, Хабр!

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

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


Как модель обучать?


Модель

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

Поскольку перед нами стояла задача построить модель, которая будет осуществлять перенос только по одному изображению человека, мы сразу отмели многие существующие решения. Из тех, что нам показались интересными были модели FaceShifter, SimSwap и One Shot Face Swapping on Megapixels.

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

В итоге, за основу мы взяли здесь подход из статьи FaceShifter, а именно его часть — AEI-Net, но также взяли несколько интересных идей из других статей. На всякий случай напомним, что работает AEI-Net так:


  1. C помощью некоторой предобученной модели ArcFace (Identity Encoder на рисунке 1) из $X_s$ извлекается вектор, который хорошо кодирует «identity» человека, которого мы хотим перенести
  2. С помощью некоторой U-NET-like архитектуры (синяя часть на рисунке 1) мы извлекаем признаки из $X_t$ — $z^{1}_{att}, z^{2}_{att}, ..., z^{n}_{att}$
  3. Последовательно смешиваем атрибутивные и identity вектора признаков
  4. Получаем на выходе изображение $\hat{Y}_{s,t}$, которое сохраняет «identity» человека из $X_s$ и атрибуты из $X_t$

Рисунок 1 - Схема AEI-Net - модели переноса лиц на изображениях
Рисунок 1 — Схема AEI-Net — модели переноса лиц на изображениях

На верхнем уровне данная архитектура состоит из трех частей:


  1. Identity Encoder
  2. Энкодер для $X_t$ — синяя часть на схеме, которая извлекает признаки из target изображения
  3. Генератор, который по признакам из $X_t$ и identity вектору генерирует итоговое изображение

Мы попробовали модифицировать каждую из данных частей, сохранив общую структуру модели. Рассмотрим эксперименты, которые мы провели.

Эксперименты с Identity Encoder

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


  1. Использование разных версий ArcFace в качестве identity энкодера
  2. Использование в качестве вектора $z_{id}$ среднего из векторов, которые получены на выходе разных версий модели ArcFace

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

Эксперименты с энкодером для $X_t$

Чтобы понять, как можно поменять энкодер, давайте подумаем —, а какая вообще перед нами стоит задача? Мы хотим доставать признаки из $X_t$, причем они должны содержать данные с разных уровней детализации, чтобы у генератора была подробная информация про $X_t$ — иначе он просто не сможет сгенерировать изображение, достаточно на $X_t$ похожее.

Мы попробовали следующие подходы:


  1. U-NET
    Модель использовалась в оригинальной архитектуре FaceShifter. Мы пробовали добавлять Resblock в encoder части, однако улучшений это не дало. Также пробовали уменьшить число каналов в извлекаемых картах признаков, что позволило ускорить модель.
  2. LinkNET
    Архитектура похожа на U-Net, однако карты признаков в декодере складываются, а не конкатенируются, что в итоге приводит к снижению размерности признаков и повышению скорости инференса модели. За счет уменьшения размерности признаков атрибутов удалось улучшить identity на изображениях — мы стали меньше ориентироваться на target, однако на видео получаем немного нестабильную генерацию.
  3. FPN
    Данный подход лежал в основе гипотезы о том, что с его помощью можно более эффективно сохранить информацию об $X_t$, однако в ходе экспериментов эта гипотеза не подтвердилась, а использование FPN лишь привело к снижению качества работы.
  4. Resnet
    Тут мы отказались от идеи encoder-decoder архитектуры, а взяли карты признаков на разных уровнях ResNet. В результате информация об атрибутах стала протекать сильнее в сравнении с архитектурами encoder-decoder, что позволило сохранять при переносе различные детали target лица, такие как челка, уши и т.п. Архитектура получилась очень легковесной — однако работу всей модели это не сильно ускорило.

Рисунок 2 - Результаты переноса при использовании различных энкодеров для $<!-- math>$inline$X_t$inline$</math -->$» /><br />Рисунок 2 — Результаты переноса при использовании различных энкодеров для <img src=

Ниже в таблице приведены метрики для моделей с разными энкодерами с учетом применения алгоритмов постобработки (блендинга).

Все метрики рассчитываются так, как описано в статьях FaceShifter, SimSwap, Hifiswap, Smooth-Swap. Не углубляясь в математические формулировки, используемые метрики можно описать следующим образом:


  • ID retrieval и shape_ringnet — отвечают за сохранение identity (форма головы и т.д.);
  • exp_ringnet — отвечает за выражение лица и сохранение эмоций;
  • pose_ringnet и poseHN — отвечают за сохранение позы с помощью моделей Ringnet и Hopenet;
  • eye_ldmk — за сохранение направления взгляда.

Эксперименты с итоговым генератором

В случае с самим генератором, нам необходимо смешивать признаки $X_t$ разного уровня, вектор $z_{id}$ и то, что получилось на этапе генерации на предыдущем шаге. В оригинальной архитектуре в основе всей работы генератора лежит AAD ResBlk, состоящий из нескольких AAD блоков, внутри каждого из которых мы забираем информацию из id-вектора с помощью AdaIN, а из $z^{k}_{att}$ с помощью Spade блока (рисунок 3).

Рисунок 3 - AAD генератор
Рисунок 3 — AAD генератор

В ходе экспериментов мы также проверили следующие гипотезы:


  1. Заменить смешивание id-вектора AdaIN на Adaconv, Attention в AAD блоке.

    Использование AdaConv и Attention вместо AdaIN не дало нам прироста по качеству, так что мы остановились на тех же AdaIN и Spade блоках, которые были изначально.


  2. Изменить число AAD блоков внутри AAD ResBlk

    Использование меньшего количества блоков позволило значительно ускорить модель, поскольку самой долгой составляющей является работа именно генератора, однако это привело к небольшому снижению качества работы модели. Чтобы облегчить обучение более легкой версии модели мы прибегли к использованию knowledge distillation — добавили loss, который бы требовал, чтобы выход более легкой версии модели был близок к тому, что выдает обычная модель — здесь нам помог обычный l2-loss, а также perceptual loss.


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

Рисунок 4 - Перенос с разным числом AAD блоков
Рисунок 4 — Перенос с разным числом AAD блоков

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

Видно, что более тяжелая модель лучше переносит «identity» человека, однако приводит к снижению других метрик. Мы в качестве основной модели используем модель с 2 блоками.


Функции потерь

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

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

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

В ходе экспериментов мы остановились на первом подходе, так как именно он давал значительные улучшения в плане реалистичности и сохранения направления взгляда, как у ${X}_{t}$.

Помимо этого, мы пробовали довольно много лоссов для улучшения переноса identity:


  • Предварительно обучали классификатор на VggFace2, который по фотографии определял, что за человек изображен. Использовали его для предсказания меток класса $\hat{Y}_{s,t}$ — хотелось, чтобы после переноса классификатор угадывал нового человека на изображении.
  • Пробовали обучить классификатор, но уже для пары изображений — он должен был понять, один ли человек на двух картинках или нет. Затем используя обученный классификатор хотели, чтобы для генератора $({X}_{s1},\hat{Y}_{s,t})$ → 1 и $({X}_{t},\hat{Y}_{s,t})$ → 0.
  • Пробовали обучить identity дискриминатор — идея похожая на предыдущую, однако тут дискриминатор обучался в процессе.

К сожалению, все эти подходы практически не дали улучшения в качестве.


Данные

В качестве данных мы использовали несколько датасетов:


  1. VGG Face 2

    Наш основной датасет, состоящий из 3.31 млн изображений 9131 личности. Содержит большую вариацию в позах, возрасте, этнической принадлежности людей и качестве самих изображений.


  2. CelebA-HQ

    Датасет, содержащий 30.000 изображений лиц в высоком качестве.


Для повышения качества данных мы предварительно убираем все изображения с размерностью ниже 250×250 пикселей. Затем мы обрезаем лицо на каждом изображении с помощью нашего детектора — важно использовать одну и ту же версию кропа как на стадии обучения, так и на инференсе. Также важно, чтобы детектор был хороший — если детекция лиц на соседних кадрах видео не будет близка — перенос будет содержать артефакты.


Как модель использовать?

После того, как модель обучена, она умеет переносить личность человека с одного изображения на другое. Однако, это только половина дела — наша модель умеет осуществлять перенос только для кропнутых кадров головы. Как же осуществить перенос личности человека с одного изображения на видео и что для этого нужно?

В первую очередь необходимо понять, какие вообще есть составляющие у переноса с изображения на видео:


  1. Вырезаем лицо человека, которого хотим перенести
  2. Вырезаем лицо человека с каждого кадра видео, куда будем переносить
  3. Осуществляем перенос моделью на каждый кадр
  4. Вставляем перенос обратно в видео и сохраняем видео

Тут возникает сразу несколько вопросов, например — как правильно вырезать лицо с видео, если там больше одного человека? Как вставить перенос обратно? А можно ли улучшать качество переноса уже после того, как отработала наша основная модель? На эти вопросы мы и постараемся ответить.


Маски и все, что с ними связано

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

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

image-loader.svg

Рисунок 5 - Как устроен перенос на основе масок лица
Рисунок 5 — Как устроен перенос на основе масок лица

Также при использовании указанного выше подхода мы столкнулись со следующей проблемой — иногда у $\hat{Y}_{s,t}$ и $X_t$ не совпадают пропорции лица, так как модель старается сохранить форму $X_s$. Если итоговое лицо на $\hat{Y}_{s,t}$ сильно шире, чем то, которое было на $X_t$, то перенос будет лишь частичный, и мы не сохраним форму лица как на $X_s$.

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


Можно ли улучшить качество?

Помимо хорошего переноса identity нам необходимо учитывать и качество самой картинки. Поэтому было решено добавить постобработку в виде этапа суперразрешения (SuperResolution). Так как модель должна обрабатывать каждый кадр видео, скорость инференса в данном случае была крайне важна. Решено было остановиться на подходе, схожим с face-renovation у HiFaceGAN — на стадии обучения мы предварительно ухудшаем изображения за счет понижения размерности, небольшого размытия и деформации и подаем на вход модели. Затем мы используем модель для восстановления изначального качества изображения. Использование SuperResolution модели действительно позволило улучшить качество итогового переноса.

image-loader.svg

Рисунок 6 - Пример работы super resolution модуля
Рисунок 6 — Пример работы super resolution модуля


Как ускорить пайплайн?

Скорость работы всего пайплайна зависит, как мы уже обсудили, от нескольких вещей — как быстро мы находимо лица на видео, скорость работы модели, скорость получения масок и вставки результата переноса в оригинальные кадры, скорость добавления звука. Добавив обработку кадров батчами, переписав некоторые функции через библиотеку kornia, поместив большую часть операций на CUDA и проведя оптимизацию добавления звука, нам удалось ускорить общий пайплайн всего переноса в два раза — до менее чем 20 секунд на 15 секундное видео на Tesla V100 (что является близким к квазиоптимальному решению в смысле вычислительной эффективности). Использование SuperResolution модели также может добавить несколько дополнительных секунд на весь пайплайн.


Перенос на видео

Последнее, о чем мы еще не поговорили — это детали переноса на видео. Например, что будет, если на видео показаны несколько человек? Как сделать так, чтобы на каждом кадре мы переносили лицо только на одного человека? И как этого человека выбрать?
Допустим, на видео присутствует несколько человек:


  1. Необходимо указать лицо (или несколько), на которые необходимо сделать перенос. Это можно сделав загрузив изображение человека, на которого мы хотим перенести, например, взяв его изображение с одного из кадров видео
  2. Затем на каждом кадре будут находиться лица, и перенос будет осуществляться на то, которое соответствует выбранному нами согласно близости по identity вектору

Если не указывать никакое лицо в качестве референса, то перенос будет осуществляться на случайного человека из видео, а если в кадре только один человек — перенос будет осуществляться на него (вектор ArcFace не вычисляется). Здесь можно заметить, что такой подход может существенно снизить производительность работы пайплайна, поскольку для каждого лица в кадре нам придется вычислять identity вектор с помощью ArcFace. В связи с этим, если в видео присутствует несколько человек, выполнение пайплайна может занять больше времени. На gif показан результат работы нашей модели.

image-loader.svg


Результаты

Рисунок 7 - Результат переноса нашей модели
Рисунок 7 — Результат переноса нашей модели

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

В данной таблице мы можем наблюдать сравнение нашего метода до блендинга (то есть выход самой модели) с другими методами — все значения метрик для других методов взяты из статей. Здесь в качестве сравнения мы используем нашу обычную модель с U-Net энкодером и двумя AAD блоками.

А также мы посчитали метрики самостоятельно для каждого метода, однако уже после блендинга. Для этого мы использовали предоставленные видео в датасете FaceForensics++.

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


Заключение

Мы провели достаточно большую работу и провели множество различных экспериментов, как с архитектурой модели, так и с другими составляющими обучения — в данной статье мы рассказали лишь о наиболее основных из них. В итоге, мы получили действительно хорошо работающую модель, которую можно применять для осуществления FaceSwap как на изображениях, так и на видео. Если вам понравилась данная статья и вам стало интересно попробовать работу наших моделей самостоятельно, то вы можете изучить код подробнее в репозитории GitHub или попробовать осуществить перенос с помощью Google Colab.

Команда: Даниил Чесаков, Анастасия Мальцева, Александр Грошев, Андрей Кузнецов, Денис Димитров

Ссылки:


  1. FaceShifter: Towards High Fidelity And Occlusion Aware Face Swapping https://arxiv.org/abs/1912.13457
  2. SimSwap https://arxiv.org/abs/2106.06340
  3. FaceForensics++: Learning to Detect Manipulated Facial Images https://arxiv.org/abs/1901.08971
  4. Adaconv https://studios.disneyresearch.com/app/uploads/2021/04/Adaptive-Convolutions-for-Structure-Aware-Style-Transfer.pdf
  5. Attention https://openaccess.thecvf.com/content/ICCV2021/papers/Liu_AdaAttN_Revisit_Attention_Mechanism_in_Arbitrary_Neural_Style_Transfer_ICCV_2021_paper.pdf

© Habrahabr.ru