[人工智能]Pytorch学习系列之十二：OpenCV（C++）加载onnx模型进行推理

1 前言

博客Pytorch学习系列之九：OpenCV加载onnx模型进行推理_thequitesunshine007的博客-CSDN博客_opencv 加载onnx记载了python版本opencv加载onnx模型进行推理的方法。这里记录的是C++版本OpenCV加载onnx模型进行推理。

2 代码

#include<opencv2/highgui/highgui.hpp>
#include<opencv2/imgproc/imgproc.hpp>
#include<opencv2/dnn/dnn.hpp>
#include <opencv2/core/utils/logger.hpp>
#include<string>
#include<iostream>
#include<fstream>
#include<vector>
using namespace cv;
using namespace std;

class ONNXClassifier
{
public:
	ONNXClassifier(const std::string& model_path, const std::string& label_path, cv::Size _input_size);
	void Classify(const cv::Mat& input_image, std::string& out_name, double& confidence);
private:
	void preprocess_input(cv::Mat& image);
	bool read_labels(const std::string& label_paht);
private:
	cv::Size		input_size;
	cv::dnn::Net	net;
	cv::Scalar		default_mean;
	cv::Scalar		default_std;
	std::vector<std::string> labels;
};

ONNXClassifier::ONNXClassifier(const std::string& model_path, const std::string& label_path, cv::Size _input_size) 
	:default_mean(0.485, 0.456, 0.406),
	default_std(0.229, 0.224, 0.225), input_size(_input_size)
{
	if (!read_labels(label_path))
	{
		throw std::runtime_error("label read fail!");
	}
	net = cv::dnn::readNet(model_path);
	net.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
	net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
}

bool ONNXClassifier::read_labels(const std::string& label_path)
{
	std::ifstream ifs(label_path);
	assert(ifs.is_open());
	std::string line;
	while (std::getline(ifs, line))
	{
		std::size_t index = line.find_first_of(':');
		labels.push_back(line.substr(index + 1));
	}
	if (labels.size() > 0)
		return true;
	else
		return false;
}

//数据规范化Normalization
void ONNXClassifier::preprocess_input(cv::Mat& image)
{
	image.convertTo(image, CV_32F, 1.0 / 255.0);
	cv::subtract(image, default_mean, image);
	cv::divide(image, default_std, image);
}

void ONNXClassifier::Classify(const cv::Mat& input_image, std::string& out_name, double& confidence)
{
	out_name.clear();
	cv::Mat image = input_image.clone();
	preprocess_input(image);
	cv::Mat input_blob = cv::dnn::blobFromImage(image, 1.0, input_size, cv::Scalar(0, 0, 0), true);
	net.setInput(input_blob);
	const std::vector<cv::String>& out_names = net.getUnconnectedOutLayersNames();
	cv::Mat out_tensor = net.forward(out_names[0]);
	cv::Point maxLoc;
	double minV;
	cv::minMaxLoc(out_tensor, &minV, &confidence, (cv::Point*)0, &maxLoc);
	out_name = labels[maxLoc.x];
}


int main(int argc, char* argv[])
{
	cv::utils::logging::setLogLevel(cv::utils::logging::LogLevel::LOG_LEVEL_SILENT);
	
	//std::string model_path("../model/classifier.onnx");
	//std::string label_path("../model/labels.txt");
	//cv::Size input_size(224, 224);


	std::vector<std::string> imgVec;
	cv::glob("./defect_test/", imgVec);
	std::string model_path("../model/surface_defect_model.onnx");
	std::string label_path("../model/labels_defect.txt");
	cv::Size input_size(300, 300);

	for (size_t i = 0; i < imgVec.size(); i++)
	{
		cv::Mat test_image = cv::imread(imgVec[i], cv::IMREAD_COLOR);
		ONNXClassifier classifier(model_path, label_path, input_size);
		std::string result;
		double confidence = 0;
		classifier.Classify(test_image, result, confidence);
		std::cout << imgVec[i] <<",预测结果为:"<< result<<",confidence: " << confidence << std::endl;
	}

	
	return 0;
}