YOLOP可检测类别数量及推理边界框标签显示问题咨询

一、YOLOP可检测的类别数量及范围 #

YOLOP基于BDD100K数据集训练，检测头支持10类自动驾驶相关目标，除车辆类（car、bus、truck）外，还包括：

• 行人（pedestrian）
• 骑行者（rider）
• 自行车（bicycle）
• 摩托车（motorbike）
• 交通灯（traffic light）
• 交通标志（traffic sign）
• 火车（train）

二、修改推理代码显示检测标签 #

要让demo.py在边界框旁显示类别标签，需要修改代码中绘制检测框的逻辑，步骤如下：

1. 定位核心绘图代码
   打开tools/demo.py，找到遍历检测结果并绘制框的代码段，或是调用plot_one_box函数的位置（该函数通常在utils/plots.py中）。

2. 添加标签参数
   原代码仅绘制框未传入标签信息，修改plot_one_box的调用，传入类别名称和置信度组合的标签。示例：

# 原代码（仅绘制框）
plot_one_box(xyxy, im0, color=colors[int(cls)], line_thickness=3)

# 修改后代码（添加标签）
# 定义BDD100K的类别名称列表
names = ['car', 'bus', 'truck', 'pedestrian', 'rider', 'bicycle', 'motorbike', 'traffic light', 'traffic sign', 'train']
# 拼接类别名与置信度，保留两位小数
label = f'{names[int(cls)]} {conf:.2f}'
plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)

3. 验证修改效果
   重新运行推理命令：

python tools/demo.py --source /content/YOLOP/inference/images --device 0

输出图像的边界框旁会显示对应类别名称和置信度。

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