将模型转为onnx
选用pytorch框架，训练resnet18二分类，将二分类模型转为onnx模型
from typing_extensions import dataclass_transform
import torch
import torch.nn as nn
from torchvision import models, transforms
import cv2 as cv
import numpy as np
from PIL import Image
class_name = ['ant', 'bee']
# device = torch.device('cuda' if torch.cuda.is_available else 'cpu')
img = Image.open(r'D:\Projects\PythonProjects\ResnetClassifier\data\hymenoptera_data\val\ants\800px-Meat_eater_ant_qeen_excavating_hole.jpg')
img = img.convert('RGB')
img = img.resize((224, 224))
img = np.array(img, np.float32)
preprocess_img = img/255.0
img_data = np.expand_dims(np.transpose(preprocess_img, (2, 0, 1)), 0)
device = torch.device('cpu')
model = models.resnet18()
model_l = model.fc.in_features
model.fc = nn.Linear(model_l, 2)
model_x = model.to(device)
model_x.eval()
model_x.load_state_dict(torch.load("D:\\Projects\\PythonProjects\\Datasets\\resnet18_bee_and_ant.pth", map_location=torch.device('cpu')))
tensor_img = torch.from_numpy(img_data)
out = model_x(tensor_img)
_, pred = torch.max(out, 1)
pred_name = class_name[pred]
print(pred_name)
input = torch.rand(1, 3, 224, 224)
# torch.save(model_x, 'new_resnet18_v3.pth')
torch.onnx.export(model_x, input, 'new_resnet18_trtpost_v3.onnx', verbose=True)
模型部署
将转换好的onnx模型采用c++进行部署，定义一个ONNXClassifier类，初始化构造函数和分类函数。并将标准化和读标签声明为私有函数，并将重要参数设置为私有变量
class ONNXClassifier {
public:
ONNXClassifier(const string& model_path, const string& label_path, Size input_size);
void Classify(const Mat& input_image, string& out_name, double& confidence);
private:
void preprocess_input(Mat& image);
bool read_labels(const string& label_path);
private:
Size input_size;
cv::dnn::Net net;
cv::Scalar default_mean;
cv::Scalar default_std;
std::vector<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);
}
定义读标签函数，首先声明ifstream对象用于读文件
bool ONNXClassifier::read_labels(const std::string& label_path) {
ifstream ifs(label_path);
string line;
while (getline(ifs, line)) {
size_t index = line.find_first_of(":");
labels.push_back(line.substr(index + 1));
}
if (labels.size() > 0)
return true;
else
return false;
}
进行图片标准化，先将图片像素点进行标准化，并转为32位浮点型
void ONNXClassifier::preprocess_input(Mat& image) {
image.convertTo(image, CV_32F, 1.0 / 255.0); //scale:比例因子
subtract(image, default_mean, image); //图像相减
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);
//图像预处理 1、整体像素减去均值 2、通过放缩系数对图片像素进行放缩
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]);
cout << out_tensor << endl;
cv::Point maxLoc;
double minV;
cv::minMaxLoc(out_tensor, &minV, &confidence, (cv::Point*)0, &maxLoc); // 寻找矩阵或一维向量中的最大、最小值
cout << "maxLoc.x:" << maxLoc.x<<"\tmaxLoc.y:"<<maxLoc.y<<endl;
out_name = labels[maxLoc.x];
}
定义主函数
int main(int argc, char** argv) {
cv::utils::logging::setLogLevel(cv::utils::logging::LogLevel::LOG_LEVEL_SILENT);
vector<string> imgVec;
cv::glob("D:\\Projects\\PythonProjects\\ResnetClassifier\\data\\test_data\\", imgVec);
string model_path = ("D:\\Projects\\PythonProjects\\ResnetClassifier\\new_resnet18_trtpost_v3.onnx");
string label_path = ("D:\\Projects\\PythonProjects\\ResnetClassifier\\ant_and_bee.txt");
Size input_size(300, 300);
for (size_t i = 0; i < imgVec.size(); i++) {
Mat test_image = imread(imgVec[i], cv::IMREAD_COLOR);
ONNXClassifier classifier(model_path, label_path, input_size);
string result;
double confidence = 0;
classifier.Classify(test_image, result, confidence);
cout << imgVec[i] << "\n"<<"预测结果为:" << result << "\tconfidence:" << confidence << endl;
}
return 0;
}
