【TensorRT】Yolov5-DeepSORT 目标跟踪

573 阅读2分钟

Yolov5-DeepSORT-TensorRT

本项目是 Yolo-DeepSORT 的 C++ 实现,使用 TensorRT 进行推理。

🚀🚀🚀 开源地址:Yolov5_DeepSORT_TensorRT,求 star⭐ ~

引言

  • ⚡ 推理速度可达25-30FPS,可以落地部署
  • 🐳 提供了 dockerfile 以快速搭建开发环境,不用浪费时间配置环境了;
  • 📥 下载并使用本文提供的 onnx 文件(见正文),在程序运行时会自动解析 onnx 并序列化出 engine 文件(*.trtmodel)到 workspace 目录下,无需自行转换

我的另一个 PyTorch 版本的实现,含撞线检测,可对行人进行计数,开源地址:Yolov5_Deepsort_Person_Count

一、快速使用

💻 环境配置

参考项目内的 README 使用Docker容器,或参考下方自行配置:

  • python: 3.8
  • cuda: 11.2
  • cudnn: 8.2.2.26
  • tensorRT: 8.0.3.4
  • protobuf: 3.11.4

📥 下载视频&模型

创建工作目录:

mkdir workspace

下载文件到 workspace 目录下:

文件链接
yolov5s.onnx下载 (提取码: tg42)
deepsort.onnx下载 (提取码: iyms)
test.mp4下载 (提取码: vatx)

yolo 的 onnx 导出自 yolov5-6.0,deepsort 的 onnx 导出自 deep_sort_pytorch,可参考脚本 exportOnnx.py

🏃‍ 运行

修改 MakeFile 中的相关头文件和库文件路径(若使用本项目提供的 docker 则不需要),然后执行:

make run

项目运行时将显示推理结果,按 ESC 退出。

目前在 GeForce RTX 2060 上,推理 test.mp4 的速度约为 40 ms/帧,满足实时性要求。

二、文件说明

  • Infer、Yolo 和 DeepSORT 使用接口模式和 RAII 进行封装:

    • infer.hyolo.hdeepsort.h 仅暴露 create_*推理接口
    • 使用 create_* 创建对象实例,将自动解析 onnx 文件,生成 engine 并加载
  • infer.cpp: 分四个线程,两两之间为生产者-消费者关系:

在这里插入图片描述

三、项目参考

本项目构建过程中参考了以下优秀的开源项目,在此对这些项目的作者表示感谢: