Раскрашиваем черно-белые изображения разными инструментами: подробный разбор

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

Определяем набор данных

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

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

Далее был разработан скрипт для обработки набора данных. Он выполнял следующие шаги:

• Преобразовывал изображения в черно-белый формат (градации серого);
• Разделял данные на обучающий и тестовый наборы в виде пар «оригинал — чёрно-белая версия»;
• Обработал более 1700 картин, сохранив их в нужном формате.

Для работы с изображениями использовалась библиотека PIL (Python Imaging Library). Несмотря на то, что она устарела и была заменена на Pillow, её функционала оказалось достаточно для нашей задачи, так как требовалось лишь преобразовать изображения в чёрно-белый формат.

Ниже представлен код, который выполняет преобразование изображений:

			# loop through all the files in the input folder
for images in os.listdir(folder_path):
    # get the file path
    input_file = os.path.join(folder_path, images)
    # load image using matplotlib
    image = plt.imread(os.path.join(folder_path, images))
    if count < 5:
      print(images)
      # display the image
      plt.imshow(image)
      plt.show()
      # open the image
    with Image.open(input_file) as img:
        # convert the image to black and white
        bw_img = img.convert('L')
        # get the output file path
        output_file = os.path.join(output_folder, images)
        # save the image
        bw_img.save(os.path.join(output_folder, "bw-"+images))
        if count < 5:
          plt.imshow(bw_img)
          plt.show()
          count += 1

		

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

Выборка из полученного датасета

Выбираем инструмент

Сравним более чем 5 различных подходов к решению задачи. У каждого из них есть свои преимущества и недостатки.

*Спойлер: Наилучший результат показали GAN и Autoencoder + ResNet.

OpenCV

Пример оригинально изображения и 3-х изображений в разных палитрах (слева направо): Rainbow, Plasma, Turbo

OpenCV — open source библиотека компьютерного зрения, которая предназначена для анализа, классификации и обработки изображений.

Преимущества:

• Позволяет гибко настраивать параметры алгоритмов и экспериментировать с различными методами раскрашивания;

• Отличается высокой скоростью работы, поэтому удобен для обработки больших объемов данных в реальном времени;

• Имеет понятный интерфейс для работы с изображениями, что упрощает применение методов машинного зрения.

Недостатки:

• Может требовать больших вычислительных ресурсов при обработке сложных изображений;

• Подбор цветового профиля и параметров раскрашивания затруднен, так как они сильно влияют на конечный результат;

• Возможны искажения и ошибки, особенно если на изображении присутствуют сложные текстуры или объекты необычной формы.

MATPLOTLIB

Пример оригинально изображения и 3-х изображений в разных палитрах (слева направо): Terrain, HSV, Turbo

Matplotlib — популярная Python-библиотека для визуализации данных. Она используется для создания любых видов графиков: линейных, круговых диаграмм, построчных гистограмм и других — в зависимости от задач.

Преимущества:

• Позволяет быстро и просто изменять схему, не требуя детальной работы с цветами;

• Обеспечивает удобство визуализации цветных изображений благодаря встроенным палитрам.

Недостатки:

• Выбор цветовой схемы — сложный процесс, требующий внимательного подхода. Автоматическое использование get_cmap не всегда приводит к оптимальному результату;

• В некоторых случаях может потребоваться гибкое управление цветами, что выходит за рамки стандартных возможностей Matplotlib.

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

CNN

Один из самых простых способов автоматического раскрашивания чёрно-белых изображений — применение свёрточных нейронных сетей (CNN) в архитектуре автокодировщика. Для реализации использовался готовый код с Kaggle.

Преимущества:

• Выявляет закономерности оттенков, заполняя изображение наиболее естественными цветами;

• Имеет высокое качество результатов в сравнении с традиционными методами;

• Тратит мало времени на обработку, что особенно важно при работе с большими объемами изображений.

Недостатки:

• Не может точно предсказать цвета объектов — например, модель может раскрасить траву в желтоватый, а не в зелёный цвет;

• Требует большого количества данных для обучения, иначе точность раскрашивания будет низкой;

• Имеет ограничение на размер изображения — слишком большие изображения могут потребовать значительных вычислительных ресурсов;

• Имеет риск переобучения, особенно если модель плохо настроена.

