Всем привет! Меня зовут Денис Бобков, я сейчас обучаюсь на совместной магистерской программе ВШЭ и ШАД под названием «Современные компьютерные науки», а также работаю исследователем в AIRI в команде Controllable Generative AI лаборатории FusionBrain. Область моих исследований касается методов редактирования изображений.

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

Совсем недавно нашу статью приняли на одну из топ‑конференций по компьютерному зрению CVPR 2024 (эта конференция недавно стала самой цитируемой!). Наша статья про то, как можно редактировать лица в высоком качестве с помощью генеративной модели StyleGAN. Почитать её целиком можно на архиве, а здесь же я хотел кратко рассказать о том, что именно мы сделали.

Недавно, а именно 10 июня, исполнилось ровно 10 лет с момента выхода статьи Яна Гудфеллоу и его коллег, в котором были впервые описаны генеративно‑состязательные нейросети или просто GANs — generative adversarial networks. Эта такая архитектура, в которой две независимые сети — генератор и дискриминатор — бьются друг с другом ради того, чтобы качество генерации было на высшем уровне. О них уже написано довольно много (например, рекомендую вот эту статью моего руководителя Айбека Аланова), поэтому останавливаться на них я не стану.

В конце 2017 года спецы из NVIDIA открыли новую страницу в использовании GAN, представив модель StyleGAN, которая в своё время стала фундаментальной базой для задания SOTA‑планки в области генерации изображений, в особенности человеческих лиц. Чем же StyleGAN так хороша?

Кратко про StyleGAN

Для начала стоит отметить, что в основе генератора StyleGAN лежат свёрточные нейронные сети. Генерация начинается с блока, отвечающего за разрешение 4 х 4, потом увеличивается до 8 х 8, потом до 16 х 16 и так далее до 1024 х 1024. Таким образом формируется 9 блоков, каждый из которых состоит из 2 специальных свёрточных слоёв.

В обычных GAN на вход в генератор подаётся нормальный шум, выбранный из так называемого латентного или, по‑другому, Z‑пространства в форме вектора некоторого размера, например, 512. В StyleGAN всё хитрее: z‑вектор сначала прогоняется через специальную MappingNetwork, и становится w‑вектором того же размера — 512. W‑пространство нужно, поскольку пространство нормальных шумов Z слишком простое. В нём сложно найти такое смещение (направление), которое, будучи добавленным к исходному латенту, добавляет на изображение, например, улыбку (см. Рис. 3). W‑пространство, напротив, выучивается быть достаточно сложным, за счёт чего искать в нём направления для редактирования сильно легче.

Далее этот w‑вектор попадает в каждый из блоков StyleGAN с помощью специальной операции AdaIN (в более поздних версиях StyleGAN немного по другому, но суть такая же). Ключевое здесь то, что в обычных GAN вектор входит только в первый слой сети, в то время как в StyleGAN он скармливается каждому слою, поэтому поздние слои не забывают информацию об атрибутах синтезируемого изображения, которая содержится в w‑векторе.

А что с реальными картинками?

Мы в общих чертах разобрали, как работает StyleGAN, и понимаем, как генерируется изображение, и как управлять его атрибутами — круто! Но как же работать с реальными картинками? Тут‑то и вступает в дело задача инверсии — когда у нас есть реальная картинка и мы хотим найти какие то внутренние представления StyleGAN (латенты), из которых предобученный генератор синтезирует эту самую картинку.

В самых первых работах по инверсии авторы искали именно w‑латенты. Это удобно, т.к., чтобы редактировать изображения, достаточно соответствующим образом поменять найденные w‑вектора (а для этого мы знаем кучу направлений из W‑пространства, каждое из которых отвечает за свой атрибут), и синтезировать результат из изменённых латентов.

Однако всё не так просто. Если наша картинка довольно большая — скажем, 1024 x 1024 x 3, — то при попытке запихать её в вектор размерности 512 потеряется много информации. Чтобы уменьшить этот негативный эффект, предлагается работать в так называемом W⁺‑пространстве. В нём для каждого свёрточного слоя StyleGAN ищется свой собственный w‑вектор. Так как в StyleGAN у нас 9 х 2 = 18 слоёв, то размерность W⁺ будет 18 x 512, что в 18 раз увеличивает количество информации о картинке, которую мы можем сохранить.

