百行程式碼構建神經網路黑白圖片自動上色系統

機器之心發表於2017-10-31

使用神經網路對圖片進行風格化渲染是計算機視覺領域的熱門應用之一。本文將向你介紹一種簡單而有效的黑白圖片上色方法,僅需 100 行程式碼,你也可以搭建自己的神經網路,幾秒鐘內讓計算機自動完成手動工作需要幾個月的任務。

今年 7 月,Amir Avni 用神經網路向 Reddit 的 Colorization 社群宣戰——那是一個為歷史黑白照片上色的版面,使用者們通常使用的工具是 Photoshop。

社群使用者們驚訝於 Amir 的深度學習系統。以前人們需要幾個月手工上色的任務在計算機那裡只需要幾秒鐘就完成了。我也被 Amir 的神經網路迷住了,所以在看到結果之後我開始試著復現它。首先,讓我們看看我們自動上色的成果/失敗品。

百行程式碼構建神經網路黑白圖片自動上色系統

目前,圖片上色的工作主要由人們通過 Photoshop 手工完成,這需要相當多的工序和時間。

簡而言之,一張照片可能需要一個月的時間來上色。對於畫師來說,這還意味著很多研究工作,一張人臉可能需要 20 張粉色、綠色和藍色圖層,經過不斷調整最終才能獲得正確的效果。

本文面向初學者,如果你認為自己對於深度學習理解不足,可先看看:

https://blog.floydhub.com/my-first-weekend-of-deep-learning/

https://blog.floydhub.com/coding-the-history-of-deep-learning/

以及 Andrej Karpathy 的 CS231 課程視訊:https://www.youtube.com/watch?v=LxfUGhug-iQ

本文將展示如何三步構建自己的神經網路上色系統。

第一部分先解釋核心邏輯。我們會構建一個 40 行的神經網路作為「Alpha」上色機器人,其中沒有應用到太多深度技術,這個過程會讓我們熟悉模型的核心機制。

第二部是建立一個可以泛化的神經網路——我們的「beta」版本。它可以讓我們對此前未見的圖片進行上色。

在「正式版」中,我們會將神經網路與分類器相結合。我們會使用訓練過 120 萬張圖片的 Inception ResNet V2,為了讓上色風格更符合現代審美,我們使用 Unsplash 上的圖片來訓練神經網路。

該專案的 Jupyter Notebook:https://www.floydhub.com/emilwallner/projects/color/43/code/Alpha-version/alpha_version.ipynb

FloydHub:https://www.floydhub.com/emilwallner/projects/color/43/code

GitHub:https://github.com/emilwallner/Coloring-greyscale-images-in-Keras

以及 FloydHub 上的雲 GPU 實驗:https://www.floydhub.com/emilwallner/projects/color/jobs

核心邏輯

在本節中,我會介紹如何渲染影象,數字色彩的機理以及神經網路的主要邏輯。黑白圖片可以在畫素網格中表示。每個畫素由對應其亮度的數值,範圍在 0-255 之間,對應從黑到白。

百行程式碼構建神經網路黑白圖片自動上色系統

彩色影象(RGB)由三個層組成:紅色層、綠色層與藍色層。若將白色背景上一片綠葉的圖片分析為三個通道。直觀地,你可能認為樹葉只存在於綠色層。但是正如你所見到的,樹葉在三個通道中都出現了,這是因為圖層不僅決定顏色,還決定亮度。

百行程式碼構建神經網路黑白圖片自動上色系統

為了實現白色,你需要讓三種顏色的分佈平均,通過為紅色和藍色加入相同數值——結果就是綠色變得更亮了。所以,彩色影象使用三個通道的數值對比對顏色進行編碼。

百行程式碼構建神經網路黑白圖片自動上色系統

對於黑白圖片,每個圖層的數值都是 0-255 之間的相同值。如果所有通道數值為 0,則這個畫素為黑色。你可能知道神經網路會在輸入與輸出之間構建聯絡。將這個思路歸納到我們的著色任務中——神經網路需要找到灰度影象和彩色影象之間的聯絡。

再準確一點,我們在尋找的是將灰度值連結到三色圖層數值的方法。

百行程式碼構建神經網路黑白圖片自動上色系統

Alpha 版本

我們將從神經網路的簡單版本開始,先來為這張圖片中的人臉進行著色。這樣你就可以熟悉模型的核心機制了。

只需 40 行程式碼,我們就可以做到下圖中的著色效果。下圖中間是神經網路的著色,而右側是原圖。神經網路也經過了這張圖的訓練。在 Beta 版本的段落中我們將對此進行詳細解釋。

百行程式碼構建神經網路黑白圖片自動上色系統


色彩空間

首先,我們需要用演算法來改變顏色通道,從 RGB 到 Lab。L 代表亮度,而 a 和 b 代表顏色光譜綠-紅和藍-黃。

正如下圖所示,Lab 編碼的圖片有一個灰度層,而顏色層由三個減少為兩個。這意味著我們可以在最終的預測中使用原來的灰度圖片,同時只需要預測兩個通道。

百行程式碼構建神經網路黑白圖片自動上色系統

從實驗上來說,人類眼睛裡大約 94% 的細胞是用來探測亮度的,只有剩下 6% 是用來感應顏色的。正如你在上圖中所看到的,灰度圖比顏色圖看起來清楚多了,這也是我們需要在神經網路輸出中保留灰度圖的另一個原因。

從黑白到彩色

神經網路的最終預測就像這樣:我們有一張灰度圖的輸入,我們希望以此來預測兩個色彩圖層,Lab 中的 ab。最終的影象包括了輸入的灰度圖層 L 和預測的兩個圖層,最終組成 Lab 圖片。

百行程式碼構建神經網路黑白圖片自動上色系統

為了將一層轉換為兩層,我們需要使用卷積濾波器。你可以把他們想象為 3D 眼鏡中的藍/紅色偏振鏡。每個偏振鏡都會決定我們看到的圖片是什麼樣的,它可以突出或刪除圖片中的部分資訊。神經網路也可以從一個/多個過濾器組合之上創造新的影象。

對於卷積神經網路而言,每一個濾波器都會自動調整以幫助達到預期的結果。我們先要從堆疊數百個濾波器,並將它們塞進兩個顏色圖層 a、b 中做起。不過首先,我們來看看程式碼。


在 FloydHub 中部署程式碼

如果你不太瞭解 FloydHub,你可以先看看它的安裝指南:

https://www.floydhub.com/

或我的上手指南:

https://blog.floydhub.com/my-first-weekend-of-deep-learning/

Alpha 版本

安裝完 FloydHub 後,鍵入下列命令:

git clone https://github.com/emilwallner/Coloring-greyscale-images-in-Keras

開啟資料夾,啟動 FloydHub。

cd Coloring-greyscale-images-in-Keras/floydhub
floyd init colornet

FloydHub 的網路控制皮膚就在網頁裡。系統將提示你建立一個名為 colornet 的新 FloydHub 專案。一旦完成,回到你的終端並執行相同的 init 命令。

floyd init colornet

好了,開始執行我們的程式:

floyd run --data emilwallner/datasets/colornet/2:data --mode jupyter --tensorboard

注意:我們已經在 FloydHub 上安裝了一個公開資料集(我已經上傳了),資料集的目錄在這裡:

--dataemilwallner/datasets/colornet/2:data

你可以直接在 FloydHub 上使用這個資料集(其他的也是)並且:

  • 通過--tensorboard 指令開啟 Tensorboard
  • 通過--mode jupyter 指令執行 Jupyter Notebook 模式
  • 如果你有可用 GPU,你可以在命令列中加入 GPU flag: --gpu,這樣訓練的速度可以大約快上 50 倍。

在 FloydHub 網站的 Jobs 標籤下,點選 Jupyter Notebook 連結,並找到這個檔案:

floydhub/Alpha version/working_floyd_pink_light_full.ipynb

開啟它,並執行所有程式碼塊,隨著逐漸增加的 epoch 值,我們可以體會神經網路是如何學習的。

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

首先從 epoch 值為 1 開始,逐漸增加到 10、100。Epoch 的數值表示神經網路從訓練集影象中學習的次數。在訓練完神經網路後,你可以在的 img_result.png 主資料夾中找到結果圖。

# Get images
image = img_to_array(load_img('woman.png'))
image = np.array(image, dtype=float)
# Import map images into the lab colorspace
X = rgb2lab(1.0/255*image)[:,:,0]
Y = rgb2lab(1.0/255*image)[:,:,1:]
Y = Y / 128
X = X.reshape(1, 400, 400, 1)
Y = Y.reshape(1, 400, 400, 2)
model = Sequential()
model.add(InputLayer(input_shape=(None, None, 1)))
# Building the neural network
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'))
# Finish model
model.compile(optimizer='rmsprop',loss='mse')
#Train the neural network
model.fit(x=X, y=Y, batch_size=1, epochs=3000)
print(model.evaluate(X, Y, batch_size=1))
# Output colorizations
output = model.predict(X)
output = output * 128
canvas = np.zeros((400, 400, 3))
canvas[:,:,0] = X[0][:,:,0]
canvas[:,:,1:] = output[0]
imsave("img_result.png", lab2rgb(cur))
imsave("img_gray_scale.png", rgb2gray(lab2rgb(cur)))

執行神經網路的 FloydHub 命令:

floyd run --data emilwallner/datasets/colornet/2:data --mode jupyter --tensorboard


技術說明

簡單來說,輸入就是以網格表示的黑白影象,輸出是兩個具有顏色值的網格。在輸入和輸出之間用濾波器將它們連線起來。這是一個卷積神經網路。
百行程式碼構建神經網路黑白圖片自動上色系統

從左側開始,分別是 B&W 輸入、濾波器和神經網路的預測。


我們需要在相同的區間將預測值和真實值建立對映,從而將值進行比較。區間的範圍在-1 到 1 之間。為了對映預測值,我們使用了 tanh 啟用函式。tanh 函式的任意輸入的輸出值都在-1 到 1 之間。

真實的顏色值分佈在-128 到 128 的區間內。這是 Lab 色彩空間的預設區間。用 128 除這些值就能獲得-1 到 1 的區間分佈。這種「標準化」操作使我們能比較預測的誤差。

計算除最後的誤差之後,網路將更新濾波器以減少總誤差。網路會持續迭代這個過程直到誤差儘可能的小。接下來我們將描述程式碼片段中的一些想法與表達。

X = rgb2lab(1.0/255*image)[:,:,0]
Y = rgb2lab(1.0/255*image)[:,:,1:]

1.0/255 表示我們使用了一個 24 位元的 RGB 色彩空間,意味著在每個色彩通道中使用 0-255 之間的一個數字,總共有 1670 萬個色彩組合。

由於人類只能感知 200 萬到 1000 萬種顏色,使用這麼大的色彩空間不太必要。

Y = Y / 128

Lab 色彩空間的分佈範圍和 RGB 不同。Lab 的色譜範圍從-128 到 128。通過用 128 除輸出層的值,可以將範圍轉換為-1 到 1。

我們的神經網路的輸出也是在這個範圍內,因此可以互相匹配。

在使用 rgb2lab() 函式轉換色彩空間後,我們用 [ : , : , 0] 選擇灰度層,這是神經網路的輸入,[ : , : , 1: ] 選擇的是綠-紅、藍-黃這兩個層。在訓練了神經網路之後,根據最後預測結果生成影象。

output = model.predict(X)
output = output * 128

在這裡,我們使用灰度影象作為輸入,並用於訓練神經網路。我們對所有的位於-1 到 1 範圍內輸出乘上 128,從而得到 Lab 色彩空間的真實顏色。

canvas = np.zeros((400, 400, 3))
canvas[:,:,0] = X[0][:,:,0]
canvas[:,:,1:] = output[0]

最後,通過三層全為 0 值的網格建立一個黑色 RGB 畫布,然後從測試影象中複製灰度影象,並將這兩個顏色層新增到 RGB 畫布上去,最後將這個畫素值陣列轉換為一張影象。


研究經驗

  • 閱讀研究文獻是個難關:一旦我總結了每一篇論文中的核心特徵之後,瀏覽論文就變得容易多了,我還學會將細節放入特定的背景中理解。
  • 從簡單的開始很關鍵:我在網上找到的大多數實現都是 2000 到 10000 行的長度的程式碼。這使得理解問題的核心邏輯變得很困難。一旦獲得主要思維框架之後,閱讀程式碼實現和研究文獻就變得容易多了。
  • 探索開源專案:為了粗略理解寫程式碼的要領,我在 GitHub 上瀏覽了 50-100 個著色專案。
  • 過程不總是很順利:在剛開始的時候,網路只能生成紅色和黃色。我最初用 ReLU 函式作為最後一層的啟用函式,由於它只能將數字對映為正值,而無法輸出負值,即藍色和綠色的色譜。通過改為使用 tanh 為啟用函式解決了這個問題。
  • 從理解到加速:我見過的大多數的實現工作得很快但很難使用。而我更注重模型革新的速度,而不是程式碼執行速度。

beta 版本

如果想理解 alpha 版本的弱點,可以嘗試為未訓練過的圖片上色,可以看到,它通常會給出糟糕的結果,因為網路只對資訊進行了記憶,從而無法為沒見過的影象上色。我們將以這個方向在 beta 版本做出改進,即讓它學會泛化。

百行程式碼構建神經網路黑白圖片自動上色系統

特徵提取器

我們的網路能找到連線灰度影象和其色彩版本的特徵。

假設你需要對黑白影象上色,但因為一些限制你一次只能看到 9 個畫素。你可以從左上方到右下方掃描每一張影象並嘗試預測每一個畫素對應的顏色。

百行程式碼構建神經網路黑白圖片自動上色系統

例如,這 9 個畫素是下方所示的女人的鼻孔的邊緣。可以想象,這幾乎不可能做出正確的上色。因此需要分成幾步來完成。

首先,尋找簡單的模式:對角線、所有的黑色畫素,等等。在每一個方塊中尋找相同的模式並移除不匹配的畫素。從 64 個小型濾波器中生成 64 張新影象。

百行程式碼構建神經網路黑白圖片自動上色系統

當再次掃描影象的時候,你能找到與已經檢測到的相同的小範圍模式。為了獲得對影象的高層次的理解,需要將影象尺寸減半。

百行程式碼構建神經網路黑白圖片自動上色系統

目前仍然只有一個 3x3 濾波器用於掃描每一張影象。但通過用低階的濾波器組合成新的 9 畫素濾波器,可以探測更加複雜的模式。每個畫素組合可能用於探測一個半圓、一個點或一條線。你可以重複地在影象中提取相同的特徵。這一次總共生成了 128 張新的過濾影象。

經過幾步的過濾之後,過濾影象的樣子可能如下圖所示:

百行程式碼構建神經網路黑白圖片自動上色系統

之前提到過,這個過程從低階的特徵開始,比如邊緣,靠近輸出的層將其組合成模式,然後再組合成面部細節,最後轉換成一張臉。這個視訊教程中提供了更詳細的解釋:https://www.youtube.com/watch?v=AgkfIQ4IGaM

整個過程和大部分處理視覺的神經網路很相似。這裡用的神經網路的型別是卷積神經網路。這些網路組合多個濾波器以理解影象中的語境。

從特徵提取到色彩

神經網路以試驗和除錯誤差的方式訓練。它首先為每一個畫素給出一個隨機預測,基於每一個畫素計算出的誤差,通過反向傳播提升特徵提取的效能。

在我們的案例中,調整網路的目標是顏色的匹配度和目標的定位。我們的網路從將所有目標轉換為褐色開始。褐色是和所有其它顏色相差最小的顏色,因此產生的誤差最小。

由於大多數的訓練資料都很相似,網路在分辨不同的目標時遇到了困難。這將使網路在生成更加細緻的顏色時失敗。這將是我們在完整的版本中要探索的問題。

以下是 beta 版本的程式碼,隨後是程式碼的技術解釋。

# Get images
X = []
for filename in os.listdir('../Train/'):
    X.append(img_to_array(load_img('../Train/'+filename)))
X = np.array(X, dtype=float)
# Set up training and test data
split = int(0.95*len(X))
Xtrain = X[:split]
Xtrain = 1.0/255*Xtrain
#Design the neural network
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)))
# Finish model
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 = 50
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)
# Train model
TensorBoard(log_dir='/output')
model.fit_generator(image_a_b_gen(batch_size), steps_per_epoch=10000, epochs=1)
# Test images
Xtest = rgb2lab(1.0/255*X[split:])[:,:,:,0]
Xtest = Xtest.reshape(Xtest.shape+(1,))
Ytest = rgb2lab(1.0/255*X[split:])[:,:,:,1:]
Ytest = Ytest / 128
print model.evaluate(Xtest, Ytest, batch_size=batch_size)
# Load black and white images
color_me = []
for filename in os.listdir('../Test/'):
        color_me.append(img_to_array(load_img('../Test/'+filename)))
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 = output * 128
# Output colorizations
for i in range(len(output)):
        cur = np.zeros((256, 256, 3))
        cur[:,:,0] = color_me[i][:,:,0]
        cur[:,:,1:] = output[i]
        imsave("result/img_"+str(i)+".png", lab2rgb(cur))

下面是 FloydHub 指令,用來執行 Beta 神經網路:

floyd run --data emilwallner/datasets/colornet/2:data --mode jupyter --tensorboard


技術說明

該網路與其他視覺神經網路的主要不同是畫素位置的重要性。在上色網路中,影象大小或比例在網路中保持不變。在其他型別的網路中,影象距離最後一層越近就越失真。

分類網路中的最大池化層增加了資訊密度,同時也使影象失真。它只對資訊估值,而不是影象的佈局。在上色網路中,我們使用步幅 2(帶 padding)以減少一半的寬度和高度。這也會增加資訊密度,但不會使影象失真。

百行程式碼構建神經網路黑白圖片自動上色系統

兩個進一步的區別是:上取樣層和影象比例的保持問題。分類網路只關心最後的分類。因此,整個網路架構將持續降低影象的大小和質量。

上色網路保持影象比率的穩定。這通過新增空白填充(如上圖所示)來完成。此外,每個卷積層剪下影象。它通過引數 padding='SAME'來完成。

為了使影象大小翻倍,上色網路使用上取樣層https://keras.io/layers/convolutional/#upsampling2d)。

for filename in os.listdir('/Color_300/Train/'):
X.append(img_to_array(load_img('/Color_300/Test'+filename)))

該迴圈首先計算目錄中所有檔名稱,然後通過影象目錄進行迭代,將影象轉換成畫素陣列。最後,將它們連線成為巨型向量(giant vector)。

datagen = ImageDataGenerator(
shear_range=0.2,
zoom_range=0.2,
rotation_range=20,
horizontal_flip=True)

我們可以使用 ImageDataGenerator(https://keras.io/preprocessing/image/)調整影象生成器的設定。這樣,每個影象都不一樣,從而提高學習率。shear_range 使影象向左或向右傾斜,另一種設定是放大、旋轉和水平翻轉。

batch_size = 50
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)

我們使用資料夾 Xtrain 中的影象,來生成基於以上設定的影象。然後我們為 X_batch 提取黑色層和白色層,併為兩個色彩層(color layer)提取兩個顏色值。

model.fit_generator(image_a_b_gen(batch_size), steps_per_epoch=1, epochs=1000)

GPU 越強大,可以擬合的影象就越多。通過這個設定,你可以每次迭代中使用 50-100 張影象。steps_per_epoch 通過將訓練的影象數量除以批次大小(batch size)來計算。

例如,如果有 100 張影象,而批次大小為 50,則 steps_per_epoch 為 2。epoch 表示訓練全部影象的次數。在一個 Tesla K80 GPU 上用 10000 張影象訓練 21 個 epoch 需要大約 11 個小時。

