基于c++推理yolov5的代码（三）全部代码

yzZ_here

于 2024-08-15 00:04:42 发布

阅读量47

点赞数

分类专栏：基于tensorrt的yolov5推理实现文章标签： c++ YOLO 开发语言 python tensorrt

本文链接：https://blog.csdn.net/qq_39333636/article/details/141203142

版权

基于tensorrt的yolov5推理实现专栏收录该内容

8 篇文章 0 订阅

订阅专栏

目录树结构：

├── build
├── include
│   └── parse_yolov5_output.h
├── input
│   ├── image
│   └── video
├── src
│   └── parse_yolov5_output.cpp
├── CMakeLists.txt
├── cpp_tensorrt_infer.cpp
├── yolov5s_wsl.onnx
└── yolov5s.wsl.trt

CMakeLists.txt：

坦白来说，这一块还没有完全搞明白，但这一份针对我的设备是ok的。这个地方留个坑。

cmake_minimum_required(VERSION 3.12)
project(tensorrt_cpp_infer)

set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Ofast -DNDEBUG -Wno-deprecated-declarations")
set(OpenCV_DIR /usr/lib/x86_64-linux-gnu/cmake/opencv4)

# 设置 TensorRT、CUDA、cuDNN 和 OpenCV 的安装路径
set(TensorRT_DIR /home/zhangyaozu/download/TensorRT-8.6.1.6)
set(CUDA_TOOLKIT_ROOT_DIR /usr/local/cuda-11.3)
set(CUDNN_DIR /home/zhangyaozu/download/cudnn-linux-x86_64-8.9.5.30_cuda11-archive)
set(OpenCV_DIR /usr/lib/x86_64-linux-gnu/cmake/opencv4)
# set(Torch_DIR /home/zhangyaozu/download/libtorch)

# 添加自定义 CMake 模块路径
set(CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} "${CMAKE_SOURCE_DIR}/cmake")

# 设置CMAKE_PREFIX_PATH以包含LibTorch
list(APPEND CMAKE_PREFIX_PATH "${Torch_DIR}")

# 查找 TensorRT、CUDA 和 OpenCV
find_package(TensorRT REQUIRED)
find_package(CUDA REQUIRED)
find_package(OpenCV REQUIRED)
# find_package(Torch REQUIRED)
find_package(OpenCV REQUIRED)

include_directories(
    ${OpenCV_INCLUDE_DIRS}
    ${CUDA_TOOLKIT_ROOT_DIR}/include
    ${TensorRT_INCLUDE_DIRS}
    ${CUDNN_DIR}/include
    # ${TORCH_DIR}/include
)

link_directories(
    ${CUDA_TOOLKIT_ROOT_DIR}/lib64
    ${TensorRT_DIR}/lib
    ${CUDNN_DIR}/lib64
    # ${TORCH_DIR}/lib
)

set(SOURCE_FILES cpp_tensorrt_infer.cpp
                src/parse_yolov5_output.cpp
                                        )

set(INCLUDE_DIRS
  include)

add_executable(tensorrt_cpp_infer ${SOURCE_FILES})

target_link_libraries(tensorrt_cpp_infer
    ${CUDA_LIBRARIES}
    ${TensorRT_LIBRARIES}
    cudart
    # ${TORCH_LIBRARIES}
    ${OpenCV_LIBS}
)

maincpp

也就是cpp_tensorrt_infer.cpp

#include <NvInfer.h> // TensorRT 主要接口头文件
#include <cuda_runtime_api.h> // CUDA 运行时 API
#include <opencv2/opencv.hpp> // OpenCV，用于图像处理
#include <iostream> // 输入输出流
#include <fstream> // 文件流
#include <memory> // 智能指针
#include <vector> // STL 向量容器
#include <sys/stat.h>
#include "include/parse_yolov5_output.h"
#include <random>
#include <chrono>
#include <filesystem>
// #include <torch/torch.h>

const int H = 640;
const int W = 640;

bool fileExists(const std::string& filename) {
    struct stat buffer;
    return (stat(filename.c_str(), &buffer) == 0);
}

std::vector<std::string> listFiles(const std::string& directory);
std::string getFileExtension(const std::string& filename);
std::string getFileName(const std::string& filepath);
void processImage(YoLov5TRT& yolov5_trt, const std::string& img_path);
void processVideo(YoLov5TRT& yolov5_trt, const std::string& video_path);

// 自定义的 Logger 类，继承自 nvinfer1::ILogger
class MyLogger : public nvinfer1::ILogger {
public:
    void log(Severity severity, const char* msg) noexcept override {
        // 如果日志级别是错误或更严重的情况，则输出日志消息
        if (severity <= Severity::kERROR) {
            std::cout << msg << std::endl;
        }
    }
};

// 获取 Logger 实例
nvinfer1::ILogger& getCustomLogger() {
    static MyLogger logger;
    return logger;
}

// 构造函数实现
YoLov5TRT::YoLov5TRT(const std::string& engine_file_path) {
    loadEngine(engine_file_path); // 调用私有方法加载引擎
}

// 析构函数实现
YoLov5TRT::~YoLov5TRT() {
    // 销毁 CUDA 流
    cudaStreamDestroy(stream_);
    // 释放 GPU 内存
    cudaFree(buffers_[0]);
    cudaFree(buffers_[1]);
    // 销毁执行上下文
    context_->destroy();
    // 销毁引擎
    engine_->destroy();
    // 销毁运行时
    runtime_->destroy();
}

// 加载 TensorRT 引擎的实现
void YoLov5TRT::loadEngine(const std::string& engine_file_path) {

    nvinfer1::ILogger& logger = getCustomLogger(); // 获取 ILogger 实例
    runtime_ = nvinfer1::createInferRuntime(logger); // 创建 TensorRT 运行时  指针
    std::ifstream engine_file(engine_file_path, std::ios::binary);
    if (!engine_file) {
        std::cerr << "Failed to open engine file: " << engine_file_path << std::endl;
        return;
    }

    engine_file.seekg(0, std::ifstream::end); // 移动文件指针到末尾
    size_t fsize = engine_file.tellg(); // 获取文件大小
    if (fsize == static_cast<size_t>(-1)) {
        std::cerr << "Failed to get file size: " << engine_file_path << std::endl;
        return;
    }

    engine_file.seekg(0, std::ifstream::beg); // 移动文件指针到开头
    std::vector<char> engine_data(fsize); // 分配缓冲区
    engine_file.read(engine_data.data(), fsize); // 读取文件数据到缓冲区
    engine_ = runtime_->deserializeCudaEngine(engine_data.data(), fsize, nullptr); // 反序列化引擎
    context_ = engine_->createExecutionContext(); // 创建执行上下文
    input_size_ = 1 * 3 * H * W; // 设置输入大小
    output_size_ = 1 * 25200 * 85; // 设置输出大小
    // 分配输入和输出的 GPU 内存
    cudaMalloc(&buffers_[0], input_size_ * sizeof(float));
    cudaMalloc(&buffers_[1], output_size_ * sizeof(float));
    cudaStreamCreate(&stream_); // 创建 CUDA 流
}

// 图像前处理函数实现
void YoLov5TRT::preprocessImage(const cv::Mat& image, float* gpu_input) {
     cv::Mat Image = image;

    // preImage函数中对图片进行详细前处理
    cv::Mat risezed_image = preImage(Image, H, W);

    size_t image_size = risezed_image.total() * risezed_image.channels(); // 计算输入大小，用来分配缓存区
    std::vector<float> input_data(image_size); // 分配输入数据缓冲区

    // 5. 使用 memcpy 复制数据
    std::memcpy(input_data.data(), risezed_image.data, image_size * sizeof(float));

    // 6. 异步复制数据到 GPU
    cudaMemcpyAsync(gpu_input, input_data.data(), image_size * sizeof(float), cudaMemcpyHostToDevice, stream_);
}

// 结果后处理函数实现
std::vector<float> YoLov5TRT::postprocessResults(float* gpu_output) {
    std::vector<float> output_data(output_size_); // 分配输出数据缓冲区
    cudaMemcpyAsync(output_data.data(), gpu_output, output_size_ * sizeof(float), cudaMemcpyDeviceToHost, stream_); // 异步复制数据到主机
    cudaStreamSynchronize(stream_); // 同步 CUDA 流，确保数据复制完成
    return output_data; // 返回结果数据
};

