Нейросеть получила удобный интерфейс редактирования изображений

Ученые из США и Германии создали инструмент для редактирования объектов на изображении DragGAN. Он основан на нейросети, которая на лету генерирует новые изображения. Пользователь может отмечать мышкой разные элементы на фотографии, и они реалистично перемещаются, поворачиваются, растягиваются и уменьшаются — похожий подход нередко используют в видеоиграх для редактирования внешности персонажа. Препринт доступен на arXiv.org.

В последние пару лет случился бум диффузионных нейросетей для генерации изображений по текстовым запросам — например, Stable Diffusion, DALL-E и Midjourney. Диффузионная модель умеет генерировать реалистичные изображения из случайного шума, а текстовое описание подсказывает модели, что именно должно быть на изображении. Но есть проблема: диффузионные нейросети недостаточно точны и не подходят, если нужно выполнить очень конкретную задачу при редактировании: сдвинуть объект на несколько пикселей в определенную сторону или повернуть его на определенное количество градусов.

Задачу точного редактирования изображений можно решить с помощью других нейросетевых архитектур. Например, до появления диффузионных нейросетей часто использовали GAN — генеративно-состязательные сети. Такая модель состоит из двух нейросетей: генеративной и состязательной. По сути две нейросети соревнуются друг с другом: состязательная сеть учится отличать реальные изображения от сгенерированных, а генеративная, в свою очередь, пытается сгенерировать максимально реалистичные изображения из случайного шума (на этой идее основаны и диффузионные модели). Состязательная сеть посылает сигнал генеративной модели — сообщает ей, насколько реалистичное получилось изображение. И так до тех пор, пока генеративная сеть не научится обманывать состязательную.

Группа ученых под руководством Кристиана Теобальта (Christian Theobalt) из Института информатики Общества Макса Планка придумала, как научить модель GAN делать сложные точечные изменения в изображении. Новый алгоритм называется DragGAN. Если пользователь хочет отредактировать какой-то объект на изображении, ему нужно отметить точками, где некоторые части объекта находятся до изменений и куда они должны переместиться после изменений. Например, чтобы повернуть мордочку кота на фотографии, нужно указать где сейчас находится нос и в какой точке он должен оказаться после поворота. Также можно отмечать область изображения, которая будет меняться.

Главная задача DragGAN — реалистично трансформировать объект на изображении, опираясь на отмеченные пользователем точки до и после изменений. Она решается алгоритмом как задача оптимизации. Путь от каждой точки «до» к точке «после» разбивается на множество маленьких шагов. На каждом шаге генерируется новое изображение, которое совсем немного отличается от предыдущего. После каждого шага алгоритм определяет новые позиции точек «до», которые понемногу меняются в ходе оптимизации. Когда они совпадут с позициями «после», алгоритм завершит работу и пользователь получит последнее сгенерированное изображение. Обычно для одной операции редактирования требуется от 30 до 200 маленьких шагов, которые в сумме занимают несколько секунд.

Модель обучали на нескольких датасетах с людьми (FFHQ, SHHQ), животными (AFHQCat), автомобилями (LSUN Car), пейзажами (LSUN, Landscapes HQ) и объектами под микроскопом (microscrope). Ее качество сравнили с похожей нейросетью UserControllableLT, которую авторы считают одной из лучших моделей для редактирования изображений с перемещением точек. На примерах видно, что DragGAN лучше определяет новое положение объектов и не делает лишних изменений.

Также DragGAN сравнили с UserControllableLT на классической задаче по генерации изображений — трансформации ключевых точек лица. На изображении лица всегда можно выделить координаты точек, которые определяют его строение, мимику и положение на фото, в том числе контуры глаз, носа, губ. Модели получили изображения двух разных лиц и должны были изменить первое изображение так, чтобы ключевые точки на нем совпали или хотя бы максимально приблизились к ключевым точкам на втором изображении. Чем лучше совпадали два набора ключевых точек, тем точнее модели удавалось повторить мимику и положение второго лица, сохранив черты первого. Разница в координатах ключевых точек у DragGAN оказалась в 3 раза меньше, чем у UserControllableLT.

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

Раньше мы рассказывали, как другая нашумевшая нейросеть GPT-4 научилась работать с изображениями. Модель от компании OpenAI может понимать и изображения, и текст, хотя ответы по-прежнему выдает только в текстовом виде.