Этот метод — хорошая отправная точка, но для улучшения качества раскрашивания его часто приходится комбинировать с более сложными архитектурами, такими как Pix2Pix GAN или Autoencoder + ResNe.

AUTOENCODER + RESNET

Пример покраски за 8 и 12 эпох

Один из наиболее эффективных методов автоматической раскраски изображений — использование автокодировщика (Autoencoder) в сочетании с ResNet. Этот подход позволяет не только предсказывать цвета, но и сохранять детали и текстуру изображения. Код реализации был взят с Kaggle.

Преимущества:

• Глубокое понимание структуры изображения — модель способна распознавать границы объектов и корректно раскрашивать их, сохраняя детали;

• Автоматизированный процесс — после обучения модель сама выбирает подходящие цвета, сводя к минимуму вмешательство человека;

• Хорошие результаты за короткое время — всего за 8 эпох удалось получить качественные раскрашенные изображения.

Недостатки:

• Высокие вычислительные затраты — обучение заняло 5–6 часов;

• Требовательность к данным — для обучения необходима большая и разнообразная выборка изображений;

• Сложность реализации — настройка модели и подбор гиперпараметров требуют знаний в области глубокого обучения.

Несмотря на эти ограничения, комбинация Autoencoder + ResNet — один из самых перспективных методов раскрашивания чёрно-белых изображений. Она превосходит базовые архитектуры CNN и традиционные алгоритмы, обеспечивая высокое качество цветопередачи.

GAN

Пример покраски изображения(слева-направо): черно белая фотография, сгенерированное изображение, оригинал

Эти сети состоят из двух моделей: генератора, который пытается создать реалистичное цветное изображение, и дискриминатора, который оценивает, насколько оно похоже на настоящее. Такой механизм позволяет GAN создавать изображения, максимально приближенные к оригиналу. Код реализации был взят с Kaggle.

Преимущества:

• Высокое качество изображений — благодаря состязательной природе модель генерирует цветные изображения с реалистичной текстурой;

• Учет контекста — GAN анализирует структуру изображения, определяя объекты и их окружение, что помогает точнее предсказывать цвета;

• Гибкость и детализация — сеть адаптирует цветопередачу к особенностям изображения, улучшая конечный результат.

Недостатки:

• Проблемы с переобучением — дискриминатор может просто запоминать входные данные, а генератор — создавать цветные изображения, которые не соответствуют оригинальному контенту;

• Сложность обучения — GAN требует больших вычислительных ресурсов и времени на обучение;

• Не всегда стабильные результаты — иногда сеть может генерировать цвета, не соответствующие реальности, особенно при недостатке обучающих данных.

DeOldify

DeOldify — проект с открытым исходным кодом, доступный на GitHub. Он основан на трёх моделях, одна из которых — NoGAN. Репозиторий с github.

NoGAN представляет собой усовершенствованный метод обучения генеративно-состязательных сетей (GAN), который сочетает в себе преимущества GAN-обучения (реалистичная и качественная цветокоррекция) с минимизацией побочных эффектов (мерцание объектов в видео).

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

Преимущества:

• Высокая реалистичность — модель обучена на огромном количестве изображений, что позволяет ей точно предсказывать цвета;

• Стабильные результаты — минимизация артефактов и мерцания, особенно при обработке видео;

• Эффективность — за счет NoGAN-обучения достигается высокое качество колоризации без многодневного обучения;

• Доступность — проект можно легко запустить, используя доступные в интернете контрольные точки модели.

Пример покраски

Недостатки:

• Ограниченная кастомизация — модель хорошо работает с большинством изображений, но иногда цветопередача может отличаться от реальной;

• Требования к ресурсам — обработка изображений, особенно видео, требует мощного оборудования.

Есть сайт, на котором можно проверить работу этой модели.

Раскраска, которую мы будем имплементировать в телеграм бота + сам бот

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

У нее есть две версии 2016 и 2017 года.

Пример

Телеграм-бот реализован с помощью библиотеки pyTelegramBotAPI. Он принимает фотографии от пользователя, раскрашивает их с помощью моделей, автоматически отправляет обратно пользователю и удаляет загруженные файлы с сервера. Для работы с кодом требуется установка пакетов argparse и matplotlib, а также наличие файлов с весами моделей eccv16 и siggraph17, которые будут использоваться для раскраски фотографий.

Код реализации на гитхабе.

Что у нас получилось