python图片自动上色_自己实现黑白图片自动上色AI（一）

前言

我们将用两篇文章来从零实现一个自动上色的AI。我们知道人工智能能够学习到图片里面的概率模型，也某一类物体与另外一类物体交错存在的概率，那么有没有可能实现一种AI，看遍上千种物体之后，开始学会知道：皮肤是黄色的，头发是黑色的，牙齿是白色的，水是天蓝的，花是红的，草是绿的，天空是蓝的呢？答案是存在的，也是很多时候我们梦寐以求去寻觅的。那么我们接下来就是真正从零来写一个深度学习应用，去实现这么一个AI。在动手之前，让我们来看一下，我们最终要达到的效果是怎样的：

当我们自己能够实现这个之后， 你就可以。。。自己开发一套系统专门用人工智能给黑白图片着色，甚至给素描着色，或者做一个APP，专门自动给图片上色，给设计师用，上一张图片五毛钱。。。想想还是非常激动，升级CEO，迎娶白富美，走向人生巅峰就在眼前…… 好了梦想是非常美好的但是不能贪杯哦，让我们直接开始吧！

我们做一个简单版本的自动上色AI之后的效果：

简直美哭有没有。这是400x400的上色AI，比你以前看到的256x256都要大。闲话不多说，直接开始咯。

1. 从最简单的版本开始

本篇博客是阅读了很多论文、源码之后才开始写作的，重复造轮子是个bad idea，所以你可能会看到一些borrow的代码(其实几乎是自己从头写的)，但是即使是同一行代码我都会做出自己的修改，写上自己的注释，所以大家有任何疑问，欢迎从文章末尾的博客找到我与我进行互动。闲话不多说，在正式开始之前，我简要列举一下图片彩色的几种思路(说实话非常不喜欢一些人写博客把论文直接照搬，甚至把模型结构直接复制粘贴，我敢说他们自己都没有理解论文的真实含义)：我们知道图片正常情况下是RGB的，这样的话图片的颜色信息就被分散在了R，G，B三个通道里，我们要做AI上色，首先你肯定要有黑白图片吧，黑白图片都没有那还上个毛？而黑白图片是gray也就是灰度图，灰度图只有一个通道，这和我们熟悉的mnist是一样的，用这个输入，输出是啥？那么这里就引出了一个新的图片制式，LAB，也就是明度+颜色A+颜色B，这里的颜色A是红绿，颜色B是黄蓝，而明度里面并没有包含颜色的信息，因此使用LAB来做自动上色简直是再合适不过；

上面说的办法用的还是x是灰度图，y是原图对应的ab通道，那么有没有办法用生成对抗网络来生成一张彩色的图呢？答案是有的，这种方法思路是训练一个G，专门从灰度图生成彩色图，训练一个D，专门辨别这张图和原图的差距是多少，通过训练来生成一个强大无比的G，专门生成彩色图。

还有一些其他的思路，比如Recursive PixelColor，就是用RNN来对图片像素级别的学习，从而生成彩色图，但是本质上也和第一种方法类似，只不过用的网络模型不同。

好了，概览了人类所能搞出的一些方法后，我们就得评估一下了。用哪种方法好呢？如果你现在正在看这篇文章，那么你算是走了狗*运，我们不会实现所有的方法(摊手)。但是我们会通过三个阶段来从0制作一个自动上色程序，分为 simple, advance, complex三个版本。当然啦，代码是从少到多了。而我们采用的思路呢，就是第一种，大家别小看第一种方法，很多专业的上色程序说不定就是第一种做出来的，只要你的模型够复杂，你的训练数据足够多，你的网络可以收敛，上色的效果绝对是杠杠的！

OK，让我们直接从最简单的版本开始吧！

开始之前，我们得新建一个colorize的文件夹，然后 touch simple.py ,新建一个python文件，我们依旧采用python3，当然也可以兼容python2啦，如果你有pycharm的话就更好啦。开始写一个简单的函数把：

"""

this file is the simple version of colorize

you are going to need skimage and keras for it

any version should be compatible

"""

import keras

import tensorflow as tf

from skimage.io import imread, imsave

from skimage.color import rgb2gray, gray2rgb, rgb2lab, lab2rgb

