【OpenCV】133 OpenCV DNN 图像颜色化模型使用

133 OpenCV DNN 图像颜色化模型使用

代码

import numpy as np
import cv2 as cv

W_in = 224
H_in = 224
modelTxt = "../models/color/colorization_deploy_v2.prototxt";
modelBin = "../models/color/colorization_release_v2.caffemodel";
pts_txt = "../models/color/pts_in_hull.npy";

# Select desired model
net = cv.dnn.readNetFromCaffe(modelTxt, modelBin)
pts_in_hull = np.load(pts_txt) # load cluster centers

# populate cluster centers as 1x1 convolution kernel
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)]

frame = cv.imread("../images/test1.png")
h, w = frame.shape[:2]
img_rgb = (frame[:,:,[2, 1, 0]] * 1.0 / 255).astype(np.float32)

img_lab = cv.cvtColor(img_rgb, cv.COLOR_RGB2Lab)
img_l = img_lab[:,:,0] # pull out L channel
(H_orig,W_orig) = img_rgb.shape[:2] # original image size

# resize image to network input size
img_rs = cv.resize(img_rgb, (W_in, H_in))
img_lab_rs = cv.cvtColor(img_rs, cv.COLOR_RGB2Lab)
img_l_rs = img_lab_rs[:,:,0]
img_l_rs -= 50 # subtract 50 for mean-centering

# run network
net.setInput(cv.dnn.blobFromImage(img_l_rs))
ab_dec = net.forward()[0,:,:,:].transpose((1,2,0))

(H_out,W_out) = ab_dec.shape[:2]
ab_dec_us = cv.resize(ab_dec, (W_orig, H_orig))
img_lab_out = np.concatenate((img_l[:,:,np.newaxis],ab_dec_us),axis=2)
img_bgr_out = np.clip(cv.cvtColor(img_lab_out, cv.COLOR_Lab2BGR), 0, 1)

frame = cv.resize(frame, (w, h))
cv.imshow('origin', frame)
cv.imshow('gray', cv.cvtColor(frame, cv.COLOR_RGB2GRAY))
# fix 4.0 imshow issue
cv.normalize(img_bgr_out, img_bgr_out, 0, 255, cv.NORM_MINMAX)
cv.imshow('colorized', cv.resize(np.uint8(img_bgr_out), (w, h)))
cv.waitKey(0)
cv.destroyAllWindows()

实验结果

解释

OpenCV DNN在4.0还支持灰度图像的彩色化模型，是根据2016年ECCV的论文而来，基于卷积神经网络模型，通过对Lab色彩空间进行量化分割，映射到最终的CNN输出结果，最后转换为RGB彩色图像。
相关论文详见：
https://arxiv.org/pdf/1603.08511.pdf

模型下载地址

OpenCV DNN使用该模型时候，除了正常的Caffe模型与配置文件之外，还需要一个Lab的量化表。

---

所有内容均来源于贾志刚老师的知识星球——OpenCV研习社，本文为个人整理学习，已获得贾老师授权，有兴趣、有能力的可以加入贾老师的知识星球进行深入学习。