[Deep learning - Object Detection] mAP(mean Average Precision) 계산 함수 구현

객체 탐지 모델의 성능은 mAP(mean Average Precision) 값을 통해서 평가할 수 있다. 
본문에서는 mAP을 구하는 함수를 직접 구현한다.

Ground-Truth Bounding Box 하나는 (xmin, ymin, xmax, ymax, label)로 구성되고,
Prediction Bounding Box 하나는 (xmin, ymin, xmax, ymax, confidence score, label)로 구성되는 것을 전제로 코드를 작성한다.

mAP의 함수에 입력되는 GT-BBOX, Pred-BBOX는 모든 이미지에 대한 BBOX 정보를 포함한다.

 

1. 패키지 로드

import numpy as np

2. 함수 정의

2.1. Label 획득 함수

모델이 추론할 수 있는 Label의 경우에 사전에 정의되어 있는 경우가 대부분이지만, 별도의 Label Map이 존재하지 않고 GT-BBOX 데이터만 존재하는 경우도 있다. 따라서 GT-BBOX로부터 Label을 획득하는 함수를 작성한다.

"""
all_gt_bboxes : 모든 이미지에 대한 ground-truth bounding boxes
"""
def get_labels(all_gt_boxes):
    labels = set()
    for gt_boxes in all_gt_boxes:
        for gt_box in gt_boxes:
            labels.add(gt_box[-1])
        
    return labels

2.2. IOU 계산 함수

