Перенос головы с картинки — сложно ли это? Модель GHOST-2.0. GAN.. GAN. head swap.. GAN. head swap. sberai.. GAN. head swap. sberai. аватар.. GAN. head swap. sberai. аватар. генеративные модели.. GAN. head swap. sberai. аватар. генеративные модели. смешивание изображений.

В последнее время технологии замены лиц находят все больше применений. Помимо использования в развлекательных целях, они стали особенно важны для индустрии фильмов и рекламы, позволяя существенно ускорить и удешевить производство. Несколько лет назад мы так же не остались в стороне от этой темы и выпустили свою модель переноса лиц с фото на фото или видео (в простонародье, face swap) — GHOST. Основная фишка модели была в отсутствии необходимости обучать её под каждый новый источник переноса лица — весь процесс буквально происходил по одной фотке (отсюда и буквы “OS=one shot” в названии модели)

Однако в таком подходе, где мы заменяем лишь область лица, есть несколько существенных недостатков. Например, эти подходы основываются на предположении, что формы лица source и target одинаковы. Если это условие не соблюдается, то это сильно повлияет на восприятие лица после переноса. Именно для того, чтобы эта проблема не мешала, мы начали смотреть в сторону создания технологии переноса головы целиком (~~чувствуем себя трансплантологами~~).

Что и следовало ожидать, процесс переноса головы ни разу не легче задачи переноса лица, а даже обрастает дополнительными серьёзными челленджами, которые надо обязательно учитывать при разработке модели. Например, очень важно, чтобы сгенерированная голова органично смотрелась в исходном фоне target и не отличалась по контрасту и цвету кожи от остальных частей тела и окружения. Получается, мы не только генерируем source голову с позой и выражением лица target, но и буквально адаптируем новую голову под сцену.

После переноса головы в новый фон требуется дополнительный inpainting фона и адаптация цветов

Примечательно, что для данной задачи, в отличии от переноса лица, существует мало решений, а точнее даже одно (тут мы конечно не учитываем спектр работ по 3D аватарам и целый пласт подходов на сплатах). Той самой единственной моделью является HeSer, состоящая из двух модулей. Первый модуль Aligner используется для реконструкции головы с помощью генеративно-состязательных сетей (GAN). Второй модуль Blender отвечает за органичную вставку полученной головы в target изображение. Данная работа была продолжена в статье Generalist FaceX , где функции обоих модулей реализованы через Stable Diffusion. К сожалению, к обеим статьям нет официального кода, так что данная задача до недавнего времени не имела опенсорс решений. Настало время исправить несправедливость в мире опенсорс комьюнити:)

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

Aligner

Процесс переноса source головы начинается с реконструкции целевых позы и выражения лица. С этой задачей призван справляться модуль Aligner, при разработке которого мы отталкивались от оригинальной статьи HeSer. Условно процесс генерации можно разделить на 2 этапа:

считывание необходимой информации с source и target изображений
последующая генерация нового изображения

Схема модуля Aligner

Для считывания информации мы используем три энкодера:

портретный энкодер $E_{por}$ на архитектуре ResNet. Он используется для кодирования source айдентити по фотографии целой головы, таким образом учитывая не только лицо, но также прическу и форму головы
энкодер source айдентити $E_{id}$ , в качестве которого мы берем замороженную сеть для распознавания лиц, предобученную с Arcface лоссом. На вход этому энкодеру подается кроп лица
энкодер движения $E_{motion}$ , о котором мы подробнее расскажем ниже

