8

YOLOv5+ByteTrack统计人流、车流

 1 year ago
source link: https://www.51cto.com/article/743776.html
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.
neoserver,ios ssh client

YOLOv5+ByteTrack统计人流、车流

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

哈喽,大家好。

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

以人流统计为例,实现这样的系统,需要两个基本的步骤,第一步是识别道路上的行人,可以采用目标检测模型。第二步是给识别出的每个行人分配一个唯一的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。

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


About Joyk


Aggregate valuable and interesting links.
Joyk means Joy of geeK