YOLOv8模型训练后，如何在PyTorch中计算IOU分数

方法一：用YOLOv8自带验证命令批量计算 #

YOLOv8内置了完整的评估工具，能直接计算测试集的各类指标（包括IOU），步骤简单：

• 确保你的测试集对应标注文件存放在dataset/labels/test/目录下，格式符合YOLO的class_id x_center y_center width height要求
• 运行以下命令：

!yolo val model=/content/runs/detect/train/weights/best.pt data=你的数据集配置.yaml split=test

• data参数需指定你的数据集配置yaml文件，文件里要包含测试集的图像、标注路径等信息
• 执行完成后，/content/runs/detect/val目录会生成评估报告，里面包含各类别的IoU@0.5、IoU@0.5:0.95以及平均IOU等数据

方法二：手动计算单张图片的IOU #

如果只需要计算单张测试图的IOU，按以下步骤操作：

1. 加载模型并获取预测框 #

from ultralytics import YOLO

# 加载训练好的模型
model = YOLO('/content/runs/detect/train/weights/best.pt')
# 对目标图片做预测
results = model('/content/dataset/images/test/004250005.jpg')
# 提取预测框坐标（格式：x1, y1, x2, y2，左上角和右下角）
pred_boxes = results[0].boxes.xyxy

2. 加载真实标注框 #

把YOLO格式的标注转换为x1,y1,x2,y2的坐标格式：

def load_gt_boxes(label_path, img_shape):
    gt_boxes = []
    img_h, img_w = img_shape
    with open(label_path, 'r') as f:
        for line in f.readlines():
            parts = line.strip().split()
            x_center = float(parts[1]) * img_w
            y_center = float(parts[2]) * img_h
            width = float(parts[3]) * img_w
            height = float(parts[4]) * img_h
            # 转换为x1,y1,x2,y2
            x1 = x_center - width/2
            y1 = y_center - height/2
            x2 = x_center + width/2
            y2 = y_center + height/2
            gt_boxes.append([x1, y1, x2, y2])
    return gt_boxes

# 获取图像尺寸
img_shape = results[0].orig_img.shape[:2]
# 加载对应标注文件
gt_boxes = load_gt_boxes('/content/dataset/labels/test/004250005.txt', img_shape)

3. 计算IOU #

可以用PyTorch内置工具，也可以手动实现：

方式A：用torchvision内置函数

from torchvision.ops import box_iou
import torch

# 转换为张量格式
pred_tensor = torch.tensor(pred_boxes)
gt_tensor = torch.tensor(gt_boxes)

# 计算IOU矩阵（每行对应一个预测框，每列对应一个真实框）
iou_matrix = box_iou(pred_tensor, gt_tensor)
# 获取每个预测框的最大匹配IOU
max_iou_per_pred = iou_matrix.max(dim=1)[0]

print("每个预测框的最大IOU值：", max_iou_per_pred.numpy())
print("单张图片的平均IOU：", max_iou_per_pred.mean().item())

方式B：手动实现IOU计算

def calculate_iou(box1, box2):
    # 计算交集区域的坐标
    x1 = max(box1[0], box2[0])
    y1 = max(box1[1], box2[1])
    x2 = min(box1[2], box2[2])
    y2 = min(box1[3], box2[3])
    
    # 计算交集面积
    intersection = max(0, x2 - x1) * max(0, y2 - y1)
    # 计算并集面积
    box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1])
    box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1])
    union = box1_area + box2_area - intersection
    
    # 避免除以0的情况
    return intersection / union if union > 0 else 0.0

# 遍历所有预测框，计算每个框的最大匹配IOU
for pred_box in pred_boxes.numpy():
    max_iou = 0.0
    for gt_box in gt_boxes:
        current_iou = calculate_iou(pred_box, gt_box)
        if current_iou > max_iou:
            max_iou = current_iou
    print(f"该预测框的最大IOU：{max_iou:.4f}")

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