Получение информации о целевой позе головы и выражении лица происходит через энкодер $E_{motion}$ , которому на вход передается полное target изображение. Однако при прямолинейном обучении этого энкодера вероятна ситуация, что он просто будет кодировать target изображение, а не выучивать нужную информацию о движении. Поскольку на этапе инференса мы подаем разных людей в качестве source и target, это приведет к “протеканию” внешности target в сгенерированное изображение и смешению внешности двух людей. Чтобы этого избежать, вводятся дополнительные аугментации, которые заставляют энкодер учить именно информацию о позе и выражении, исключая внешность. Эти аугментации независимо по горизонтальной и вертикальной осям растягивают или сужают target изображение, фактически меняя его внешность, но оставляя прежними его позу и выражение. Кроме того, чтобы дополнительно сфокусировать энкодер на получение только нужной информации, специально берется сеть с небольшим количеством параметров на архитектуре MobileNetV2.

Обучение энкодера движения без аугментаций приводит к смешению внешности

Выходом энкодеров являются отдельные вектора, которые затем конкатенируются в один глобальный эмбеддинг, который используется для обусловливания генератора. Сам генератор по архитектуре напоминает StyleGAN. Входом в него служит обучаемый тензор шума , который проходит через свёрточные ResNet блоки с AdaIN нормализацией. На каждом слое этот тензор обуславливается на глобальный эмбеддинг через AdaIN и в 2 раза увеличивается в разрешении. На выходе мы получаем финальное изображение разрешением 512х512, а также бинарную маску головы. Предсказание маски помогает стабилизировать процесс обучения.

Ключевая модификация нашей модели по сравнению с HeSer — получение представлений позы и выражения с помощью единого энкодера $E_{motion}$ В оригинальной статье для этого использовались два отдельных энкодера. Однако один из них был избыточен и кодировал информацию о внешности target, в результате чего портилась финальная генерация. Особенно проблема была заметна при большой разнице в прическе. Например, если у target головы была чёлка, а у source – нет, то на генерации она появлялась (точнее даже сказать “проявлялась”). Как показано на примере ниже, наше решение с объединением энкодеров помогло уменьшить такое протекание и значительно улучшило качество генерации.

Сравнение Aligner модулей моделей GHOST 2.0 и HeSer

Данная проблема с избыточностью энкодеров выявилась в ходе следующего эксперимента. Мы решили выяснить, какую информацию выучивают энкодеры позы и выражения. Для этого мы обнуляли эмбеддинг соответствующего энкодера и смотрели, как изменится генерация (эдакий ablation study). Оказалось, что при обнулении энкодера позы всё равно генерируется изображение с нужной позой и выражением. При обнулении энкодера выражения лицо логично приобретает нейтральное выражение, однако поза головы перестает соответствовать целевой позе. Это указывает на то, что все движения выучиваются через энкодер выражения, а энкодер позы кодирует преимущественно детали внешности target.
Вывод: энкодер позы избыточен!

Иллюстрация представлений, выучиваемых энкодерами

Следующим шагом мы решили попробовать разъединить представления так, чтобы каждый энкодер выучивал только свою часть информации. Для этого мы провели несколько экспериментов по изменению архитектуры и процесса обучения. В частности, мы пробовали заимствовать идеи генератора и циклических лоссов из статьи DPE: Disentanglement of Pose and Expression или брать предобученные энкодеры, например из модели DECA: Detailed Expression Capture and Animation. В какой-то момент мы решили полностью изменить архитектуру и взять за основу метод варпинга (отображение пикселей из одного изображения в другое без изменения цветов) по ключевым точкам лица, для большей абстракции взаимодействия с target изображением. Однако все эти методы либо не давали нужного разъединения представлений, либо ухудшали качество генерации по сравнению с HeSer.

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

На данном этапе может возникнуть логичный вопрос: почему бы не попробовать другие модели в качестве Aligner’а? Действительно, мы протестировали довольно большое количество опенсорс методов, основанных на 2D представлениях. Большинство таких предобученных моделей обучалось на узких кропах лица, которые не захватывают голову полностью и следовательно не подходят для нашей задачи. Следовательно, их пришлось бы повторно обучать с нуля для использования в задачах head swap. Среди протестированных методов есть довольно свежие работы:

