如何解决YOLOv5同时使用两个权重文件检测视频时标签异常问题

问题原因 #

你训练的两个单类别模型（pen.pt、cap.pt），在单独训练时各自的类别索引都是0，对应的类别名分别是pen和cap。但同时加载多个模型时，YOLOv5默认会采用第一个模型的类别名称列表，导致所有检测结果都被标注为cap。

无需重新训练的解决方法 #

方法1：修改模型类别名称映射 #

直接修改detect.py中模型加载后的类别名称配置，让每个模型对应自己的标签：

1. 打开detect.py，找到模型加载的代码段（通常在models = attempt_load(weights, map_location=device)附近）
2. 在加载模型后添加以下代码，手动指定每个模型的类别名：

   models = attempt_load(weights, map_location=device)
   # 按加载顺序指定模型类别：第一个是cap.pt，第二个是pen.pt
   models[0].names = ['cap']
   models[1].names = ['pen']
   

3. 保存修改后，重新执行原检测命令即可。

方法2：分开检测后合并结果 #

分别用两个模型单独检测视频，再将结果合并可视化：

1. 用cap.pt检测视频，保存结果到单独目录：

   ! python detect.py --weights cap.pt --img 640 --conf 0.50 --source VID_20220727_185703.mp4 --name cap_detect
   

2. 用pen.pt检测视频，保存结果到另一目录：

   ! python detect.py --weights pen.pt --img 640 --conf 0.50 --source VID_20220727_185703.mp4 --name pen_detect
   

3. 编写简单脚本，读取两个检测结果的视频或标签文件，在原视频帧上同时绘制两个类别的检测框和对应标签。

方法3：给模型预测结果添加类别偏移 #

通过修改检测代码，让第二个模型的类别索引偏移，匹配自定义的类别列表：

1. 打开detect.py，找到预测处理的代码段
2. 在每个模型预测后，给第二个模型的类别索引加1，并设置全局类别名称：

   # 先定义全局类别列表
    names = ['cap', 'pen']
    # 遍历模型进行预测
    for i, model in enumerate(models):
        pred = model(im, augment=augment)[0]
        # 第二个模型（pen.pt）的类别索引偏移为1
        if i == 1:
            pred[:, 5] += 1
        # 非极大值抑制
        pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
        # 后续结果处理...
   

3. 确保后续可视化时使用自定义的names列表，即可正确显示两个标签。

内容的提问来源于stack exchange，提问作者sourab maity