Колоризация изображения — это процесс получения черно-белого изображения в качестве входных данных и возврата автоматически окрашенного изображения в качестве выходного. Предыдущие методы раскрашивания изображений основывались на ручной аннотации человека и часто давали ненасыщенные результаты, которые не были правдоподобны как истинные раскраски.
В этом посте мы будем использовать модель 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)
Большой! Попробуйте другие черно-белые фотографии и убедитесь, насколько удивительным может быть искусственный интеллект!