Раскрашиваем ч/б фото с помощью Python
Привет Хабр.
Одной из интересных и популярных (особенно перед разными юбилеями) задач является «раскрашивание» старых черно-белых фотографий и даже фильмов. Тема это достаточно интересная, как с математической, так и с исторической точки зрения. Мы рассмотрим реализацию этого процесса на Python, который любой желающий сможет запустить на своем домашнем ПК.
Результат работы на фото.
Для тех кому интересно, принцип работы, исходники и примеры под катом.
Принцип работы
Итак, у нас есть черно-белые фото, которые делались в то время, когда цветной печати еще не было. Что это значит для нас — что информация о цвете в этих фото физически отсутствует, взять её оттуда мы не можем. Есть два варианта — либо раскрашивать вручную (чем занимаются некоторые специально обученные люди, что может включать в себя не только раскрашивание, но и работу с архивами, музеями и пр), либо призвать всю мощь машинного обучения, и пусть черновую работу за нас делает компьютер.
Основной подход по сути прост. Первым шагом мы переводим изображение в формат Lab, который содержит раздельно яркостную (L) и цветовую (аb) составляющие. L-канал это фактически и есть искомое черно-белое фото.
Вторым шагом мы обучаем нейросеть, которая по каналу L сможет «угадывать» для нас недостающие цветовые каналы. Для этого нужно много цветных изображений, которые будут использоваться для обучения.
И наконец, когда все готово, мы можем использовать обученную нейросеть для раскрашивания черно-белых фотографий.
Мы рассмотрим два подхода — реализацию нейросети с нуля, и использование «профессиональной» предобученной нейросети. Первый способ скорее учебный, и желающие посмотреть на практические результаты могут сразу перейти ко второму.
Способ 1. Реализация «с нуля» в Keras
Приступим к коду. Рассмотрим пример реализации от Emil Wallner, его код на github можно посмотреть по ссылке.
Загрузка фотографий
Тут ничего необычного. Загружаем массив и приводим цвета к диапазону 0…1. Как можно видеть, нейросеть работает с изображениями 256×256. Немного, но даже это уже на пределе возможностей современных «домашних» видеокарт — при обучении сети я периодически получал сообщение 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, код целиком приведен под спойлером.
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. Хотя если этого не знать, то заметить цветовые артефакты не так и просто, цветовая разрешающая способность глаза в несколько раз ниже яркостной:
На КДПВ принцип тот же.
Для желающих поэкспериментировать самостоятельно с разными картинками, исходник под спойлером.
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, приведен под спойлером.
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 кадров/с.
Как обычно, всем удачных экспериментов.
Можно объявить в комментариях конкурс на лучшее цветное ретро-фото ;)
