极智项目 | 基于YOLOv5+DeepSort实现的车流量统计算法

news2024/9/20 22:00:20

欢迎关注我的公众号「极智视界」，获取我的更多技术分享

大家好，我是极智视界，本文分享一下基于YOLOv5+DeepSort实现的车流量统计算法。

车流量统计的实现一般原理是跟踪画面上检测出来的车辆，并给每个框标上了序号，当有一个方框跨过检测线时，计数器就会+1，这其实跟人流量统计的算法实现原理是一样的。所以这种类型的项目一般就会用到两种算法，首先是目标检测算法，先把待统计流量的目标给检测出来，然后是跟踪算法，如果不使用跟踪算法的话，就会产生重复计数的问题。

这个项目使用主流的深度学习框架 Pytorch + YOLOv5 目标检测算法 + DeepSort 目标跟踪算法来实现，项目的特点是支持自定义多条检测线、支持高密车流场景的流量计数、能够一键执行推理，能够适应复杂的车流路况。项目提供完整的代码，包括推理代码、一键执行脚本、训练好的模型权重、项目三方依赖库 (requirements.txt)、、待检测的测试视频、检测后的效果视频等。

在拿到项目代码后，进行开发环境的配置，老习惯用 Anaconda3 管理 py 环境，如下，

# 安装 Anaconda3，过程略 ~

# 创建conda py环境
conda create -n aipro_py39 python=3.9

# 激活conda py环境
conda activate aipro_py39

然后自行安装 pytorch，然后直接运行 run.sh一键执行脚本即可进行三方依赖安装 + 执行训练 + 执行推理测试。run.sh 的内容很简单，如下：

#!/bin/bash

# 安装三方依赖，先自己安装好pytorch
pip install -r requirements.txt

# 执行车流量统计推理
python count.py

可以看到，这里的 count.py 是主要的执行函数，在这个py脚本中，支持一些客制化的配置，具体可以设置以下参数：

source_dir : 要打开的视频文件。若要调用摄像头，需要设置为字符串'0'，而不是数字0，按q退出播放
output_dir : 要保存到的文件夹
show_video : 运行时是否显示
save_video : 是否保存运行结果视频
save_text :  是否保存结果数据到txt文件中，将会保存两个文本文件：result.txt和number.txt。result.txt的格式是(帧序号,框序号,框到左边距离,框到顶上距离,框横长,框竖高,-1,-1,-1,-1)，number.txt的内容是统计到第几帧时每条线沿两个方向的跨线物体数

class_list : 要检测的类别序号，在coco_classes.txt中查看（注意是序号不是行号），可以有一个或多个类别

lines : 定义检测线的两个端点的xy坐标、颜色、粗细，可以定义多条检测线
point_idx : 方框的检测点位置(0, 1, 2, 3, 4)，看下边的图，当一个方框的检测点跨过检测线时，统计数会+1

执行的检测结果保存在 inference/output 文件夹下，会输出三个文件，分别是 number.txt、results.txt 和结果视频 test.mp4。完整的车辆统计效果可以直接观看 `inference/output/rest.mp4`，下面展示了部分效果截图。

预测结果保存在 res 文件夹下，下面展示医学图像分割效果，分割结果以 mask 二值图的形式展示，部分分割结果如下 (左边为原图，右边为分割图)，可以看到即使是边界不太清楚的结构也是可以分割的比较清楚，且项目在增加 + 丰富训练数据集后，分割效果会进一步提升。