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<opencv2/highgui/highgui.hpp>
#include<opencv2/imgproc/imgproc.hpp>
#include<opencv2/dnn/dnn.hpp>
#include<string>
#include<iostream>
#include<fstream>
#include<vector>
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++代码和模型文件