// PthONNX.cpp : 基于OpenCV dnn、 onnx 的cat、dog二分类程序
// Created by -牧野- 2019年10月29日
//
#include <iostream>
#include <opencv2/highgui.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/dnn.hpp>
#include <fstream>
//ONNX 执行推理类
class PthONNX {
public:
//@model_path ONNX模型路径
//@classes_file_path 分类信息文件
//@input_size 网络输入大小
PthONNX(const std::string &model_path, const std::string &classes_file_path, cv::Size input_size);
//@input_image 输入图片，BGR格式
//@classification_output 网络输出的分类名称 0：cat 1：dog 1：None
void Classify(const cv::Mat &input_image, std::string &classification_output);
private:
void ClassifyImplement(const cv::Mat &image, std::string &classification_output);
private:
cv::Size input_size_;
cv::dnn::Net net_classify_;
std::vector<std::string> classes_;
};
// 构造函数
PthONNX::PthONNX(const std::string &model_path, const std::string &classes_file_path,
cv::Size input_size) : input_size_(input_size) {
std::ifstream ifs(classes_file_path.c_str());
assert(ifs.is_open());
std::string line;
while (getline(ifs, line)) {
line = line;
classes_.push_back(line);
}
net_classify_ = cv::dnn::readNetFromONNX(model_path);
net_classify_.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
net_classify_.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
}
// ONNX推理入口函数
void PthONNX::Classify(const cv::Mat &input_image, std::string &classification_results) {
assert(input_image.data);
cv::Mat image = input_image.clone();
cv::resize(image, image, cv::Size(90, 90));
cv::cvtColor(image, image, cv::COLOR_BGR2GRAY);
ClassifyImplement(image,classification_results);
}
//ONNX推理主函数
void PthONNX::ClassifyImplement(const cv::Mat &image,std::string &classification_results) {
classification_results.clear();
//***********前处理***********
cv::Scalar mean_value(0, 0, 0);
cv::Mat input_blob = cv::dnn::blobFromImage(image, 1, input_size_, mean_value, false, false, CV_32F);
//***********前处理***********
net_classify_.setInput(input_blob);
const std::vector<cv::String> &out_names = net_classify_.getUnconnectedOutLayersNames();
cv::Mat out_tensor = net_classify_.forward(out_names[0]);
//***********后处理***********
double minVal;
double maxVal;
cv::Point minIdx;
cv::Point maxIdx; // minnimum Index, maximum Index
cv::minMaxLoc(out_tensor, &minVal, &maxVal, &minIdx, &maxIdx);
int index_class = maxIdx.x;
classification_results = (index_class <= 1) ? classes_[index_class] : "None";
//***********后处理***********
}
int main()
{
const std::string img_path = "D:/1/1/SimpleNet-master/images/train/dog/dog.9312.jpg";
const std::string onnx_model_path = "D:/software/VS2019_Test/PthONNX/x64/cat_dog_classify.onnx";
const std::string class_names_file_path = "D:/software/VS2019_Test/PthONNX/x64/class_names.class";
const cv::Size net_input_size(90, 90);
cv::Mat img = cv::imread(img_path);
std::string classify_output; // 分类结果
PthONNX classifier(onnx_model_path, class_names_file_path, net_input_size);
classifier.Classify(img, classify_output);
std::cout << "图片类别：" << classify_output << std::endl << std::endl;
cv::putText(img, classify_output, cv::Point(20,20), 2, 1.2, cv::Scalar(0, 0, 255));
cv::imshow("classify", img);
cv::waitKey();
}
- 1
- 2
前往页