Раскрашиваем ч/б фото с помощью Python

Моя цель - предложение широкого ассортимента товаров и услуг на постоянно высоком качестве обслуживания по самым выгодным ценам.

Привет Хабр.

Одной из интересных и популярных (особенно перед разными юбилеями) задач является «раскрашивание» старых черно-белых фотографий и даже фильмов. Тема это достаточно интересная, как с математической, так и с исторической точки зрения. Мы рассмотрим реализацию этого процесса на Python, который любой желающий сможет запустить на своем домашнем ПК.

Результат работы на фото.



Для тех кому интересно, принцип работы, исходники и примеры под катом.

Принцип работы


Итак, у нас есть черно-белые фото, которые делались в то время, когда цветной печати еще не было. Что это значит для нас — что информация о цвете в этих фото физически отсутствует, взять её оттуда мы не можем. Есть два варианта — либо раскрашивать вручную (чем занимаются некоторые специально обученные люди, что может включать в себя не только раскрашивание, но и работу с архивами, музеями и пр), либо призвать всю мощь машинного обучения, и пусть черновую работу за нас делает компьютер.

Основной подход по сути прост. Первым шагом мы переводим изображение в формат Lab, который содержит раздельно яркостную (L) и цветовую (аb) составляющие. L-канал это фактически и есть искомое черно-белое фото.



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

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

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

Способ 1. Реализация «с нуля» в Keras


Приступим к коду. Рассмотрим пример реализации от Emil Wallner, его код на github можно посмотреть по ссылке.

Загрузка фотографий

Тут ничего необычного. Загружаем массив и приводим цвета к диапазону 0..1. Как можно видеть, нейросеть работает с изображениями 256x256. Немного, но даже это уже на пределе возможностей современных «домашних» видеокарт — при обучении сети я периодически получал сообщение Out Of Memory.
import numpy as np
import cv2

IMG_SIZE = 256

def load_train_images(folder):
    x = []
    for filename in os.listdir(folder):
        x.append(img_to_array(load_img(folder + os.sep + filename)))
    x = np.array([cv2.resize(i, (IMG_SIZE, IMG_SIZE)) for i in x], dtype=float)/255.0
    return x

Подготовка данных

Здесь используется класс ImageDataGenerator, который позволяет из одного изображения получить несколько. Как можно видеть, в качестве параметров передаются zoom_range=0.2, rotation_range=20, horizontal_flip=True, т.е. из одного изображения будет создано несколько с разным поворотом, увеличением и отражением по горизонтали.

from skimage.color import rgb2lab, lab2rgb, rgb2gray, xyz2lab
from keras.preprocessing.image import ImageDataGenerator, array_to_img, img_to_array, load_img
from keras.layers import Conv2D, UpSampling2D, InputLayer, Conv2DTranspose
from keras.layers import Activation, Dense, Dropout, Flatten

datagen = ImageDataGenerator(shear_range=0.2, zoom_range=0.2, rotation_range=20, horizontal_flip=True)

# Generate training data
def image_a_b_gen(batch_size):
    for batch in datagen.flow(train_images, batch_size=batch_size):
        lab_batch = rgb2lab(batch)
        x_batch = lab_batch[:, :, :, 0]
        y_batch = lab_batch[:, :, :, 1:] / 128
        yield (x_batch.reshape(x_batch.shape + (1,)), y_batch)

Внутри самого генератора происходит вызов функции rgb2lab, и берутся соответствующие каналы.

Обучение нейронной сети

Как можно видеть, здесь используется достаточно многослойная сверточная сеть, которая настроена так, что результирующее изображение имеет такой же размер, что и входное.

