Реалистичные пейзажи из рисунков от NVIDIA: глубокое погружение

Многие уже слышали про нейросеть GauGAN от NVIDIA, которая умеет создавать потрясающие пейзажи из примитивных набросков. Мы решили детально разобраться, какие методы использовались для её создания.

Раньше подобные фотореалистичные изображения генерировались путём «сшивания» между собой фрагментов из готовой базы данных с фотографиями. Современные же алгоритмы на основе нейросетей изучают непосредственно отображения — связи между исходными и целевыми данными. Они, как правило, работают гораздо быстрее и не используют внешние наборы данных.

Механизм, реализованный NVIDIA — это преобразование маски, состоящей из нескольких различных по содержанию сегментов, в фотореалистичный снимок. Они назвали его «синтез семантического образа». Традиционные сетевые архитектуры обычно используют несколько алгоритмов для различных слоёв, но имеют тенденцию «вымывать» информацию из входной маски. Чтобы решить эту проблему, разработчики использовали специальный пространственно-адаптивный слой, который может эффективно распространять семантическую информацию по всей сети. С его помощью компактная сеть синтезирует гораздо лучшие результаты, чем другие современные методы. Кроме того, она может генерировать изображения разного стиля, как показано на рисунке:

Синтез семантического образа

Напомним, что для создания уникальных образцов, в данном случае изображений, обычно используются глубокие генеративно-состязательные сети (GAN). GAN состоят из двух компонентов: генератора и дискриминатора. Генератор пытается создать реалистичные изображения таким образом, чтобы дискриминатор не смог отличить их от настоящих. Модель NVIDIA основана на GAN, но она синтезирует не совсем случайные, а условно случайные изображения, поскольку входные данные задаются пользователем (карта пейзажа).

Важный компонент глубоких нейронных сетей — слои нормализации. Они бывают безусловные (не зависят от внешних данных) и условные. Пространственно-адаптивный слой, используемый NVIDIA, является слоем условной нормализации и называется SPADE (SPatially-Adaptive (DE)normalization). По принципу работы он похож на Batch normalization и фактически обобщает несколько существующих слоёв нормализации. Со SPADE нет необходимости отправлять карту сегментов маски на первый слой генератора, поскольку он заранее изучает параметры модуляции и закодированную информацию о расположении меток. Поэтому энкодер генератора, обычно использующийся в архитектурах GAN, отбрасывается. Так сеть получается более простой и «лёгкой».

На рисунке показана архитектура генератора, которая включает в себя несколько блоков ResNet со слоями, повышающими разрешение изображения. Поскольку каждый остаточный блок работает с участками разного масштаба, SPADE уменьшает семантическую маску для соответствия пространственному разрешению.

Почему же SPADE так хорошо работает? Если вкратце — потому что он лучше сохраняет семантическую информацию по сравнению с обычными слоями нормализации, которые часто «вымывают» её при работе с однородными или плоскими сегментными масками.

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

Это ограничение присутствует у широкого диапазона архитектур генераторов, включая pix2pixHD. На рисунке можно увидеть, как выглядит «вымывание»:

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

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

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

Детали

Ко всем слоям генератора и дискриминатора применяется спектральная норма. Скорости обучения установлены на 0.0001 для генератора и 0.0004 для дискриминатора. В качестве оптимизатора используется Adam с параметрами β1 = 0, β2 = 0.999. Все эксперименты проводятся на NVIDIA DGX-1 с 8 графическими процессорами V100.

Наборы данных

Эксперименты проводятся на нескольких датасетах, содержащих в себе фотографии с различными пейзажами: COCO-Stuff (118 тыс. обучающих и 5 тыс. тестовых изображений), ADE20K (20120 / 2000), ADE20K-outdoor (часть предыдущего набора), Cityscapes (3000 / 500), Flickr Landscapes (41000 фотографий из Flickr и 1000 тестовых образцов).

Метрики эффективности

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

Для измерения точности сегментации используется метрика среднего пересечения над объединением (Intersection-over-Union, mIoU) и подлинность пикселей (accu). Для сегментации каждого набора данных используются современные нейросети: DeepLabV2 для COCO-Stuff, UperNet101 для ADE20K и DRN-D-105 для Cityscapes. Помимо этого, для измерения расстояния между распределениями синтезированных и реальных изображений используется расстояние Фреше (Frechet Inception Distance, FID) — чем оно меньше, тем лучше модель.

Исходные данные

Метод сравнивается с тремя наиболее популярными моделями семантического синтеза: pix2pixHD, сеть каскадного уточнения (cascaded refinement network, CRN) и алгоритм “Semi-parametric Image Synthesis” (SIMS). pix2pixHD синтезирует изображения с помощью GAN. CRN использует глубокую нейросеть, которая многократно уточняет выходной снимок, увеличивая его разрешение. SIMS же объединяет сегменты фотографий из обучающей выборки и повышает качество границ между ними. Чтобы сравнение было честным, модели CRN и pix2pixHD обучены с использованием параметров, предоставленных их авторами. Архитектура SIMS очень требовательна к наборам данных, и её нельзя обучить на больших датасетах, таких как COCO-Stuff и ADE20K. Поэтому там, где это возможно, использовались только результаты предварительно обученной модели.

Результаты

Количественные сравнения

Как показано в таблице, метод NVIDIA (помечен как Ours) превосходит все остальные с большим отрывом. Модель SIMS дала более низкий показатель FID на наборе Cityscapes, но её характеристики сегментации оказались хуже. Это связано с тем, что SIMS «сшивает» патчи изображений из обучающей выборки, и полученный снимок может лучше соответствовать реальному. Но при этом нет гарантии, что в наборе найдутся фотографии с точно совпадающими фрагментами (например, человек в необычной позе вряд ли окажется сразу на двух изображениях), поэтому SIMS часто копирует объекты с несовпадающими сегментами.

