Нейросеть pix2pix обучили автоматически раскрашивать старые черно-белые фотографии. Модель позволит специалистам по колоризации сократить время работы, затрачиваемое на подбор цветов и ручную раскраску изображений.
Помимо ручного раскрашивания оригинальной черно-белой фотографии, процесс колоризации сопровождается обширными исследованиями аналогичных цветных изображений и исторических данных, требуемыми для точного восстановления цветов. Условная генеративно-состязательная нейронная сеть pix2pix позволила упростить этот процесс. Модель состоит из двух частей (рис. 1): генератора, обученного раскрашивать черно-белые изображения, и дискриминатора. Для обучения использовался датасет, состоящий из цветных изображений. Путем удаления цвета из каждого изображения генерировались пары цветных и черно-белых изображений. Задача генератора состояла в том, чтобы раскрасить черно-белые изображения так, чтобы заставить дискриминатор думать, что это изначально цветные изображения, и чтобы исходные цвета восстановлены как можно точнее. Дискриминатор обучается отличать исходно цветные фотографии от раскрашенных генератором. Благодаря этому конкурирующему процессу модель с каждой итерацией начинает более достоверно раскрашивать изображения.
Обучение нейросети производилось в два этапа. На первом этапе использовался датасет COCO, состоящий из 35 000 изображений объектов, для обучения нейросети типичной распространенности цветов в окружающей среде. На втором этапе использовался датасет из 2 000 портретов, взятых из Unsplash, для обучения модели раскрашивать лица людей. Поскольку старые фотографии более размытые по сравнению с цифровыми фотографиями, к каждому изображению из датасетов было добавлено случайное размытие. Результаты восстановления приведены на рис. 2 совместно с вручную скорректированными цветами. Несмотря на видимые на рис. 2 недостатки восстановления, а именно окраски некоторых частей изображения в слишком яркие цвета, нейросеть корректно распознает цвета на всем изображении целиком. Это сокращает время, затрачиваемое специалистами по колористике на подбор цветов и ручное раскрашивание, сводя их работу к быстрой цветокоррекции. Особенностью модели состоит в отсутствии необходимости распознавания объектов на изображении – метода, используемого в большинстве state-of-the-art-подходов к раскрашиванию изображений и требующего больших размеченных датасетов.