English | 简体中文

🚀 TensorRT YOLO

🚀 TensorRT-YOLO 是一款专为 NVIDIA 设备设计的易用灵活、极致高效的YOLO系列推理部署工具。项目不仅集成了 TensorRT 插件以增强后处理效果，还使用了 CUDA 核函数以及 CUDA 图来加速推理。TensorRT-YOLO 提供了 C++ 和 Python 推理的支持，旨在提供📦开箱即用的部署体验。包括 目标检测、实例分割、图像分类、姿态识别、旋转目标检测、视频分析等任务场景，满足开发者多场景的部署需求。

✨ 主要特性

🎯 多样化的 YOLO 支持

• 全面兼容：支持 YOLOv3 至 YOLOv11 全系列模型，以及 PP-YOLOE 和 PP-YOLOE+，满足多样化需求。
• 灵活切换：提供简洁易用的接口，支持不同版本 YOLO 模型的快速切换。
• 多场景应用：提供丰富的示例代码，涵盖Detect、Segment、Classify、Pose、OBB等多种应用场景。

🚀 性能优化

• CUDA 加速：通过 CUDA 核函数优化前处理流程，并采用 CUDA 图技术加速推理过程。
• TensorRT 集成：深度集成 TensorRT 插件，显著加速后处理，提升整体推理效率。
• 多 Context 推理：支持多 Context 并行推理，最大化硬件资源利用率。
• 显存管理优化：适配多架构显存优化策略（如 Jetson 的 Zero Copy 模式），提升显存效率。

🛠️ 易用性

• 开箱即用：提供全面的 C++ 和 Python 推理支持，满足不同开发者需求。
• CLI 工具：内置命令行工具，支持快速模型导出与推理，提升开发效率。
• Docker 支持：提供 Docker 一键部署方案，简化环境配置与部署流程。
• 无第三方依赖：全部功能使用标准库实现，无需额外依赖，简化部署流程。
• 部署便捷：提供动态库编译支持，方便调用和部署。

🌐 兼容性

• 多平台支持：全面兼容 Windows、Linux、ARM、x86 等多种操作系统与硬件平台。
• TensorRT 兼容：完美适配 TensorRT 10.x 版本，确保与最新技术生态无缝衔接。

🔧 灵活配置

• 预处理参数自定义：支持多种预处理参数灵活配置，包括 通道交换 (SwapRB)、归一化参数、边界值填充。

🔮 文档教程

• 安装指南
  • 📦 快速编译安装
• 使用示例
  • 目标检测 示例
  • 实例分割 示例
  • 图像分类 示例
  • 姿态识别 示例
  • 旋转目标检测 示例
  • 📹视频分析 示例
  • 多线程多进程 示例
• API文档
  • Python API文档（⚠️ 未实现）
  • C++ API文档（⚠️ 未实现）
• 常见问题
  • ⚠️ 收集中 ...
• 模型支持列表
  • 🖥️ 模型支持列表

💨 快速开始

1. 前置依赖

• CUDA：推荐版本 ≥ 11.0.1
• TensorRT：推荐版本 ≥ 8.6.1
• 操作系统：Linux (x86_64 或 arm)（推荐）；Windows 亦可支持

2. 安装

• 参考 📦 快速编译安装 文档。

3. 模型导出

• 参考 🔧 模型导出 文档，导出适用于该项目推理的ONNX模型并构建为TensorRT引擎。

4. 推理示例

Note

ClassifyModel、DetectModel、OBBModel、SegmentModel 和 PoseModel 分别对应于图像分类（Classify）、检测（Detect）、方向边界框（OBB）、分割（Segment）、姿态估计（Pose）和模型。

