Cockpit-cpp/ZjQt/Qt/YOLOv5Detector.h

#pragma once

#include <algorithm> 
#include <iostream>
#include <opencv2/opencv.hpp>
#include <onnxruntime_cxx_api.h>
#include <vector>
#include <string>
#include<QImage>

using cv::Mat;
using std::cout;
using std::endl;
using std::string;
using std::vector;

class YOLOv5Detector
{
public:
    YOLOv5Detector(const wchar_t* model_path);
    vector<vector<float>> detect(const cv::Mat& img);
    void draw_boxes(Mat& img, const vector<vector<float>>& info);
    cv::Mat QImageToMat(const QImage& image);
    
private:
    Ort::Env env;
    Ort::MemoryInfo memory_info;
    Ort::SessionOptions session_options;
    Ort::Session session;
    bool use_cuda = false;

    const char* input_names[1] = { "images" }; // 根据模型的输入名调整
    const char* output_names[1] = { "output" }; // 根据模型的输出名调整

    static vector<vector<float>> get_info(const float* pdata, int total);
    static void info_simplify(vector<vector<float>>& info);   
    static void nms(vector<vector<float>>& info);    
       
};