帮我用c++编写动态链接库封装yolov5,导入onnx模型
时间: 2023-07-30 12:05:54 浏览: 130
好的,下面是一个简单的示例代码,仅供参考:
```c++
#include <iostream>
#include <vector>
#include <opencv2/opencv.hpp>
#include <onnxruntime_cxx_api.h>
class YOLOv5 {
public:
YOLOv5(const std::string& model_path) {
// 创建会话
Ort::SessionOptions options;
options.SetIntraOpNumThreads(1);
options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
session_ = Ort::Session(env_, model_path.c_str(), options);
}
std::vector<cv::Rect> detect(const cv::Mat& image) {
// 预处理图像
cv::Mat resized_image;
cv::resize(image, resized_image, cv::Size(input_size_, input_size_));
cv::Mat float_image;
resized_image.convertTo(float_image, CV_32FC3, 1.0 / 255.0);
float* input_data = float_image.ptr<float>();
// 准备输入张量
Ort::MemoryInfo input_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
std::vector<int64_t> input_shape = {1, 3, input_size_, input_size_};
Ort::Value input_tensor = Ort::Value::CreateTensor<float>(input_info, input_data, input_shape.data(), input_shape.size());
// 准备输出张量
Ort::MemoryInfo output_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
std::vector<const char*> output_names = {"output"};
std::vector<int64_t> output_shape = {1, 25200, 85};
std::vector<float> output_data(25200 * 85);
// 运行模型
Ort::RunOptions run_options;
session_.Run(run_options, &input_names_[0], &input_tensor, 1, &output_names[0], 1, &output_info, &output_data[0], output_data.size());
// 解析输出张量
std::vector<cv::Rect> detections;
for (int i = 0; i < 25200; i++) {
float* data = &output_data[i * 85];
int class_id = std::max_element(data + 5, data + 85) - data - 5;
float confidence = data[class_id + 5];
if (confidence >= confidence_threshold_) {
int x1 = static_cast<int>((data[0] - data[2] / 2.0) * image.cols);
int y1 = static_cast<int>((data[1] - data[3] / 2.0) * image.rows);
int x2 = static_cast<int>((data[0] + data[2] / 2.0) * image.cols);
int y2 = static_cast<int>((data[1] + data[3] / 2.0) * image.rows);
detections.emplace_back(x1, y1, x2 - x1, y2 - y1);
}
}
return detections;
}
private:
Ort::Env env_{ORT_LOGGING_LEVEL_ERROR, "yolov5"};
Ort::Session session_;
int input_size_ = 640;
float confidence_threshold_ = 0.5;
std::vector<const char*> input_names_ = {"input"};
};
```
这个类使用了 ONNX Runtime C++ API,可以加载 ONNX 模型并进行推理。在这个示例中,模型的输入张量名为 "input",输出张量名为 "output",模型的输入尺寸为 640x640,置信度阈值为 0.5。
在使用时,只需要创建 YOLOv5 对象并调用 detect 方法即可:
```c++
YOLOv5 yolo("yolov5.onnx");
cv::Mat image = cv::imread("test.jpg");
std::vector<cv::Rect> detections = yolo.detect(image);
```
阅读全文