Merge pull request #10777 from berak:dnn_colorize_cpp

* dnn: add colorization.cpp

* Update arguments of dnn/colorization.py

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 <opencv2/dnn.hpp>
+#include <opencv2/imgproc.hpp>
+#include <opencv2/highgui.hpp>
+using namespace cv;
+using namespace cv::dnn;
+
+#include <iostream>
+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<string>("proto");
+    String modelBin = parser.get<string>("model");
+    String imageFile = parser.get<String>("image");
+    if (parser.get<bool>("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<dnn::Layer> class8_ab = net.getLayer("class8_ab");
+    class8_ab->blobs.push_back(pts_in_hull);
+
+    Ptr<dnn::Layer> conv8_313_rh = net.getLayer("conv8_313_rh");
+    conv8_313_rh->blobs.push_back(Mat(1, 313, CV_32F, 2.606f));
+
+    if (parser.get<bool>("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;
+}

samples/dnn/colorization.py

@@ -1,6 +1,6 @@
-# Script is based on https://github.com/richzhang/colorization/blob/master/colorize.py
+# 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/models
+# 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/resources/pts_in_hull.npy
+# 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('--prototxt', help='Path to colorization_deploy_v2.prototxt', required=True)
-    parser.add_argument('--caffemodel', help='Path to colorization_release_v2.caffemodel', default='./models/colorization_release_v2.caffemodel')
+    parser.add_argument('--caffemodel', help='Path to colorization_release_v2.caffemodel', required=True)
-    parser.add_argument('--kernel', help='Path to pts_in_hull.npy', default='./resources/pts_in_hull.npy')
+    parser.add_argument('--kernel', help='Path to pts_in_hull.npy', required=True)
 
     args = parser.parse_args()
     return args