Главная » Софт » [Перевод] DeOldify: программа для раскрашивания чёрно-белых изображений

[Перевод] DeOldify: программа для раскрашивания чёрно-белых изображений

Если вкратце, задача этого проекта — раскрашивать и восстанавливать старые снимки. Я немного углублюсь в детали, но сначала давайте посмотрим фотографии! Кстати, большинство исходных изображений взято из подреддита r/TheWayWeWere, благодарю всех за такие качественные большие снимки.

Это лишь немногие примеры, и они вполне типичные!

Мария Андерсон в роли Маленькой феи и её паж Любовь Рябцова в балете «Спящая красавица» в Императорском театре, Санкт-Петербург, Россия, 1890


Женщина расслабляется в своей гостиной (1920, Швеция)

Студенты-медики позируют возле трупа, около 1890

Сёрфер на Гавайях, 1890

Крутящаяся Лошадь, 1898

Интерьер бара Миллера и Шумейкера, 1899

Париж в 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)

Я бы сохранил размер около 500px, плюс-минус, если запускаете программу на GPU с большим количеством памяти (например, GeForce 1080Ti 11 GB). Если памяти меньше, то вам придётся уменьшить размер картинок или попробовать запустить на CPU. Я на самом деле пытался сделать последнее, но по какой-то причине модель работала очень, абсурдно медленно, и я не нашел времени, чтобы исследовать проблему. Знатоки рекомендовали собрать Pytorch из исходников, тогда получится большой прирост производительности. Мда… В тот момент было не до того.
Визуализацию сгенерированных изображений по мере обучения можно выполнить и в Jupyter: нужно лишь установить значение true при создании экземпляра этого хука визуализации:

GANVisualizationHook(TENSORBOARD_PATH, trainer, 'trainer', jupyter=True, visual_iters=100

Поверьте, вам тоже захочется сделать именно так. Я предпочитаю оставить false и просто использовать Tensorboard. Кроме того, если оставить его работать слишком долго, Jupyter съест много памяти с такими изображениями.

По умолчанию, они сохраняются каждые 1000 итераций (это дорогостоящая операция). Веса модели тоже автоматически сохраняются во время тренировочных прогонов GANTrainer. Веса сохраняются отдельно для каждого размера тренировки. Они хранятся в корневой папке, которую вы указали для обучения, а название соответствует save_base_name, указанному в расписании обучения.

Я отношусь к этим заметкам просто как к удобному интерфейсу для прототипирования и визуализации, всё остальное пойдёт в файлы .py, как только я найду для них место. Я бы рекомендовал перемещаться по коду сверху вниз, начиная с Jupyter Notebook. У меня уже есть примеры визуализации, которые можно удобно включить и посмотреть: просто откройте xVisualization в Notebook, там указаны включённые в проект тестовые изображения (они лежат в test_images).

Если увидите GAN Schedules, то это самая уродливая вещь в проекте, всего лишь моя версия реализации прогрессивного обучения GAN, подходящая для генератора Unet.

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

Обычно при обучении вы увидите первые хорошие результаты на полпути, то есть с размера 192px (если используете предоставленные учебные примеры).

Уверен, что я где-то напортачил, так что пожалуйста, дайте знать, если это так.

  • Придётся немного поиграть с размером изображения, чтобы получить лучший результат. Модель явно страдает некоторой зависимостью от соотношения сторон и размера при генерации изображений. Раньше было намного хуже, но ситуация значительно улучшилась с увеличением освещения/контраста и внедрением прогрессивного обучения. Я хочу полностью устранить эту проблему и сосредоточусь на ней, но пока не отчаивайтесь, если изображение выглядит чрезмерно насыщенным или со странными глюками. Скорее всего, всё станет нормально после небольшого изменения размера. Как правило, для перенасыщенных картинок нужно увеличить размер.
  • В дополнение к вышесказанному: получение лучших изображений действительно сводится к искусству выбора оптимальных параметров. Да, результаты выбираются вручную. Я очень доволен качеством, и модель работает вполне надёжно, но не идеально. Проект ещё продолжается! Думаю, инструмент можно использовать как «ИИ художника», но он пока не готов для широкой публики. Просто ещё не время.
  • Чтобы усложнить ситуацию: в настоящее время модель зверски жрёт память, поэтому на моей карте 1080Ti получается обработать картинки максимум 500-600px. Держу пари, что здесь есть много вариантов оптимизации, но я пока не сделал этого.
  • Я добавил нулевое заполнение (zero padding) в генератор Unet для всего, что не соответствует ожидаемым размерам (именно так я могу загрузить изображение произвольного размера). Это был очень быстрый хак, и он приводит к глупым правым и нижним границам на выходе тестовых изображений произвольного размера. Уверен, что есть лучший способ, но пока не нашёл его.
  • Модель любит синюю одежду. Не совсем уверен, почему, решение в поиске!

Буду выкладывать новые результаты в твиттере.

Дополнение от переводчика.
Из последнего в твиттере:

Представители народности сами у своей землянки, 1880


(оригинал)

Строительство лондонского метро, 1860


(оригинал)

Трущобы Балтимора, 1938

Тренажёрный зал на «Титанике», 1912


(оригинал)


Оставить комментарий

Ваш email нигде не будет показан
Обязательные для заполнения поля помечены *

*

x

Ещё Hi-Tech Интересное!

Получаем музыку Вк через сторонний API

В этот раз дело начиналось после закрытия методов audio в методе execute. Я решил посмотреть, как получают музыку сайты, которые предоставляют возможность ее скачать. Меня заинтересовал сайт vrit.me. Я залез во вкладку network и увидел интересный запрос: То есть, можно ...

Все, что нужно знать о стрессе и сильных эмоциях

Под катом — цикл из 3-х видео и текстовых версий к ним. Причина внутри нас. Главное — стресс вызван не внешними событиями. Мы можем управлять этим процессом, потому что он основан на нашем опыте. Внутри каждого существует процесс, который отвечает ...