【源码】基于Python和TensorRT的DETR目标检测模型部署

项目简介

本项目是一个基于Python和TensorRT的目标检测模型部署项目，主要实现了DETR（Detection Transformer）目标检测模型的训练、测试、导出为ONNX格式，以及在C++环境中使用TensorRT进行推理部署。

项目的主要特性和功能

1. 模型训练与测试：使用PyTorch框架实现DETR模型的训练和测试，支持自定义数据集的训练和评估。
2. 模型导出：支持将训练好的DETR模型导出为ONNX格式，便于跨平台部署。
3. 模型推理部署：在C++环境中使用TensorRT加载ONNX模型，进行高效的推理部署，支持模型的序列化和反序列化，便于模型的保存和加载。
4. 数据预处理：提供数据预处理功能，包括图像的大小调整、颜色空间转换等，以满足模型输入要求。
5. 性能评估：提供计算模型性能（如mAP）的功能，支持在COCO数据集上的性能评估。

安装使用步骤

1. 环境准备

• Python环境：安装Python 3.x，安装所需的Python库：torch==1.7.1、torchvision==0.8.2、numpy==1.24.4、onnx==1.15.0、onnxruntime==1.18.0、opencv-python==4.6.0.66、tqdm==4.66.4、scipy==1.10.1。
• C++环境：安装CMake 3.16.3，安装TensorRT 8.2.1.8，安装OpenCV 4.8.0，安装yaml-cpp。

2. 数据准备

• 准备训练所需的数据集，按照VOC2007格式准备，包括图像和对应的标注信息。
• 下载预训练权重和模型文件，链接：百度网盘，提取码：6syh。

3. 训练模型

• 修改python/train.py中的数据集路径和参数，进行模型训练。
• 训练结束后，权重文件将保存在log文件夹下。

4. 导出ONNX模型

• 修改python/export_onnx.py中的resume和onnx_path参数。
• 运行python export_onnx.py，导出的ONNX模型将保存在onnx文件夹下。

5. C++环境部署

• 构建项目：
• 进入CPP目录，配置config.yaml文件，设置isBuild为true，isSaveEngine为true，指定onnx_file_path和classes_txt路径。
• 运行以下命令构建项目： bash cmake -B build cmake --build build
• 运行build/DETR生成detr.engine文件。
• 推理部署：
• 修改config.yaml文件，将isBuild设置为false，指定enginefile路径。
• 运行build/DETR进行推理，预测图片生成结果。

注意事项

• 确保C++环境配置正确，包括模型文件路径、库文件路径等。
• 根据实际需求，可能需要对模型输入和输出进行进一步处理，如图像读取、结果可视化等。
• 部署时可能需要调整模型参数以适应不同环境或任务需求。

下载地址

点击下载 【提取码: 4003】【解压密码: www.makuang.net】