智能黑白图像自动上色——C++实现

前言

《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 
#include 
#include 


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