TensorRT-YOLO：灵活易用的 YOLO 部署工具

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

Detect	Segment
Pose	OBB

项目地址：TensorRT-YOLO GitHub:https://github.com/laugh12321/TensorRT-YOLO

✨ 主要特性

• 多样化的YOLO支持：全面兼容YOLOv3至YOLOv11以及PP-YOLOE和PP-YOLOE+，满足不同版本需求。
• 多场景应用：提供Detect、Segment、Classify、Pose、OBB等多样化场景的示例代码。
• 模型优化与推理加速：
  • ONNX支持：支持ONNX模型的静态和动态导出，包括TensorRT自定义插件支持，简化模型部署流程。
  • TensorRT集成：集成TensorRT插件，包括自定义插件，加速Detect, Segment, Pose, OBB等场景的后处理，提升推理效率。
  • CUDA加速：利用CUDA核函数优化前处理，CUDA图技术加速推理流程，实现高性能计算。
• 语言支持：支持C++和Python（通过Pybind11映射，提升Python推理速度），满足不同编程语言需求。
• 部署便捷性：
  • 动态库编译：提供动态库编译支持，方便调用和部署。
  • 无第三方依赖：全部功能使用标准库实现，无需额外依赖，简化部署流程。
• 快速开发与部署：
  • CLI工具：提供命令行界面（CLI）工具，实现快速模型导出和推理。
  • 跨平台支持：支持Windows、Linux、ARM、x86等多种设备，适应不同硬件环境。
  • Docker部署：支持Docker一键部署，简化环境配置和部署流程。
• TensorRT兼容性：兼容TensorRT 10.x版本，确保与最新技术兼容。

🔮 文档教程

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

💨 快速开始

🔸 前置依赖

• 推荐 CUDA 版本 >= 11.0.1 （CUDA 最低版本 11.0.1）
• 推荐 TensorRT 版本 >= 8.6.1 （TensorRT 最低版本 8.6.1）
• OS: Linux x86_64 (推荐) arm / Windows /

🎆 快速安装

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

在进行推理之前，请参考🔧 CLI 导出模型文档，导出适用于该项目推理的ONNX模型并构建为TensorRT引擎。

Python SDK快速开始

Python CLI 推理示例

使用 --cudaGraph 选项可以显著提升推理速度，但需知此功能仅适用于静态模型。

通过 -m, --mode 参数可以选择不同的模型类型，其中 0 代表检测（Detect）、1 代表旋转边界框（OBB）、2 代表分割（Segment）、3 代表姿态估计（Pose）、4 代表图像分类（Classify）。

1. 使用 tensorrt_yolo 库的 trtyolo 命令行工具进行推理。运行以下命令查看帮助信息：

   trtyolo infer --help 

2. 运行以下命令进行推理：

   trtyolo infer -e models/yolo11n.engine -m 0 -i images -o output -l labels.txt --cudaGraph 

   推理结果将保存至 output 文件夹，并生成可视化结果。

Python 推理示例

DeployDet、DeployOBB、DeploySeg、DeployPose 和 DeployCls 分别对应于检测（Detect）、方向边界框（OBB）、分割（Segment）、姿态估计（Pose）和图像分类（Classify）模型。

对于这些模型，CG 版本利用 CUDA Graph 来进一步加速推理过程，但请注意，这一功能仅限于静态模型。

import cv2 from tensorrt_yolo.infer import DeployDet, generate_labels, visualize  # 初始化模型 model = DeployDet("yolo11n-with-plugin.engine") # 加载图片 im = cv2.imread("test_image.jpg") # 模型预测 result = model.predict(cv2.cvtColor(im, cv2.COLOR_BGR2RGB)) print(f"==> detect result: {result}") # 可视化 labels = generate_labels("labels.txt") vis_im = visualize(im, result, labels) cv2.imwrite("vis_image.jpg", vis_im) 

C++ SDK快速开始

DeployDet、DeployOBB、DeploySeg、DeployPose 和 DeployCls 分别对应于检测（Detect）、方向边界框（OBB）、分割（Segment）、姿态估计（Pose）和图像分类（Classify）模型。

对于这些模型，CG 版本利用 CUDA Graph 来进一步加速推理过程，但请注意，这一功能仅限于静态模型。

#include <opencv2/opencv.hpp> // 为了方便调用，模块除使用 CUDA、TensorRT 其余均使用标准库实现 #include "deploy/vision/inference.hpp" #include "deploy/vision/result.hpp"  int main() {     // 初始化模型     auto model = deploy::DeployDet("yolo11n-with-plugin.engine");     // 加载图片     cv::Mat cvim = cv::imread("test_image.jpg");     cv::cvtColor(cvim, cvim, cv::COLOR_BGR2RGB);     deploy::Image im(cvim.data, cvim.cols, cvim.rows);     // 模型预测     deploy::DetResult result = model.predict(im);     // 可视化（代码省略）     // ...     return 0; } 

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

🖥️ 模型支持列表

Detect	Segment
Pose	OBB

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

☕ 请作者喝杯咖啡

开源不易，如果本项目有帮助到你的话，可以考虑请作者喝杯咖啡，你的支持是开发者持续维护的最大动力~