Делайте черно-белые фотографии красочными с помощью Generative Adversarial Network (GAN) от Fastai

Введение

Когда я вижу историческое черно-белое фото, мне всегда интересно, какой был настоящий цвет, что увидел фотограф, когда делал снимок.

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

Другой подход состоит в том, чтобы научить компьютер понимать цвета различных объектов, предоставив большое количество цветных фотографий, а затем попросить компьютер раскрасить, определив правдоподобный цвет для обнаруженного возражения. Благодаря развитию глубокого обучения это можно сделать с помощью Generative Adversarial Network (GAN).

Генеративная состязательная сеть (GAN)

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

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

В этой работе я взял DeOldify Джейсона Антика в качестве образца (https://github.com/jantic/DeOldify) и обучил GAN с помощью fastai.

Набор данных

Я использовал ImageNet в качестве набора данных для обучения. Однако из-за ограничения вычислений (ограничение времени работы 12 часов в Colab) было непросто использовать все изображения, поэтому в итоге я использовал только 5000 изображений из него.

Делайте черно-белые фотографии

Сначала мне нужно сделать цветное изображение черно-белым для тренировки. Я также добавил немного шума, чтобы придать изображениям ощущение старой фотографии.

Предпусковой генератор

Теперь у меня готов набор данных, я могу приступить к обучению. Я создал обучающую программу U-Net для генератора и предварительно обучил ее перед GAN. Это легко сделать с помощью библиотеки fastai.

Я начал с небольшого размера (64 пикселя), поэтому я могу обучать его быстрее с большим размером пакета.

Затем я увеличил размер до 128 и 256. Я начал видеть некоторые хорошие результаты с большим количеством эпох.

Предварительно подготовленный дискриминатор

После предварительной подготовки генератора я использовал его для создания цветных фотографий из набора черно-белых данных. Затем я использовал эти сгенерированные фотографии и исходные фотографии для обучения дискриминатора.

Через 10 эпох дискриминатор смог определить, создано ли изображение с точностью 85%.

GAN

Пришло время собрать их вместе. Fastai предоставляет очень полезный инструмент для создания GAN, объединяющий предварительно обученный генератор и дискриминатор. После 50-летней тренировки я получил свои предварительные результаты.

Полученные результаты

Я скачал для тестирования несколько черно-белых фотографий, которых модель еще не видела. Вот некоторые результаты.

Модель умеет раскрашивать дерево и траву. Он также хорошо раскрашивал человеческую кожу.

Некоторые цвета не совсем реалистичны, но все же достаточно убедительны.

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

Веб приложение

Прототип веб-приложения можно быстро создать с помощью Streamlit. Для этого не требуется много навыков HTML и Java, поэтому он действительно полезен для таких веб-манекенов, как я.



Благодарим за чтение, предложения и отзывы приветствуются.

Еще примеры: