本文主要介紹採用python中的cv2模組實現圖片上色。

首先需要下載預訓練的model。

1、建立getModels.sh檔案，並輸出如下內容：

mkdir models
wget https://github.com/richzhang/colorization/blob/master/colorization/resources/pts_in_hull.npy?raw=true -O ./pts_in_hull.npy
wget https://raw.githubusercontent.com/richzhang/colorization/master/colorization/models/colorization_deploy_v2.prototxt -O ./models/colorization_deploy_v2.prototxt
wget http://eecs.berkeley.edu/~rich.zhang/projects/2016_colorization/files/demo_v2/colorization_release_v2.caffemodel -O ./models/colorization_release_v2.caffemodel
wget http://eecs.berkeley.edu/~rich.zhang/projects/2016_colorization/files/demo_v2/colorization_release_v2_norebal.caffemodel -O ./models/colorization_release_v2_norebal.caffemodel
 

2、

chmod +x getModels.sh

3、

sh getModels.sh

其次，執行如下程式碼：

import numpy as np
import cv2
import os.path

# 讀取圖片
frame = cv2.imread('./greyscaleImage.png')

# 指定model所在路徑
protoFile = './models/colorization_deploy_v2.prototxt'
weightFile = './models/colorization_release_v2.caffemodel'

# 載入聚類中心
pts_in_hull = np.load('./pts_in_hull.npy')

# 讀取網路
net = cv2.dnn.readNetFromCaffe(protoFile, weightFile)

# 將聚類中心填充為1x1卷積核
pts_in_hull = pts_in_hull.transpose().reshape(2, 313, 1, 1)
net.getLayer(net.getLayerId('class8_ab')).blobs = [pts_in_hull.astype(np.float32)]
net.getLayer(net.getLayerId('conv8_313_rh')).blobs = [np.full([1, 313], 2.606, np.float32)]

W_in = 224
H_in = 224

img_rgb = (frame[:, :, [2, 1, 0]] * 1.0 / 255).astype(np.float32)
img_lab = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2Lab)
img_l = img_lab[:, :, 0]  # 提取出L通道

# 將L通道的影象重新設定大小為network的輸出大小
img_l_rs = cv2.resize(img_l, (W_in, H_in))
img_l_rs -= 50 # 從mean-centering減去50
net.setInput(cv2.dnn.blobFromImage(img_l_rs))
ab_dec = net.forward()[0, :, :, :].transpose((1, 2, 0))  # 結果

(H_orig, W_orig) = img_rgb.shape[:2]  # 原始的圖片大小
ab_dec_us = cv2.resize(ab_dec, (W_orig, H_orig))
img_lab_out = np.concatenate((img_l[:, :, np.newaxis], ab_dec_us), axis=2)  # 與原始圖片L連線
img_bgr_out = np.clip(cv2.cvtColor(img_lab_out, cv2.COLOR_Lab2BGR), 0, 1)

outputFile = 'colorized.png'
cv2.imwrite(outputFile, (img_bgr_out * 255).astype(np.uint8))
print('Done!')
 

測試圖片：

測試結果：

最後，附上github連結。