diff --git a/samples/dnn/colorization.cpp b/samples/dnn/colorization.cpp new file mode 100644 index 0000000000..1c4ae07aac --- /dev/null +++ b/samples/dnn/colorization.cpp @@ -0,0 +1,125 @@ +// +// This program is based on https://github.com/richzhang/colorization/blob/master/colorization/colorize.py +// download the caffemodel from: http://eecs.berkeley.edu/~rich.zhang/projects/2016_colorization/files/demo_v2/colorization_release_v2.caffemodel +// and the prototxt from: https://github.com/richzhang/colorization/blob/master/colorization/models/colorization_deploy_v2.prototxt +// +#include +#include +#include +using namespace cv; +using namespace cv::dnn; + +#include +using namespace std; + + +// the 313 ab cluster centers from pts_in_hull.npy (already transposed) +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) +{ + CommandLineParser parser(argc, argv, + "{ help | false | print this help message }" + "{ proto | colorization_deploy_v2.prototxt | model configuration }" + "{ model | colorization_release_v2.caffemodel | model weights }" + "{ image | space_shuttle.jpg | path to image file }" + "{ opencl | false | enable OpenCL }" + ); + + String modelTxt = parser.get("proto"); + String modelBin = parser.get("model"); + String imageFile = parser.get("image"); + if (parser.get("help") || modelTxt.empty() || modelBin.empty() || imageFile.empty()) + { + cout << "A sample app to demonstrate recoloring grayscale images with dnn." << endl; + parser.printMessage(); + return 0; + } + + // fixed input size for the pretrained network + int W_in = 224; + int H_in = 224; + + Net net = dnn::readNetFromCaffe(modelTxt, modelBin); + + // setup additional layers: + int sz[] = {2, 313, 1, 1}; + Mat pts_in_hull(4, sz, CV_32F, hull_pts); + Ptr class8_ab = net.getLayer("class8_ab"); + class8_ab->blobs.push_back(pts_in_hull); + + Ptr conv8_313_rh = net.getLayer("conv8_313_rh"); + conv8_313_rh->blobs.push_back(Mat(1, 313, CV_32F, 2.606f)); + + if (parser.get("opencl")) + { + net.setPreferableTarget(DNN_TARGET_OPENCL); + } + + Mat img = imread(imageFile); + if (img.empty()) + { + std::cerr << "Can't read image from the file: " << imageFile << std::endl; + exit(-1); + } + + // 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("class8_ab"); + + // 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); + + namedWindow("color", WINDOW_NORMAL); + namedWindow("original", WINDOW_NORMAL); + imshow("color", color); + imshow("original", img); + waitKey(); + return 0; +} diff --git a/samples/dnn/colorization.py b/samples/dnn/colorization.py index 5cdbfd6224..d55ac45093 100644 --- a/samples/dnn/colorization.py +++ b/samples/dnn/colorization.py @@ -1,6 +1,6 @@ -# Script is based on https://github.com/richzhang/colorization/blob/master/colorize.py -# To download the caffemodel and the prototxt, see: https://github.com/richzhang/colorization/tree/master/models -# To download pts_in_hull.npy, see: https://github.com/richzhang/colorization/blob/master/resources/pts_in_hull.npy +# Script is based on https://github.com/richzhang/colorization/blob/master/colorization/colorize.py +# To download the caffemodel and the prototxt, see: https://github.com/richzhang/colorization/tree/master/colorization/models +# To download pts_in_hull.npy, see: https://github.com/richzhang/colorization/blob/master/colorization/resources/pts_in_hull.npy import numpy as np import argparse import cv2 as cv @@ -8,9 +8,9 @@ import cv2 as cv def parse_args(): parser = argparse.ArgumentParser(description='iColor: deep interactive colorization') parser.add_argument('--input', help='Path to image or video. Skip to capture frames from camera') - parser.add_argument('--prototxt', help='Path to colorization_deploy_v2.prototxt', default='./models/colorization_release_v2.prototxt') - parser.add_argument('--caffemodel', help='Path to colorization_release_v2.caffemodel', default='./models/colorization_release_v2.caffemodel') - parser.add_argument('--kernel', help='Path to pts_in_hull.npy', default='./resources/pts_in_hull.npy') + parser.add_argument('--prototxt', help='Path to colorization_deploy_v2.prototxt', required=True) + parser.add_argument('--caffemodel', help='Path to colorization_release_v2.caffemodel', required=True) + parser.add_argument('--kernel', help='Path to pts_in_hull.npy', required=True) args = parser.parse_args() return args