object-detection-metric

import numpy as np
from collections import defaultdict

def calculate_iou(box1, box2):
    # IoU 계산 함수
    x1, y1, w1, h1 = box1
    x2, y2, w2, h2 = box2

    inter_x1 = max(x1, x2)
    inter_y1 = max(y1, y2)
    inter_x2 = min(x1 + w1, x2 + w2)
    inter_y2 = min(y1 + h1, y2 + h2)

    inter_area = max(0, inter_x2 - inter_x1) * max(0, inter_y2 - inter_y1)
    union_area = w1 * h1 + w2 * h2 - inter_area

    return inter_area / union_area if union_area > 0 else 0


def compute_ap_pascal_voc(recall, precision):
    # PASCAL VOC 스타일: 11-point 보간법
    ap = 0.
    for t in np.arange(0., 1.1, 0.1):
        if np.sum(recall >= t) == 0:
            p = 0
        else:
            p = np.max(precision[recall >= t])
        ap = ap + p / 11.
    return ap

def compute_ap_coco(recall, precision):
    # COCO 스타일: 모든 포인트를 사용한 적분
    mrec = np.concatenate(([0.], recall, [1.]))
    mpre = np.concatenate(([0.], precision, [0.]))

    for i in range(mpre.size - 1, 0, -1):
        mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])

    i = np.where(mrec[1:] != mrec[:-1])[0]
    ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
    return ap

def calculate_ap(gt_boxes, pred_boxes, iou_threshold, ap_method):
    # AP 계산 (공통 로직)
    num_gt = len(gt_boxes)
    num_pred = len(pred_boxes)

    pred_boxes = sorted(pred_boxes, key=lambda x: x[4], reverse=True)

    tp = np.zeros(num_pred)
    fp = np.zeros(num_pred)

    matched = [False] * num_gt

    for i, pred in enumerate(pred_boxes):
        best_iou = 0
        best_gt_idx = -1

        for j, gt in enumerate(gt_boxes):
            if not matched[j]:
                iou = calculate_iou(pred[:4], gt[:4])
                if iou > best_iou:
                    best_iou = iou
                    best_gt_idx = j

        if best_iou >= iou_threshold:
            if not matched[best_gt_idx]:
                tp[i] = 1
                matched[best_gt_idx] = True
            else:
                fp[i] = 1
        else:
            fp[i] = 1

    tp_cumsum = np.cumsum(tp)
    fp_cumsum = np.cumsum(fp)

    recalls = tp_cumsum / num_gt
    precisions = tp_cumsum / (tp_cumsum + fp_cumsum)

    if ap_method == 'pascal_voc':
        return compute_ap_pascal_voc(recalls, precisions)
    elif ap_method == 'coco':
        return compute_ap_coco(recalls, precisions)
    else:
        raise ValueError("Unknown AP method")

def calculate_map(gt_boxes_all, pred_boxes_all, iou_thresholds, ap_method):
    aps = defaultdict(list)

    for image_id in gt_boxes_all.keys():
        gt_boxes = gt_boxes_all[image_id]
        pred_boxes = pred_boxes_all[image_id]

        for class_id in set([box[4] for box in gt_boxes + pred_boxes]):
            gt_class = [box[:4] for box in gt_boxes if box[4] == class_id]
            pred_class = [box for box in pred_boxes if box[4] == class_id]

            for iou_threshold in iou_thresholds:
                ap = calculate_ap(gt_class, pred_class, iou_threshold, ap_method)
                aps[class_id].append(ap)

    mean_ap = np.mean([np.mean(ap_list) for ap_list in aps.values()])
    return mean_ap

def pascal_voc_map(gt_boxes_all, pred_boxes_all):
    return calculate_map(gt_boxes_all, pred_boxes_all, [0.5], 'pascal_voc')

def coco_map(gt_boxes_all, pred_boxes_all):
    return calculate_map(gt_boxes_all, pred_boxes_all, np.arange(0.5, 1.0, 0.05), 'coco')

# 사용 예시
gt_boxes_all = {
    'image1': [[10, 10, 100, 100, 0], [150, 150, 50, 50, 1]],
    'image2': [[30, 30, 80, 80, 0]]
}

pred_boxes_all = {
    'image1': [[15, 15, 95, 95, 0, 0.9], [155, 155, 45, 45, 1, 0.8]],
    'image2': [[35, 35, 75, 75, 0, 0.95]]
}

pascal_map_score = pascal_voc_map(gt_boxes_all, pred_boxes_all)
print(f"PASCAL VOC mAP: {pascal_map_score}")

coco_map_score = coco_map(gt_boxes_all, pred_boxes_all)
print(f"COCO mAP: {coco_map_score}")