自己實現黑白圖片自動上色AI（一）

01-23

前言

我們將用兩篇文章來從零實現一個自動上色的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 colorizeyou are going to need skimage and keras for itany version should be compatible"""import kerasimport tensorflow as tffrom skimage.io import imread, imsavefrom skimage.color import rgb2gray, gray2rgb, rgb2lab, lab2rgbfrom keras.models import Sequentialfrom keras.layers import Conv2D, UpSampling2D, InputLayer, Conv2DTransposeimport numpy as npdef 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 colorizeyou are going to need skimage and keras for itany version should be compatible"""import kerasimport tensorflow as tffrom skimage.io import imread, imsavefrom skimage.color import rgb2gray, gray2rgb, rgb2lab, lab2rgbfrom keras.models import Sequentialfrom keras.layers import Conv2D, UpSampling2D, InputLayer, Conv2DTransposefrom keras.preprocessing.image import img_to_array, load_imgimport numpy as npfrom keras.preprocessing.image import ImageDataGeneratorimport osdef 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 modeldef 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()

讓我們看看這些圖片的結果咋樣：