• 使用 Python 进行推理：

  import cv2
  from tensorrt_yolo.infer import InferOption, DetectModel, generate_labels, visualize
  
  # 配置推理选项
  option = InferOption()
  option.enable_swap_rb()
  
  # 初始化模型
  model = DetectModel("yolo11n-with-plugin.engine", option)
  
  # 加载图片
  im = cv2.imread("test_image.jpg")
  
  # 模型预测
  result = model.predict(im)
  print(f"==> detect result: {result}")
  
  # 可视化检测结果
  labels = generate_labels("labels.txt")
  vis_im = visualize(im, result, labels)
  cv2.imwrite("vis_image.jpg", vis_im)
  
  # 克隆模型并进行预测
  clone_model = model.clone()
  clone_result = clone_model.predict(im)
  print(f"==> detect clone result: {clone_result}")

• 使用 C++ 进行推理：

  #include <memory>
  #include <opencv2/opencv.hpp>
  
  // 为了方便调用，模块除使用CUDA、TensorRT外，其余均使用标准库实现
  #include "deploy/model.hpp"  // 包含模型推理相关的类定义
  #include "deploy/option.hpp"  // 包含推理选项的配置类定义
  #include "deploy/result.hpp"  // 包含推理结果的定义
  
  int main() {
      // 配置推理选项
      deploy::InferOption option;
      option.enableSwapRB();  // 启用通道交换（从BGR到RGB）
  
      // 初始化模型
      auto model = std::make_unique<deploy::DetectModel>("yolo11n-with-plugin.engine", option);
  
      // 加载图片
      cv::Mat cvim = cv::imread("test_image.jpg");
      deploy::Image im(cvim.data, cvim.cols, cvim.rows);
  
      // 模型预测
      deploy::DetResult result = model->predict(im);
  
      // 可视化（代码省略）
      // ...  // 可视化部分代码未提供，可根据需要实现
  
      // 克隆模型并进行预测
      auto clone_model = model->clone();
      deploy::DetResult clone_result = clone_model->predict(im);
  
      return 0;  // 程序正常结束
  }

更多部署案例请参考模型部署示例 .

🖥️ 模型支持列表

Detect	Segment
Pose	OBB

符号说明: (1) ✅ : 已经支持; (2) ❔: 正在进行中; (3) ❎ : 暂不支持; (4) 🟢 : 导出自行实现，即可推理.

任务场景	模型	CLI 导出	推理部署
Detect	ultralytics/yolov3	✅	✅
Detect	ultralytics/yolov5	✅	✅
Detect	meituan/YOLOv6	❎ 参考官方导出教程	✅
Detect	WongKinYiu/yolov7	❎ 参考官方导出教程	✅
Detect	WongKinYiu/yolov9	❎ 参考官方导出教程	✅
Detect	THU-MIG/yolov10	✅	✅
Detect	ultralytics/ultralytics	✅	✅
Detect	PaddleDetection/PP-YOLOE+	✅	✅
Segment	ultralytics/yolov3	✅	✅
Segment	ultralytics/yolov5	✅	✅
Segment	meituan/YOLOv6-seg	❎ 参考tensorrt_yolo/export/head.py 自行实现	🟢
Segment	WongKinYiu/yolov7	❎ 参考tensorrt_yolo/export/head.py 自行实现	🟢
Segment	WongKinYiu/yolov9	❎ 参考tensorrt_yolo/export/head.py 自行实现	🟢
Segment	ultralytics/ultralytics	✅	✅
Classify	ultralytics/yolov3	✅	✅
Classify	ultralytics/yolov5	✅	✅
Classify	ultralytics/ultralytics	✅	✅
Pose	ultralytics/ultralytics	✅	✅
OBB	ultralytics/ultralytics	✅	✅

📄 许可证

TensorRT-YOLO采用 GPL-3.0许可证，这个OSI 批准的开源许可证非常适合学生和爱好者，可以推动开放的协作和知识分享。请查看LICENSE 文件以了解更多细节。

感谢您选择使用 TensorRT-YOLO，我们鼓励开放的协作和知识分享，同时也希望您遵守开源许可的相关规定。

📞 联系方式

对于 TensorRT-YOLO 的错误报告和功能请求，请访问 GitHub Issues！

🙏 致谢