Detection precision,recall,f1-score 계산 (Feat. 헝가리안 알고리즘)

헝가리안 알고리즘을 활용하여 Detection에서 Precision, Recall, F1-score를 계산하는 코드를 작성할 수 있습니다. 헝가리안 알고리즘은 매칭 문제에 사용되며, 여기서는 Detection 결과와 Ground Truth 간의 IoU를 기준으로 최적의 매칭을 찾아서 Precision, Recall, F1-score를 계산할 것입니다. 아래는 예시 코드입니다:

import numpy as np
from scipy.optimize import linear_sum_assignment
from sklearn.metrics import precision_score, recall_score, f1_score

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])

    intersection = max(0, x2 - x1) * max(0, y2 - y1)
    area_box1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
    area_box2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
    union = area_box1 + area_box2 - intersection

    iou = intersection / union if union > 0 else 0
    return iou

def hungarian_algorithm(prediction_boxes, gt_boxes, iou_threshold=0.5):
    num_preds = len(prediction_boxes)
    num_gts = len(gt_boxes)

    cost_matrix = np.zeros((num_preds, num_gts))
    for i in range(num_preds):
        for j in range(num_gts):
            cost_matrix[i, j] = 1 - calculate_iou(prediction_boxes[i][1:], gt_boxes[j][1:])

    row_ind, col_ind = linear_sum_assignment(cost_matrix)

    matched_pairs = []
    for i, j in zip(row_ind, col_ind):
        iou = calculate_iou(prediction_boxes[i][1:], gt_boxes[j][1:])
        if iou >= iou_threshold:
            matched_pairs.append((i, j))

    return matched_pairs

def calculate_precision_recall_f1(prediction_boxes, gt_boxes, iou_threshold=0.5):
    matched_pairs = hungarian_algorithm(prediction_boxes, gt_boxes, iou_threshold)

    true_positives = len(matched_pairs)
    false_positives = len(prediction_boxes) - true_positives
    false_negatives = len(gt_boxes) - true_positives

    precision = true_positives / (true_positives + false_positives) if true_positives + false_positives > 0 else 0
    recall = true_positives / (true_positives + false_negatives) if true_positives + false_negatives > 0 else 0
    f1 = 2 * (precision * recall) / (precision + recall) if precision + recall > 0 else 0

    return precision, recall, f1

# 예시 사용
prediction_boxes = [
    [1, 10, 10, 20, 20, 0.9],
    [2, 30, 30, 40, 40, 0.8],
    # ...
]

gt_boxes = [
    [1, 8, 8, 18, 18],
    [2, 32, 32, 42, 42],
    # ...
]

precision, recall, f1 = calculate_precision_recall_f1(prediction_boxes, gt_boxes)
print(f'Precision: {precision}, Recall: {recall}, F1-score: {f1}')

이 코드는 헝가리안 알고리즘을 사용하여 IoU가 일정 임계값 이상인 것들을 매칭시킨 후, Precision, Recall, F1-score를 계산합니다. IoU 임계값은 iou_threshold 매개변수를 통해 조절할 수 있습니다.