基于OnnxRuntime在C++端部署YOLOV5

// 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;


}