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

В этом посте мы будем использовать модель CNN (Convolutional Neural Network) для раскрашивания старинных черно-белых фотографий. Этот подход представляет собой предварительно обученную модель Ричарда Чжана, в которой он реализовал подход искусственного интеллекта (ИИ) как проход с прямой связью в CNN во время тестирования, а затем обучил более миллиона цветных изображений. Причина выбора этой модели для обучения нашей собственной модели заключается в том, что автор обучил свою модель, используя миллионы изображений. Из-за отсутствия такого большого количества изображений и нехватки вычислительной мощности для обучения такого большого количества изображений мы используем его предварительно обученную модель.

Вот фото канадского танка, наступающего во время битвы при Кане 10 июля 1944 года:

И его раскрашенная версия с использованием моделей, разработанных с использованием этих методов ИИ:

Итак, как работает эта система?

Теперь, когда мы получаем интенсивность в качестве входных данных и нам нужно угадывать цвета, мы не можем использовать цветовое пространство RGB, поскольку оно не содержит информации об освещении. У нас есть два других варианта: либо цветовое пространство YCbCr, либо цветовое пространство Lab, поскольку оба канала Y и L кодируют информацию об освещении. В этом блокноте мы будем использовать цветовое пространство Lab.

Что такое лабораторное цветовое пространство? Подобно цветовому пространству RGB, цветовое пространство Lab имеет три канала. Но в отличие от цветового пространства RGB, Lab кодирует информацию о цвете по-другому, т.е.

L-канал: кодирует только интенсивность света.

канал: кодирует зеленый-красный.

b-канал: кодирует сине-красный

Поскольку канал L кодирует только интенсивность, мы можем использовать его в качестве входных данных для сети в градациях серого. Оттуда сеть должна научиться предсказывать каналы a и b. Имея входной канал L и предсказанные каналы ab, мы можем затем сформировать окончательное выходное изображение.

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

Весь (упрощенный) процесс можно резюмировать следующим образом:

  • Загрузите модель и точки свертки/ядра
  • Преобразуйте входное изображение в цветовое пространство Lab.
  • Используйте L-канал в качестве входных данных для сети и обучите сеть предсказывать каналы ab.
  • Объедините входной L-канал с предсказанными ab-каналами
  • Преобразование лабораторного изображения в RGB

Чтобы перейти к дальнейшему объяснению раскраски черно-белых изображений, нам нужно загрузить эти 3 файла:

  • "colorization_release_v2.caffemodel": предварительно обученная модель, хранящаяся в формате платформы Caffe, которую можно использовать для прогнозирования новых невидимых данных.
  • colorization_deploy_v2.prototxt: состоит из различных параметров, определяющих сеть, а также помогает в развертывании модели Caffe.
  • pts_in_hull.npy: файл NumPy, в котором хранятся центральные точки кластера в формате NumPy. Он состоит из 313 ядер кластера.

Вы можете скачать файлы colorization_deploy_v2.prototxt и pts_in_hull.npy с моего GitHub, а затем загрузить их в свою папку.

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

# import necessary libraries
import cv2
import numpy as np
import matplotlib.pyplot as plt

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

# path to the Caffe prototxt file
prototxt = 'colorization_deploy_v2.prototxt'
# path to the Caffe pre-trained model
model = 'colorization_release_v2.caffemodel'
# path to a NumPy cluster center points file
points = 'pts_in_hull.npy'
# path to our input black & white image
bw_image = 'test_image.jpg'

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

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

Теперь мы загрузили нашу модель Caffe непосредственно из значений аргументов командной строки. OpenCV может читать модели Caffe через функцию cv2.dnn.readNetFromCaffe. Функция cv2.dnn.readNetFromCaffe() принимает два параметра:

  • prototxt — путь к файлу «.prototxt»
  • caffe_model — путь к файлу «.caffemodel»

Затем мы загрузили центральные точки кластера непосредственно из командной строки в файл точек с помощью NumPy.

Следующим шагом является получение идентификаторов слоев из последнего слоя модели кафе с помощью функции «.getLayerId()». «.getLayerId()» принимает один параметр.

layer1 = net.getLayerId('class8_ab')
print(layer1)
layer2 = net.getLayerId('conv8_313_rh')
print(layer2)

Чтобы понять, как мы получили два вышеупомянутых имени вывода «class8_ab» и «conv8_313_rh», мы получили их из последней части кода в файле «.prototxt». Как мы видим ниже, количество выходов равно двум и выше, которые являются именами выходов.

# read the last 20 lines of colorization_deploy_v2.prototxt file
with open('colorization_deploy_v2.prototxt', 'r') as file:
    for line in (file.readlines() [-20:]):
        print(line)

Затем мы транспонируем наш файл NumPy и изменим форму центров кластеров, хранящихся в них, в виде матрицы 1 × 1, а затем добавим ее в нашу модель.