研究經驗

  • 在大規模執行之前先進行大量小批次實驗。即使實驗了二三十次,還是能發現錯誤。因為執行不代表就一定奏效。神經網路中的 bug 通常比傳統程式設計錯誤更加細微。
  • 更多樣的資料集使影象更加呈褐色。如果你的影象很相似,那麼無需過於複雜的架構,你就可以得到不錯的結果。缺點是網路的泛化效果不好。
  • 形狀、形狀、形狀。每個影象的大小必須確定且與網路比例相符。一開始,我用的是一個大小 300 的影象。把這個影象分割了三次,得到 150、75 和 35.5 的影象。結果是丟了一半的畫素!這導致很多不好的結果,直到我意識到最好使用 2 的冪次方:2、8、16、32、64、256 等。
  • 建立資料集:a)停用.DS_Store 檔案。b)有創造性。我最後使用了 Chrome console 指令碼和擴充套件程式 ImageSpark 來下載檔案。c)複製抓取的原始檔案,使你的清理指令碼結構化。

完整版

我們最終版的上色神經網路有四個組成部分。我們將之前的網路分割成編碼器和解碼器。在二者中間,使用一個融合層。如果你對分類網路不是很瞭解,推薦學習該教程:http://cs231n.github.io/classification/

與編碼器並行,輸入影象還在最強大的分類器之一 Inception ResNet v2 中執行。這是一個在 120 萬張影象上訓練的神經網路。我們提取分類層,然後將它與編碼器的輸出融合起來。

百行程式碼構建神經網路黑白圖片自動上色系統

原始論文詳細視覺化:https://github.com/baldassarreFe/deep-koalarization

通過從分類器到上色網路的遷移學習,該網路可以瞭解影象中有什麼。這樣就可以使該網路通過上色機制匹配目標表示(object representation)。

這裡是部分驗證影象,僅使用 20 張影象用於網路訓練。

百行程式碼構建神經網路黑白圖片自動上色系統

大部分影象的輸出結果效果不好。但是我在一個較大的驗證集(包含 2500 個影象)上找到了一些不錯的影象。在更多影象上訓練模型可以給出更一致的結果,但是大多數影象經過處理後呈褐色。我使用多個影象(包括驗證影象)進行了一些實驗,全部實驗列表:https://www.floydhub.com/emilwallner/projects/color

這裡是之前研究中最常見的架構:

常見色彩空間:Lab、YUV、HSV 和 LUV

常見損失函式:均方差、分類、加權分類

我選擇了「融合層」架構(上述列表的第五個)。

因為融合層的輸出結果最好,而且在 Keras 中理解和復現都更加容易。儘管它不是最強大的上色網路,但很適合初學者,而且最適合理解上色問題的動態。

我使用 Federico Baldassarre 等人論文中設計的神經網路https://github.com/baldassarreFe/deep-koalarization/blob/master/report.pdf),在 Keras 中進行了自己的操作。

技術說明

