File size: 2,959 Bytes
979472c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | #include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/imgcodecs.hpp>
#include <iostream>
#include <string>
#include <vector>
static std::string argVal(int argc, char** argv, const std::string& key, const std::string& def)
{
for (int i = 1; i + 1 < argc; ++i)
if (key == argv[i]) return argv[i + 1];
return def;
}
// Run one pyramid head and rebuild a colour image from its luma output.
// Mirrors dnn_superres' preprocess_YCrCb / reconstruct_YCrCb: the network only
// ever sees the Y channel, and Cr/Cb are bicubically upscaled and merged back.
static cv::Mat upsample(cv::dnn::Net& net, const cv::Mat& img, const std::string& nodeName, int scale)
{
cv::Mat ycrcb;
cv::cvtColor(img, ycrcb, cv::COLOR_BGR2YCrCb);
ycrcb.convertTo(ycrcb, CV_32F, 1.0 / 255.0);
cv::Mat ch[3];
cv::split(ycrcb, ch);
cv::Mat Y = ch[0];
if (!Y.isContinuous()) Y = Y.clone();
// This model takes NHWC [1, H, W, 1] even though its outputs are NCHW.
int blobShape[] = {1, Y.rows, Y.cols, 1};
cv::Mat blob(4, blobShape, CV_32F, Y.data);
net.setInput(blob);
cv::Mat outBlob = net.forward(nodeName);
// Output blob is NCHW [1, 1, H*scale, W*scale] -> wrap as a single-channel Mat.
cv::Mat yHr(outBlob.size[2], outBlob.size[3], CV_32F, outBlob.ptr<float>());
cv::Mat crHr, cbHr;
cv::resize(ch[1], crHr, cv::Size(), scale, scale);
cv::resize(ch[2], cbHr, cv::Size(), scale, scale);
std::vector<cv::Mat> merged = {yHr, crHr, cbHr};
cv::Mat hr;
cv::merge(merged, hr);
hr.convertTo(hr, CV_8U, 255.0);
cv::cvtColor(hr, hr, cv::COLOR_YCrCb2BGR);
return hr;
}
int main(int argc, char** argv)
{
std::string model = argVal(argc, argv, "--model", "lapsrn_x4_2026sep.onnx");
std::string image = argVal(argc, argv, "--image", "example_outputs/input_image.png");
std::string outputDir = argVal(argc, argv, "--output-dir", "example_outputs");
cv::Mat img = cv::imread(image);
if (img.empty())
{
std::cerr << "could not read image: " << image << std::endl;
return 1;
}
cv::dnn::Net net = cv::dnn::readNetFromONNX(model);
// The two heads of the Laplacian pyramid, and the scale each one produces.
const std::vector<std::pair<std::string, int> > outputs =
{{"NCHW_output_2x", 2}, {"NCHW_output_4x", 4}};
std::cout << "lapsrn_x4 input " << img.cols << "x" << img.rows << std::endl;
for (size_t i = 0; i < outputs.size(); ++i)
{
const std::string& nodeName = outputs[i].first;
int scale = outputs[i].second;
cv::Mat hr = upsample(net, img, nodeName, scale);
std::string path = cv::format("%s/output_image_%dx.png", outputDir.c_str(), scale);
cv::imwrite(path, hr);
std::cout << " " << nodeName << " x" << scale
<< " -> " << hr.cols << "x" << hr.rows << " " << path << std::endl;
}
return 0;
}
|