Однако большинство этих методов хорошо работает на небольших (< 45 градусов) углах поворота. Нам же хотелось иметь более устойчивый метод, который поддерживает перенос при повороте target головы до 90 градусов. Наша модель успешно справляется с этой задачей. При этом мы не использовали в обучении каких-то специальных датасетов с экстремальными позами, вроде EFHQ. Кроме того, на некоторых методах возникают проблемы с фликерингом (заметное мерцание из-за разницы в цветах соседних кадров) при генерации видео, хотя генерации отдельных кадров довольно качественные.

В статье HeSer модуль Aligner обучался со следующими лоссами:

Adversarial loss
Discriminator perceptual loss, который притягивает признаки дискриминатора от реальных и синтетических данных, и используется на этапе обучения генератора
L1 лосс на реконструкцию
VGG perceptual loss, который притягивает признаки нейросети VGG от реальных и синтетических данных
Аналогичный ID perceptual loss, использующий модель Arcface для распознавания лиц
ID лосс, сравнивающий эмбеддинги от Arcface по косинусному расстоянию. Мы хотим, чтобы эмбеддинги от генерации и target на этапе тренировки были как можно ближе, чтобы лучше обучить энкодер $E_{por}$ кодировать source внешность.
Dice лосс для обучения бинарных масок головы

Однако, если обучать модель с объединенным энкодером позы и выражения $E_{motion}$ и данным набором лоссов, то качество мимики на генерации сильно падает. Для решения проблемы мы добавили 3 новых лосса:

Lip closure loss. Считается расстояние между губами на генерации, и сравнивается с аналогичным расстоянием между губами на target изображении. Таким образом, этот лосс позволяет контролировать степень открытия и закрытия рта.
Аналогичный ему eye closure loss для закрытия глаз
Emotion loss, который использует ResNet модель распознавания эмоций для сравнения выражения лица генерации и target

По аналогии с предыдущей версией модели GHOST, мы также ввели gaze loss для переноса направления взгляда. Это позволяет существенно улучшить качество генераций. Этот лосс работает на основе сопоставления эмбеддингов из соответствующих моделей. Мы пробовали использовать такие модели, как L2CS, GazeTR и сопоставление хитмапов, полученных по ключевым точкам глаз. Однако в итоге остановились на RT-GENE, который обеспечил лучшее качество генерации глаз и переноса взгляда. В ходе этих экспериментов мы также выявили, что добавление данного лосса с большим весом может испортить source айдентити на генерации или вызвать протекание target внешности при добавлении на ранних этапах обучения. Поэтому мы включаем его с небольшим весом ближе к концу обучения, когда энкодеры уже хорошо выучили нужные представления.

Датасет

Для обучения наших моделей используется датасет VoxCeleb2 в разрешении 512×512. Он состоит из порядка 200 000 видео-интервью более 5500 различных знаменитостей. Однако не все изображения в нем высокого качества — одна часть смазана, другая часть в низком разрешении — поэтому мы дополнительного его фильтруем. Для этого используем две модели для image quality assessment (IQA): HyperIQA и CLIP-IQA. Эти модели принимают на вход изображение и выдают оценку его качества от 0 до 100. После фильтрации у нас осталось около 140 000 видео. Их мы разделили на обучающий (4850 личностей) и тестовый (360 личностей) подсеты.

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

Также для работы нашей модели необходимо иметь качественную сегментационную модель — она будет использоваться в модели Blender, а также на переходных этапах, чтобы выделять голову человека. Есть два основных требования к модели: она должна уметь выделять волосы на голове и на лице, как отдельные классы, а также быть аддитивной, то есть выделять сегменты таким образом, чтобы мы могли при их объединении получить целую голову. Для обучения модели был собран и размечен датасет из 20 тысячи FullHD изображений, и на его основе обучена сегментационная модель Segformer-B5 (за обучение выражаем благодарность команде Vision из SberDevices).

Сравнение нашей сегментационной модели с популярной сетью BiSeNet

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

