1// This file is part of OpenCV project.
2// It is subject to the license terms in the LICENSE file found in the top-level directory
3// of this distribution and at http://opencv.org/license.html
4// To download the onnx model, see: https://storage.googleapis.com/ailia-models/colorization/colorizer.onnx
5
6#include <opencv2/dnn.hpp>
7#include <opencv2/imgproc.hpp>
8#include <opencv2/imgcodecs.hpp>
9#include "common.hpp"
10#include <opencv2/highgui.hpp>
11#include <iostream>
12
13using namespace cv;
14using namespace std;
15using namespace cv::dnn;
16
17
18int main(int argc, char** argv) {
19 const string about =
20 "This sample demonstrates recoloring grayscale images with dnn.\n"
21 "This program is based on:\n"
22 " http://richzhang.github.io/colorization\n"
23 " https://github.com/richzhang/colorization\n"
24 "To download the onnx model:\n"
25 " https://storage.googleapis.com/ailia-models/colorization/colorizer.onnx\n";
26
27 const string param_keys =
28 "{ help h | | Print help message. }"
29 "{ input i | baboon.jpg | Path to the input image }"
30 "{ onnx_model_path | | Path to the ONNX model. Required. }";
31
32 const string backend_keys = format(
33 "{ backend | 0 | Choose one of computation backends: "
34 "%d: automatically (by default), "
35 "%d: Intel's Deep Learning Inference Engine (https://software.intel.com/openvino-toolkit), "
36 "%d: OpenCV implementation, "
37 "%d: VKCOM, "
38 "%d: CUDA, "
39 "%d: WebNN }",
40 cv::dnn::DNN_BACKEND_DEFAULT, cv::dnn::DNN_BACKEND_INFERENCE_ENGINE, cv::dnn::DNN_BACKEND_OPENCV,
41 cv::dnn::DNN_BACKEND_VKCOM, cv::dnn::DNN_BACKEND_CUDA, cv::dnn::DNN_BACKEND_WEBNN);
42 const string target_keys = format(
43 "{ target | 0 | Choose one of target computation devices: "
44 "%d: CPU target (by default), "
45 "%d: OpenCL, "
46 "%d: OpenCL fp16 (half-float precision), "
47 "%d: VPU, "
48 "%d: Vulkan, "
49 "%d: CUDA, "
50 "%d: CUDA fp16 (half-float preprocess) }",
51 cv::dnn::DNN_TARGET_CPU, cv::dnn::DNN_TARGET_OPENCL, cv::dnn::DNN_TARGET_OPENCL_FP16,
52 cv::dnn::DNN_TARGET_MYRIAD, cv::dnn::DNN_TARGET_VULKAN, cv::dnn::DNN_TARGET_CUDA,
53 cv::dnn::DNN_TARGET_CUDA_FP16);
54
55 const string keys = param_keys + backend_keys + target_keys;
56 CommandLineParser parser(argc, argv, keys);
57 parser.about(about);
58
59 if (parser.has("help")) {
60 parser.printMessage();
61 return 0;
62 }
63
64 string inputImagePath = parser.get<string>("input");
65 string onnxModelPath = parser.get<string>("onnx_model_path");
66 int backendId = parser.get<int>("backend");
67 int targetId = parser.get<int>("target");
68
69 if (onnxModelPath.empty()) {
70 cerr << "The path to the ONNX model is required!" << endl;
71 return -1;
72 }
73
74 Mat imgGray = imread(samples::findFile(inputImagePath), IMREAD_GRAYSCALE);
75 if (imgGray.empty()) {
76 cerr << "Could not read the image: " << inputImagePath << endl;
77 return -1;
78 }
79
80 Mat imgL = imgGray;
81 imgL.convertTo(imgL, CV_32F, 100.0/255.0);
82 Mat imgLResized;
83 resize(imgL, imgLResized, Size(256, 256), 0, 0, INTER_CUBIC);
84
85 // Prepare the model
86 EngineType engine = ENGINE_AUTO;
87 if (backendId != 0 || targetId != 0){
88 engine = ENGINE_CLASSIC;
89 }
90 dnn::Net net = dnn::readNetFromONNX(onnxModelPath, engine);
91 net.setPreferableBackend(backendId);
92 net.setPreferableTarget(targetId);
93 //! [Read and initialize network]
94
95 // Create blob from the image
96 Mat blob = dnn::blobFromImage(imgLResized, 1.0, Size(256, 256), Scalar(), false, false);
97
98 net.setInput(blob);
99
100 // Run inference
101 Mat result = net.forward();
102 Size siz(result.size[2], result.size[3]);
103 Mat a(siz, CV_32F, result.ptr(0,0));
104 Mat b(siz, CV_32F, result.ptr(0,1));
105 resize(a, a, imgGray.size());
106 resize(b, b, imgGray.size());
107
108 // merge, and convert back to BGR
109 Mat color, chn[] = {imgL, a, b};
110
111 // Proc
112 Mat lab;
113 merge(chn, 3, lab);
114 cvtColor(lab, color, COLOR_Lab2BGR);
115
116 imshow("input image", imgGray);
117 imshow("output image", color);
118 waitKey();
119
120 return 0;
121}
