Skip to content

Commit 56bcdfd

Browse files
berakalalek
berak
authored and
alalek
committed
Merge pull request opencv#10777 from berak:dnn_colorize_cpp
* dnn: add colorization.cpp * Update arguments of dnn/colorization.py
1 parent 14089b1 commit 56bcdfd

File tree

2 files changed

+131
-6
lines changed

2 files changed

+131
-6
lines changed

‎samples/dnn/colorization.cpp‎

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
//
2+
// This program is based on https://github.com/richzhang/colorization/blob/master/colorization/colorize.py
3+
// download the caffemodel from: http://eecs.berkeley.edu/~rich.zhang/projects/2016_colorization/files/demo_v2/colorization_release_v2.caffemodel
4+
// and the prototxt from: https://github.com/richzhang/colorization/blob/master/colorization/models/colorization_deploy_v2.prototxt
5+
//
6+
#include <opencv2/dnn.hpp>
7+
#include <opencv2/imgproc.hpp>
8+
#include <opencv2/highgui.hpp>
9+
using namespace cv;
10+
using namespace cv::dnn;
11+
12+
#include <iostream>
13+
using namespace std;
14+
15+
16+
// the 313 ab cluster centers from pts_in_hull.npy (already transposed)
17+
float hull_pts[] = {
18+
-90., -90., -90., -90., -90., -80., -80., -80., -80., -80., -80., -80., -80., -70., -70., -70., -70., -70., -70., -70., -70.,
19+
-70., -70., -60., -60., -60., -60., -60., -60., -60., -60., -60., -60., -60., -60., -50., -50., -50., -50., -50., -50., -50., -50.,
20+
-50., -50., -50., -50., -50., -50., -40., -40., -40., -40., -40., -40., -40., -40., -40., -40., -40., -40., -40., -40., -40., -30.,
21+
-30., -30., -30., -30., -30., -30., -30., -30., -30., -30., -30., -30., -30., -30., -30., -20., -20., -20., -20., -20., -20., -20.,
22+
-20., -20., -20., -20., -20., -20., -20., -20., -20., -10., -10., -10., -10., -10., -10., -10., -10., -10., -10., -10., -10., -10.,
23+
-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.,
24+
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.,
25+
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.,
26+
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.,
27+
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.,
28+
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.,
29+
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.,
30+
90., 90., 90., 90., 90., 90., 90., 90., 90., 100., 100., 100., 100., 100., 100., 100., 100., 100., 100., 50., 60., 70., 80., 90.,
31+
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.,
32+
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.,
33+
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.,
34+
-40., -30., -20., -10., 0., 10., 20., 30., 40., 50., 60., 70., 80., 90., 100., -60., -50., -40., -30., -20., -10., 0., 10., 20.,
35+
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.,
36+
100., -80., -70., -60., -50., -40., -30., -20., -10., 0., 10., 20., 30., 40., 50., 60., 70., 80., 90., -80., -70., -60., -50.,
37+
-40., -30., -20., -10., 0., 10., 20., 30., 40., 50., 60., 70., 80., 90., -90., -80., -70., -60., -50., -40., -30., -20., -10.,
38+
0., 10., 20., 30., 40., 50., 60., 70., 80., 90., -100., -90., -80., -70., -60., -50., -40., -30., -20., -10., 0., 10., 20., 30.,
39+
40., 50., 60., 70., 80., 90., -100., -90., -80., -70., -60., -50., -40., -30., -20., -10., 0., 10., 20., 30., 40., 50., 60., 70.,
40+
80., -110., -100., -90., -80., -70., -60., -50., -40., -30., -20., -10., 0., 10., 20., 30., 40., 50., 60., 70., 80., -110., -100.,
41+
-90., -80., -70., -60., -50., -40., -30., -20., -10., 0., 10., 20., 30., 40., 50., 60., 70., 80., -110., -100., -90., -80., -70.,
42+
-60., -50., -40., -30., -20., -10., 0., 10., 20., 30., 40., 50., 60., 70., -110., -100., -90., -80., -70., -60., -50., -40., -30.,
43+
-20., -10., 0., 10., 20., 30., 40., 50., 60., 70., -90., -80., -70., -60., -50., -40., -30., -20., -10., 0.
44+
};
45+
46+
47+
int main(int argc, char **argv)
48+
{
49+
CommandLineParser parser(argc, argv,
50+
"{ help | false | print this help message }"
51+
"{ proto | colorization_deploy_v2.prototxt | model configuration }"
52+
"{ model | colorization_release_v2.caffemodel | model weights }"
53+
"{ image | space_shuttle.jpg | path to image file }"
54+
"{ opencl | false | enable OpenCL }"
55+
);
56+
57+
String modelTxt = parser.get<string>("proto");
58+
String modelBin = parser.get<string>("model");
59+
String imageFile = parser.get<String>("image");
60+
if (parser.get<bool>("help") || modelTxt.empty() || modelBin.empty() || imageFile.empty())
61+
{
62+
cout << "A sample app to demonstrate recoloring grayscale images with dnn." << endl;
63+
parser.printMessage();
64+
return 0;
65+
}
66+
67+
// fixed input size for the pretrained network
68+
int W_in = 224;
69+
int H_in = 224;
70+
71+
Net net = dnn::readNetFromCaffe(modelTxt, modelBin);
72+
73+
// setup additional layers:
74+
int sz[] = {2, 313, 1, 1};
75+
Mat pts_in_hull(4, sz, CV_32F, hull_pts);
76+
Ptr<dnn::Layer> class8_ab = net.getLayer("class8_ab");
77+
class8_ab->blobs.push_back(pts_in_hull);
78+
79+
Ptr<dnn::Layer> conv8_313_rh = net.getLayer("conv8_313_rh");
80+
conv8_313_rh->blobs.push_back(Mat(1, 313, CV_32F, 2.606f));
81+
82+
if (parser.get<bool>("opencl"))
83+
{
84+
net.setPreferableTarget(DNN_TARGET_OPENCL);
85+
}
86+
87+
Mat img = imread(imageFile);
88+
if (img.empty())
89+
{
90+
std::cerr << "Can't read image from the file: " << imageFile << std::endl;
91+
exit(-1);
92+
}
93+
94+
// extract L channel and subtract mean
95+
Mat lab, L, input;
96+
img.convertTo(img, CV_32F, 1.0/255);
97+
cvtColor(img, lab, COLOR_BGR2Lab);
98+
extractChannel(lab, L, 0);
99+
resize(L, input, Size(W_in, H_in));
100+
input -= 50;
101+
102+
// run the L channel through the network
103+
Mat inputBlob = blobFromImage(input);
104+
net.setInput(inputBlob);
105+
Mat result = net.forward("class8_ab");
106+
107+
// retrieve the calculated a,b channels from the network output
108+
Size siz(result.size[2], result.size[3]);
109+
Mat a = Mat(siz, CV_32F, result.ptr(0,0));
110+
Mat b = Mat(siz, CV_32F, result.ptr(0,1));
111+
resize(a, a, img.size());
112+
resize(b, b, img.size());
113+
114+
// merge, and convert back to bgr
115+
Mat color, chn[] = {L, a, b};
116+
merge(chn, 3, lab);
117+
cvtColor(lab, color, COLOR_Lab2BGR);
118+
119+
namedWindow("color", WINDOW_NORMAL);
120+
namedWindow("original", WINDOW_NORMAL);
121+
imshow("color", color);
122+
imshow("original", img);
123+
waitKey();
124+
return 0;
125+
}