Результаты

Мы сравнивали наш Aligner с оригинальным подходом HeSer и моделью StyleHeat, основанной на инверсии StyleGAN. Все модели работают в разрешении 512. Следует отметить, что в настоящее время существует довольно мало опенсорсных решений для 2D реконструкции лиц в данном разрешении.

Сравнение проводилось в двух сценариях: когда в качестве source и target берется один и тот же человек (self-reenactment) и разные люди (cross-reenactment). Для сравнения в сценарии self-reenactment использовались следующие метрики:

LPIPS, оценивающая, насколько по человеческому восприятию похожи изображения (чем ниже, тем лучше)
SSIM, также сравнивающая изображения по восприятию (чем выше, тем лучше)
PSNR, оценивающая попиксельное различие в изображениях (чем выше, тем лучше)
VGG ID, оценивающая сохранение source айдентити с помощью сети VGG Face (чем выше, тем лучше)

На cross-reenactment мы считаем VGG ID и метрику FID, оценивающую сходство выборки сгенерированных изображений с выборкой настоящих (чем ниже, тем лучше).

	ID self	LPIPS	PSNR	SSIM	ID cross	FID
HeSer	0.780	0.139	22.99	0.856	0.612	35.33
StyleHeat	0.789	0.574	8.791	0.582	0.673	75.856
GHOST 2.0	0.754	0.142	22.04	0.848	0.628	29.57

Результаты реконструкции голов

StyleHEAT по метрикам лучше других методов сохраняет source внешность, но плохо справляется с поворотами головы больше 15 градусов и в целом генерирует довольно размытые изображения. HeSer показывает наилучшие метрики на self-reenactment из-за протекания target изображения, однако на cross-reenactment сильно искажает source внешность. Наша модель показывает хорошие результаты генерации в различных позах головы, при этом устраняя протекание target.

Blender

После того, как мы реконструировали голову, нужно вставить ее в target фон так, чтобы итоговое изображение смотрелось естественно. Это делается с помощью модуля Blender. Наша реализация представляет собой немного модифицированный подход, описанный в статье HeSer. В оригинальной модели Blender состоит из двух основных модулей:

Color Reference Creator: модуль создает цветной референс для раскрашивания черно-белой реконструкции, полученной из Aligner, в соответствии с цветами target. Это позволяет перенести тон кожи, и делает реконструкцию более соответствующей остальному фону по контрасту. Для создания такого референса используются маски сегментации обоих портретов. По ним считаются попарные корреляции признаков соответствующих семантических регионов. Признаки извлекаются с помощью Feature Pyramid Network (FPN). Затем эти корреляции используются для определения итоговых цветов на референсе.
Blending UNet: принимает на вход набор изображений и масок, описывающих фон и само лицо, и “сшивает” их в итоговую генерацию. Этими входными изображениями являются:
- цветной референс головы $I^{HR}_{T rightarrow A}$ , полученный из Color Reference Creator
- маска головы
- черно-белая версия реконструкции головы $I_A^{GH}$
- inpainting маска , описывающая регион фона вокруг головы, требующий заполнения.
- inpainting референс $I^{IR}_{T rightarrow A}$ , полученный c помощью inpainting модели по маске вокруг головы
- фон $I_T^{BG}$ из target изображения

Схема модуля Blender

Однако в процессе реализации и тестирования оригинальной модели из HeSer обнаружилось несколько недостатков.

Проблема 1. Если на target изображении не найдены семантические регионы, присутствующие на реконструкции, на генерации они будут окрашены в серый цвет, как показано на примере ниже. Такая ситуация может возникнуть из-за ошибок сегментации или особенности фотографии (например, спрятаны уши или не был обнаружен нос).

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

При отсутствии референсов на target изображении некоторые области на реконструированном лице не окрашиваются

Проблема 2. Появление контрастной границы вокруг перенесенной головы из-за полупрозрачности волос.