from keras.models import Sequential

from keras.layers import Conv2D, UpSampling2D, InputLayer, Conv2DTranspose

import numpy as np

def image_color_test():

# this method to let's know something about RGB and LAB color space

test_image_gray = rgb2gray(imread('data/Train/11Se02.jpg'))

test_image_rgb = imread('data/Train/11Se02.jpg')

print(test_image_gray.shape)

print(test_image_gray)

print(test_image_rgb.shape)

if __name__ == '__main__':

image_color_test()

这个函数的是读入了一张彩色图片，注意是彩色图片，然后我们打印转成gray之后和原始类型，输出是这样的：

(256, 256)

[[ 0.47397569 0.33166039 0.19745608 ..., 0.73957725 0.73565569

0.73565569]

[ 0.4322451 0.27283686 0.12351176 ..., 0.73957725 0.73565569

0.73565569]

[ 0.25969608 0.18656235 0.1291149 ..., 0.73957725 0.73565569

0.73565569]

...,

[ 0.49407098 0.49407098 0.50191412 ..., 0.45594353 0.42678824

0.39679176]

[ 0.48425569 0.48425569 0.48425569 ..., 0.44023451 0.40795333

0.36139961]

[ 0.47641255 0.48033412 0.48033412 ..., 0.41465765 0.36416745

0.31229333]]

(256, 256, 3)

经过rgb2gray之后，原来的RGB变成了一个通道了，没毛病。后面我们要输入到模型里面的也就是这个gray的数据。这时我们需要用到lab这个色彩空间，至于rgb怎么转换到lab我们先不管，当然需要先把rgb转xyz，在xyz转lab。具体转换公式是人们根据人类视觉来调试的(如果是公式推导出来的我吃翔)。好，那么接下来我们的思路明确了，我们要做这些准备工作：设计网络，网路的输入是(256, 256)的灰度图，哦不对，不是灰度图，输入的图片是灰度图没有错，但是要将灰度图转成lab，然后以l这个通道作为输入，原始图片的ab通道作为label，来训练模型；

我们要自己做一个数据集，数据集的x和label都是图片本身，这就非常简单了，我们只要下载一些图片即可，然后预测处理一下，把x和y一个个喂入模型训练即可。

好了，上面只是给大家一个大体 的认识，我们按照这个计划来。我们先对一张图片进行处理，对它进行训练，训练1000次，看看能不能把它自己的黑白版本编程彩色版本。

先get一下我们的训练数据：

"""

this file is the simple version of colorize

you are going to need skimage and keras for it

any version should be compatible

"""

import keras

import tensorflow as tf

from skimage.io import imread, imsave

from skimage.color import rgb2gray, gray2rgb, rgb2lab, lab2rgb

from keras.models import Sequential

from keras.layers import Conv2D, UpSampling2D, InputLayer, Conv2DTranspose

from keras.preprocessing.image import img_to_array, load_img

import numpy as np

from keras.preprocessing.image import ImageDataGenerator

import os

def get_train_data(img_file):

image = img_to_array(load_img(img_file))

image_shape = image.shape

image = np.array(image, dtype=float)

x = rgb2lab(1.0 / 255 * image)[:, :, 0]

y = rgb2lab(1.0 / 255 * image)[:, :, 1:]

x = x.reshape(1, image_shape[0], image_shape[1], 1)

y = y.reshape(1, image_shape[0], image_shape[1], 2)

return x, y, image_shape

这个函数相当简单，但是我还得解释一下：

我们把图片读取出来了，这是一张彩色图片，训练用的，然后我们将其转成了LAB，并且把L作为X，把AB作为Y。

具体来说我们在load数据的时候做了以下事情：首先keras api load_image会把图片以RGB的形式load成array，然后我们将其除以255再转成LAB，为什么要除以255？加入你不除以的话就会发现预测的图片明度会变很暗，所以为了不重蹈覆辙，大家可以按照这个来，其实除以255也就是对RGB进行归一化，LAB转换的时候好像还真的要对图片进行归一化；

