前言
《Colorful Image Colorization》是加里福利亚大学Richard Zhang发表在ECCV 2016上的文章,论文主要解决的问题是给灰度图的自动着色,算法并不是为恢复灰度图的真实颜色,而是用灰度图中物体的纹理、语义等信息作为线索,来预测可能的上色,最后的上色结果只要真实即可。这不仅降低了上色的难度,而且也符合人们的认知:比如一个苹果,给它上青色,上红色都是正常的,不限于某一个颜色,只要不是紫色黑色等奇怪的颜色。
算法模型处理步骤是,输入图片的L通道,使用一个CNN预测对应的ab通道取值的概率分布,最后转化为RGB图像结果。
官方给的测试效果图:
算法代码:https://github.com/richzhang/colorization
C++ 模型推理
1.这里使用的开发环境是win10,显卡RTX3080,cuda10.2,cudnn7.1,OpenCV4.5,NCNN,IDE 是Vs2019。
2.模型是从onnx转成ncnn的模型,转换方法可以参考ncnn的官网。
3.代码步骤就是加载模型后,输入一张灰度图像,模型给图像上色,输出一张彩色图像。
4.C++源码
#pragma once
#include <ncnn/net.h>
#include <ncnn/layer.h>
#include <opencv2/opencv.hpp>
class Color
{
public:
Color();
~Color();
int read_models(std::string param_path,
std::string model_path, bool use_gpu);
int image_color(cv::Mat& cv_src, cv::Mat& cv_dst);
private:
ncnn::Net net;
int width = 256;
int height = 256;
};
#include "Color.h"
class Sig17Slice : public ncnn::Layer
{
public:
Sig17Slice()
{
one_blob_only = true;
}
virtual int forward(const ncnn::Mat& bottom_blob, ncnn::Mat& top_blob, const ncnn::Option& opt) const
{
int w = bottom_blob.w;
int h = bottom_blob.h;
int channels = bottom_blob.c;
int outw = w / 2;
int outh = h / 2;
//int outc = channels * 4;
int outc = channels;
top_blob.create(outw, outh, outc, 4u, 1, opt.blob_allocator);
if (top_blob.empty())
return -100;
#pragma omp parallel for num_threads(opt.num_threads)
for (int p = 0; p < outc; p++)
{
const float* ptr = bottom_blob.channel(p % channels).row((p / channels) % 2) + ((p / channels) / 2);
float* outptr = top_blob.channel(p);
for (int i = 0; i < outh; i++)
{
for (int j = 0; j < outw; j++)
{
*outptr = *ptr;
outptr += 1;
ptr += 2;
}
ptr += w;
}
}
return 0;
}
};
DEFINE_LAYER_CREATOR(Sig17Slice)
Color::Color()
{
}
Color::~Color()
{
}
int Color::read_models(std::string param_path,
std::string model_path, bool use_gpu)
{
net.register_custom_layer("Sig17Slice", Sig17Slice_layer_creator);
bool has_gpu = false;
#if NCNN_VULKAN
ncnn::create_gpu_instance();
has_gpu = ncnn::get_gpu_count() > 0;
#endif
bool to_use_gpu = has_gpu && use_gpu;
net.opt.use_vulkan_compute = to_use_gpu;
int rp = net.load_param(param_path.c_str());
int rb = net.load_model(model_path.c_str());
if (rp < 0 || rb < 0)
{
return 1;
}
return 0;
}
int Color::image_color(cv::Mat& cv_src, cv::Mat& cv_dst)
{
cv::Mat cv_base, cv_lab, cv_channel, cv_input;
cv_base = cv_src.clone();
cv_base.convertTo(cv_base, CV_32F, 1.0 / 255);
cvtColor(cv_base, cv_lab, cv::COLOR_BGR2Lab);
cv::extractChannel(cv_lab, cv_channel, 0);
resize(cv_channel, cv_input, cv::Size(width, height));
ncnn::Mat in(cv_input.cols, cv_input.rows, 1, (void*)cv_input.data);
in = in.clone();
ncnn::Extractor ex = net.create_extractor();
ex.input("input", in);
ncnn::Mat out;
ex.extract("out_ab", out);
cv::Mat colored_LAB(out.h, out.w, CV_32FC2);
memcpy((uchar*)colored_LAB.data, out.data,
out.w * out.h * 2 * sizeof(float));
//get separsted LAB channels a&b
cv::Mat a(out.h, out.w, CV_32F, (float*)out.data);
cv::Mat b(out.h, out.w, CV_32F, (float*)out.data + out.w * out.h);
//Resize a, b channels to origina image size
cv::resize(a, a, cv_base.size());
cv::resize(b, b, cv_base.size());
cv::Mat chn[] = { cv_channel, a, b };
cv::merge(chn, 3, cv_lab);
cvtColor(cv_lab, cv_dst, cv::COLOR_Lab2BGR);
cv_dst.convertTo(cv_dst, CV_8UC3, 255);
return 0;
}
5.运行结果
源码配置
1.源码地址:https://download.csdn.net/download/matt45m/87374018
2.配置运行环境:
3.配置包含库:
4.配置lib路径:
4.添加链接:
这个里添加的工程目录下lib下的所的.lib的名字。
GenericCodeGen.lib
glslang.lib
MachineIndependent.lib
ncnn.lib
OGLCompiler.lib
opencv_world450.lib
OSDependent.lib
SPIRV.lib
VkLayer_utils.lib
vulkan-1.lib