如何从YOLOv5中提取检测标签并存储到数组中

1. 仓库克隆与环境配置 #

# 克隆YOLOv5仓库
!git clone https://github.com/ultralytics/yolov5
%cd yolov5
# 安装依赖
%pip install -qr requirements.txt
%pip install -q roboflow

import torch
import os
from IPython.display import Image, clear_output

print(f"Setup complete. Using torch {torch.__version__} ({torch.cuda.get_device_properties(0).name if torch.cuda.is_available() else 'CPU'})")

2. 训练命令 #

!python train.py --img 597 --batch 16 --epochs 100 --data {dataset.location}/data.yaml --weights yolov5s.pt --cache

3. 原始检测命令 #

!python detect.py --weights runs/train/exp/weights/best.pt --img 416 --conf 0.1 --source /content/IMG-20221220-WA0031.jpg

方法一：通过YOLOv5 Python API直接获取 #

这种方式无需依赖命令行输出，可直接在代码中获取检测结果并提取标签，适合程序化处理：

# 加载自定义训练好的模型
model = torch.hub.load('ultralytics/yolov5', 'custom', path='runs/train/exp/weights/best.pt', force_reload=True)

# 设置检测参数
model.conf = 0.1  # 置信度阈值
model.imgsz = 416  # 输入图像尺寸

# 执行检测
img_path = '/content/IMG-20221220-WA0031.jpg'
results = model(img_path)

# 提取标签到数组
detected_labels = []
for pred in results.pred:
    # 遍历每个检测框的类别索引，转换为标签名称
    for cls_idx in pred[:, 5]:
        label = model.names[int(cls_idx)]
        detected_labels.append(label)

# 输出结果
print("检测标签数组：", detected_labels)

方法二：修改detect.py脚本实现 #

如果需要在原detect.py执行流程中提取标签，可修改脚本中的结果处理逻辑：

1. 打开yolov5目录下的detect.py文件，定位到run()函数内处理检测结果的代码段。
2. 添加以下代码提取并存储标签：

# 在结果处理区域添加
detected_labels = []
for pred in results.pred:
    for cls_idx in pred[:, 5]:
        label = model.names[int(cls_idx)]
        detected_labels.append(label)

# 打印标签数组
print("检测到的标签：", detected_labels)
# 可选：保存到文件
with open('detected_labels.txt', 'w') as f:
    f.write('\n'.join(detected_labels))

修改后运行原检测命令，即可在控制台看到标签数组，或在当前目录下的detected_labels.txt文件中查看结果。

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