pts = pts.transpose().reshape(2, 313, 1, 1)
net.getLayer(layer1).blobs = [pts.astype('float32')]
net.getLayer(layer2).blobs = [np.full([1, 313], 2.606, dtype = 'float32')]

Следующим шагом является использование OpenCV для чтения нашего тестового изображения из пути, его масштабирования и преобразования в лабораторное: обратите внимание, что сначала мы преобразуем изображение в оттенки серого. Этот шаг на самом деле не является необходимым, но я понимаю, что некоторые черно-белые фотографии, особенно старые, могут подвергаться некоторой обработке с годами, поэтому лучше их немного почистить.

# read image from the path
test_image = cv2.imread(bw_image)
# convert image into gray scale
test_image = cv2.cvtColor(test_image, cv2.COLOR_BGR2GRAY)
# convert image from gray scale to rgb format
test_image = cv2.cvtColor(test_image, cv2.COLOR_GRAY2RGB)
# check image using matplotlib
plt.imshow(test_image)
plt.rcParams['figure.figsize'] = [16, 9]
plt.title('Original Image')
plt.axis('off')
plt.show()

Далее мы выполним операцию масштабирования, нормализовав пиксели изображения в диапазоне от 0 до 1. Затем преобразуйте формат изображения в лабораторное пространство. Затем измените размер изображения до размера 224 × 224, поскольку предварительно обученная модель принимает входное изображение размером (224 224). Функция cv2.split() разбивает изображение на три канала: L, a и b.

# normalise the image
normalised= test_image.astype("float32") / 255.0
# convert the image into LAB
lab_image = cv2.cvtColor(normalised, cv2.COLOR_RGB2LAB)
# resize the image
resized = cv2.resize(lab_image, (224, 224))

Мы можем извлечь L-канал из изображения Lab, используя его порядковый номер. После этого мы выполняем вычитание среднего для канала L изображения.

# extract the value of L for Lab image
L = cv2.split(resized)[0]
L -= 50

Теперь мы передаем L-канал в качестве входных данных для нашей модели, а затем прогнозируем значения a и b из модели.

# set the input
net.setInput(cv2.dnn.blobFromImage(L))
# findi the values of a and b
ab = net.forward()[0, :, :, :].transpose((1, 2, 0))
# resize
ab = cv2.resize(ab, (test_image.shape[1], test_image.shape[0]))

Затем снова извлекается L-канал, но уже из исходного изображения Lab, поскольку размеры всех трех плоскостей (L, a, b) должно быть одинаково. Затем мы объединяем L-канал с a и b с помощью Numpy, чтобы получить цветное изображение Lab. Наконец, мы используем Matplotlib для отображения изображения.

L = cv2.split(lab_image)[0]
# combining L, a, b
Lab_coloured = np.concatenate((L[:, :, np.newaxis], ab), axis = 2)
# check the Lab image
plt.imshow(Lab_coloured)
plt.rcParams['figure.figsize'] = [16, 9]
plt.title('Lab image')
plt.axis('off')
plt.show()

Теперь мы получили цветное изображение Lab, но изображение непонятное. Итак, нам нужно преобразовать его (изображение Lab) в формат RGB и посмотреть результат.

# convert Lab image to RGB_colored
RGB_coloured = cv2.cvtColor(Lab_coloured, cv2.COLOR_LAB2RGB)

Затем мы будем использовать np.clip() для обрезки изображения RGB между 0 и 1, потому что RGB работает именно так. Отсечение означает, что если интервал равен [0,1], то все значения меньше нуля станут равны нулю, а все значения больше единицы станут единицей. Если мы помним, ранее мы нормализовали пиксели изображения между 0–1. Теперь мы меняем пиксели изображения обратно в диапазоне от 0 до 255.

# limits the values in array
RGB_coloured = np.clip(RGB_coloured, 0, 1)
# change the pixel intensity back to [0,255]
RGB_coloured = (255 * RGB_coloured).astype('uint8')

Теперь, после построения изображения RGB с помощью Matplotlib, мы получим идеально окрашенное изображение для нашего черно-белого тестового изображения.

# check the final coloured image
plt.imshow(RGB_coloured)
plt.rcParams['figure.figsize'] = [16, 9]
plt.title('Coloured Image')
plt.axis('off')
plt.show()

Чтобы сохранить цветное изображение, сначала нам нужно преобразовать его из формата RGB в формат BGR, а затем с помощью OpenCV сохранить изображение по пути к каталогу. Как мы видим, cv2.imwrite() принимает аргументы, то есть путь (место, где должен быть сохранен файл) и RGB_BGR (файл).

# converting RGB to BGR
RGB_BGR = cv2.cvtColor(RGB_coloured, cv2.COLOR_RGB2BGR)
# save the image in desired path
# cv2.imwrite('../'+bw_image, RGB_BGR)
cv2.imwrite('coloured_image.jpg', RGB_BGR)

Большой! Попробуйте другие черно-белые фотографии и убедитесь, насколько удивительным может быть искусственный интеллект!