我们把x取为了读取的图片的第一个通道，也就是这里的[:, :, 0], 这个数组切片应该知道，但是为什么还要讲x和y reshape一下呢？因为你要转成一个batch，这个batch是直接输入到网络的，神经网络在图片上默认的输入至少是4维度，如果因为图片是三维的，加上batch就是4维，除此之外，如果要继续深入为什么，那么就要涉及到深度学习框架内部的实现了，一般这个都是确定的，所以你必须即使只有一张图片也要指定batch为1.

现在将其提供给网络训练，看看我们的模型构建，一切训练的脚本：

def build_model():

model = Sequential()

model.add(InputLayer(input_shape=(None, None, 1)))

model.add(Conv2D(8, (3, 3), activation='relu', padding='same', strides=2))

model.add(Conv2D(8, (3, 3), activation='relu', padding='same'))

model.add(Conv2D(16, (3, 3), activation='relu', padding='same'))

model.add(Conv2D(16, (3, 3), activation='relu', padding='same', strides=2))

model.add(Conv2D(32, (3, 3), activation='relu', padding='same'))

model.add(Conv2D(32, (3, 3), activation='relu', padding='same', strides=2))

model.add(UpSampling2D((2, 2)))

model.add(Conv2D(32, (3, 3), activation='relu', padding='same'))

model.add(UpSampling2D((2, 2)))

model.add(Conv2D(16, (3, 3), activation='relu', padding='same'))

model.add(UpSampling2D((2, 2)))

model.add(Conv2D(2, (3, 3), activation='tanh', padding='same'))

model.compile(optimizer='rmsprop', loss='mse')

return model

def train():

x, y, img_shape = get_train_data('./data/test.jpg')

model = build_model()

num_epochs = 6000

batch_size = 6

model_file = 'simple_model.h5'

model.fit(x, y, batch_size=1, epochs=1000)

model.save(model_file)

整个过程非常简单，也不奎是我们的simple版本，相信你应该一目了然。如果你觉得过于简单，那么后面会有更加复杂的版本。我们训练6000次，同时把模型保存一下，后面再预测的时候会load整个模型。

看到了吗，我们分分钟就搭建起来训练框架，对图片训练个1000次，看看效果怎么样？直接开始跑？

OK，所有事情都没有错了，让我们run一下把：

Run玩之后，发现尼玛这个loss。。。甚至都不收敛是什么鬼。难道是我的网络写的太垃圾？？？？先不管那么多了，先跑它个100000000个epoch再说。。。。开个玩笑，这么训练下去感觉也没有什么卵用，不多说，改一下优化器把？？我们把这行代码：

model.compile(optimizer='rmsprop', loss='mse')

改为adam来试一下。然而我并没有改优化器，我猛然发现，这个问题出现的原因是我好想没有对y进行归一化。那就对y归一化一下把。我们修改一下getdata的代码，对y除以一个128，为什么要除以128呢？因为可以将像素点的值归一化，但是不一定是0-1：

def get_train_data(img_file):

image = img_to_array(load_img(img_file))

image_shape = image.shape

image = np.array(image, dtype=float)

x = rgb2lab(1.0 / 255 * image)[:, :, 0]

y = rgb2lab(1.0 / 255 * image)[:, :, 1:]

y /= 128

x = x.reshape(1, image_shape[0], image_shape[1], 1)

y = y.reshape(1, image_shape[0], image_shape[1], 2)

return x, y, image_shape

前面import那些就一贴了。好接下来看一下我们的loss：

简直牛逼啊有木有。虽然说好想这个loss也并没有降低的样子但是至少看上去没有那么恐怖，哦嚯嚯。好了，接下来要做什么呢?上个厕所，喝杯咖啡。。上厕所喝咖啡？这里说明一下，这里我是运算了12张图片，所以你会看到一个mnibatch，如果你只是训练一张图片，batchsize就是1，因此也就没有minibatch。好了，你的模型应该一下就运行完了。让我们实现一下预测的脚本：

def colorize():

x, y, image_shape = get_train_data('./data/test.jpg')

model = build_model()

model.load_weights('simple_model.h5')

output = model.predict(x)

output *= 128

tmp = np.zeros((400, 400, 3))

tmp[:, :, 0] = x[0][:, :, 0]

tmp[:, :, 1:] = output[0]

