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

TensorRT YOLO

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

Detect	Segment
Pose	OBB

主要特性

• 多样化的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) : 导出自行实现，即可推理.

任务场景	模型	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

请作者喝杯咖啡

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