TensorRT 部署加速基于 Deepsort + YOLOv5 的目标跟踪

一. 项目简介

  • 基于 TensorRT-v8.2 ,加速YOLOv5-v5.0 + deepsort 的目标跟踪;
  • Jetson nano 上进行部署;
  • Linux x86_64 系统上也是可行的,需要修改本项目中所有的 MakefileCMakeLists.txt ,把这 2 种文件里 TensorRTOpenCV 头文件、库文件的路径换成对应系统的即可,CUDA的路径一般不用修改,最好也确认下。

二. 项目效果

  • 下面是从一段测试视频当中选取的两帧跟踪的结果
    在这里插入图片描述
    在这里插入图片描述

三. 环境配置

  • Jetson nano 上烧录的系统镜像为 Jetpack 4.6.1,该jetpack 原装环境如下:
CUDAcuDNNTensorRTOpenCV
10.28.28.2.14.1.1

关于如何在 Jetson nano 上烧录镜像,网上资料还是很多的,这里就不赘述了,注意下载 Jetpack镜像时选择 4.6.1 版本,该版本对应的 TensorRT v8 版本

  • 安装Eigen
apt install libeigen3-dev

记得在之前将 apt 源更换为国内源

四. 项目代码

五. 模型转换

YOLO检测模型,和ReID特征提取模型,转换成TensorRT的序列化文件,后缀 .plan(作者的习惯,也可以是.engine或其他)

5.1 原模型下载

  • 链接:https://pan.baidu.com/s/1YG-A8dXL4zWvecsD6mW2ug
  • 提取码:y2oz

下载并解压后,

目录中的文件说明如下:

Deepsort模型文件
    |
    ├── ReID模型  # 该目录中存放的是 ReID 特征提取网络的模型
    │   ├── ckpt.t7  # 官方 PyTorch 格式的模型文件
    │   └── deepsort.onnx  # 根据 ckpt.t7 导出的 onnx 格式模型文件
    |
    └── YOLOv5-v5.0  # 该目录中存放的是 YOLOv5 目标检测网络的模型
        ├── yolov5s.pt  # 官方 PyTorch 格式的模型文件
        └── para.wts  # 根据 yolov5s.pt 导出的 wts 格式模型文件

5.2 YOLO模型转换

  • 将上述 yolov5s.pt 转为 model.plan,或 para.wts转为 model.plan
  • 具体的方法步骤参考 TensorRT-YOLOv5,也是作者自己发布的一个项目

注意:使用时,记得把 MakefileTensorRTOpenCV 头文件、库文件所在的路径,换成自己 Jetson nano 上的路径

完成之后便可得到 model.plan ,为检测网络的 TensorRT 序列化模型文件。

5.3 ReID网络模型转换

  • 将上述 ckpt.t7 转为 deepsort.plan,或 deepsort.onnx 转为 deepsort.plan
  • 项目代码链接在下方
  • 按如下步骤运行
# 进入本项目的 reid_torch2trt 目录
cd reid_torch2trt

# 执行命令
# 如果jetson nano上没有安装pytorch,这一步随便在任何有pytorch上的环境运行都可以
python onnx_export.py
# 此命令运行后,ckpt.t7 转为 deepsort.onnx
# 不过作者已给出 deepsort.onnx,因此可以跳过这步,从下面的开始

# 依次执行
make
./trt_export
# 此命令运行后,deepsort.onnx 转为 deepsort.plan

完成之后便可得到 deepsort.plan ,为特征提取网络的 TensorRT 序列化模型文件。

六. 运行项目

  • 开始编译并运行目标跟踪的代码
  • 按如下步骤运行
# 进入本项目的 yolov5-deepsort-tensorrt 目录
cd yolov5-deepsort-tensorrt
mkdir resources
# 把上面转换得到的 2 个 plan 文件复制到目录resources中
cp {TensorRT-v8-YOLOv5-v5.0}/model.plan ./resources
cp ../reid_torch2trt/deepsort.plan ./resources

mkdir test_videos  # 向其中放入测试的视频文件,并命名为 demo.mp4

vim src/main.cpp  # 可以根据自己的需求修改一些配置信息

mkdir build
cd build
cmake ..
make
./yolosort  # 运行后即看到跟踪效果

七. 项目参考

7.1 参考链接

主要参考了下面的项目:

  • https://github.com/RichardoMrMu/yolov5-deepsort-tensorrt
  • https://blog.csdn.net/weixin_42264234/article/details/120152117

7.2 本项目的改动

  • 本项目的代码大部分参考了上方项目的内容,给予了作者很多启发
  • 但作者自己也做了不少变更,具体如下:
  1. 所参考项目ReID网络模型输出为 512 维度的向量,但代码中是 256 维,作者做了修正;

  2. 对于YOLOv5-v5.0目标检测网络:

    2.1 TensorRT 模型转换使用的是作者自己的项目,进行了预处理的 CUDA 加速

    2.2 对 YOLOv5 的 TensorRT 推理做了类的封装,并编译为动态链接库,在其他项目中使用也非常方便;

  3. 实现了类别筛选功能,可以在 main.cpp 中设置想要跟踪的类别,忽略掉那些不关心的种类;

  4. 为了适配 TensorRT 8 版本,对模型推理的地方做了一定的改动。

感谢阅毕,编码整理不易,若有帮助,望不吝star,欢迎点赞收藏。

============================= END ================================

  • 3
    点赞
  • 13
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值