// 推理函数实现
std::vector<float> YoLov5TRT::infer(const cv::Mat& image) {
    preprocessImage(image, static_cast<float*>(buffers_[0])); // 调用图像前处理
    context_->enqueueV2(buffers_, stream_, nullptr); // 执行推理
    return postprocessResults(static_cast<float*>(buffers_[1])); // 调用结果后处理并返回结果
}

//warm_up
void warm_up(YoLov5TRT& warmup_yolo,int iterations){
    cv::Mat image_zero = cv::Mat::ones(640,640,CV_8UC3);
    // std::vector<Detection> warmup_out;
    std::cout<<"start warm up…………"<<std::endl;
    auto start = std::chrono::high_resolution_clock::now();

    for (int i = 0; i < iterations; ++i) {
        warmup_yolo.infer(image_zero);
    }
    std::cout<<"warm up end"<<std::endl;

    auto end = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();

    std::cout << "Warm up end. Total time: " << duration << " ms" << std::endl;

}

//box
void draw_detections(cv::Mat& img, const std::vector<Detection>& detections, const std::vector<std::string>& categories) {
    for (const auto& det : detections) {

        // std::cout << "原图上坐标是 " << det.x1 << " " << det.y1 << " " << det.x2 << " " << det.y2 << std::endl;

        // 绘制检测框
        cv::rectangle(img, cv::Point(det.x1, det.y1), cv::Point(det.x2, det.y2), cv::Scalar(0, 255, 0), 2);

        // 准备标签文本
        std::string label = categories[det.class_id] + ": " + std::to_string(det.confidence).substr(0, 4);

        // 获取文本大小
        int baseline;
        cv::Size textSize = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseline);

        // 确保文本不超出图像边界
        int text_x = det.x1;
        int text_y = det.y1;

        // 绘制文本背景框
        cv::rectangle(img, cv::Point(text_x, text_y - textSize.height), 
                      cv::Point(text_x + textSize.width, text_y + baseline), 
                      cv::Scalar(255, 255, 255), cv::FILLED);

        // 绘制文本
        cv::putText(img, label, cv::Point(text_x, text_y), 
                    cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 0), 1);
    }
}

// // 主函数
int main() {

    //文件路径
    std::string trtfilepath = "/mnt/d/code/cpp_tensorrt_infer/yolov5s_wsl.trt";
    std::string img_path = "/mnt/d/code/cpp_tensorrt_infer/input/image/team.jpg";

    //检查文件是否存在
    if (!fileExists(trtfilepath)) {
        std::cerr << "File does not exist: ./yolov5s_wsl.trt" << std::endl;
        return -1;
        }

    // 创建 YOLOv5 TensorRT 对象，并指定引擎文件路径
    YoLov5TRT yolov5_trt(trtfilepath);

    warm_up(yolov5_trt,50);
    
    //读取image
    cv::Mat original_image = cv::imread(img_path);
    int ori_H = original_image.rows;
    int ori_W = original_image.cols;

    // 执行推理并获取结果
    std::vector<float> output = yolov5_trt.infer(original_image);

    //postProcess
    std::vector<Detection> yolo_out = parse_yolo(output,  ori_H, ori_W);

    // 打印结果
    if(yolo_out.size() == 0){
        std::cout<<"yolo_out is empty！！！"<< std::endl;
        return -1;
    }

    // 绘制检测结果
    draw_detections(original_image, yolo_out, categories);

    // 显示图像
    cv::imshow("Detections", original_image);
    cv::waitKey(0);

    // 保存图像
    cv::imwrite("output_image.jpg", original_image);

    return 0;
}

void processImage(YoLov5TRT& yolov5_trt, const std::string& img_path) {
    cv::Mat original_image = cv::imread(img_path);
    int ori_H = original_image.rows;
    int ori_W = original_image.cols;

    // 执行推理并获取结果
    std::vector<float> output = yolov5_trt.infer(original_image);

    // postProcess
    std::vector<Detection> yolo_out = parse_yolo(output, ori_H, ori_W);

    // 打印结果
    if (yolo_out.empty()) {
        std::cout << "yolo_out is empty for image: " << img_path << std::endl;
        return;
    }

    // 绘制检测结果
    draw_detections(original_image, yolo_out, categories);

    cv::imshow("Detections", original_image);
    cv::waitKey(0);

    // 保存图像
    std::string output_path = "output_" + getFileName(img_path);
    cv::imwrite(output_path, original_image);

    std::cout << "Processed image: " << img_path << " -> " << output_path << std::endl;
}

void processVideo(YoLov5TRT& yolov5_trt, const std::string& video_path) {
    cv::VideoCapture cap(video_path);
    if (!cap.isOpened()) {
        std::cerr << "Error opening video file: " << video_path << std::endl;
        return;
    }

    int frame_width = cap.get(cv::CAP_PROP_FRAME_WIDTH);
    int frame_height = cap.get(cv::CAP_PROP_FRAME_HEIGHT);
    double fps = cap.get(cv::CAP_PROP_FPS);

    std::string output_path = "output_" + getFileName(video_path);
    cv::VideoWriter video_writer(output_path, cv::VideoWriter::fourcc('M','J','P','G'), fps, cv::Size(frame_width, frame_height));

    cv::Mat frame;
    int frame_count = 0;

    while (cap.read(frame)) {
        std::vector<float> output = yolov5_trt.infer(frame);
        std::vector<Detection> yolo_out = parse_yolo(output, frame_height, frame_width);

        draw_detections(frame, yolo_out, categories);

        video_writer.write(frame);

        // 显示当前帧
        cv::imshow("Video", frame);

        // 等待一小段时间，并检查是否按下 'q' 键来退出
        char c = (char)cv::waitKey(1);
        if (c == 'q' || c == 'Q')
            break;

        frame_count++;
        if (frame_count % 100 == 0) {
            std::cout << "Processed " << frame_count << " frames of " << video_path << std::endl;
        }
    }

    cap.release();
    video_writer.release();

    std::cout << "Processed video: " << video_path << " -> " << output_path << std::endl;
}

std::string getFileExtension(const std::string& filename) {
    size_t pos = filename.find_last_of(".");
    if (pos != std::string::npos) {
        return filename.substr(pos + 1);
    }
    return "";
}

std::string getFileName(const std::string& filepath) {
    size_t pos = filepath.find_last_of("/\\");
    if (pos != std::string::npos) {
        return filepath.substr(pos + 1);
    }
    return filepath;
}

std::vector<std::string> listFiles(const std::string& directory) {
    std::vector<std::string> files;
    for (const auto& entry : std::filesystem::directory_iterator(directory)) {
        files.push_back(entry.path().string());
    }
    return files;
}

前后处理文件

parse_yolov5_output.cpp

#include <vector>
#include <algorithm>
#include <cmath>
#include <iostream>
#include "../include/parse_yolov5_output.h"
// #include <torch/torch.h>

const int NUM_ANCHORS = 25200;
const int NUM_CLASSES = 80;
const int NUM_OUTPUTS = 85;
const float CONFIDENCE_THRESHOLD = 0.5f;
const float NMS_THRESHOLD = 0.4f;

void nms_process(std::vector<Detection>& detection, std::vector<Detection>& result, float iou_threshold);

//calculate iou function
float calculate_iou(const Detection& a,const Detection& b){
    float i_x1 = std::max(a.x1,b.x1);
    float i_y1 = std::max(a.y1,b.y1);
    float i_x2 = std::min(a.x2,b.x2);
    float i_y2 = std::min(a.y2,b.y2);

    float i_area = std::max(0.0f, i_x2 - i_x1) * 
                   std::max(0.0f, i_y2 - i_y1);
    float a_area = (a.x2 - a.x1) * (a.y2 - a.y1);
    float b_area = (b.x2 - b.x1) * (b.y2 - b.y1); 
    float iou = i_area / (a_area + b_area - i_area);
    return iou;
}

// 预测框校准函数
std::vector<Detection> correct_boxes(const std::vector<Detection>& detections, const cv::Size& input_shape, const cv::Size& image_shape) {
    // 构建存放结果的结构体
    std::vector<Detection> corrected_boxes;
    // 得到训练缩放的尺寸 640
    float input_w = input_shape.width;     //640
    float input_h = input_shape.height;    //640

    // 计算原图的w和h
    float image_w = image_shape.width;     //810
    float image_h = image_shape.height;    //1080

    // 计算缩放因子
    float scale = std::min(input_w / image_w, input_h / image_h);  // 0.592593

    float new_w = image_w * scale;     //480
    float new_h = image_h * scale;     //640

    float pad_w = (input_w - new_w) / 2.0f;
    float pad_h = (input_h - new_h) / 2.0f;

    // std::cout<<"校准框之后的的边框坐标分别是"<<std::endl;
    for (const auto& detection : detections) {
        float x1 = (detection.x1 - pad_w) / scale;
        float y1 = (detection.y1 - pad_h) / scale;
        float x2 = (detection.x2 - pad_w) / scale;
        float y2 = (detection.y2 - pad_h) / scale;
        corrected_boxes.push_back({x1, y1, x2, y2, detection.class_id, detection.confidence});
    }
    return corrected_boxes;
}

std::vector<Detection> parse_yolo(const std::vector<float>& yoloout,int img_height,int img_width){
    
    // std::cout << "进入parse_yolo"<<std::endl;
    std::vector<Detection> detections;

    // std::cout<<"图片的高和宽    "<<img_height<<img_width<<std::endl;

    for (int i = 0; i < NUM_ANCHORS; ++i) {
        const float* data = &yoloout[i * NUM_OUTPUTS];

        //边界框坐标
        float x = data[0];
        float y = data[1];
        float w = data[2];
        float h = data[3];
        // std::cout<<"x is:"<<x<<", y is:"<<y<<", w is :"<<w<<", h is:"<<h<<std::endl;

        //计算obj score
        float obj_score = data[4];

        float class_score = 0;
        int class_id = -1;
        //遍历每一类的score，取最大值并记录该类id
        for(int j = 0; j < NUM_CLASSES; ++j){
            float score = data[5 + j];
            if (score > class_score){
                class_score = score;
                class_id = j;
            }
        };

        //计算最后的置信度
        float confidence = class_score * obj_score;
        // std::cout << i << "   confidence is "<< confidence <<std::endl;

        //设置置信度阈值进行过滤
        if(confidence > CONFIDENCE_THRESHOLD){
            // ----------------------------------------------------------
            // yolov5推理输出的坐标框已经是反归一化之后的，不需要再反归一化
            // ----------------------------------------------------------
            // float x_center = x * img_width;
            // float y_center = y * img_height;
            // float width = w * img_width;
            // float height = h * img_height;
            
            //图像坐标系和笛卡尔坐标系的y轴是反的
            float x1 = x - w / 2;
            float y1 = y - h / 2;
            float x2 = x + w / 2;
            float y2 = y + h / 2;

            //push倒detection这个vector的末尾
            detections.push_back({x1, y1, x2, y2, class_id, confidence});

        }

    }

    // std::cout << "开始NMS"<<std::endl;
    //NMS 
    std::vector<Detection> nms_result;
    nms_process(detections,nms_result,NMS_THRESHOLD);
    // std::cout << "结束NMS"<<std::endl;

    // return nms_result;
    // std::cout<<"nms之后的边框坐标分别是"<<std::endl;
    // for(std::vector<Detection>::const_iterator iter = nms_result.cbegin(); iter != nms_result.cend(); iter++){
    //     std::cout<<"x1是"<<(*iter).x1<<",y1是"<<(*iter).y1<<",x2是"<<(*iter).x2<<",y2是"<<(*iter).y2<<std::endl;
    // }

    // 校准预测框
    std::vector<Detection> corrected_boxes = correct_boxes(nms_result, cv::Size(640, 640), cv::Size(img_width,img_height));
    return corrected_boxes;

}

void nms_process(std::vector<Detection>& detection, 
                 std::vector<Detection>& result, float nms_threshold){
    
    //sort函数默认升序，这边使用confidence进行降序排序
    std::sort(detection.begin(),detection.end(),
              //Lambda表达式作为第三个参数输入。
              [](const Detection& a,const Detection& b){return a.confidence > b.confidence;}
              );

    // 创建一个布尔类型的std::vector is_suppressed，其长度等于detections的大小。
    // 所有元素被初始化为false，表示初始时所有检测结果都没有被抑制。
    std::vector<bool> is_suppressed(detection.size(),false);

    for(size_t i = 0; i < detection.size(); ++i){
        if(is_suppressed[i]) continue;

        //依次遍历整个vector，先选择i，然后依次和（i+1）开始对比iou
        //每次只从第一个循环中写值
        result.push_back(detection[i]);

        for(size_t j = i + 1; j < detection.size(); ++j){

            if(is_suppressed[j]) continue;

            if(detection[i].class_id == detection[j].class_id){

               float iou = calculate_iou(detection[i],detection[j]);
               if(iou > nms_threshold){
                is_suppressed[j] = true;
               }
            }
        }
    }
}

//输入image处理函数
//------------------------------------------------------------------------
// 首先不保证传进来的image是正方形，保证输入的图像为640*640
// 输入图像需要转换到BGR-->RGB；resized（优先保证长边缩放到位，然后填充短边）;
// 然后归一化；HWC-->CHW; CHW-->BCHW; 转为行存储（连续内存）；
//------------------------------------------------------------------------
cv::Mat preImage(cv::Mat& rawBGRImage,int inputH, int inputW){
    
    // BGR --> RGB
    cv::Mat image;
    cv::cvtColor(rawBGRImage,image,cv::COLOR_BGR2RGB);

    // cv::Mat oriImage = rawBGRImage.clone();
    int ori_h = rawBGRImage.rows;
    int ori_w = rawBGRImage.cols;

    // 取最小的因子，以保证长的一边被正确缩放，短边过度缩放可在后面使用padding填充
    // 这两个结果肯定小于1，使用整除会将结果截断为0
    double r = std::min(static_cast<double>(inputH) / ori_h, static_cast<double>(inputW) / ori_w);  

    int new_h = std::round(ori_h * r);
    int new_w = std::round(ori_w * r);

    //计算padding的大小,
    int dh = inputH - new_h;
    int dw = inputW - new_w;

    //判断缩放后的图像和原图长和宽不同的话进行缩放，
    if(dh > 0 || dw > 0){
        cv::resize(image,image,cv::Size(new_w,new_h),0, 0, cv::INTER_LINEAR);
        
        // Pad the short side with (128,128,128)
        int top = dh / 2;
        int bottom = dh - top;
        int left = dw / 2;
        int right = dw - left;
        //扩充边界
        cv::copyMakeBorder(image, image, top, bottom, left, right, cv::BORDER_CONSTANT, cv::Scalar(128, 128, 128));
    }

    //归一化
    image.convertTo(image, CV_32F); //自动转，如果原图是32f，只拷贝，不做转换。
    image /= 255.0; // 8u转32f之后，原始像素值还是不变的，只是可表示范围变大了。

    // HWC to NCHW format:并处理成cv中的dnn多维格式数据，类似于tensor
    cv::Mat imageNCHW;
    cv::dnn::blobFromImage(image, imageNCHW, 1.0, cv::Size(), cv::Scalar(), false, false);//这种缩放会影响图像精度
    //该函数默认bs维度=1。cv::dnn::blobFromImages函数可以自由设定bs维度的值
    // std::cout << "NCHW shape: " << imageNCHW.size << std::endl;

    // Convert the image to row-major order, also known as "C order":
    // 复制矩阵 转换位行主序（C order）存储。意味着数据在内存中是按行连续存储的
    // 某些深度学习框架或推理引擎可能要求输入数据是连续存储的
    cv::Mat imageRowMajor;
    imageNCHW.copyTo(imageRowMajor);

    return imageRowMajor;

}

头文件

主要包含了yoloTRT的类


#pragma once
#include <vector>
#include <string>
#include <NvInfer.h> // TensorRT 主要接口头文件
#include <cuda_runtime_api.h> // CUDA 运行时 API
#include <opencv2/opencv.hpp> // OpenCV，用于图像处理

const std::vector<std::string> categories = {"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat",
                  "traffic light",
                  "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
                  "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase",
                  "frisbee",
                  "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard",
                  "surfboard", "tennis racket", "bottle",
                  "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange",
                  "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", "potted plant", "bed",
                  "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave",
                  "oven",
                  "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier",
                  "toothbrush"};

