如何用Python从YOLO V5预测图片中提取标签以验证模型精度？

下面分两种常见场景实现标签提取：

场景1：从已生成的预测结果txt文件中提取标签 #

如果你已经通过YOLOv5官方的predict.py完成预测，程序会在runs/detect/exp（或自定义路径）下生成labels文件夹，里面的每个txt文件对应一张图片的预测结果，每行格式为class_id x_center y_center width height confidence。

实现代码 #

import yaml
import os

# 加载训练时的data.yaml，获取类别名称映射
with open('data.yaml', 'r', encoding='utf-8') as f:
    data_config = yaml.safe_load(f)
class_names = data_config['names']  # 索引对应class_id

# 遍历labels文件夹下的所有预测结果文件
labels_dir = 'runs/detect/exp/labels'  # 替换为你的实际labels文件夹路径
for filename in os.listdir(labels_dir):
    if not filename.endswith('.txt'):
        continue
    img_name = filename.replace('.txt', '.jpg')  # 根据你的图片格式调整后缀
    print(f"【{img_name}】的预测标签：")
    with open(os.path.join(labels_dir, filename), 'r') as f:
        for line in f.readlines():
            line = line.strip()
            if not line:
                continue
            parts = line.split()
            class_id = int(parts[0])
            confidence = float(parts[-1])
            print(f"- {class_names[class_id]} (置信度: {confidence:.2f})")

场景2：推理时直接提取标签（无需依赖生成的txt文件） #

如果需要在推理过程中实时提取标签，可直接调用YOLOv5的模型接口获取结构化结果：

实现代码 #

import torch
import os

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

# 单张图片推理并提取标签
img_path = 'test_images/sample.jpg'  # 替换为你的测试图片路径
results = model(img_path)

# 转换为DataFrame格式，方便提取信息
pred_df = results.pandas().xyxy[0]

print(f"【{os.path.basename(img_path)}】的预测标签：")
for _, row in pred_df.iterrows():
    print(f"- {row['name']} (置信度: {row['confidence']:.2f})")

# 可选：将标签保存为YOLO格式的txt文件
save_dir = 'predicted_labels'
os.makedirs(save_dir, exist_ok=True)
save_filename = os.path.splitext(os.path.basename(img_path))[0] + '.txt'
with open(os.path.join(save_dir, save_filename), 'w') as f:
    img_width, img_height = results.imgs[0].shape[1], results.imgs[0].shape[0]
    for _, row in pred_df.iterrows():
        class_id = row['class']
        # 转换为YOLO归一化坐标格式
        x_center = (row['xmin'] + row['xmax']) / 2 / img_width
        y_center = (row['ymin'] + row['ymax']) / 2 / img_height
        width = (row['xmax'] - row['xmin']) / img_width
        height = (row['ymax'] - row['ymin']) / img_height
        confidence = row['confidence']
        f.write(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f} {confidence:.6f}\n")

注意事项 #

• 所有路径需替换为你的实际文件/文件夹路径
• 类别名称映射必须与训练时data.yaml中的names字段完全一致，否则会出现类别名称不匹配
• 批量处理图片时，只需循环遍历图片文件夹，重复上述推理和提取逻辑即可

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