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

알겠습니다. 클래스별로 Precision, Recall, F1-score를 계산하기 위해서는 각 클래스에 대한 True Positive (TP), False Positive (FP), False Negative (FN) 값을 추적하고, 이를 이용하여 Precision, Recall, F1-score를 계산해야 합니다. 아래는 클래스별로 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):
    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):
            if prediction_boxes[i][0] == gt_boxes[j][0]:  # Match classes
                cost_matrix[i, j] = 1 - calculate_iou(prediction_boxes[i][1:], gt_boxes[j][1:])
            else:
                cost_matrix[i, j] = 1  # Assign a high cost for different classes

    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_per_class(prediction_boxes, gt_boxes, iou_threshold=0.5):
    class_ids = set([box[0] for box in prediction_boxes] + [box[0] for box in gt_boxes])
    precision_per_class, recall_per_class, f1_per_class = {}, {}, {}

    for class_id in class_ids:
        # Filter prediction and ground truth boxes by class
        pred_boxes_cls = [box for box in prediction_boxes if box[0] == class_id]
        gt_boxes_cls = [box for box in gt_boxes if box[0] == class_id]

        matched_pairs = hungarian_algorithm(pred_boxes_cls, gt_boxes_cls, iou_threshold)

        true_positives = len(matched_pairs)
        false_positives = len(pred_boxes_cls) - true_positives
        false_negatives = len(gt_boxes_cls) - 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

        precision_per_class[class_id] = precision
        recall_per_class[class_id] = recall
        f1_per_class[class_id] = f1

    return precision_per_class, recall_per_class, f1_per_class

# 예시 사용
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_per_class, recall_per_class, f1_per_class = calculate_precision_recall_f1_per_class(prediction_boxes, gt_boxes)
for class_id in precision_per_class:
    print(f'Class {class_id}: Precision: {precision_per_class[class_id]}, Recall: {recall_per_class[class_id]}, F1-score: {f1_per_class[class_id]}')

이 코드는 클래스별로 True Positive, False Positive, False Negative를 추적하고, 각 클래스에 대한 Precision, Recall, F1-score를 계산합니다. 결과는 클래스 아이디를 키로 하는 딕셔너리로 반환됩니다.