Но даже этого мало, чтобы учесть все детали сложных изображений. Поэтому было придумано расширить W⁺‑пространство с помощью дополнительного F‑пространства. Идея следующая: давайте помимо маленьких w‑векторов, мы также будем искать большой F‑тензор, составленный из выходов k‑го слоя генератора. То есть мы начнём синтез в генераторе не с константы размера 512 х 4 х 4, а прямо с предсказанного тензора (в нашей работе мы использовали размерность 512 x 64 x 64). Этот тензор настолько большой, что в него уже помещается достаточно много информации об исходной картинке, и это позволяет сохранить много деталей.

Однако при решении задачи о редактировании стиля важно научиться правильно менять этот тензор желаемым образом. И если тому, как менять w‑латенты, чтобы получилось желаемое редактирование, посвящено много работ, то манипуляции в F‑пространстве далеко не так очевидны. Именно этой проблеме и посвящена наша статья.

А как инвертировать?

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

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

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

Какие латенты искать?

Как я уже писал выше, существуют разные латентные пространства, в которых мы можем искать представления для наших картинок, в частности W⁺ и F. Ещё при создании StyleGAN W⁺‑пространство подразумевалось как пространство, в котором будут производиться все редактирования, поэтому с этим там всё хорошо. Однако маленький размер пространства не помогает сохранить все детали большого исходного изображения.

F, напротив, за счёт предсказания большого тензора позволяет перенести все детали в инверсию. Но тут возникает проблема: бóльшая часть информации об изображении находится в F‑тензоре, а во время редактирования мы знаем, как изменять только w‑латенты. Если менять w, но не менять F, то редактирования просто перестанут работать, либо создадут кучу артефактов. И тогда мы сталкиваемся с ещё одним tradeoff'ом: между качество инверсии (количеством сохраняемых деталей исходного изображения) и возможностями к редактированию полученной инверсии.

Как получить всё и сразу

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

Нам удалось получить несколько профитов сразу:

• быстрое время работы,

• высокое качество реконструкции,

• хорошую редактируемость.

Для этого мы придумали фреймворк, в котором обучение энкодера делится на две фазы. Сначала мы просто учим наш энкодер выучивать представления в F‑пространстве — модуль, отвечающий за это, мы назвали Inverter (ведь он занимается инверсией). При этом теряется редактируемость: чтобы её восстановить, во второй фазе мы вводим и обучаем новый модуль под названием Feature Editor, чья задача — модифицировать F‑тензор таким образом, чтобы добиться желаемого изменения стиля.

На словах просто, но как учить такой модуль? Это можно было бы сделать, будь у нас множество пар с исходным и модифицированным изображениями. Например: реальное лицо и ровно оно же, но с улыбкой, или ровно оно же, но с другой причёской и т. п.. Однако на практике у нас нет таких данных, а собирать подобный датасет было бы ООООЧЕНЬ затратно и долго.

В нашей статье для решения этой проблемы мы предлагаем трюк: а давайте для получения таких пар сгенерируем данные. Основная проблема w‑энкодеров в том, что они не сохраняют деталей. Однако если мы посчитаем, что инверсия w‑энкодера — это реальная картинка, то её редактирование w‑энкодером выглядят вполне себе реалистично (см. Рис. 9). Так давайте использовать инверсию, полученную от w‑энкодера, как вход для нашего пайплайна, а какое‑то случайное редактирование, полученное также от этого w‑энкодера, как таргет для нашего пайплайна. Внутри обучаемого пайплайна как раз будем проводить то же самое редактирование.

Таким способом мы выучим FeatureEditor, чтобы его редактирования работали хотя бы на не слишком сложных данных. Однако реальные картинки обычно действительно содержат сильно больше деталей, чем те, что используются у нас в обучении. Чтобы избежать потери деталей (которые Inverter успевает собрать), мы предложили параллельно с задачей редактирования обучать Feature Editor на задачу инверсии, но уже с реальными картинками. Таким образом, мы показываем Feature Editor как редактировать простые (синтетические) данные, как инвертировать сложные (реальные) данные, и, о чудо, он сам начинает понимать как достаточно хорошо редактировать сложные данные!