model = Sequential()
model.add(InputLayer(input_shape=(IMG_SIZE, IMG_SIZE, 1)))
model.add(Conv2D(64, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(64, (3, 3), activation='relu', padding='same', strides=2))
model.add(Conv2D(128, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(128, (3, 3), activation='relu', padding='same', strides=2))
model.add(Conv2D(256, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(256, (3, 3), activation='relu', padding='same', strides=2))
model.add(Conv2D(512, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(256, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(128, (3, 3), activation='relu', padding='same'))
model.add(UpSampling2D((2, 2)))
model.add(Conv2D(64, (3, 3), activation='relu', padding='same'))
model.add(UpSampling2D((2, 2)))
model.add(Conv2D(32, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(2, (3, 3), activation='tanh', padding='same'))
model.add(UpSampling2D((2, 2)))
model.compile(optimizer='rmsprop', loss='mse')

model.fit_generator(image_a_b_gen(batch_size=50), steps_per_epoch=steps_per_epoch, epochs=epochs)

Собственно раскрашивание состоит в подготовке L-канала (ч/б фото), вызову нейронной сети и объединении каналов вместе.

output = model.predict(data_in)

for i in range(len(output)):
    cur = np.zeros((256, 256, 3))
    cur[:,:,0] = color_me[i][:,:,0]
    cur[:,:,1:] = output[i] * 128
    img_rgb = lab2rgb(cur)*brightness_corr
    imsave(folder_dst + os.sep + "img_%d.png" % i, img_rgb.astype(np.uint8))

В общем, и вся «магия».

Увы, добиться стабильной работы от этой программы так и не удалось, иногда она показывает неплохие результаты, иногда вообще ничего, так что этот пример можно лишь рассматривать как учебный. Желающие могут посмотреть результаты на github, исходники можно взять там же. Для тех кому лень идти на github, код целиком приведен под спойлером.

colorize1.py
import os
# os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # Force CPU
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'  # 0 = all messages are logged, 3 - INFO, WARNING, and ERROR messages are not printed

from keras.layers import Conv2D, Conv2DTranspose, UpSampling2D
from keras.layers import Activation, Dense, Dropout, Flatten, InputLayer
from keras.layers.normalization import BatchNormalization
from keras.callbacks import TensorBoard
from keras.models import Sequential
from keras.preprocessing.image import ImageDataGenerator, array_to_img, img_to_array, load_img
from skimage.color import rgb2lab, lab2rgb, rgb2gray
from skimage.io import imsave
import numpy as np
import random
import tensorflow as tf
import cv2


folder_train = "Faces"
folder_src = "Input"
folder_dst = "Output"
model_file = "faces1.h5"
brightness_corr = 250
do_train = True

# Get train images
X = []
for filename in os.listdir(folder_train):
    X.append(img_to_array(load_img(folder_train + os.sep + filename)))
X = np.array([cv2.resize(i, (256, 256)) for i in X], dtype=float)/255.0
# X = np.array(X, dtype=float)
Xtrain = X

# Model
model = Sequential()
model.add(InputLayer(input_shape=(256, 256, 1)))
model.add(Conv2D(64, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(64, (3, 3), activation='relu', padding='same', strides=2))
model.add(Conv2D(128, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(128, (3, 3), activation='relu', padding='same', strides=2))
model.add(Conv2D(256, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(256, (3, 3), activation='relu', padding='same', strides=2))
model.add(Conv2D(512, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(256, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(128, (3, 3), activation='relu', padding='same'))
model.add(UpSampling2D((2, 2)))
model.add(Conv2D(64, (3, 3), activation='relu', padding='same'))
model.add(UpSampling2D((2, 2)))
model.add(Conv2D(32, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(2, (3, 3), activation='tanh', padding='same'))
model.add(UpSampling2D((2, 2)))
model.compile(optimizer='rmsprop', loss='mse')

# Image transformer
datagen = ImageDataGenerator(shear_range=0.2, zoom_range=0.2, rotation_range=20, horizontal_flip=True)

# Generate training data
batch_size = 10
def image_a_b_gen(batch_size):
    for batch in datagen.flow(Xtrain, batch_size=batch_size):
        lab_batch = rgb2lab(batch)
        X_batch = lab_batch[:,:,:,0]
        Y_batch = lab_batch[:,:,:,1:] / 128
        yield (X_batch.reshape(X_batch.shape+(1,)), Y_batch)

if do_train:
    # Train model
    model.fit_generator(image_a_b_gen(batch_size), epochs=1, steps_per_epoch=400)

    # Save model
    model.save_weights(model_file)

# Load model
model.load_weights(model_file)

# Process images
color_me = []
for filename in os.listdir(folder_src):
    color_me.append(img_to_array(load_img(folder_src + os.sep + filename)))
color_me = np.array([cv2.resize(i, (256, 256)) for i in color_me], dtype=float)
# color_me = np.array(color_me, dtype=float)
color_me = rgb2lab(1.0/255*color_me)[:,:,:,0]
color_me = color_me.reshape(color_me.shape+(1,))

# Test model
output = model.predict(color_me)

# Output colorizations
for i in range(len(output)):
    cur = np.zeros((256, 256, 3))
    cur[:,:,0] = color_me[i][:,:,0]
    cur[:,:,1:] = output[i] * 128
    img_rgb = lab2rgb(cur)*brightness_corr
    imsave(folder_dst + os.sep + "img_%d.png" % i, img_rgb.astype(np.uint8))
    print("img_%d.png saved" % i)


Способ 2. Предобученная сеть


В качестве второго метода рассмотрим уже готовую нейронную сеть от Rich Zhang. Это гораздо более серьезный проект, нейросеть была обучена на 1.3млн изображений, и только лишь файлы сохраненной модели для этой сети занимают около 300Мбайт.

Шаг 1. Скачиваем файлы модели:

Загружаем файлы colorization_deploy_v2.prototxt, colorization_release_v2.caffemodel и pts_in_hull.npy, кладем их в папку models.

Шаг 2. Загружаем модель нейронной сети

def load_model() -> Any:
    # Load serialized black and white colorizer model and cluster
    # The L channel encodes lightness intensity only
    # The a channel encodes green-red.
    # And the b channel encodes blue-yellow
    print("Loading model...")

    prototxt = "models/colorization_deploy_v2.prototxt"
    model = "models/colorization_release_v2.caffemodel"
    points = "models/pts_in_hull.npy"

    net = cv2.dnn.readNetFromCaffe(prototxt, model)
    pts = np.load(points)

    # Add the cluster centers as 1x1 convolutions to the model:

    class8 = net.getLayerId("class8_ab")
    conv8 = net.getLayerId("conv8_313_rh")
    pts = pts.transpose().reshape(2, 313, 1, 1)
    net.getLayer(class8).blobs = [pts.astype("float32")]
    net.getLayer(conv8).blobs = [np.full([1, 313], 2.606, dtype="float32")]
    return net

Шаг 3. Раскрашиваем изображения. Принцип такой же, как и в нашем «учебном» примере: получаем каналы в Lab, конвертируем их в RGB и записываем в файл.

def colorize_image(net: Any, image_in: str, image_out: str):
    # Load the input image, scale it and convert it to Lab:
    image = cv2.imread(image_in)
    height, width, channels = image.shape
    # image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)

    # Extracting "L"
    scaled = image.astype("float32") / 255.0
    lab = cv2.cvtColor(scaled, cv2.COLOR_RGB2LAB)

    # Resize to network size
    resized = cv2.resize(lab, (224, 224))
    L = cv2.split(resized)[0]
    L -= 50

    # Predicting "a" and "b"
    net.setInput(cv2.dnn.blobFromImage(L))
    ab = net.forward()[0, :, :, :].transpose((1, 2, 0))

    # Creating a colorized Lab photo (L + a + b)
    L = cv2.split(lab)[0]
    ab = cv2.resize(ab, (width, height))
    colorized = np.concatenate((L[:, :, np.newaxis], ab), axis=2)

    # Convert to RGB
    colorized = cv2.cvtColor(colorized, cv2.COLOR_LAB2RGB)
    colorized = np.clip(colorized, 0, 1)
    colorized = (255 * colorized).astype("uint8")

    cv2.imwrite(image_out, cv2.cvtColor(colorized, cv2.COLOR_RGB2BGR))
    print("Image %s saved" % image_out)

Результаты


В качестве исходных файлов я взял следующие фотографии:



Результат неидеален, но в принципе не так уж плох:



Важно иметь в виду, что нейросеть работает с изображениями 224х224, так что High-Res получить увы, не удастся — изображение на выходе будет иметь оригинальное разрешение (L-канал), но канал цвета будет масштабирован из 224х224. Хотя если этого не знать, то заметить цветовые артефакты не так и просто, цветовая разрешающая способность глаза в несколько раз ниже яркостной:



На КДПВ принцип тот же.

Для желающих поэкспериментировать самостоятельно с разными картинками, исходник под спойлером.

colorize2.py
import numpy as np
from typing import *
import cv2
import os

# Sources:
# https://www.learnopencv.com/convolutional-neural-network-based-image-colorization-using-opencv/
# https://www.pyimagesearch.com/2019/02/25/black-and-white-image-colorization-with-opencv-and-deep-learning/
# http://richzhang.github.io/colorization/

# Models:
# http://eecs.berkeley.edu/~rich.zhang/projects/2016_colorization/files/demo_v2/colorization_release_v2.caffemodel
# http://eecs.berkeley.edu/~rich.zhang/projects/2016_colorization/files/demo_v2/colorization_release_v2_norebal.caffemodel
# http://eecs.berkeley.edu/~rich.zhang/projects/2016_colorization/files/demo_v1/colorization_release_v1.caffemodel


prototxt = "models/colorization_deploy_v2.prototxt"
model = "models/colorization_release_v2.caffemodel"
points = "models/pts_in_hull.npy"
folder_in = "Input"
folder_out = "Output"


def load_model() -> Any:
    # Load serialized black and white colorizer model and cluster
    # The L channel encodes lightness intensity only
    # The a channel encodes green-red.
    # And the b channel encodes blue-yellow
    print("Loading model...")

    net = cv2.dnn.readNetFromCaffe(prototxt, model)
    pts = np.load(points)

    # Add the cluster centers as 1x1 convolutions to the model:

    class8 = net.getLayerId("class8_ab")
    conv8 = net.getLayerId("conv8_313_rh")
    pts = pts.transpose().reshape(2, 313, 1, 1)
    net.getLayer(class8).blobs = [pts.astype("float32")]
    net.getLayer(conv8).blobs = [np.full([1, 313], 2.606, dtype="float32")]
    return net


def colorize_image(net: Any, image_in: str, image_out: str):
    # Load the input image, scale it and convert it to Lab:
    image = cv2.imread(image_in)
    height, width, channels = image.shape
    # image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)

    # Extracting "L"
    scaled = image.astype("float32") / 255.0
    lab = cv2.cvtColor(scaled, cv2.COLOR_RGB2LAB)

    # Resize to network size
    resized = cv2.resize(lab, (224, 224))
    L = cv2.split(resized)[0]
    L -= 50

    # Predicting "a" and "b"
    net.setInput(cv2.dnn.blobFromImage(L))
    ab = net.forward()[0, :, :, :].transpose((1, 2, 0))

    # Creating a colorized Lab photo (L + a + b)
    L = cv2.split(lab)[0]
    ab = cv2.resize(ab, (width, height))
    colorized = np.concatenate((L[:, :, np.newaxis], ab), axis=2)

    # Convert to RGB
    colorized = cv2.cvtColor(colorized, cv2.COLOR_LAB2RGB)
    colorized = np.clip(colorized, 0, 1)
    colorized = (255 * colorized).astype("uint8")

    cv2.imwrite(image_out, cv2.cvtColor(colorized, cv2.COLOR_RGB2BGR))
    print("Image %s saved" % image_out)


if __name__ == '__main__':

    net = load_model()
    for filename in os.listdir(folder_in):
       colorize_image(net, folder_in + os.sep + filename, folder_out + os.sep + filename.replace(".", "-out."))


Для запуска программы навыков программирования не требуется, достаточно уметь запустить программу из консоли, и положить исходные файлы в папку Input, цветные фото будет сгенерированы в папке Output (также необходимо скачать файлы модели из вышеприведенных ссылок в папку models).

Заключение


Как можно видеть, задача раскрашивания фото является весьма непростой с вычислительной точки зрения. Даже обучение нейросети, обрабатывающей файлы 224х224пкс, может требовать мощной видеокарты и нескольких часов или даже дней обучения. Существуют ли доступные нейросети, способные сделать хотя бы FullHD-качество, мне неизвестно.

Желающие могут изучить проект от Rich Chang более подробно, там же есть инструкции по обучению сети на разных датасетах, ну а готовый вышеприведенной код с уже обученной сетью можно просто запустить, он работает быстро и почти не требует ресурсов ПК.

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

И бонус для тех, кто дочитал до сюда. С помощью аналогичного подхода можно раскрашивать и видео, принцип точно такой же. Исходный код, принимающий и сохраняющий MP4, приведен под спойлером.
video.py
def colorize_video(net: Any, video_in: str, video_out: str):
    vid_in = cv2.VideoCapture(video_in)
    frames, fps = int(vid_in.get(cv2.CAP_PROP_FRAME_COUNT)), int(round(vid_in.get(cv2.CAP_PROP_FPS)))
    width, height = int(vid_in.get(cv2.CAP_PROP_FRAME_WIDTH)), int(vid_in.get(cv2.CAP_PROP_FRAME_HEIGHT))
    print("Video {}: {}x{}, {} frames, {} fps".format(video_in, width, height, frames, fps))
    size_out = (width, height)
    vid_out = cv2.VideoWriter(video_out, cv2.VideoWriter_fourcc(*"mp4v"), fps, size_out)  # .mp4
    # vid_out = cv2.VideoWriter(video_out, cv2.VideoWriter_fourcc(*'DIVX'), fps, size_out)  # .avi

    count = 0
    while True:
        success, frame = vid_in.read()
        if frame is None or success is False:
            break

        if (count % 10) == 0:
            print("Frame {} of {}".format(count, frames))

        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB)
        scaled = frame.astype("float32") / 255.0
        lab = cv2.cvtColor(scaled, cv2.COLOR_RGB2LAB)
        resized = cv2.resize(lab, (224, 224))
        L = cv2.split(resized)[0]
        L -= 50

        net.setInput(cv2.dnn.blobFromImage(L))
        ab = net.forward()[0, :, :, :].transpose((1, 2, 0))
        ab = cv2.resize(ab, (frame.shape[1], frame.shape[0]))
        L = cv2.split(lab)[0]
        colorized = np.concatenate((L[:, :, np.newaxis], ab), axis=2)
        colorized = cv2.cvtColor(colorized, cv2.COLOR_LAB2BGR)
        colorized = np.clip(colorized, 0, 1)
        colorized = (255 * colorized).astype("uint8")
        frame_out = colorized  # cv2.cvtColor(colorized, cv2.COLOR_RGB2BGR)

        vid_out.write(frame_out)

        count += 1

    vid_in.release()
    vid_out.release()

    print("File %s saved" % video_out)

Его достаточно вставить в вышеприведенную программу. Процесс обработки видео не быстрый, скорость порядка 1-2 кадров/с.

Как обычно, всем удачных экспериментов.
Можно объявить в комментариях конкурс на лучшее цветное ретро-фото ;)
Источник: https://habr.com/ru/post/471656/

Интересные статьи

Интересные статьи

Каждый, кому хоть раз приходилось строить диаграммы в draw.io или Google Diagrams, помнит всю утомительность и медлительность этого процесса. Сегодня делимся с вами материалом, в ко...
Дисциплина Computational fluid dynamics(CFD) или, на русском языке, Вычислительная гидродинамика изучает поведение различных потоков, в том числе вихревых. Это и моделирование цунами,...
Наверное каждому разработчику хотя бы раз в жизни приходила идея что-нибудь автоматизировать. Ведь если есть возможность избавиться от рутины, то грех ей не воспользоваться. Для меня эта идея...
Существует традиция, долго и дорого разрабатывать интернет-магазин. :-) Лакировать все детали, придумывать, внедрять и полировать «фишечки» и делать это все до открытия магазина.
Несмотря на то, что “в коробке” с Битриксом уже идут модули как для SOAP (модуль “Веб сервисы” в редакции “Бизнес” и старше), так и для REST (модуль “Rest API” во всех редакциях, начиная с...