目标检测 AP 计算实例 python

以下是使用 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}")