Решение. Чтобы избавиться от этого, мы создаем дополнительное изображение target фона $I_{ext}$ , в котором экстраполируем фон на маску человека. Далее, мы смешиваем по маске изображение , полученное из Aligner, с этим изображением фона, передаем смешанное изображение $hat{I_A}$ в Blender и получаем итоговую генерацию $hat{I_B}$ .

Пример с экстраполяцией маски.

– результат без экстраполяции, $hat{I_B}$ – с экстраполяцией

При обучении Blender мы берем одно и то же изображение в качестве входа в модуль и ground truth. Однако, чтобы модель не выучилась просто копировать пиксели, вводятся дополнительные аугментации, например, изменение цвета входного изображения. Также в изначальном подходе входное и ground truth изображения горизонтально отражались относительно друг друга. Мы решили отказаться от такой аугментации, так как это приводило к появлению артефактов при инференсе (хотя при обучении никаких артефактов не наблюдалось).

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

Примеры аугментаций для масок

При этом качество самого inpainting, особенно на больших зонах, требовалось улучшить. Для этого мы внедрили внешнюю inpainting модель. Мы протестировали несколько диффузионных подходов, но в итоге остановились на LaMa. Эта модель используется для создания inpainting референса $I^{IR}_{T rightarrow A}$ , который далее передается в Blending UNet. Таким образом получилось немного улучшить метрики и качество inpainting.

Оценку качества блендинга мы производим с помощью метрик LPIPS, PSNR и SSIM. Также мы подсчитываем ошибку реконструкции отдельно для фона ( $text{PNSR}_{back}$ ) и головы ( $text{PNSR}_{head}$ ). Как видно из таблицы ниже, наше решение превосходит бейзлайн HeSer по всем метрикам.

	LPIPS	SSIM		$text{PSNR}_{back} uparrow$	$text{PSNR}_{head} uparrow$
HeSer	0.034	0.972	33.86	35.13	42.13
GHOST 2.0	0.018	0.984	36.65	37.98	45.83

Финальные результаты после блендинга головы с фоном

Перенос на реальные данные

Схема дополнительной постобработки с Кандинским 2.2

При переходе к реальным данным необходимо учитывать, что наша модель работает с обрезанными изображениями. В отличие от задачи замены лица, где лицо человека всегда находится внутри сгенерированной области, в задаче замены головы волосы могут выходить за границы рассматриваемой области. В таком случае необходимо удалить избыточные волосы из изображения. Для этой цели мы предлагаем алгоритм постобработки с использованием нашей диффузионной модели Kandinsky 2.2. Финальное изображение получается путем восстановления области, выделенной маской $M_{inpainting}$ для выхода UNet , как показано на схеме выше. Стоит отметить, что этот шаг является опциональным и применяется только в случае, если у человека на target изображении длинные волосы.

Заключение

В этой статье мы познакомились с задачей переноса головы и рассказали про наше новое опенсорс решение в этой области. Напомним, что подход состоит из двух этапов: реконструкции source головы с target движением, и блендинг сгенерированной головы с target фоном. Метрики позволяют оценить положение модели в ряду немногочисленных существующих альтернатив. Так, наша модель позволяет получить более реалистичные реконструкции головы, а также обеспечивает более естественные результаты при блендинге. Далее вы найдёте список важных и полезных ссылок на демо, страницу проекта и technical report:

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

Команда: Александр Грошев*, Павел Парамонов*, Анастасия Ященко*, Денис Димитров, Андрей Кузнецов

* — равный вклад

По всем возникающим вопросам и предложениям можно писать Денису Димитрову или Андрею Кузнецову.

Каналы авторов:

Автор: NastasiaY

Источник

Запись добавлена: 27.02.2025 в 08:59
Оставлено в

Перенос головы с картинки — сложно ли это? Модель GHOST-2.0

Меню навигации

На главную

Главное

Рубрики

Методики

Информация

Из архивов