‎samples/dnn/colorization.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,16 @@
1-
# Script is based on https://github.com/richzhang/colorization/blob/master/colorize.py
2-
# To download the caffemodel and the prototxt, see: https://github.com/richzhang/colorization/tree/master/models
3-
# To download pts_in_hull.npy, see: https://github.com/richzhang/colorization/blob/master/resources/pts_in_hull.npy
1+
# Script is based on https://github.com/richzhang/colorization/blob/master/colorization/colorize.py
2+
# To download the caffemodel and the prototxt, see: https://github.com/richzhang/colorization/tree/master/colorization/models
3+
# To download pts_in_hull.npy, see: https://github.com/richzhang/colorization/blob/master/colorization/resources/pts_in_hull.npy
44
import numpy as np
55
import argparse
66
import cv2 as cv
77

88
def parse_args():
99
parser = argparse.ArgumentParser(description='iColor: deep interactive colorization')
1010
parser.add_argument('--input', help='Path to image or video. Skip to capture frames from camera')
11-
parser.add_argument('--prototxt', help='Path to colorization_deploy_v2.prototxt', default='./models/colorization_release_v2.prototxt')
12-
parser.add_argument('--caffemodel', help='Path to colorization_release_v2.caffemodel', default='./models/colorization_release_v2.caffemodel')
13-
parser.add_argument('--kernel', help='Path to pts_in_hull.npy', default='./resources/pts_in_hull.npy')
11+
parser.add_argument('--prototxt', help='Path to colorization_deploy_v2.prototxt', required=True)
12+
parser.add_argument('--caffemodel', help='Path to colorization_release_v2.caffemodel', required=True)
13+
parser.add_argument('--kernel', help='Path to pts_in_hull.npy', required=True)
1414

1515
args = parser.parse_args()
1616
return args

0 commit comments

Comments
 (0)