Качественные сравнения

На рисунках видно, что модель NVIDIA показывает гораздо лучшее визуальное качество. Когда размер обучающей выборки невелик, модель SIMS тоже генерирует хорошие изображения, хотя некоторые детали могут отклоняться от заданной маски (например, форма бассейна на втором рисунке).

Для получения фотографий с разным стилем из одной и той же маски авторы протестировали модель на наборе данных Flickr Landscape:

При этом пользователь сам может выбрать стиль и внешний вид изображения, поэтому процесс синтеза является управляемым.

Исходный код

Не так давно NVIDIA опубликовала исходный код своей нейросети, и теперь каждый желающий может обучить её и проверить на практике. Код доступен по лицензии CC BY-NC-SA 4.0 (Creative Commons Attribution-NonCommercial-ShareAlike 4.0), то есть для использования только в некоммерческих целях. Модель написана на языке Python с помощью фреймворка PyTorch.

Чтобы достичь тех же результатов обучения, что и NVIDIA, вам понадобится система не хуже, чем NVIDIA DGX-1 с 8 графическими процессорами V100. Например, можно использовать наши серверы с GPU NVIDIA Tesla V100.

Установка

Клонируйте репозиторий:

Перед использованием вам понадобятся PyTorch 1.0 и Python 3+. Вы можете установить все необходимые зависимости с помощью команды:

Также вам потребуется репозиторий Synchronized-BatchNorm-PyTorch:

Подготовка данных

Для работы нужен один из датасетов: COCO-Stuff, Cityscapes или ADE20K.

Подготовка набора COCO-Stuff

Набор можно скачать здесь. В частности, вам понадобится загрузить архивы train2017.zip, val2017.zip, stuffthingmaps_trainval2017.zip, и annotations_trainval2017.zip. Изображения, метки и карты объектов должны иметь ту же структуру файлов и папок, как в datasets/coco_stuff/. Карты объектов можно разбить на «карты вещей» ("things instance map") и «карты меток» ("stuff label map"). Для этого необходимо установить pycocotools с помощью pip install pycocotools и запустить скрипт datasets/coco_generate_instance_map.py.

Подготовка набора ADE20K

Скачайте набор отсюда. После его распаковки поместите jpg-файлы изображений из папки ADEChallengeData2016/images/ и png-файлы меток из ADEChallengeData2016/annotatoins/ в одну директорию.

Есть разные режимы загрузки изображений (--preprocess_mode), зависящие от параметров --load_size и --crop_size. Чтобы увидеть все возможные режимы, запустите python train.py --help и посмотрите на data/base_dataset.py. По умолчанию на этапе обучения изображения случайным образом поворачиваются по горизонтали. Чтобы предотвратить это, используйте параметр no_flip.

Генерация изображений с предварительно обученной моделью

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

Как только вы подготовили набор данных, сделайте следующее:

1. Загрузите tar-архив с моделью из гугл-диска, сохраните его в папке checkpoints/ и запустите:

2. Сгенерируйте изображения:

[type]_pretrained — это каталог с файлом контрольной точки, загруженный на шаге 1. Он должен быть одним из этих файлов: coco_pretrained, ade20k_pretrained или cityscapes_pretrained. [dataset] — набор данных, coco, ade20k или cityscapes, а [path_to_dataset] — путь к нему. Если вы запускаете модель на CPU, установите --gpu_ids на -1.

3. Выходные изображения по умолчанию сохраняются в ./results/[type]_pretrained/. Вы можете просмотреть их с помощью автоматически сгенерированного HTML-файла в этой папке.

Обучение новых моделей

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

1. Подготовьте набор данных. Если вы хотите использовать один из датасетов, указанных в статье, можете скачать его и применить опцию --dataset_mode, которая выберет нужный подкласс BaseDataset. Для пользовательских наборов данных самый простой способ — использовать ./data/custom_dataset.py, указав параметр --dataset_mode custom вместе с --label_dir [path_to_labels] и --image_dir [path_to_images]. Вам также нужно указать такие параметры как --label_nc для количества меток классов в наборе, --contain_dontcare_label — имеется ли неизвестная метка, или --no_instance  — указывает, что в наборе нет карт объектов.
2. Начните обучение:

Для просмотра списка доступных параметров используйте команду python train.py --help. Число используемых GPU можно указать с помощью --gpu_ids. Не забудьте, что --gpu_ids 0 — первый процессор, --gpu_ids 1 — второй, и так далее.

Тестирование

Процесс аналогичен тому, что используется для предварительно обученной модели:

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

Структура файлов

— train.py, test.py: скрипты для обучения и тестирования

— trainers/pix2pix_trainer.py: управляет процессом обучения

— models/pix2pix_model.py: создаёт сети и вычисляет потери

— models/networks/: определяет архитектуру всех моделей

— options/: создаёт списки параметров, используя пакет argparse

— data/: определяет класс для загрузки изображений и меток

Обучение с возможностью управлять стилем

Если вы хотите добавить возможность выбирать стиль изображения, используйте параметр --use_vae. Модель создаст и обучит ещё одну сеть netE в дополнение к netG и netD.

Новая нейросеть поможет архитекторам, художникам, дизайнерам и разработчикам игр быстро создавать и редактировать эскизы и модели. Не забывайте делиться своими результатами и предположениями, насколько далеко искусственный интеллект сможет продвинуться в живописи и творчестве. С оригинальной статьёй можно ознакомиться на сайте arxiv.org.