YOLOv5+ByteTrack统计人流、车流

哈喽，大家好。

今天给大家分享一个人流、车流统计系统。

以人流统计为例，实现这样的系统，需要两个基本的步骤，第一步是识别道路上的行人，可以采用目标检测模型。第二步是给识别出的每个行人分配一个唯一的id，这样我们可以跟踪到每一个行人，计算他的行走方向，行走速度等。

源码和数据集获取方式放在文末，大家自行获取即可。

1. YOLOv5 训练行人识别模型

之前分享过很多篇 YOLOv5 训练目标检测模型的方式，这里再简单说一下。

如果使用的数据没有标注，可以用用LabelImage进行标注

我提供的训练数据集是已经标注好了的，大家可以直接用。

标注好的数据集，按照图片和标注文件分别存放。

数据集目录

下载YOLOv5源码到本地

git clone https://github.com/ultralytics/yolov5.git

在yolov5目录下，修改data/coco128.yaml配置文件

path: ../datasets/people  # 数据集目录
train: images/train  # 训练集
val: images/train  # 验证集

# Classes
names:
  0: person

这里我们只训练行人模型，所以names只有1个。

在yolov5目录下，创建weights目录，下载预训练模型。

yolov5预训练模型

下载后，放到weights文件夹中，这里我用的是 yolov5s.pt。

修改models/yolov5s.yaml文件中的分类数量。

# YOLOv5  by Ultralytics, GPL-3.0 license

# Parameters
nc: 1  # number of classes

我们只识别行人，设置成1即可。

执行一下命令进行训练

python ./train.py --data ./data/coco128.yaml --cfg ./models/yolov5s.yaml --weights ./weights/yolov5s.pt --batch-size 30 --epochs 120 --workers 8 --name base_s --project yolo_people

训练完成后，可以查看训练效果

yolov5训练效果

生成的模型放存放在yolo_people/base_s/weights/best.pt中，后面可以直接用它来做推理。

1. ByteTrack 行人跟踪

识别出行人，我们可以利用多目标跟踪技术(MOT)技术来跟踪行人，并给每个行人分配唯一的ID。

算法思路为：

算法流程

关键思路是用卡尔曼滤波预测当前帧的跟踪轨迹在下一帧的位置，预测框和实际框之间通过匈牙利算法，用 IoU 进行快速相似度匹配。

MOT的方案有很多，如：SORT、DeepSORT、ByteTrack、BoT-SORT等等。

与DeepSORT不同的是，ByteTrack没有使用 ReID 特征计算表观相似度，这样做的目的，一是尽可能做到简单高速，二是检测结果足够好的情况下，卡尔曼滤波的预测准确性已经比较高了，能够代替ReID。

所以，ByteTrack比较依赖目标检测的准确度的。

下面是在VisDrone2019-MOT数据集训练约 10 epochs, 采用YOLO v7 w6结构, COCO预训练模型基础上训练。对比几个常见的MOT方案的效果

MOT方案效果对比

指标解释如下：

• MOTA：多目标跟踪准确度，数值越高代表跟踪精确度越好
• IDF1：被检测和跟踪的目标中，获取正确ID的比例，综合考虑准召，是 F1 score
• IDS：id switch次数
• fps：帧率

权衡准确度和性能，我选择了 ByteTrack 作为本项目的多目标追踪方案。

各种MOT追踪的API大致类似，先准备目标检测框

box_list = yolo_pd.to_numpy()
detections = []
for box in box_list:
    l, t = int(box[0]), int(box[1])
    r, b = int(box[2]), int(box[3])

    conf = box[4]

    detections.append([l, t, r, b, conf])

这里将识别出的行人检测框，转为numpy结构。

sys.path.append('../../github/ByteTrack/')
from yolox.tracker.byte_tracker import BYTETracker, STrack

@dataclass(frozen=True)
class BYTETrackerArgs:
    track_thresh: float = 0.25
    track_buffer: int = 30
    match_thresh: float = 0.8
    aspect_ratio_thresh: float = 3.0
    min_box_area: float = 1.0
    mot20: bool = False
    
byte_tracker = BYTETracker(BYTETrackerArgs())

tracks = byte_tracker.update(
        output_results=np.array(detections, dtype=float),
        img_info=frame.shape,
        img_size=frame.shape
    )

调用ByteTrack的update函数进行匹配，匹配后会给每一个检测框一个唯一的ID。

主要思路和核心代码就是这些，基于此在做一些工程编码就可以实现文章开头的效果。