// OnnxRuntimeInference.cpp : 此文件包含 "main" 函数。程序执行将在此处开始并结束。
//
#pragma once
#include <iostream>
#include <assert.h>
#include <vector>
#include <onnxruntime_cxx_api.h>
#include <string>
#include <opencv2/opencv.hpp>
#include "postProcess.h"
using namespace std;
using namespace cv;
using namespace postProcess;
void preDataDet(cv::Mat& matSrc, int& model_width, int& model_height, cv::Mat& matDst)
{
cv::cvtColor(matSrc, matDst, cv::COLOR_BGR2RGB);
cv::resize(matDst, matDst, cv::Size(model_width, model_height));
matDst.convertTo(matDst, CV_32FC3);
matDst = matDst / 255.0f;
}
int main()
{
Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "ONNX_DETECTION"); //设置图优化类型
Ort::SessionOptions session_options;
session_options.SetIntraOpNumThreads(1);
session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_BASIC);
#ifdef _WIN32
const wchar_t* model_path = L"/OnnxRuntimeInference/weights/yolov5_model/best.onnx";
#else
const char* model_path = "/OnnxRuntimeInference/weights/yolov5_model/best.onnx";
#endif
Ort::Session session(env, model_path, session_options); // 创建会话,把模型加载到内存中 CPU
// print model input layer (node names, types, shape etc.)
Ort::AllocatorWithDefaultOptions allocator;
size_t num_input_nodes = session.GetInputCount();
size_t numInputNodes = session.GetInputCount(); //输入输出节点数量
size_t numOutputNodes = session.GetOutputCount();
vector<char*> input_names; // 定义一个字符指针vector
vector<char*> output_names; // 定义一个字符指针vector
vector<vector<int64_t>> input_node_dims; // >=1 outputs ,二维vector
vector<vector<int64_t>> output_node_dims; // >=1 outputs ,int64_t C/C++标准
for (int i = 0; i < numInputNodes; i++)
{
input_names.push_back(session.GetInputName(i, allocator)); // 内存
Ort::TypeInfo input_type_info = session.GetInputTypeInfo(i); // 类型
auto input_tensor_info = input_type_info.GetTensorTypeAndShapeInfo(); //
auto input_dims = input_tensor_info.GetShape(); // 输入shape
input_node_dims.push_back(input_dims); // 保存
std::cout << "InputName: " << session.GetInputName(i, allocator)<< " InputShape: ";
std::copy(input_dims.begin(), input_dims.end(), std::ostream_iterator<int64_t>(std::cout, " "));
std::cout << std::endl;
}
for (int i = 0; i < numOutputNodes; i++)
{
output_names.push_back(session.GetOutputName(i, allocator));
Ort::TypeInfo output_type_info = session.GetOutputTypeInfo(i);
auto output_tensor_info = output_type_info.GetTensorTypeAndShapeInfo();
auto output_dims = output_tensor_info.GetShape();
output_node_dims.push_back(output_dims);
std::cout << "OutputName: " << session.GetOutputName(i, allocator) << " OutputShape: ";
std::copy(output_dims.begin(), output_dims.end(), std::ostream_iterator<int64_t>(std::cout, " "));
std::cout << std::endl;
}
int nModelwidth = 640; //模型input width
int nModelheigh = 640; //模型input height
// 定义一个输入矩阵,int64_t是下面作为输入参数时的类型
array<int64_t, 4> input_shape_{ 1, 3, nModelheigh, nModelwidth }; //NCHW 1,3,640,640
//std::array<float, nModelheigh * nModelwidth * 3> input_image_{}; //输入图片,HWC,不建议这样写,在栈上创建较大的数组容易出现栈溢出(Stack Overflow)的问题
std::vector<float> input_image_(nModelheigh * nModelwidth * 3, 0.0f);
std::string imgPath = "\\OnnxRuntimeInference\\weights\\yolov5_model\\test.bmp";
cv::Mat srcMat = cv::imread(imgPath);
int srch = srcMat.rows, srcw = srcMat.cols;
cv::Mat matNormImage;
preDataDet(srcMat, nModelwidth, nModelheigh, matNormImage); // 减均值除方差
for (int i = 0; i < nModelheigh; ++i) {
for (int j = 0; j < nModelwidth; ++j) {
for (int c = 0; c < srcMat.channels(); ++c) {
//input_image_[c * nModelheigh * nModelwidth + i * nModelwidth + j]= matNormImage.at<cv::Vec<float, 3>>(i, j)[c];
input_image_[c * nModelheigh * nModelwidth + i * nModelwidth + j]= matNormImage.at<float>(i, j * srcMat.channels() + c);
}
}
}
//创建输入tensor
/*
这一行代码的作用是创建一个指向CPU内存的分配器信息对象(AllocatorInfo),用于在运行时分配和释放CPU内存。它调用了CreateCpu函数并传递两个参数:OrtDeviceAllocator和OrtMemTypeCPU。
其中,OrtDeviceAllocator表示使用默认的设备分配器,OrtMemTypeCPU表示在CPU上分配内存。通过这个对象,我们可以在运行时为张量分配内存,并且可以保证这些内存在计算完成后被正确地释放,避免内存泄漏的问题。
*/
auto allocator_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
//使用Ort库创建一个输入张量,其中包含了需要进行目标检测的图像数据。
Ort::Value input_tensor_ = Ort::Value::CreateTensor<float>(allocator_info, input_image_.data(), input_image_.size(), input_shape_.data(), input_shape_.size());
// 开始推理
vector<Ort::Value> ort_outputs = session.Run(Ort::RunOptions{ nullptr }, &input_names[0], &input_tensor_, 1, output_names.data(), output_names.size()); // 开始推理
// generate proposals
std::cout<<"ort_outputs_size: "<<ort_outputs.size()<<std::endl;
const float* prob = ort_outputs[0].GetTensorMutableData<float>(); // 数组,存放预测数据 [bs,anchor'classes,anchor'number,pos+conf+ num'classes]
int img_w = 8;
int img_h = 25200;
cv::Mat probImg = cv::Mat(cv::Size(img_w, img_h), CV_32FC1);
//单张图
memcpy(probImg.data, (float*)prob, (size_t)img_w * img_h * sizeof(float));
std::vector<int> candi;
for (size_t y = 0; y < probImg.rows; ++y)
{
float tmpVal = probImg.at<float>(y, 4);
//std::cout << y << " " << tmpVal << std::endl;
if (tmpVal > 0.5)
{
//std::cout << y << " " << tmpVal << std::endl;
candi.push_back(y);
}
}
//Detections matrix nx6 (xyxy, conf, cls)
std::vector<TDetectBbox> all_box;
for (size_t g = 0; g < candi.size(); ++g)
{
float* ptr = probImg.ptr<float>(candi[g]); //候选行的首地址
TDetectBbox candi_single;
std::vector<float> score;
for (size_t m = 0; m < 3; ++m) //3个类别 获取类别得分
{
float _temVal = ptr[m + 5];
score.push_back(_temVal);
}
//返回score的最大值和index
float _score{ 0.0 };
int classIndx = findVectorMax(score, _score); // 获取得分最高的类别
candi_single.x0 = ptr[0] - ptr[2] / 2;
candi_single.y0 = ptr[1] - ptr[3] / 2;
candi_single.x1 = ptr[0] + ptr[2] / 2;
candi_single.y1 = ptr[1] + ptr[3] / 2;
float _confThresh = ptr[4];
candi_single.confidence = ptr[4] * _score;
candi_single.classIndx = classIndx;
all_box.push_back(candi_single);
}
//计算nms
nms(all_box, 0.4);
//计算坐标映射关系
int shift_x{ 0 }, shift_y{ 0 };
float scale_x{ .0f }, scale_y{ .0f };
//当前是直接resize
scale_x = (float)nModelwidth / srcw;
scale_y = (float)nModelheigh / srch;
for (size_t i = 0; i < all_box.size(); i++)
{
if (all_box[i].confidence > 0.1)
{
int x1 = max(0, (int)(all_box[i].x0 - shift_x) / scale_x);
int y1 = max(0, (int)(all_box[i].y0 - shift_y) / scale_y);
int x2 = min(srcMat.cols, (int)(all_box[i].x1 - shift_x) / scale_x);
int y2 = min(srcMat.rows, (int)(all_box[i].y1 - shift_y) / scale_y);
cv::rectangle(srcMat, cv::Rect(x1, y1, x2 - x1, y2 - y1), cv::Scalar(255, 255, 255));
std::cout << "confidence: " << all_box[i].confidence << std::endl;
}
}
//cv::imwrite("\\OnnxRuntimeInference\\weights\\yolov5\\output.bmp", srcMat);
cv::imshow("Image", srcMat);
cv::waitKey(0);
std::cout << "Hello World!\n";
return 0;
}
基于OnnxRuntime在C++端部署YOLOV5
最新推荐文章于 2024-06-20 11:23:41 发布