我們想級聯或融合多個模型,Keras 的功能性 API 非常合適。(API 地址:https://keras.io/getting-started/functional-api-guide/

百行程式碼構建神經網路黑白圖片自動上色系統

首先,我們下載 Inception ResNet v2 神經網路,載入權重。由於我們將並行使用這兩個模型(Inception ResNet v2 和編碼器),我們需要明確要使用的模型。這兩個模型可以使用帶 TensorFlow 後端的 Keras 實現:

inception = InceptionResNetV2(weights=None, include_top=True)
inception.load_weights('/data/inception_resnet_v2_weights_tf_dim_ordering_tf_kernels.h5')
inception.graph = tf.get_default_graph()

我們使用影象微調來建立批次,然後將它們轉換成黑白影象,在 Inception ResNet 模型中執行。

grayscaled_rgb = gray2rgb(rgb2gray(batch))
embed = create_inception_embedding(grayscaled_rgb)

首先,我們必須調整影象的大小,使之適應 Inception 模型。然後使用前處理器根據模型對畫素和顏色值進行格式處理。最後,在 Inception 模型中執行該影象,並提取模型的最後一層。

def create_inception_embedding(grayscaled_rgb):
grayscaled_rgb_resized = []
for i in grayscaled_rgb:
i = resize(i, (299, 299, 3), mode='constant')
grayscaled_rgb_resized.append(i)
grayscaled_rgb_resized = np.array(grayscaled_rgb_resized)
grayscaled_rgb_resized = preprocess_input(grayscaled_rgb_resized)
with inception.graph.as_default():
embed = inception.predict(grayscaled_rgb_resized)
return embed

我們回過頭來看一下生成器。每一個批次都生成 20 個影象,格式如下。這在 Tesla K80 GPU 上需要花費大約一個小時。在沒有記憶體問題的情況下,該模型最多可以一次性生成 50 個影象。

yield ([X_batch, create_inception_embedding(grayscaled_rgb)], Y_batch)

這和我們的 colornet 模型的格式相匹配。

model = Model(inputs=[encoder_input, embed_input], outputs=decoder_output)

將 encoder_input 饋送至編碼器模型,然後在融合層中把編碼器模型的輸出和 embed_input 融合起來;融合層的輸出作為解碼器模型的輸入,然後返回最終輸出 decoder_output。

fusion_output = RepeatVector(32 * 32)(embed_input) 
fusion_output = Reshape(([32, 32, 1000]))(fusion_output)
fusion_output = concatenate([fusion_output, encoder_output], axis=3)
fusion_output = Conv2D(256, (1, 1), activation='relu')(fusion_output)

在融合層中,我們首先將帶有 1000 個類別的層級乘 1024 (32 * 32)。這樣,我們就使用 Inception 模型的最後一層得到了 1024 個單元。

然後將它們從 2D 重塑為 3D,即將維度更改為 32 x 32x1000 的張量。然後把它們和編碼器模型的輸出連線起來。我們使用 256 個 1X1 卷積核的卷積網路,饋送到 ReLU 啟用函式後作為融合層的最終輸出。

下一步

為影象上色是一個非常有趣迷人的問題。它是一個科學問題,也是一個藝術問題。我寫這篇文章的目的就是使大家可以瞭解影象上色,並且繼續開發相關技術。以下是一些初學建議:

  • 使用另一個預訓練模型實現它
  • 嘗試一個不同的資料集
  • 使用更多影象來提高網路的準確率
  • 在 RGB 色彩空間內構建一個放大器。建立一個與上色網路類似的模型,該模型使用高飽和彩色影象作為輸入,正確的彩色影象作為輸出。
  • 實現加權分類
  • 將該網路應用到視訊。不用擔心上色,試圖使影象之間的轉換更加連貫。你還可以對較大的影象進行一些類似的操作。

你還可以使用 FloydHub,用這三個版本的上色神經網路為黑白影象上色。

  • 對於 alpha 版本,只需要將 woman.jpg 檔案替換成你自己的同名檔案即可(影象大小 400x400)。
  • 對於 beta 版本和完整版本,將你的影象新增至 Test 資料夾,然後執行 FloydHub 命令。如果 Notebook 正在執行的話,你還可以將 Notebook 中的影象直接上傳至 Test 資料夾。注意:這些影象必須是 256x256 畫素。還有,你可以上傳彩色影象作為測試影象,因為系統可以自動將它們轉換成黑白影象。

程式碼

注意:使用下列程式碼時我從 Keras 序列模型轉向了它們的功能性 API。(文件:https://keras.io/getting-started/functional-api-guide/

# Get images
X = []
for filename in os.listdir('/data/images/Train/'):
    X.append(img_to_array(load_img('/data/images/Train/'+filename)))
X = np.array(X, dtype=float)
Xtrain = 1.0/255*X
#Load weights
inception = InceptionResNetV2(weights=None, include_top=True)
inception.load_weights('/data/inception_resnet_v2_weights_tf_dim_ordering_tf_kernels.h5')
inception.graph = tf.get_default_graph()
embed_input = Input(shape=(1000,))
#Encoder
encoder_input = Input(shape=(256, 256, 1,))
encoder_output = Conv2D(64, (3,3), activation='relu', padding='same', strides=2)(encoder_input)
encoder_output = Conv2D(128, (3,3), activation='relu', padding='same')(encoder_output)
encoder_output = Conv2D(128, (3,3), activation='relu', padding='same', strides=2)(encoder_output)
encoder_output = Conv2D(256, (3,3), activation='relu', padding='same')(encoder_output)
encoder_output = Conv2D(256, (3,3), activation='relu', padding='same', strides=2)(encoder_output)
encoder_output = Conv2D(512, (3,3), activation='relu', padding='same')(encoder_output)
encoder_output = Conv2D(512, (3,3), activation='relu', padding='same')(encoder_output)
encoder_output = Conv2D(256, (3,3), activation='relu', padding='same')(encoder_output)
#Fusion
fusion_output = RepeatVector(32 * 32)(embed_input) 
fusion_output = Reshape(([32, 32, 1000]))(fusion_output)
fusion_output = concatenate([encoder_output, fusion_output], axis=3) 
fusion_output = Conv2D(256, (1, 1), activation='relu', padding='same')(fusion_output)
#Decoder
decoder_output = Conv2D(128, (3,3), activation='relu', padding='same')(fusion_output)
decoder_output = UpSampling2D((2, 2))(decoder_output)
decoder_output = Conv2D(64, (3,3), activation='relu', padding='same')(decoder_output)
decoder_output = UpSampling2D((2, 2))(decoder_output)
decoder_output = Conv2D(32, (3,3), activation='relu', padding='same')(decoder_output)
decoder_output = Conv2D(16, (3,3), activation='relu', padding='same')(decoder_output)
decoder_output = Conv2D(2, (3, 3), activation='tanh', padding='same')(decoder_output)
decoder_output = UpSampling2D((2, 2))(decoder_output)
model = Model(inputs=[encoder_input, embed_input], outputs=decoder_output)
#Create embedding
def create_inception_embedding(grayscaled_rgb):
    grayscaled_rgb_resized = []
    for i in grayscaled_rgb:
        i = resize(i, (299, 299, 3), mode='constant')
        grayscaled_rgb_resized.append(i)
    grayscaled_rgb_resized = np.array(grayscaled_rgb_resized)
    grayscaled_rgb_resized = preprocess_input(grayscaled_rgb_resized)
    with inception.graph.as_default():
        embed = inception.predict(grayscaled_rgb_resized)
    return embed
# Image transformer
datagen = ImageDataGenerator(
        shear_range=0.4,
        zoom_range=0.4,
        rotation_range=40,
        horizontal_flip=True)
#Generate training data
batch_size = 20
def image_a_b_gen(batch_size):
    for batch in datagen.flow(Xtrain, batch_size=batch_size):
        grayscaled_rgb = gray2rgb(rgb2gray(batch))
        embed = create_inception_embedding(grayscaled_rgb)
        lab_batch = rgb2lab(batch)
        X_batch = lab_batch[:,:,:,0]
        X_batch = X_batch.reshape(X_batch.shape+(1,))
        Y_batch = lab_batch[:,:,:,1:] / 128
        yield ([X_batch, create_inception_embedding(grayscaled_rgb)], Y_batch)
#Train model      
tensorboard = TensorBoard(log_dir="/output")
model.compile(optimizer='adam', loss='mse')
model.fit_generator(image_a_b_gen(batch_size), callbacks=[tensorboard], epochs=1000, steps_per_epoch=20)
#Make a prediction on the unseen images
color_me = []
for filename in os.listdir('../Test/'):
    color_me.append(img_to_array(load_img('../Test/'+filename)))
color_me = np.array(color_me, dtype=float)
color_me = 1.0/255*color_me
color_me = gray2rgb(rgb2gray(color_me))
color_me_embed = create_inception_embedding(color_me)
color_me = rgb2lab(color_me)[:,:,:,0]
color_me = color_me.reshape(color_me.shape+(1,))
# Test model
output = model.predict([color_me, color_me_embed])
output = output * 128
# Output colorizations
for i in range(len(output)):
    cur = np.zeros((256, 256, 3))
    cur[:,:,0] = color_me[i][:,:,0]
    cur[:,:,1:] = output[i]
    imsave("result/img_"+str(i)+".png", lab2rgb(cur))

以下是執行完整版神經網路的 FloydHub 命令:

floyd run --data emilwallner/datasets/colornet/2:data --mode jupyter --tensorboard

相關文章