OpenCV4.x可实现加载模型与执行推断

概述

OpenCV4.x发布以后，有很多新的特性与黑科技支持，无论是支持OpenVINO加速、图计算模块、二维码识别，还是DNN中新增加的人脸检测与识别模型，作为OpenCV开发者的我深深被吸引，几乎只要有时间就会一个一个的去发现与之前的不同之处。OpenCV DNN模块，不仅支持图像分类、对象检测、人脸检测、图像分割等操作除外，还支持对灰度图像的自动彩色化转换，而且效果十分靠谱，亲测有效！

着色模型(colorization model)

该模型是在2016发表在ECCV上面的，该模型与之前的基于CNN模型的不同之处在于，它是一个无监督的学习过程，不会把着色对象与训练生成看成是一个回归问题、而且是使用CIE Lab色彩空间，使用L分量作为输入，输出为颜色分量a，b，通过对颜色分量进行量化，把网络作为一个分类问题对待， 对得到输出结果，最终加上L分量之后，得到着色之后的图像，模型架构如下：

其中卷积层每一个block是有几个重复的conv卷积操作与ReLU + BN层构成！其中蓝色部分，是a，b颜色的313对ab量化表示。最终学习到的就是WxHx313输出，进一步转换为Color ab的输出, 加上L分量之后就是完整的图像输出！313对ab色彩空间量化表示如下：

针对自然场景下，ab值较低导致生成图像的失真问题，作者通过分类再平衡技术依靠训练阶段，通过对损失函数调整像素权重，实现了比较好的效果。作者的github上可以查看该模型的实现源码。

模型下载地址如下

https://github.com/e-lab/ENet-training

论文地址

https://arxiv.org/abs/1606.02147

OpenCV中使用

下载ENet预训练模型，通过OpenCV DNN支持，可以实现加载模型与执行推断，对大多数的灰度图像实现自然着色，毫无违和感！步骤如下：

加载模型

 

modelTxt = "D:/projects/models/color/colorization_deploy_v2.prototxt";
modelBin = "D:/projects/models/color/colorization_release_v2.caffemodel";
pts_txt = "D:/projects/models/color/pts_in_hull.npy";

# 加载网络
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("D:/images/yuan_test.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_BGR2Lab)
img_l = img_lab[:,:,0] # pull out L channel
(H_orig,W_orig) = img_rgb.shape[:2] # original image size

# resize为输入网络图像大小
img_rs = cv.resize(img_rgb, (W_in, H_in))
img_lab_rs = cv.cvtColor(img_rs, cv.COLOR_BGR2Lab)

# 选择一L通道分量
img_l_rs = img_lab_rs[:,:,0]
img_l_rs -= 50 # subtract 50 for mean-centering

# 输入L分量，开始操作
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)))

 

运行结果：

输入彩色图像，转为灰度图像，然后自动着色对比一下！

直接输入灰度图像，着色：

看效果，从此以后再也不担心灰度图像无法自动上色啦！

OpenCV成功解锁！　　

      审核编辑：彭静