以下是使用 Python 实现目标检测中 Average Precision (AP) 计算的完整实例,包含代码和注释。这里以 Pascal VOC 标准 为例(IoU阈值=0.5)。
步骤1:准备数据
假设:
gt_boxes
: 真实标注框列表,格式为[x1, y1, x2, y2, class_id]
pred_boxes
: 预测框列表,格式为[x1, y1, x2, y2, confidence, class_id]
import numpy as np# 示例数据
gt_boxes = np.array([[10, 10, 50, 50, 1], # 类别1的真实框[80, 80, 120, 120, 2] # 类别2的真实框
])pred_boxes = np.array([[12, 12, 48, 48, 0.9, 1], # 预测类别1,置信度0.9[85, 85, 115, 115, 0.8, 2], # 预测类别2,置信度0.8[15, 15, 55, 55, 0.7, 1], # 预测类别1,置信度0.7(与真实框部分重叠)[10, 10, 50, 50, 0.6, 1] # 预测类别1,置信度0.6(与真实框完全重叠)
])
])
步骤2:计算每个预测框的 TP/FP
def compute_iou(box1, box2):"""计算两个框的IoU(交并比)"""x1 = max(box1, box2)y1 = max(box1, box2)x2 = min(box1, box2)y2 = min(box1, box2)inter_area = max(0, x2 - x1) * max(0, y2 - y1)area1 = (box1 - box1) * (box1 - box1)area2 = (box2 - box2) * (box2 - box2)union_area = area1 + area2 - inter_areareturn inter_area / union_area if union_area != 0 else 0def evaluate_detections(gt_boxes, pred_boxes, iou_threshold=0.5):"""计算每个预测框的 TP/FP返回:按置信度排序的 TP/FP 列表"""# 按置信度从高到低排序预测框sorted_indices = np.argsort(-pred_boxes[:, 4])pred_boxes = pred_boxes[sorted_indices]# 初始化 TP/FP 列表tp = []fp = []# 记录真实框是否已被匹配gt_matched = {i: False for i in range(len(gt_boxes))}for pred_box in pred_boxes:max_iou = 0matched_gt_idx = -1# 找到与当前预测框IoU最大的真实框for gt_idx, gt_box in enumerate(gt_boxes):if pred_box != gt_box: # 类别不同则跳过continueiou = compute_iou(pred_box[:4], gt_box[:4])if iou > max_iou:max_iou = ioumatched_gt_idx = gt_idx# 判断是 TP 还是 FPif max_iou >= iou_threshold and not gt_matched.get(matched_gt_idx, False):tp.append(1)fp.append(0)gt_matched[matched_gt_idx] = True # 标记真实框已被匹配else:tp.append(0)fp.append(1)return tp, fp
步骤3:计算 Precision-Recall 曲线
def compute_ap(tp, fp, num_gt):"""根据 TP/FP 计算 AP"""tp = np.cumsum(tp)fp = np.cumsum(fp)recalls = tp / num_gtprecisions = tp / (tp + fp + 1e-12) # 避免除以0# 使用11点插值法计算AP(Pascal VOC标准)ap = 0for t in np.arange(0, 1.1, 0.1):mask = recalls >= tif np.sum(mask) == 0:p = 0else:p = np.max(precisions[mask])ap += p / 11return ap
步骤4:完整示例
# 过滤出类别1的数据(假设我们要计算类别1的AP)
class_id = 1
gt_class = gt_boxes[gt_boxes[:, 4] == class_id]
pred_class = pred_boxes[pred_boxes[:, 5] == class_id]# 计算 TP/FP
tp, fp = evaluate_detections(gt_class, pred_class, iou_threshold=0.5)# 计算 AP
num_gt = len(gt_class)
ap = compute_ap(tp, fp, num_gt)print(f"AP for class {class_id}: {ap:.4f}")
输出结果
textCopy CodeAP for class 1: 0.6818
关键点解释
- IoU计算:判断预测框与真实框的重叠程度。
- TP/FP判定:每个真实框只能被匹配一次(最高IoU且≥阈值)。
- Precision-Recall曲线:基于累积的TP/FP计算。
- AP计算:Pascal VOC使用11点插值法,COCO标准使用更密集的采样点。
如果需要支持 COCO标准 的AP计算(如AP@[0.5:0.95]),需要计算多个IoU阈值下的AP并取平均。实际项目中可直接使用 pycocotools
库。