使用Opencv运行onnx模型-从训练到部署

1. 首先使用pytorch训练一个简单的猫狗分类模型
具体代码参考
pytorch实现kaggle猫狗识别（超详细）
记得先下载猫狗数据集，然后改一下代码里的路径

2. 将保存的模型转为onnx格式

model = torch.load('model.pt',map_location=lambda storage, loc: storage)
dummy = torch.randn(1,3,224,224)
out = model(dummy)
torch.onnx.export(model,dummy,"classifier.onnx",opset_version=12,input_names=['input'],output_names=['output'],
                  dynamic_axes={'input':{0:'batch_size'},'output':{0:'batch_size'}})

对这部分代码不熟的话，可以参考Exporting a Model from PyTorch to ONNX and Running it using ONNX Runtime
3. 准备一个label.txt，表明label和数字的关系,内容如下

0:cat
1:dog

4.opencv 的C++代码

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


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);
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;
}
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)
{
	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;
        cv::minMaxLoc(out_tensor,(double*)0,(double*)0,(cv::Point*)0,&maxLoc);
        out_name = labels[maxLoc.x];
}

int main(int argc, char* argv[])
{
	if (argc != 2)
	{
		std::cout << "input a image file path" << std::endl;
		return -1;
	}
	std::string model_path("../model/classifier.onnx");
	std::string label_path("../model/labels.txt");
	cv::Size input_size(224, 224);
	cv::Mat test_image = cv::imread(argv[1]);
	ONNXClassifier classifier(model_path, label_path, input_size);
	std::string result;
	classifier.Classify(test_image, result);
        std::cout<<"result: "<<result<<std::endl;
	return 0;
}

注意

opencv必须含有dnn模块，推荐opencv4.5

C++完整功能包下载
C++代码和模型文件