c++写Yolov5TRTContext推理

Yolov5TRTContext 是一个基于 TensorRT 的 YOLOv5 推理框架，它可以在 NVIDIA GPU 上进行高效的目标检测。下面是使用 C++ 编写 Yolov5TRTContext 推理的步骤：

1. 安装 TensorRT。

2. 下载 YOLOv5 模型文件。

3. 使用 TensorRT 的 Python API 将 YOLOv5 模型转换为 TensorRT 引擎。这个过程可以在 Python 中完成，示例代码如下：

import tensorrt as trt
import sys

def build_engine(onnx_file_path, engine_file_path):
    TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

    with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
        builder.max_workspace_size = 1 << 30 # Your workspace size
        builder.max_batch_size = 1 # Your max batch size
        with open(onnx_file_path, 'rb') as model:
            parser.parse(model.read())
        engine = builder.build_cuda_engine(network)
        with open(engine_file_path, "wb") as f:
            f.write(engine.serialize())
        return engine

if __name__ == '__main__':
    onnx_file_path = sys.argv[1]
    engine_file_path = sys.argv[2]
    build_engine(onnx_file_path, engine_file_path)

4. 在 C++ 中加载 TensorRT 引擎文件，创建 Yolov5TRTContext 对象，然后进行推理。示例代码如下：

#include "yolov5_trt.h"
#include <iostream>
#include <fstream>

int main(int argc, char** argv) {
    if (argc != 4) {
        std::cerr << "Usage: ./yolov5_trt input.bin engine.trt output.bin\n";
        return -1;
    }

    std::ifstream input(argv[1], std::ios::binary);
    std::ofstream output(argv[3], std::ios::binary);

    if (!input || !output) {
        std::cerr << "Failed to open file\n";
        return -1;
    }

    int input_size = 640 * 640 * 3;
    int output_size = 20 * 7;
    char* input_buf = new char[input_size];
    char* output_buf = new char[output_size];

    input.read(input_buf, input_size);

    Yolov5TRTContext context;
    context.load_engine(argv[2]);

    context.do_inference(input_buf, output_buf);

    output.write(output_buf, output_size);

    delete[] input_buf;
    delete[] output_buf;

    return 0;
}

其中，yolov5_trt.h 头文件包含了 Yolov5TRTContext 类的定义，示例代码如下：

#include "NvInferRuntime.h"

class Yolov5TRTContext {
public:
    Yolov5TRTContext();
    ~Yolov5TRTContext();
    bool load_engine(const char* engine_path);
    void do_inference(const char* input, char* output);
private:
    nvinfer1::ICudaEngine* engine_;
    nvinfer1::IExecutionContext* context_;
    cudaStream_t stream_;
};

具体实现可以参考 TensorRT 的官方文档和源代码。