基于 YOLOv5-DeepSort 的目标追踪器的封装
由于项目需要目标追踪,之前打算是用 OpenCV 提供的算法进行追踪,但实测下来效果不是很理想,了解到使用 YOLOv5 与 DeepSort 相结合的方式可以进行多物体追踪,在跑通作者提供的 track 示例后,将目标检测和目标追踪分别进行封装,方便以后在项目中使用。
原始 track 介绍
我们阅读 track.py 中代码,从中可以看出来代码包含以下部分
- 参数配置(YOLOv5 和 DeepSort)
- 初始化 DeepSort(DeepSort)
- 加载检测 model(YOLOv5)
- 进行推理,获得相关信息(YOLOv5)
- 进行追踪(DeepSort)
- 结果显示的处理(YOLOv5 和 DeepSort)
按照上述流程我们可以对检测和追踪分别进行封装。
参数配置
主方法中关于parser
的所有语句都是加载命令行参数【219-241 行】,我们如果要应用到项目中很少会用到命令行提供参数,所以我们需要写一个能提供参数的方法或类。
detect
方法中刚开始就是对 YOLOv5 和 DeepSort 的相关参数进行配置【47-55 行】
初始化 DeepSort
【54-62 行】是获得 DeepSort 的配置文件,并构造 deepsort 对象
加载检测 model
【77-104 行】加载相关的 model
进行推理
对以及打包的原始图像进行处理【91-95 行】(对图片信息打包)【113-118 行】(对图片进行浮点数、归一化处理),然后进行推理,之后进行非极大抑制,对所得的结果进行处理【129-150 行】
进行追踪
调用update
方法进行追踪。【154】
注意!追踪更新的返回 结果都是 tensor
最终展示结果处理
对获得的物体的 ID、类别名称(cls)、置信度(confidence)、boundingboxes(xyxy)进行最终处理【157-167】
!
对于保存视频、文件等相关代码本文未讲述到。
如何运行
-
克隆此项目
git clone https://github.com/jimyag/YOLOv5-DeepSort.git
-
安装依赖
在终端执行以下语句安装依赖包
pip install -r requirements.txt
requirements
默认安装 torch-cpu,如需安装 pytorch-gup 请移步至pytorch 下载选择相应版本进行下载 -
下载权重文件
-
下载 DeepSort 权重
点击下载,并将下载好的权重文件放在
deep_sort_pytorch/deep_sort/deep/checkpoint
下。如果要使用其他权重,需要修改
deep_sort_pytorch/configs/deep_sort.yaml
中REID_CKPT: "deep_sort_pytorch/deep_sort/deep/checkpoint/ckpt.t7"
的配置文件 -
下载 YOLOv5 权重
点击下载,并将下载好的权重文件放在
yolov5/weights
,并修改my_detect.py
中self.yolo_weights = 'yolov5/weights/basketball_robot.pt'
-
-
更改源文件
在
my_detect.py
中修改self.source = 'basketball.mp4'
-
运行
test.py