机缘

由于学校SRTP项目需要,笔者首先基于yolov8进行了模型训练,该模型只能进行单帧分析,但后续需要使用StrongSort对视频进行连续帧的分析,经过一早上的折腾终于完成了与StrongSort和训练好模型的结合。
由于网上我几乎没有发现此类教程,所以写下这篇博客,不仅方便后人,也方便以后的自己

收获

  1. 对虚拟环境的认识又加深了
  2. 本地部署能力加强了

废话不多说,开始吧


一、创建虚拟环境

使用anaconda进行虚拟环境的创建和管理,一是可以解决不同项目对环境的依赖,二是可以让每个项目的环境有区分,不那么杂乱

创建环境

conda create -n yolo_strongsort python=3.10.8

此处python的版本与你训练模型的版本保持一致即可

激活环境

conda activate yolo_strongsort

二、安装必要的依赖

1、安装pytorch

pip install torch torchvision -f https://download.pytorch.org/whl/cu124/torch_stable.html

2、安装Ultralytics和OpenCV

pip install ultralytics
pip install opencv-python

三、下载StrongSort代码

git clone https://github.com/mikel-brostrom/yolo_tracking.git
cd yolo_tracking
pip install poetry 
cd D:\Download\Steam\steamapps\common\MyDockFinder\yolo_tracking
poetry install

注意安装依赖已经不再使用requirements了,查询得知这里用的是query


四、查找track.py

我们运行

dir D:\Download\Steam\steamapps\common\MyDockFinder\yolo_tracking\examples /s /b | findstr track.py
dir D:\Download\Steam\steamapps\common\MyDockFinder\yolo_tracking\tracking /s /b | findstr track.py

检查 examples/ 和 tracking/ 目录是否有 track.py,找到后把该路径保存到记事本中(应该是在tracking目录下)


五、继续StrongSort的集成

1.复制模型

copy \path\to\runs\detect\train\weights\best.pt D:\Download\Steam\steamapps\common\MyDockFinder\yolo_tracking\weights\
mkdir D:\Download\Steam\steamapps\common\MyDockFinder\yolo_tracking\weights  # 如果目录不存在

注意这里的路径需要根据你自己的电脑调整

2.创建data.yaml

cd D:\Download\Steam\steamapps\common\MyDockFinder\yolo_tracking
echo names: > data.yaml
echo   - student >> data.yaml
echo   - electronic_device >> data.yaml
echo nc: 2 >> data.yaml

或者直接在文件夹下创建也行,保存为data.yaml

下载ReID权重

1.下载:
打开浏览器,访问:[https://drive.google.com/uc?id=1LaG1EJpHrxdAxKnSCJ_i0u-nbxSAeiFY]
下载到本地(例如 D:\Download\osnet_x1_0_imagenet.pth)。
2. 移动到指定目录

mkdir D:\Download\Steam\steamapps\common\MyDockFinder\yolo_tracking\weights
move D:\Download\osnet_x1_0_imagenet.pth D:\Download\Steam\steamapps\common\MyDockFinder\yolo_tracking\weights\

六、运行(重量级的来了 )

由于版本的更新,会有非常非常多的问题,现在假设我们不知道这些问题
我们运行

poetry run python tracking\track.py --yolo-weights "D:\SRTP\Student model\train\weights\best.pt" --reid-weights weights\osnet_x1_0_imagenet.pth --tracking-method strongsort --source "D:\path\to\classroom_video.mp4" --save-vid

会报错

track.py: error: unrecognized arguments: --yolo-weights weights\best.pt --reid-weights weights\osnet_x1_0_imagenet.pth --save-vid

这表明 tracking\track.py 不认识 --yolo-weights、--reid-weights 和--save-vid参数。原因在于 mikel-brostrom/yolo_tracking 仓库中的 track.py 脚本可能已经更新,参数名称或用法与我之前基于旧版本的建议不一致。我们需要检查脚本支持的参数并调整命令

我们进入track.py的源代码
--yolo-model:替换 --yolo-weights,并且指定 YOLO 模型路径。(找到你训练的best.pt)
--reid-model:替换 --reid-weights,并且指定 ReID 模型路径。
--save:替换 --save-vid,保存跟踪结果。
记住:所有的都要替换

同时,命令改成

poetry run python tracking/track.py --yolo-weights weights\best.pt --reid-weights weights\osnet_x1_0_imagenet.pth --tracking-method strongsort --source "D:\SRTP\东大云课堂_2.mp4" --save-vid

运行后再次报错

ERROR    | tracking.detectors:get_yolo_inferer:53 - Failed to infer inference mode from yolo model name
2025-02-22 12:05:57.116 | 
ERROR    | tracking.detectors:get_yolo_inferer:54 - Your model name has to contain either yolox, yolo_nas or yolov8= parse_opt() 
    run(opt)

为什么呢,因为get_yolo_inferer 期望模型的文件名中包含关键字,例如 yolov8、yolox 或 yolo_nas,但你的模型文件 weights\best.pt 可能没有这些标识。

我们更改文件名为yolov8_best.pt
我们再次运行
再次报错

AttributeError: 'StrongSort' object has no attribute 'plot_results'

问题出在 StrongSort tracker 上,因为它没有定义 plot_results 方法。这个方法通常用于绘制和显示跟踪结果。
有没有什么办法呢,我们在github上查询,发现StrongSort本身就不支持这个函数,好吧,只能自己写代码了
在这里插入图片描述
我们把源码中有这个函数的那一行代码替换为

for r in results:
        if args.frame_limit and frame_count >= args.frame_limit:
            break
        img = r.orig_img.copy()
        
        # 这里不再使用 Annotator 类
        if r.boxes is not None and hasattr(r.boxes, 'id') and r.boxes.id is not None:
            for box, track_id in zip(r.boxes.xyxy, r.boxes.id):
                # 直接使用 cv2 绘制边框,设置最小线宽
                cv2.rectangle(
                    img, 
                    (int(box[0]), int(box[1])),
                    (int(box[2]), int(box[3])),
                    colors(int(track_id)),
                    1,  # 固定线宽为1
                    cv2.LINE_AA  # 使用抗锯齿
                )
                
                # 如果需要显示ID,使用较小的字体大小
                label = f'{int(track_id)}'
                cv2.putText(
                    img, 
                    label,
                    (int(box[0]), int(box[1]) - 5),  # 稍微调整了文字位置
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.5,  # 更小的字体大小
                    colors(int(track_id)),
                    1  # 文字线宽也设为1
                )

        if args.show:
            cv2.imshow('BoxMOT', img)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

再次运行,终于结束了

结语

做这些前沿一些的东西没人带真的很难受,所以我写下了这篇博客,当然,里面有很多东西仍旧不准确,路径也需要读者自己去替换。但是我还是希望能对大家有一点帮助

如果有什么问题可以在评论区提出,大家可以一起交流讨论。
如果有大神能改善一下代码,让代码能只分析某个视频的前多少帧就好了

更多推荐