如何用Python手动实现YOLOv8混淆矩阵及PR曲线计算

1. 提取模型预测与真实标注数据 #

首先加载YOLOv8模型并获取测试集的预测结果与真实标注：

from ultralytics import YOLO
import numpy as np

# 加载预训练模型
model = YOLO('yolov8n.pt')

# 假设你已经有测试集的图片路径和对应的标注（xyxy格式+类别ID）
# 示例：每个样本的真实标注格式为(boxes: np.ndarray, cls: np.ndarray)
test_data = [
    ("test_img1.jpg", (np.array([[10,20,50,60], [30,40,70,80]]), np.array([0, 1]))),
    # 更多测试样本...
]

# 批量获取预测结果，整理为统一格式
preds = []
targets = []
for img_path, (target_boxes, target_cls) in test_data:
    results = model(img_path)
    pred_boxes = results[0].boxes.xyxy.numpy()
    pred_scores = results[0].boxes.conf.numpy()
    pred_cls = results[0].boxes.cls.numpy()
    preds.append((pred_boxes, pred_scores, pred_cls))
    targets.append((target_boxes, target_cls))

2. 实现IoU计算函数 #

IoU是匹配预测框与真实框的核心指标，直接用numpy实现：

def calculate_iou(box1, box2):
    # box格式：(x1, y1, x2, y2)
    x1 = max(box1[0], box2[0])
    y1 = max(box1[1], box2[1])
    x2 = min(box1[2], box2[2])
    y2 = min(box1[3], box2[3])
    
    inter_area = 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_area = box1_area + box2_area - inter_area
    
    return inter_area / union_area if union_area > 0 else 0.0

3. 按conf阈值计算混淆矩阵 #

混淆矩阵行代表真实类别，列代表预测类别，最后一行/列用于统计未匹配的FP/FN：

def compute_confusion_matrix(preds, targets, conf_thresh, iou_thresh=0.5, num_classes=80):
    conf_matrix = np.zeros((num_classes + 1, num_classes + 1), dtype=np.int32)
    
    for pred, target in zip(preds, targets):
        pred_boxes, pred_scores, pred_cls = pred
        # 筛选当前conf阈值以上的预测框
        mask = pred_scores >= conf_thresh
        pred_boxes = pred_boxes[mask]
        pred_cls = pred_cls[mask].astype(np.int32)
        
        target_boxes, target_cls = target
        target_cls = target_cls.astype(np.int32)
        target_matched = np.zeros(len(target_boxes), dtype=bool)
        
        # 按预测得分从高到低匹配，避免重复匹配
        sorted_indices = np.argsort(pred_scores[mask])[::-1]
        for idx in sorted_indices:
            box = pred_boxes[idx]
            cls = pred_cls[idx]
            # 找同类别未匹配的真实框
            same_cls_mask = (target_cls == cls) & (~target_matched)
            if not np.any(same_cls_mask):
                conf_matrix[num_classes, cls] += 1  # 未匹配的FP
                continue
            # 计算IoU并匹配
            ious = np.array([calculate_iou(box, tb) for tb in target_boxes[same_cls_mask]])
            if np.max(ious) >= iou_thresh:
                m_idx = np.argmax(ious)
                original_idx = np.where(same_cls_mask)[0][m_idx]
                target_matched[original_idx] = True
                conf_matrix[target_cls[original_idx], cls] += 1  # TP
            else:
                conf_matrix[num_classes, cls] += 1  # 低IoU的FP
        
        # 统计未被匹配的真实框（FN）
        for idx in np.where(~target_matched)[0]:
            cls = target_cls[idx]
            conf_matrix[cls, num_classes] += 1
    
    return conf_matrix

4. 计算Precision、Recall与PR曲线 #

遍历0到1的conf阈值区间，计算每个阈值对应的precision和recall：

def compute_pr_curve(preds, targets, num_classes=80, iou_thresh=0.5):
    conf_thresholds = np.linspace(0, 1, 101)  # 生成0.00到1.00的101个阈值
    precision = []
    recall = []
    
    for conf_thresh in conf_thresholds:
        tp, fp, fn = 0, 0, 0
        for pred, target in zip(preds, targets):
            pred_boxes, pred_scores, pred_cls = pred
            mask = pred_scores >= conf_thresh
            pred_boxes = pred_boxes[mask]
            pred_cls = pred_cls[mask].astype(np.int32)
            
            target_boxes, target_cls = target
            target_cls = target_cls.astype(np.int32)
            target_matched = np.zeros(len(target_boxes), dtype=bool)
            
            # 匹配预测框
            for box, cls in zip(pred_boxes, pred_cls):
                same_cls_mask = (target_cls == cls) & (~target_matched)
                if not np.any(same_cls_mask):
                    fp += 1
                    continue
                ious = np.array([calculate_iou(box, tb) for tb in target_boxes[same_cls_mask]])
                if np.max(ious) >= iou_thresh:
                    m_idx = np.argmax(ious)
                    original_idx = np.where(same_cls_mask)[0][m_idx]
                    target_matched[original_idx] = True
                    tp += 1
                else:
                    fp += 1
            
            # 统计FN
            fn += np.sum(~target_matched)
        
        # 计算当前阈值的precision和recall
        current_precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
        current_recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0
        precision.append(current_precision)
        recall.append(current_recall)
    
    return conf_thresholds, np.array(precision), np.array(recall)

绘制PR曲线：

import matplotlib.pyplot as plt

conf_thresholds, precision, recall = compute_pr_curve(preds, targets)
plt.plot(recall, precision, linewidth=2)
plt.xlabel('Recall')
plt.ylabel('Precision')
plt.title('YOLOv8 Precision-Recall Curve')
plt.grid(alpha=0.3)
plt.show()

5. 确定最优conf阈值 #

最优阈值通常选择F1分数最高的点，F1是precision和recall的调和平均数：

f1_scores = 2 * (precision * recall) / (precision + recall + 1e-8)  # 加极小值避免除零
best_idx = np.argmax(f1_scores)
best_conf_thresh = conf_thresholds[best_idx]
best_f1 = f1_scores[best_idx]

print(f"最优conf阈值: {best_conf_thresh:.2f}")
print(f"对应F1分数: {best_f1:.2f}")

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