【C++模型部署】OpenCV实现图像风格迁移模型推理

背景

利用vs+opencv实现ONNX模型推理，本文采用的模型为图像风格迁移之马赛克化
onnx模型地址 链接 没有积分的小伙伴私信我，免费发你

环境要求

我的环境配置：vs2017 + opencv4.5.1（opencv为自己编译，支持CUDA加速）
opencv地址 链接 没有积分的小伙伴私信我，免费发你
opencv也可以直接在官网下载，只要dnn模块支持ONNX模型推理就可以，建议4.0以上，官网opencv没有CUDA加速

模型推理

效果图

本模型要求输入图像尺寸224x224，所以图像在进入模型推理前要resize到224x224大小
原图

效果图

代码

话不多说，看代码，亲测可用！

#include <iostream>
#include<opencv2/opencv.hpp>

using namespace std;
using namespace cv;
//using namespace dnn;

dnn::Net net;

bool readModel(dnn::Net& net, string& netPath, bool isCuda = false) 
{
	try {
		net = dnn::readNetFromONNX(netPath);
	}
	catch (const std::exception&) {
		return false;
	}
	//cuda
	if (isCuda) {
		net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
		net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA);
	}
	//cpu
	else {
		net.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
		net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
	}
	return true;
}

int Detect(Mat SrcImg, dnn::Net& net, Mat& outImg, int netWidth, int netHeight, int channel)
{
	cv::Mat netInputImg;
	cv::resize(SrcImg, netInputImg, cv::Size(netWidth, netHeight));

	if (channel == 1)
	{
		cvtColor(netInputImg, netInputImg, COLOR_BGR2GRAY);
	}
	Mat blob;
	dnn::blobFromImage(netInputImg, blob, 1 / 255.0, cv::Size(netWidth, netHeight), cv::Scalar(0, 0, 0), true, false);
	cout << "blob.size " << blob.size << endl;
	
	net.setInput(blob);
	std::vector<cv::Mat> netOutputImg;

	net.forward(netOutputImg);

	
	float* pdata = (float*)netOutputImg[0].data;

	cout << "netOutputImg[0].size " << netOutputImg[0].size << endl;

	if (channel == 1)
	{
		cv::Mat temp1 = cv::Mat(netOutputImg[0].size[2], netOutputImg[0].size[3], CV_32FC1, pdata);
		cv::Mat normalizedImage;
		cv::normalize(temp1, normalizedImage, 0, 255, cv::NORM_MINMAX);
		normalizedImage.convertTo(outImg, CV_8UC1);
	}
	else
	{
		cv::Mat tempB = cv::Mat(netOutputImg[0].size[2], netOutputImg[0].size[3], CV_32FC1, pdata);
		cv::Mat tempG = cv::Mat(netOutputImg[0].size[2], netOutputImg[0].size[3], CV_32FC1, pdata + netWidth * netHeight);
		cv::Mat tempR = cv::Mat(netOutputImg[0].size[2], netOutputImg[0].size[3], CV_32FC1, pdata + netWidth * netHeight * 2);
	
		cv::Mat merged; // 创建一个三通道图像

		std::vector<cv::Mat> channels = { tempB, tempG, tempR };// 合并图像通道
		cv::merge(channels, merged);
		cv::Mat normalizedImage;
		cv::normalize(merged, normalizedImage, 0, 255, cv::NORM_MINMAX);//归一化
		normalizedImage.convertTo(outImg, CV_8UC3);
		//imwrite("outimg.jpg", outImg);
	}
	

	return 0;
}

int main()
{
	cout << "Built with OpenCV " << CV_VERSION << endl;
	string model_path = "mosaic-8.onnx";//加载深度学习模型
	
	bool is_cuda = 0;   //是否用CUDA加速，前提要有OPENCV CUDA支持，设为0采用cpu推理
	if (readModel(net, model_path, is_cuda) == true)
	{
		cout << "模型加载成功！" << endl;
	}
	else
	{
		cout << "模型加载失败！" << endl;
		system("pause");
	}
	
	Mat src = imread("src.jpg");
	Mat outImg;
	int netWidth = 224;  //模型输入宽度
	int netHeight = 224; //模型输入高度
	int channel = 3;     //模型输入通道数
	Detect(src, net, outImg, netWidth, netHeight, channel);

	net.empty();//模型释放
	return 0;
}