红外光和可见光融合，PIAfusion模型转ONNX，然后用Opencv推理

在github上面看到一个可见光和红外融合的神经网络模型，可见光图像就是平常拍的普通照片，比如下面这张

红外图像类似下面这种

GitHub模型链接，为了用opencv来进行推理，首先将此库里的模型文件格式转为onnx，代码如下

import torch
from models.fusion_model import PIAFusion

path = 'pretrained/fusion_model_epoch_29.pth'

model = PIAFusion()
model = model.cuda()
model.load_state_dict(torch.load(path))
model.eval()
dummy_vis = torch.randn((1,1,480,640),dtype=torch.float32).cuda()
dummy_inf = torch.randn((1,1,480,640),dtype=torch.float32).cuda()

torch.onnx.export(model,(dummy_vis,dummy_inf),"fusion_model.onnx",input_names=['vis','inf'],output_names=['output']
                  ,do_constant_folding=True
                  )

vis和inf指模型的输入名称（可见光和红外两张图像），然后模型的输出名称为output，要记住这些名称，后面推理的时候要用

有了ONNX模型，接着就可以用Opencv来推理了，推理的过程不复杂，主要是对数据的预处理和后处理部分有点麻烦，代码如下

#include
#include
#include
#include
#include
#include
#include


/*
 * input is float in range [0,1]
 * */

void RGB2YCrCb(cv::Mat &bgr,cv::Mat &Y,cv::Mat &Cr,cv::Mat &Cb)
{
    std::vector<cv::Mat> v;
    cv::split(bgr,v);
    cv::Mat b,g,r;
    v[0].copyTo(b);
    v[1].copyTo(g);
    v[2].copyTo(r);
    Y.create(b.size(),CV_32F);
    Cr.create(b.size(),CV_32F);
    Cb.create(b.size(),CV_32F);
    Y = 0.299 * r + 0.587 * g + 0.114 * b;
    Cr = (r - Y) * 0.713 + 0.5;
    Cb = (b - Y) * 0.564 + 0.5;
}

void YCrCb2RGB(const cv::Mat &Y,const cv::Mat &Cr,const cv::Mat &Cb,cv::Mat &bgr)
{
    cv::Mat r,g,b;
    r.create(Y.size(),CV_32F);
    g.create(Y.size(),CV_32F);
    b.create(Y.size(),CV_32F);

    r = Y + (Cr - 0.5) * 1.402;
    g = Y - (Cr - 0.5) * 0.714 - (Cb - 0.5) * 0.344;
    b = Y + (Cb - 0.5) * 1.733;
    cv::Mat arr[3]{b,g,r};
    cv::merge(arr,3,bgr);
}
class ONNXFusion
{
public:
    ONNXFusion(const std::string& model_path, cv::Size _input_size);
    void Fusion(const cv::Mat& vis,const cv::Mat &inf, cv::Mat &out);
private:
    void preprocess_input(cv::Mat& ,cv::Mat &,cv::Mat &,cv::Mat &,cv::Mat &);
private:
    cv::Size input_size;
    cv::dnn::Net net;
};

ONNXFusion::ONNXFusion(const std::string& model_path,cv::Size _input_size):input_size(_input_size)
{
    net = cv::dnn::readNet(model_path);
    net.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
    net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
}

void ONNXFusion::preprocess_input(cv::Mat& vis,cv::Mat &inf,cv::Mat &y,cv::Mat &cr,cv::Mat &cb)
{
    cv::resize(vis,vis,input_size);
    cv::resize(inf,inf,input_size);
    vis.convertTo(vis,CV_32F,1.0/255);
    inf.convertTo(inf,CV_32F,1.0/255);
    RGB2YCrCb(vis,y,cr,cb);
}

void ONNXFusion::Fusion(const cv::Mat& vis,const cv::Mat &inf, cv::Mat &out)
{
    assert(!vis.empty());
    assert(!inf.empty());
    cv::Mat vis_input = vis.clone();
    cv::Mat inf_input = inf.clone();
    cv::Mat y,cb,cr;
    preprocess_input(vis_input,inf_input,y,cr,cb);
    std::cout<<"y:"<<y.cols<<std::endl;
    std::cout<<"inf_input:"<<inf_input.cols<<std::endl;
    cv::Mat input_y_blob = cv::dnn::blobFromImage(y, 1.0, input_size,
                                                  cv::Scalar(0, 0, 0), false);
    cv::Mat input_inf_blob = cv::dnn::blobFromImage(inf_input, 1.0, input_size,
                                                    cv::Scalar(0, 0, 0), false);
    net.setInput(input_y_blob,"vis");
    net.setInput(input_inf_blob,"inf");
    const std::vector<cv::String>& out_names = net.getUnconnectedOutLayersNames();
    cv::Mat out_tensor = net.forward(out_names[0]);
    cv::Mat temp(input_size,CV_32F);
    memcpy(temp.data,out_tensor.data,sizeof(float) * input_size.area());
    YCrCb2RGB(temp,cr,cb,out);
    cv::convertScaleAbs(out,out,255);
}

int main(int argc, char* argv[])
{

    std::string model_path("fusion_model.onnx");//model path
    cv::Size input_size(640, 480);
    cv::Mat inf = cv::imread("inf.png",0); // inf path

    cv::Mat vis = cv::imread("vis.png"); // vis path
    ONNXFusion fusion(model_path, input_size);
    cv::Mat out;
    fusion.Fusion(vis,inf,out);
    cv::imshow("out",out);
    cv::waitKey(0);
    return 0;
}

结果如下图

效果还不错！