Итого, мы получаем:

1. Отличное качество инверсии за счёт того, что Inverter предсказывает F‑тензор достаточно большого размера;

2. Хорошее качество редактирования благодаря применению хитрой техники обучения на синтетических данных, которую описали выше;

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

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

Весь фреймворк целиком получил название StyleFeatureEditor (SFE). Мы сравнили его с существующими энкодерными (e4e, psp, StyleTransformer, ReStyle, PaddingInverter, HyperInverter, Hyperstyle, HFGI, Feature‑Style, StyleRes) и оптимизационным (PTI) методами на датасетах с лицами и автомобилями. Наш подход не только успешно справляется с добавлением новых стилевых элементов, но и бережно относится к деталям исходных изображений. Продемонстрируем это на примере лиц с большим числом необычных деталей: грим, татуировки и т. п..

Можно заметить, что для простой инверсии (первый ряд) SFE аккуратно воспроизвёл шляпку, в то время как большинство других методов её практически потеряло. Или обратите внимание на третий ряд: мы решили надеть на Ло'ака очки. С этим справилось большинство подходов, но только наш сохранил его глаза жёлтыми.

А вот как выглядят результаты эксперимента в домене с машинами:

Впрочем, количественные оценки мы тоже делали. Для измерения качества инверсии всё довольно просто — достаточно измерить похожесть оригинальной картинки и её инверсии. Для оценки похожести мы измеряли достаточно классические L2, LPIPS и MS‑SSIM. Также, чтобы проверить, насколько наши данные в целом похожи на реальные (по распределению), мы считали FID между датасетом исходных картинок, и датасетом их инверсий. По таблице сравнения видно, что наш метод с довольно неплохим отрывом побивает всех конкурентов!

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

Давайте считать, что у нас есть датасет, где для каждой фотки мы знаем какой‑то набор атрибутов — есть ли на лице очки, есть ли улыбка, молодой или пожилой человек и т. п. (такие данные уже есть, скажем, Celeba). И давайте предположим, что мы хотим проверить, насколько хорошо наш метод добавляет улыбку. С помощью разметки разделим наш датасет на две части: лица с улыбкой, и лица без неё. К лицам без улыбки применим наш метод, который добавит улыбку. Теперь мы можем сравнить распределения реальных данных с улыбкой, и тех, которые получились у нашего метода — посчитаем между ними FID.

FID подходит как нельзя лучше, т.к. не сравнивает конкретные две картинки, а старается смотреть на распределение обоих наборов данных целиком. Логично, что, если в одном датасете (реальном) везде есть улыбка, а во втором её нет (или она с артефактами), то распределения будут отличаться!

Таким, немного хитрым образом, мы и можем оценить, насколько хорошо редактирует наш метод. И тут всё также довольно неплохо! Мы редактируем практически не хуже, а иногда и лучше прошлой SOTA StyleRes, при этом за счёт инверсии явно сохраняем больше деталей. И всё это при том, что, запускаясь на TeslaV100, SFE требует всего 0.07 секунды на редактирование одного изображения! Более подробные результаты и сравнение с другими методами можно посмотреть в таблице ниже:

Возможно, кому‑то удобнее будет воспринять проведённые нами сравнения в виде диаграммы:

И ещё немного примеров работы SFE на датасете с лицами знаменитостей: 

Спасибо всем, кто дочитал до конца. Надеюсь, вам стало чуть понятнее, как работает StyleGAN, а также понравился наш фреймворк. Можете сами опробовать демо нашего метода на HuggingFace, мы также будем рады вашим звёздам и форкам к нашему репозиторию.

Список авторов: Денис Бобков, Вадим Титов, Айбек Аланов, Дмитрий Ветров

Для сотрудничества можно обращаться к руководителю лаборатории Fusion Brain Андрею Кузнецову

Следите за новостями в каналах: CompleteAI, AbstractDL, Dendi Math&AI