imsave("test_image_result.png", lab2rgb(tmp))

imsave("test_image_gray.png", rgb2gray(lab2rgb(tmp)))

这张test.jpg就是我们之前那张图片，现在可以看到灰度图和彩色图了：

这也是我们一开始看到的结果。简直神奇啊！！！结果完全和我们料想的一样，模型学会了用明度去预测颜色信息！！！

2. 用这个模型预测其他黑白图片

进行到这里，我想很多人都会有这样的疑问，这个模型如果在其他的黑白图片上预测会怎么样呢？我们必须对这个问题探索一个究竟。来，我们实现以下代码：

def pre_process_image_for_colorize(img_file, target_size=(256, 256)):

"""

we will using gray image, so that image must be gray

generate image input into keras inference model

:param img_file:

:param target_size:

:return:

"""

if os.path.exists(img_file):

img = np.asarray(imread(img_file), dtype=float)

# random crop to target size

img_shape = img.shape

assert target_size[0] <= img_shape[0] and target_size[1] <= img_shape[1], 'image file must bigger than ' \

'target_size'

crop_w = np.random.randint(0, img.shape[0] - target_size[0] + 1)

crop_h = np.random.randint(0, img.shape[1] - target_size[1] + 1)

img_random_cropped = img[crop_w:crop_w + target_size[0], crop_h:crop_h + target_size[1]]

img_rgb = gray2rgb(img_random_cropped)

# must divide 225. before rgb2lab

img_lab = rgb2lab(img_rgb/225.)

x = img_lab[:, :, 0]

x = np.reshape(x, (target_size[0], target_size[1], 1))

return x

def colorize_gray():

model = build_model()

model.load_weights('simple_model.h5')

print('model loaded..')

target_size = (256, 256)

test_img_dir = 'data/Test'

all_test_img_files = [os.path.join(test_img_dir, i) for i in os.listdir(test_img_dir) if i.endswith('jpg')]

for img_file in all_test_img_files:

f_name = os.path.basename(img_file).split('.')[0]

x_origin = pre_process_image_for_colorize(img_file=img_file)

x = np.expand_dims(x_origin, axis=0)

# y*128????????

y = model.predict(x)

y *= 128

# now the y is the AB of LAB color, we concat with x_origin

tmp = np.zeros((target_size[0], target_size[1], 3))

tmp[:, :, 0] = x_origin[:, :, 0]

tmp[:, :, 1:] = y[0]

# now the image should be LAB and colorful

imsave(os.path.join(test_img_dir, 'result_simple_{}.jpg'.format(f_name)), lab2rgb(tmp))

print('result image of {} saved.'.format(img_file))

if __name__ == '__main__':

# colorize()

colorize_gray()

让我们看看这些图片的结果咋样：

哇塞，好像上一章图片学会了很多绿色！！！

而且会把这些绿色加在带预测的黑白图片上！！！

可以说我们的实验算是成功了一半。我们实现了一个简单版本的黑白图片上色AI，除此之外，我们使用这个模型尝试着对其他黑白图片进行上色，我们惊奇的发现这个模型给所有图片加上了他所学习到的绿色！！！！

好了，这一篇就实验到这里，下一篇我们将继续以下探索：实现一个批量数据生成器，把大规模彩色图片输入网络训练；

我们将会继续尝试加深我们的模型，甚至采用encoder和decoder的架构来实现更深的上色模型！！

3. 总结

这是给黑白图片人工智能上色教程的第一篇，在这一个部分中，我们非常脚踏实地的学习了一下黑白图片上色的基本原理。当然很多人会说，为什么没有GAN，你生成图片不用GAN，真的好吗？能提现逼格吗？我18年来概率与统计岂不是白学了？莫慌，

本次列车到此结束，欢迎继续乘坐老司机的高速列车。更多的好玩有创意的深度学习文章将在我的博客继续发布，也期待大家在我的微博关注我：大魔术师金天 。

本文由在当地较为英俊的男子金天大神原创，版权所有，欢迎转载，本文首发地址 https://jinfagang.github.io 。但请保留这段版权信息，多谢合作，有任何疑问欢迎通过微信联系我交流：jintianiloveu