[Перевод] DeOldify: программа для раскрашивания чёрно-белых изображений
Если вкратце, задача этого проекта — раскрашивать и восстанавливать старые снимки. Я немного углублюсь в детали, но сначала давайте посмотрим фотографии! Кстати, большинство исходных изображений взято из подреддита r/TheWayWeWere, благодарю всех за такие качественные большие снимки.
Это лишь немногие примеры, и они вполне типичные!
Мария Андерсон в роли Маленькой феи и её паж Любовь Рябцова в балете «Спящая красавица» в Императорском театре, Санкт-Петербург, Россия, 1890 
Женщина расслабляется в своей гостиной (1920, Швеция)
Студенты-медики позируют возле трупа, около 1890
Сёрфер на Гавайях, 1890
Крутящаяся Лошадь, 1898
Интерьер бара Миллера и Шумейкера, 1899 SodaFountain
Париж в 1880-е
Эдинбург с высоты птичьего полёта в 1920-е
Техасская женщина в 1938 г.
Люди на станции Ватерлоо впервые смотрят телевизор, Лондон, 1936
Урок географии в 1850 году
Китайские курильщики опиума в 1880 году
Обратите внимание, что даже действительно старые и/или плохого качества фотографии по-прежнему выглядят довольно круто:
Дедвуд, Южная Дакота, 1877
Братья и сёстры в 1877 году (Дедвуд)
Площадь Портсмут в Сан-Франциско, 1851
Самураи, около 1860-х гг
Конечно, модель не идеальна. Эта красная рука сводит меня с ума, но в остальном она фантастически работает:
Девушка племени сенека из ирокезов, 1908
Она также может раскрашивать чёрно-белые рисунки:
Это модель на основе глубокого обучения. В частности, я совместил следующие подходы:
- Self-Attention GAN. Единственное, что в качестве генератора используется предобученная Unet и я просто изменил её для спектральной нормализации и собственно, механизма Self-Attention. Это довольно простая модификация. Скажу вам, что разница поразительная по сравнению с предыдущей версией Wasserstein GAN, которую я пытался заставить работать. Мне нравилась теория Wasserstein GAN, но на практике она не работает. Но я просто влюбился в сети Self-Attention GAN.
- Структура обучения наподобие прогрессивного роста GAN (но не в точности такая). Разница в том, что количество слоёв остаётся постоянным: я просто изменил размер входных данных и скорректировал скорость обучения, чтобы переходы между размерами проходили успешно. Похоже, она выдаёт такой же конечный результат, но быстрее обучается, сама стабильнее и лучше выполняет обобщение.
- Правило TTUR (Two Time-Scale Update Rule). Здесь довольно понятно: просто итерации один к одному генератора/дискриминатора (критика) и более высокая скорость обучения дискриминатора.
- Функция потери генератора состоит из двух частей: одна из них является основной функцией Perceptual Loss (или Feature Loss) на базе VGG16 — она просто подталкивает модель генератора для репликации входного изображения. Вторая часть — оценка потерь от дискриминатора (критика). Для любопытных: только функции Perceptual Loss недостаточно для хорошего результата. Она имеет тенденцию просто поощрять кучу коричневого/зелёного/синего — ну понимаете, обманывая тест, в чём действительно хороши нейронные сети! Ключевой момент в том, что GAN по сути сами изучают для вас функцию потерь, что на самом деле является одним большим шагом на пути к тому идеалу, к которому мы стремимся в машинном обучении. И конечно же, результаты значительно улучшатся, когда машина сама обучается тому, что вы ранее кодировали вручную. Безусловно, здесь это имеет место.
Прелесть данной модели в том, что она довольно хороша в самых разных модификациях изображения. То, что вы видите выше, — это результаты модели раскрашивания, но это лишь один компонент в конвейере, который я хочу разработать с той же моделью.
Дальше я попытаюсь довести до совершенства старые изображения, и следующий пункт на повестке дня — модель улучшения насыщенности и сочности (defade). Сейчас она находится на ранних этапах обучения. В основном это та же модель, но с некоторыми настройками контраста/яркости в качестве симуляции блеклых фотографий и снимков, сделанных со старым/плохим оборудованием. Я уже получил некоторые обнадёживающие результаты:
В чём суть этого проекта? Я просто хочу применить GAN, чтобы старые фотографии выглядели очень-очень хорошо. И что более важно, это сделает проект полезным. И да, мне определённо интересно поработать с видео, но сначала нужно разобраться, как взять эту модель под контроль по потреблению памяти (это настоящий зверь). Было бы неплохо, если бы модели не обучались от двух до трёх дней на 1080Ti (к сожалению, типично для GAN). Хотя это мой ребёнок и я буду активно обновлять и улучшать код в обозримом будущем, но я постараюсь сделать программу максимально удобной для пользователей, хотя наверняка c ней найдутся какие-то трудности.
И клянусь, что должным образом задокументирую код… когда-нибудь. По общему признанию, я один из тех людей, кто верит в «самодокументированный код» (LOL).
Проект построен на замечательной библиотеке Fast.AI. К сожалению, это старая версия, и ещё предстоит обновить её до новой (это определенно на повестке дня). Итак, предварительные требования, вкратце:
- Старая библиотека Fast.AI. Закопавшись в проект на два месяца, я немного пропустил, что с ней случилось, потому что та, которая теперь помечена как «старая», на самом деле не похожа на ту, которая у меня. Всё изменилось за последние два месяца или около того. Поэтому если ничего не работает с другими версиями, я форкнул её здесь. Опять же, обновление до последней версии стоит на повестке дня, заранее извиняюсь.
- Все зависимости Fast.AI: там есть удобные файлы requirements.txt и environment.yml.
- Pytorch 0.4.1 (требуется spectral_norm, поэтому нужен последний стабильный релиз).
- JupyterLab.
- Tensorboard (т. е. установка Tensorflow) и TensorboardX. Думаю, это не строго обязательно, но так гораздо проще. Для вашего удобства я уже предоставил все необходимые хуки/обратные вызовы в Tensorboard! Есть примеры их использования. Примечательно, что по умолчанию изображения в процессе обработки записываются в Tensorboard каждые 200 итераций, поэтому вы получите постоянный и удобный вид того, что делает модель.
- ImageNet: отличный набор данных для обучения.
- Мощная видеокарта. Я бы очень хотел иметь больше памяти, чем 11 ГБ в моей GeForce 1080Ti. Если у вас что-то послабее, то будет трудно. Unet и Critic абсурдно велики, но чем они больше, тем лучше результаты.
Если хотите самостоятельно начать обработку изображений прямо сейчас без обучения модели, то можете скачать готовые веса здесь. Затем откройте ColorizationVisualization.ipynb в JupyterLab. Убедитесь, что там есть строка со ссылкой на веса:
colorizer_path = Path('/path/to/colorizer_gen_192.h5')
Тогда нужно загрузить модель colorizer после того, как инициализируется netG:
load_model(netG, colorizer_path)
Затем просто поместите любые изображения в папку /test_images/, откуда вы запускаете программу. Можете визуализировать результаты в Jupyter Notebook такими строчками:
vis.plot_transformed_image("test_images/derp.jpg", netG, md.val_ds, tfms=x_tfms, sz=500)
Я бы сохранил размер около 500 px, плюс-минус, если запускаете программу на GPU с большим количеством памяти (например, GeForce 1080Ti 11 GB). Если памяти меньше, то вам придётся уменьшить размер картинок или попробовать запустить на CPU. Я на самом деле пытался сделать последнее, но по какой-то причине модель работала очень, абсурдно медленно, и я не нашел времени, чтобы исследовать проблему. Знатоки рекомендовали собрать Pytorch из исходников, тогда получится большой прирост производительности. Мда… В тот момент было не до того.
Визуализацию сгенерированных изображений по мере обучения можно выполнить и в Jupyter: нужно лишь установить значение true при создании экземпляра этого хука визуализации:
GANVisualizationHook(TENSORBOARD_PATH, trainer, 'trainer', jupyter=True, visual_iters=100
Я предпочитаю оставить false и просто использовать Tensorboard. Поверьте, вам тоже захочется сделать именно так. Кроме того, если оставить его работать слишком долго, Jupyter съест много памяти с такими изображениями.
Веса модели тоже автоматически сохраняются во время тренировочных прогонов GANTrainer. По умолчанию, они сохраняются каждые 1000 итераций (это дорогостоящая операция). Они хранятся в корневой папке, которую вы указали для обучения, а название соответствует save_base_name, указанному в расписании обучения. Веса сохраняются отдельно для каждого размера тренировки.
Я бы рекомендовал перемещаться по коду сверху вниз, начиная с Jupyter Notebook. Я отношусь к этим заметкам просто как к удобному интерфейсу для прототипирования и визуализации, всё остальное пойдёт в файлы .py, как только я найду для них место. У меня уже есть примеры визуализации, которые можно удобно включить и посмотреть: просто откройте xVisualization в Notebook, там указаны включённые в проект тестовые изображения (они лежат в test_images).
Если увидите GAN Schedules, то это самая уродливая вещь в проекте, всего лишь моя версия реализации прогрессивного обучения GAN, подходящая для генератора Unet.
Предобученные веса для генератора колоризатора тоже здесь. Проект DeFade пока в работе, я постараюсь выложить хорошие веса в течение нескольких дней.
Обычно при обучении вы увидите первые хорошие результаты на полпути, то есть с размера 192 px (если используете предоставленные учебные примеры).
Уверен, что я где-то напортачил, так что пожалуйста, дайте знать, если это так.
- Придётся немного поиграть с размером изображения, чтобы получить лучший результат. Модель явно страдает некоторой зависимостью от соотношения сторон и размера при генерации изображений. Раньше было намного хуже, но ситуация значительно улучшилась с увеличением освещения/контраста и внедрением прогрессивного обучения. Я хочу полностью устранить эту проблему и сосредоточусь на ней, но пока не отчаивайтесь, если изображение выглядит чрезмерно насыщенным или со странными глюками. Скорее всего, всё станет нормально после небольшого изменения размера. Как правило, для перенасыщенных картинок нужно увеличить размер.
- В дополнение к вышесказанному: получение лучших изображений действительно сводится к искусству выбора оптимальных параметров. Да, результаты выбираются вручную. Я очень доволен качеством, и модель работает вполне надёжно, но не идеально. Проект ещё продолжается! Думаю, инструмент можно использовать как «ИИ художника», но он пока не готов для широкой публики. Просто ещё не время.
- Чтобы усложнить ситуацию: в настоящее время модель зверски жрёт память, поэтому на моей карте 1080Ti получается обработать картинки максимум 500–600 px. Держу пари, что здесь есть много вариантов оптимизации, но я пока не сделал этого.
- Я добавил нулевое заполнение (zero padding) в генератор Unet для всего, что не соответствует ожидаемым размерам (именно так я могу загрузить изображение произвольного размера). Это был очень быстрый хак, и он приводит к глупым правым и нижним границам на выходе тестовых изображений произвольного размера. Уверен, что есть лучший способ, но пока не нашёл его.
- Модель любит синюю одежду. Не совсем уверен, почему, решение в поиске!
Буду выкладывать новые результаты в твиттере.
Дополнение от переводчика.
Из последнего в твиттере:
Представители народности сами у своей землянки, 1880
(оригинал)
Строительство лондонского метро, 1860
(оригинал)
Трущобы Балтимора, 1938
Тренажёрный зал на «Титанике», 1912
(оригинал)