하나의 bounding box에 대해서 ground-truth와 prediction bounding box간의 IOU를 계산한다.
(여기서 하나의 bounding box란 하나의 (xmin, ymin, xmax, ymax)을 말한다.

"""
gt_box : (xmin, ymin, xmax, ymax, label)
pred_box : (xmin, ymin, xmax, ymax, confidence score, label)
"""
def calculate_iou(gt_box, pred_box):
    # ground-truth bounding box
    gt_xmin, gt_ymin, gt_xmax, gt_ymax = gt_box[:4]
    # predicted bounding box
    pred_xmin, pred_ymin, pred_xmax, pred_ymax = pred_box[:4]
    
    # Calculate the area of intersection
    xmin = max(gt_xmin, pred_xmin)
    ymin = max(gt_ymin, pred_ymin)
    xmax = min(gt_xmax, pred_xmax)
    ymax = min(gt_ymax, pred_ymax)
    
    # non-overlap
    if xmin >= xmax or ymin >= ymax:
        return 0.0
    
    intersection_area = (xmax - xmin) * (ymax - ymin)
    
    # Calculate the area of union
    gt_area = (gt_xmax - gt_xmin) * (gt_ymax - gt_ymin)
    pred_area = (pred_xmax - pred_xmin) * (pred_ymax - pred_ymin)
    union_area = gt_area + pred_area - intersection_area
    
    # Calculate the IoU
    iou = intersection_area / union_area
    
    return iou

2.3. AP(Average Precision) 계산 함수 구현

하나의 Class(Label)에 대한 AP 값을 구한다. mAP을 구하기 위해서는 모든 Class에 대한 AP 값을 구해야 한다.

"""
gt_boxes : 모든 이미지의 하나의 class에 대한 GT-BBOX
pred_boxes : 모든 이미지의 하나의 class에 대한 Pred-BBOX
"""
def calculate_ap(gt_boxes, pred_boxes, iou_threshold=0.5):
    # Sort the prediction boxes by confidence score
    pred_boxes_sorted = sorted(pred_boxes, key=lambda x: x[4], reverse=True)
    
    # Initialize variables for TP, FP, and number of ground-truth boxes
    TP = [0] * len(pred_boxes_sorted)
    FP = [0] * len(pred_boxes_sorted)
    num_gt_boxes = len(gt_boxes)
    
    # Initialize variables for precision, recall, and AP
    precision = [0] * len(pred_boxes_sorted)
    recall = [0] * len(pred_boxes_sorted)
    AP = 0.0
    
    # Loop through each predicted box
    for i, pred_box in enumerate(pred_boxes_sorted):
        # Loop through each ground-truth box
        for j, gt_box in enumerate(gt_boxes):
            # Calculate the IoU between the predicted box and the ground-truth box
            iou = calculate_iou(gt_box, pred_box)
            
            # If the IoU is greater than the threshold and the predicted box has not been matched yet
            if iou >= iou_threshold and not TP[i] and not FP[i]:
                TP[i] = 1
                # Increase the precision for the matched box
                precision[i] = sum(TP) / (sum(TP) + sum(FP))
                
        # If the predicted box was not matched with any ground-truth box
        if not TP[i]:
            FP[i] = 1
            # Set the precision to 0 for the unmatched box
            precision[i] = 0.0
            
        # Calculate the recall for the predicted box
        recall[i] = sum(TP) / num_gt_boxes
    
    # Calculate the AP using the precision and recall arrays
    for i in range(len(precision)):
        if i == 0:
            AP += precision[i] * recall[i]
        else:
            AP += precision[i] * (recall[i] - recall[i-1])
    
    return AP

2.4. mAP(mean Average Precision) 계산 함수 구현

해당 함수는 모든 Threshold에 대한 모든 이미지의 mAP 값을 구하고, 또한 각 클래스별 mAP 값도 함께 구한다. 이를 계산하기 위해서 결과는 dictionary 형태로 저장한다.

"""
gt_boxes : 모든 이미지의 GT-BBOX
pred_boxes : 모든 이미지의 Pred-BBOX
iou_thresholds : 계산할 모든 IOU Threshold
"""
def calculate_map(gt_boxes, pred_boxes, iou_thresholds=[0.5]):
    # Calculate the AP for each IoU threshold
    result = dict()
    num_images = len(gt_boxes)
    labels = get_labels(gt_boxes)
    
    for iou_threshold in iou_thresholds:
        result[f"{iou_threshold}"] = dict()
        for label in labels:
            result[f"{iou_threshold}"]["mAP"] = float(0.0)
            result[f"{iou_threshold}"][f"{label}"] = dict()
            result[f"{iou_threshold}"][f"{label}"]["mAP"] = float()
    
    for iou_threshold in iou_thresholds:
        mAPs = []
        for label in labels:
            APs = []
            for i in range(num_images):
                # get gt_box including label
                gt_box_label = []
                for gt_box in gt_boxes[i]:
                    if gt_box[4] == label:
                        gt_box_label.append(gt_box)
                
                # get pd_box_including label
                pd_box_label = []
                for pd_box in pred_boxes[i]:
                    if pd_box[5] == label:
                        pd_box_label.append(pd_box)
                if len(gt_box_label) != 0:
                    AP = calculate_ap(gt_box_label, pd_box_label, iou_threshold)
                    APs.append(AP)
            mAP = sum(APs) / len(APs)
            result[f"{iou_threshold}"][f"{label}"]["mAP"] = mAP
            mAPs.append(mAP)
        result[f"{iou_threshold}"]["mAP"] = sum(mAPs) / len(mAPs)
    
    return result, labels

여기서 주의할 점은 GT-BBOX와 Pred-BBOX는 각 인덱스 별로 동일한 이미지에 대한 정보를 가져야 한다.
(이를 위해서는 별도의 처리가 필요하나 본문에서는 다루지 않는다)

3. 사용예시

3.1. GT-BBOX, Pred-BBOX, IOU Threshold 임의 정의

# Ground-Truth Boxes
gt_bboxes_per_images = [
    [
        [20, 30, 70, 90, "car"],
        [50, 70, 120, 350, "truck"],
        [0, 50, 200, 600, "plane"]
    ],
    [
        [70, 80, 90, 100, "person"],
        [20, 50, 60, 120, "car"]
    ]
]

# Prediction Boxes
pd_bboxes_per_images = [
    [
        [20, 30, 60, 90, 0.55, "car"],
        [50, 69, 120, 350, 0.76, "truck"],
        [0, 20, 150, 500, 0.88, "car"],
        [0, 50, 190, 550, 0.43, "plane"]
    ],
    [
        [60, 80, 93, 102, 0.66, "person"],
        [20, 50, 60, 120, 0.77, "car"],
        [20, 60, 50, 70, 0.95, "person"]
    ]
]

# Threshold : [0.5;0.95;0.05]
threshold = np.arange(0.5, 1.0, 0.05)
threshold = threshold.round(2)

3.2. mAP 계산

result, labels = calculate_map(gt_bboxes_per_images, pd_bboxes_per_images, threshold)

3.3. 결과

# 결과 표시를 위한 함수 정의
def display_result(result, labels, iou_thresholds):
    for iou_threshold in iou_thresholds:
        print(f"------ @mAP{iou_threshold} ------")
        print(f"all : {result[f'{iou_threshold}']['mAP']}")
        for label in labels:
            print(f"{label} : {result[f'{iou_threshold}'][f'{label}']['mAP']}")
            
# 결과 표시
display_result(result, labels, threshold)

------ @mAP0.5 ------
all : 0.5
plane : 0.5
person : 0.25
car : 0.75
truck : 0.5
------ @mAP0.55 ------
all : 0.5
plane : 0.5
person : 0.25
car : 0.75
truck : 0.5
------ @mAP0.6 ------
all : 0.4375
plane : 0.5
person : 0.0
car : 0.75
truck : 0.5
------ @mAP0.65 ------
all : 0.4375
plane : 0.5
person : 0.0
car : 0.75
truck : 0.5
------ @mAP0.7 ------
all : 0.4375
plane : 0.5
person : 0.0
car : 0.75
truck : 0.5
------ @mAP0.75 ------
all : 0.4375
plane : 0.5
person : 0.0
car : 0.75
truck : 0.5
------ @mAP0.8 ------
all : 0.4375
plane : 0.5
person : 0.0
car : 0.75
truck : 0.5
------ @mAP0.85 ------
all : 0.375
plane : 0.5
person : 0.0
car : 0.5
truck : 0.5
------ @mAP0.9 ------
all : 0.25
plane : 0.0
person : 0.0
car : 0.5
truck : 0.5
------ @mAP0.95 ------
all : 0.25
plane : 0.0
person : 0.0
car : 0.5
truck : 0.5