DeepStream-Yolo:在NVIDIA DeepStream中实现YOLO模型

DeepStream-Yolo项目介绍

DeepStream-Yolo是一个用于在NVIDIA DeepStream SDK中配置和运行YOLO系列目标检测模型的开源项目。该项目支持DeepStream SDK 5.1到7.0版本,适用于x86平台和Jetson平台。

主要特性

• 支持多种YOLO模型,包括Darknet、YOLOv5、YOLOv6、YOLOv7、YOLOv8等
• 支持INT8量化和动态批处理
• 提供GPU加速的后处理
• 自定义ONNX模型解析器
• 支持非方形输入模型
• 提供模型基准测试

支持的模型

项目支持以下YOLO系列模型:

• Darknet
• MobileNet-YOLO
• YOLO-Fastest
• YOLOv5
• YOLOv6
• YOLOv7
• YOLOv8
• YOLOR
• YOLOX
• DAMO-YOLO
• PP-YOLOE/PP-YOLOE+
• YOLO-NAS

环境要求

不同版本的DeepStream对环境有不同的要求,以DeepStream 7.0为例,x86平台需要:

• Ubuntu 22.04
• CUDA 12.2 Update 2
• TensorRT 8.6 GA
• NVIDIA Driver 535
• GStreamer 1.20.3

Jetson平台则需要:

• JetPack 6.0
• DeepStream SDK 7.0

基本使用方法

1. 克隆项目仓库:

git clone https://github.com/marcoslucianops/DeepStream-Yolo.git
cd DeepStream-Yolo

2. 下载YOLO模型的cfg和weights文件到项目目录

3. 编译自定义库:

make -C nvdsinfer_custom_impl_Yolo clean && make -C nvdsinfer_custom_impl_Yolo

4. 修改配置文件config_infer_primary.txt,设置模型路径等参数

5. 运行DeepStream应用:

deepstream-app -c deepstream_app_config.txt

高级配置

NMS配置

可以通过修改config_infer文件来调整NMS参数:

[class-attrs-all]
nms-iou-threshold=0.45 
pre-cluster-threshold=0.25
topk=300

INT8量化

项目支持INT8量化以提高性能,详细步骤见INT8 calibration文档。

元数据提取

可以通过Python或C/C++ API提取DeepStream的元数据,获取检测结果的标签、位置等信息。

结语

DeepStream-Yolo项目为在NVIDIA DeepStream中部署YOLO模型提供了完整的解决方案。通过该项目,开发者可以方便地将各种YOLO模型集成到DeepStream应用中,充分利用GPU加速能力实现高性能的目标检测。项目还在持续更新中,未来将提供更多教程和功能支持。