C++和神經網絡實現黑白圖片彩色化
本篇教程的環境是Ubuntu18.04和OpenCV,如果你還沒有安裝可以先配置環境,主要是安裝OpenCV,其中Windows的配置可以參見OpenCV的配置過程,最終項目會基於CMake構建,只需要配置相應依賴即可。
首先獲取caffe模型,我們需要在一個caffe的模型的基礎上,用C++導入模型來進行圖片處理,請注意這裡我們已經不需要caffe了,直接使用opencv中的DNN模塊。這是opencv非常方便的地方。
不過在所有事情之前,你可能需要先下載一下模型文件和權重。我們提供了一個腳本來下載:
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
廢話不多說,核心代碼如下:
#include
#include
#include
#include
using namespace cv;
using namespace cv::dnn;
using namespace std;
// the 313 ab cluster centers from pts_in_hull.npy (already transposed)
static float hull_pts[] = {
-90., -90., -90., -90., -90., -80., -80., -80., -80., -80., -80., -80., -80., -70., -70., -70., -70., -70., -70., -70., -70.,
-70., -70., -60., -60., -60., -60., -60., -60., -60., -60., -60., -60., -60., -60., -50., -50., -50., -50., -50., -50., -50., -50.,
-50., -50., -50., -50., -50., -50., -40., -40., -40., -40., -40., -40., -40., -40., -40., -40., -40., -40., -40., -40., -40., -30.,
-30., -30., -30., -30., -30., -30., -30., -30., -30., -30., -30., -30., -30., -30., -30., -20., -20., -20., -20., -20., -20., -20.,
-20., -20., -20., -20., -20., -20., -20., -20., -20., -10., -10., -10., -10., -10., -10., -10., -10., -10., -10., -10., -10., -10.,
-10., -10., -10., -10., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 10., 10., 10., 10., 10., 10., 10.,
10., 10., 10., 10., 10., 10., 10., 10., 10., 10., 10., 20., 20., 20., 20., 20., 20., 20., 20., 20., 20., 20., 20., 20., 20., 20.,
20., 20., 20., 30., 30., 30., 30., 30., 30., 30., 30., 30., 30., 30., 30., 30., 30., 30., 30., 30., 30., 30., 40., 40., 40., 40.,
40., 40., 40., 40., 40., 40., 40., 40., 40., 40., 40., 40., 40., 40., 40., 40., 50., 50., 50., 50., 50., 50., 50., 50., 50., 50.,
50., 50., 50., 50., 50., 50., 50., 50., 50., 60., 60., 60., 60., 60., 60., 60., 60., 60., 60., 60., 60., 60., 60., 60., 60., 60.,
60., 60., 60., 70., 70., 70., 70., 70., 70., 70., 70., 70., 70., 70., 70., 70., 70., 70., 70., 70., 70., 70., 70., 80., 80., 80.,
80., 80., 80., 80., 80., 80., 80., 80., 80., 80., 80., 80., 80., 80., 80., 80., 90., 90., 90., 90., 90., 90., 90., 90., 90., 90.,
90., 90., 90., 90., 90., 90., 90., 90., 90., 100., 100., 100., 100., 100., 100., 100., 100., 100., 100., 50., 60., 70., 80., 90.,
20., 30., 40., 50., 60., 70., 80., 90., 0., 10., 20., 30., 40., 50., 60., 70., 80., 90., -20., -10., 0., 10., 20., 30., 40., 50.,
60., 70., 80., 90., -30., -20., -10., 0., 10., 20., 30., 40., 50., 60., 70., 80., 90., 100., -40., -30., -20., -10., 0., 10., 20.,
30., 40., 50., 60., 70., 80., 90., 100., -50., -40., -30., -20., -10., 0., 10., 20., 30., 40., 50., 60., 70., 80., 90., 100., -50.,
-40., -30., -20., -10., 0., 10., 20., 30., 40., 50., 60., 70., 80., 90., 100., -60., -50., -40., -30., -20., -10., 0., 10., 20.,
30., 40., 50., 60., 70., 80., 90., 100., -70., -60., -50., -40., -30., -20., -10., 0., 10., 20., 30., 40., 50., 60., 70., 80., 90.,
100., -80., -70., -60., -50., -40., -30., -20., -10., 0., 10., 20., 30., 40., 50., 60., 70., 80., 90., -80., -70., -60., -50.,
-40., -30., -20., -10., 0., 10., 20., 30., 40., 50., 60., 70., 80., 90., -90., -80., -70., -60., -50., -40., -30., -20., -10.,
0., 10., 20., 30., 40., 50., 60., 70., 80., 90., -100., -90., -80., -70., -60., -50., -40., -30., -20., -10., 0., 10., 20., 30.,
40., 50., 60., 70., 80., 90., -100., -90., -80., -70., -60., -50., -40., -30., -20., -10., 0., 10., 20., 30., 40., 50., 60., 70.,
80., -110., -100., -90., -80., -70., -60., -50., -40., -30., -20., -10., 0., 10., 20., 30., 40., 50., 60., 70., 80., -110., -100.,
-90., -80., -70., -60., -50., -40., -30., -20., -10., 0., 10., 20., 30., 40., 50., 60., 70., 80., -110., -100., -90., -80., -70.,
-60., -50., -40., -30., -20., -10., 0., 10., 20., 30., 40., 50., 60., 70., -110., -100., -90., -80., -70., -60., -50., -40., -30.,
-20., -10., 0., 10., 20., 30., 40., 50., 60., 70., -90., -80., -70., -60., -50., -40., -30., -20., -10., 0.
};
int main(int argc, char **argv)
{
string imageFileName;
// Take arguments from commmand line
if (argc < 2)
{
cout << "Please input the greyscale image filename." << endl;
cout << "Usage example: ./colorizeImage.out greyscaleImage.png" << endl;
return 1;
}
imageFileName = argv[1];
Mat img = imread(imageFileName);
if (img.empty())
{
cout << "Can't read image from file: " << imageFileName << endl;
return 2;
}
string protoFile = "./models/colorization_deploy_v2.prototxt";
string weightsFile = "./models/colorization_release_v2.caffemodel";
//string weightsFile = "./models/colorization_release_v2_norebal.caffemodel";
double t = (double) cv::getTickCount();
// fixed input size for the pretrained network
const int W_in = 224;
const int H_in = 224;
Net net = dnn::readNetFromCaffe(protoFile, weightsFile);
// setup additional layers:
int sz[] = {2, 313, 1, 1};
const Mat pts_in_hull(4, sz, CV_32F, hull_pts);
Ptr<:layer> class8_ab = net.getLayer("class8_ab");
class8_ab->blobs.push_back(pts_in_hull);
Ptr<:layer> conv8_313_rh = net.getLayer("conv8_313_rh");
conv8_313_rh->blobs.push_back(Mat(1, 313, CV_32F, Scalar(2.606)));
// extract L channel and subtract mean
Mat lab, L, input;
img.convertTo(img, CV_32F, 1.0/255);
cvtColor(img, lab, COLOR_BGR2Lab);
extractChannel(lab, L, 0);
resize(L, input, Size(W_in, H_in));
input -= 50;
// run the L channel through the network
Mat inputBlob = blobFromImage(input);
net.setInput(inputBlob);
Mat result = net.forward();
// retrieve the calculated a,b channels from the network output
Size siz(result.size[2], result.size[3]);
Mat a = Mat(siz, CV_32F, result.ptr(0,0));
Mat b = Mat(siz, CV_32F, result.ptr(0,1));
resize(a, a, img.size());
resize(b, b, img.size());
// merge, and convert back to BGR
Mat color, chn[] = {L, a, b};
merge(chn, 3, lab);
cvtColor(lab, color, COLOR_Lab2BGR);
t = ((double)cv::getTickCount() - t)/cv::getTickFrequency();
cout << "Time taken : " << t << " secs" << endl;
string str = imageFileName;
str.replace(str.end()-4, str.end(), "");
str = str+"_colorized.png";
imwrite(str, color*255);
imshow("result_color", color * 255);
cv::waitKey(0);
cout << "Colorized image saved as " << str << endl;
return 0;
}
編譯
編譯c++,我個人比較推薦使用g++,用g++也很簡單了:
g++ `pkg-config --cflags opencv` colorizeImage.cpp -o color `pkg-config --libs opencv`
但是你運行之後會發現各種報錯,沒有辦法g++就是麻煩,改用cmake吧。 opencv的配置過程如下:
cmake_minimum_required(VERSION 3.8)
project(colorizer)
find_package(OpenCV REQUIRED)
if (NOT OpenCV_FOUND)
message(FATAL_ERROR "opencv not found.")
endif ()
file(GLOB_RECURSE source_files
"colorizeImage.cpp" "*.cc" "include/*/*.hpp" "include/*.hpp")
add_executable(colorizer ${source_files})
target_link_libraries(colorizer ${OpenCV_LIBS})
cmake的這個colorizeImage.cpp就是我們的核心代碼,就一個文件了。其他的忽視掉,然後:
mkdir build && cd build
cmake ..
make -j8
如果不出意外應該可以編譯成功了。
預測結果
吧二進制文件拷貝到當前目錄,我們可以得到預測結果。
這個結果非常不錯,比如上面這張圖片,本來應該是藍天的,但是它卻有著一種比藍天更加灰宏的感覺,誰說人工智能沒有藝術感呢?我想這就是為什麼AI如此令人陶醉的原因吧。
核心代碼講解
這部分大家可以關注我近期會推出的Youtube視頻,我的視頻鏈接為:
www.youtube.com/watch%3Fv%3D_bPfq1QrZ8k
Youtube是視頻首發場所,後面也會同步到優庫和騰訊視頻,期待大家的關注。 大家也可以添加 jintianiloveu 與我互動。
原文鏈接 https://zhuanlan.zhihu.com/p/41017675
閱讀更多 前端小學生 的文章