//COCO80的颜色
const std::vector<cv::Scalar> color80{
            cv::Scalar(128, 77, 207),cv::Scalar(65, 32, 208),cv::Scalar(0, 224, 45),cv::Scalar(3, 141, 219),cv::Scalar(80, 239, 253),cv::Scalar(239, 184, 12),
            cv::Scalar(7, 144, 145),cv::Scalar(161, 88, 57),cv::Scalar(0, 166, 46),cv::Scalar(218, 113, 53),cv::Scalar(193, 33, 128),cv::Scalar(190, 94, 113),
            cv::Scalar(113, 123, 232),cv::Scalar(69, 205, 80),cv::Scalar(18, 170, 49),cv::Scalar(89, 51, 241),cv::Scalar(153, 191, 154),cv::Scalar(27, 26, 69),
            cv::Scalar(20, 186, 194),cv::Scalar(210, 202, 167),cv::Scalar(196, 113, 204),cv::Scalar(9, 81, 88),cv::Scalar(191, 162, 67),cv::Scalar(227, 73, 120),
            cv::Scalar(177, 31, 19),cv::Scalar(133, 102, 137),cv::Scalar(146, 72, 97),cv::Scalar(145, 243, 208),cv::Scalar(2, 184, 176),cv::Scalar(219, 220, 93),
            cv::Scalar(238, 153, 134),cv::Scalar(197, 169, 160),cv::Scalar(204, 201, 106),cv::Scalar(13, 24, 129),cv::Scalar(40, 38, 4),cv::Scalar(5, 41, 34),
            cv::Scalar(46, 94, 129),cv::Scalar(102, 65, 107),cv::Scalar(27, 11, 208),cv::Scalar(191, 240, 183),cv::Scalar(225, 76, 38),cv::Scalar(193, 89, 124),
            cv::Scalar(30, 14, 175),cv::Scalar(144, 96, 90),cv::Scalar(181, 186, 86),cv::Scalar(102, 136, 34),cv::Scalar(158, 71, 15),cv::Scalar(183, 81, 247),
            cv::Scalar(73, 69, 89),cv::Scalar(123, 73, 232),cv::Scalar(4, 175, 57),cv::Scalar(87, 108, 23),cv::Scalar(105, 204, 142),cv::Scalar(63, 115, 53),
            cv::Scalar(105, 153, 126),cv::Scalar(247, 224, 137),cv::Scalar(136, 21, 188),cv::Scalar(122, 129, 78),cv::Scalar(145, 80, 81),cv::Scalar(51, 167, 149),
            cv::Scalar(162, 173, 20),cv::Scalar(252, 202, 17),cv::Scalar(10, 40, 3),cv::Scalar(150, 90, 254),cv::Scalar(169, 21, 68),cv::Scalar(157, 148, 180),
            cv::Scalar(131, 254, 90),cv::Scalar(7, 221, 102),cv::Scalar(19, 191, 184),cv::Scalar(98, 126, 199),cv::Scalar(210, 61, 56),cv::Scalar(252, 86, 59),
            cv::Scalar(102, 195, 55),cv::Scalar(160, 26, 91),cv::Scalar(60, 94, 66),cv::Scalar(204, 169, 193),cv::Scalar(126, 4, 181),cv::Scalar(229, 209, 196),
            cv::Scalar(195, 170, 186),cv::Scalar(155, 207, 148)
};

//存放结果的结构体
struct Detection {
    float x1, y1, x2, y2;
    int class_id;
    float confidence;
};

// 声明后处理函数
std::vector<Detection> parse_yolo(const std::vector<float>& output, int img_width, int img_height);

// 声明图像预处理函数
cv::Mat preImage(cv::Mat& rawBGRImage,int inputH, int inputW);

// 创建一个用于 YOLOv5 TensorRT 推理的类
class YoLov5TRT {
public:
    // 构造函数，接受引擎文件路径作为参数
    YoLov5TRT(const std::string& engine_file_path);
    // 析构函数，用于释放资源
    ~YoLov5TRT();
    // 推理函数，接受一个 OpenCV 图像并返回推理结果
    std::vector<float> infer(const cv::Mat& image);

private:
    // 加载 TensorRT 引擎
    void loadEngine(const std::string& engine_file_path);
    // 图像前处理函数
    void preprocessImage(const cv::Mat& image, float* gpu_input);
    // 结果后处理函数
    std::vector<float> postprocessResults(float* gpu_output);

    nvinfer1::IRuntime* runtime_; // TensorRT 运行时指针
    nvinfer1::ICudaEngine* engine_; // TensorRT 引擎指针
    nvinfer1::IExecutionContext* context_; // TensorRT 执行上下文指针
    cudaStream_t stream_; // CUDA 流，用于管理异步操作
    void* buffers_[2]; // 输入输出缓冲区指针
    int input_size_; // 输入大小
    int output_size_; // 输出大小
};