如何使用YOLO获取测试图像的预测值？新手技术求助

原代码存在的问题 #

• 图像路径错误：未拼接具体文件名，无法正确加载单张测试图
• 手动预处理冗余：Ultralytics YOLO模型支持直接传入图片路径/PIL图像，自动完成resize、归一化、转张量等预处理步骤，无需手动操作
• 预测结果提取方式错误：模型返回的是Results对象，不是字典，不能用["boxes"]这种键值对方式访问
• 未定义变量：print(path)中的path变量未初始化

修正后的完整代码 #

from ultralytics import YOLO
import pandas as pd

# 初始化模型
model = YOLO("runs/detect/train16/weights/best.pt")

# 假设test是你的测试数据DataFrame，若从文件读取需替换为实际路径
# test = pd.read_csv("测试数据文件路径.csv")

for index, row in test.iterrows():
    img_name = row['img_fName']
    # 拼接完整图像路径，确保test_images_phase1是测试图所在文件夹
    img_path = f'test_images_phase1/{img_name}'
    
    # 直接传入图像路径，模型自动完成所有预处理
    results = model(img_path)
    
    # 提取当前图像的预测结果
    result = results[0]
    # 获取检测框、置信度、类别ID
    boxes = result.boxes.xyxy.tolist()  # xyxy格式（左上x,左上y,右下x,右下y）
    scores = result.boxes.conf.tolist()
    cls_ids = result.boxes.cls.tolist()
    class_names = model.names  # 类别ID与名称的映射表
    
    # 打印预测信息
    print(f"Image: {img_name}")
    for box, score, cls_id in zip(boxes, scores, cls_ids):
        x1, y1, x2, y2 = box
        class_name = class_names[int(cls_id)]
        print(f"Bounding Box: {x1:.2f}, {y1:.2f}, {x2:.2f}, {y2:.2f}")
        print(f"Class: {class_name}")
        print(f"Confidence: {score:.4f}")
        print()

关键说明 #

• 路径校验：确认test_images_phase1文件夹路径正确，若测试图存放在其他位置，需修改img_path的拼接逻辑
• 自动预处理：模型会按照训练时的600像素输入尺寸，自动对测试图进行缩放、填充等适配操作，无需手动调整图像大小
• 结果过滤：若需过滤低置信度的检测结果，可在推理时添加参数，例如model(img_path, conf=0.5)仅保留置信度≥0.5的预测框
• 类别映射：model.names是训练时定义的类别名称字典，直接通过类别ID即可获取对应名称

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