metric.py

import os
import sys
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import confusion_matrix

# Usage : python metric.py test_images/labels test/model/labels

if len(sys.argv) < 3:
    print("Usage: python metric.py <GT_DIR> <PRED_DIR>")
    sys.exit(1)

GT_DIR = sys.argv[1]
PRED_DIR = sys.argv[2]

IOU_THRESHOLD = 0.50
CONF_THRESHOLD = 0.25  # Seuil de confiance fixe pour Précision/Rappel/F1/Matrice

# Normalisation et résolution des dossiers du modèle
pred_path_norm = os.path.normpath(PRED_DIR)
if os.path.basename(pred_path_norm).lower() == "labels":
    model_dir = os.path.dirname(pred_path_norm)
else:
    model_dir = pred_path_norm

model_name = os.path.basename(model_dir)


def yolo_to_corners(box):
    """Convertit (x_center, y_center, w, h) en (x1, y1, x2, y2)."""
    x_c, y_c, w, h = box
    return np.array([x_c - w / 2, y_c - h / 2, x_c + w / 2, y_c + h / 2])


def get_iou(box_a, box_b):
    """Calcule l'IoU entre deux boxes au format YOLO."""
    a = yolo_to_corners(box_a)
    b = yolo_to_corners(box_b)
    ix1 = np.maximum(a[0], b[0])
    iy1 = np.maximum(a[1], b[1])
    ix2 = np.minimum(a[2], b[2])
    iy2 = np.minimum(a[3], b[3])
    i_width = np.maximum(ix2 - ix1, 0.0)
    i_height = np.maximum(iy2 - iy1, 0.0)
    area_intersection = i_width * i_height
    area_a = max(a[2] - a[0], 0.0) * max(a[3] - a[1], 0.0)
    area_b = max(b[2] - b[0], 0.0) * max(b[3] - b[1], 0.0)
    area_union = area_a + area_b - area_intersection
    return float(area_intersection / area_union) if area_union > 0 else 0.0


def parse_yolo_file(filepath, is_gt=False):
    """Lit un .txt YOLO et extrait [class_id, xc, yc, w, h, (conf)]."""
    boxes = []
    if not os.path.exists(filepath):
        return boxes
    with open(filepath, "r") as f:
        for line in f:
            parts = line.strip().split()
            if not parts:
                continue
            class_id = int(float(parts[0]))
            box = [float(p) for p in parts[1:5]]
            conf = float(parts[5]) if len(parts) >= 6 and not is_gt else 1.0
            boxes.append({"class_id": class_id, "box": box, "conf": conf})
    return boxes


def compute_ap(recalls, precisions):
    """Calcul de l'Average Precision (AP) par enveloppe monotone (méthode VOC/COCO)."""
    mrec = np.concatenate(([0.0], recalls, [1.0]))
    mpre = np.concatenate(([0.0], precisions, [0.0]))

    # Enveloppe monotone supérieure
    for i in range(len(mpre) - 1, 0, -1):
        mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])

    # Intégration sous la courbe
    indices = np.where(mrec[1:] != mrec[:-1])[0]
    ap = np.sum((mrec[indices + 1] - mrec[indices]) * mpre[indices + 1])
    return ap, mrec, mpre


# --- Chargement de toutes les données ---
gt_files = sorted(f for f in os.listdir(GT_DIR) if f.endswith(".txt"))

all_detections = []  # [{img_id, class_id, box, conf}]
gt_by_img = {}       # {img_id: [{class_id, box, matched}]}
total_gt_count = 0

for filename in gt_files:
    img_id = os.path.splitext(filename)[0]
    gt_path = os.path.join(GT_DIR, filename)
    pred_path = os.path.join(PRED_DIR, filename)

    gts = parse_yolo_file(gt_path, is_gt=True)
    preds = parse_yolo_file(pred_path, is_gt=False)

    for gt in gts:
        gt["matched"] = False
    gt_by_img[img_id] = gts
    total_gt_count += len(gts)

    for pred in preds:
        pred["img_id"] = img_id
        all_detections.append(pred)

# --- Calcul de l'AP (Average Precision @ IoU 0.50) ---
# Tri global de TOUTES les prédictions par confiance décroissante
all_detections.sort(key=lambda x: x["conf"], reverse=True)

tps = np.zeros(len(all_detections))
fps = np.zeros(len(all_detections))
ious_matched = []

for idx, pred in enumerate(all_detections):
    img_gts = gt_by_img.get(pred["img_id"], [])
    best_iou = 0.0
    best_gt = None

    # Chercher le GT correspondant avec le meilleur IoU
    for gt in img_gts:
        if gt["class_id"] == pred["class_id"]:
            iou = get_iou(pred["box"], gt["box"])
            if iou > best_iou:
                best_iou = iou
                best_gt = gt

    if best_iou >= IOU_THRESHOLD and best_gt is not None and not best_gt["matched"]:
        tps[idx] = 1
        best_gt["matched"] = True
        ious_matched.append(best_iou)
    else:
        fps[idx] = 1

cum_tps = np.cumsum(tps)
cum_fps = np.cumsum(fps)

recalls_curve = cum_tps / total_gt_count if total_gt_count > 0 else np.zeros_like(cum_tps)
precisions_curve = cum_tps / (cum_tps + cum_fps)

ap_50, mrec, mpre = compute_ap(recalls_curve, precisions_curve)

# --- Métriques au seuil de confiance spécifié (CONF_THRESHOLD) ---
# Sélection des détections >= CONF_THRESHOLD
valid_indices = [i for i, d in enumerate(all_detections) if d["conf"] >= CONF_THRESHOLD]

if valid_indices:
    last_idx = valid_indices[-1]
    tp_fixed = int(cum_tps[last_idx])
    fp_fixed = int(cum_fps[last_idx])
else:
    tp_fixed, fp_fixed = 0, 0

fn_fixed = total_gt_count - tp_fixed
precision_fixed = tp_fixed / (tp_fixed + fp_fixed) if (tp_fixed + fp_fixed) > 0 else 0.0
recall_fixed = tp_fixed / total_gt_count if total_gt_count > 0 else 0.0
f1_fixed = (
    2 * (precision_fixed * recall_fixed) / (precision_fixed + recall_fixed)
    if (precision_fixed + recall_fixed) > 0
    else 0.0
)

# --- Préparation de la Matrice de Confusion ---
y_true_all, y_pred_all = [], []

# Réinitialisation des états de match pour la matrice à seuil fixe
for filename in gt_files:
    img_id = os.path.splitext(filename)[0]
    gts = parse_yolo_file(os.path.join(GT_DIR, filename), is_gt=True)
    preds = [p for p in parse_yolo_file(os.path.join(PRED_DIR, filename)) if p["conf"] >= CONF_THRESHOLD]

    gt_matched = [False] * len(gts)
    pred_matched = [False] * len(preds)

    # Appariement local
    for p_idx, p in enumerate(preds):
        best_iou, best_gt_idx = 0.0, -1
        for g_idx, g in enumerate(gts):
            if g["class_id"] == p["class_id"] and not gt_matched[g_idx]:
                iou = get_iou(p["box"], g["box"])
                if iou > best_iou:
                    best_iou = iou
                    best_gt_idx = g_idx

        if best_iou >= IOU_THRESHOLD and best_gt_idx != -1:
            gt_matched[best_gt_idx] = True
            pred_matched[p_idx] = True
            y_true_all.append(gts[best_gt_idx]["class_id"])
            y_pred_all.append(p["class_id"])

    # Faux Négatifs (GT non détectés)
    for g_idx, g in enumerate(gts):
        if not gt_matched[g_idx]:
            y_true_all.append(g["class_id"])
            y_pred_all.append(-1)  # Fond / Background

    # Faux Positifs (Prédictions en trop)
    for p_idx, p in enumerate(preds):
        if not pred_matched[p_idx]:
            y_true_all.append(-1)
            y_pred_all.append(p["class_id"])

# --- Écriture des Logs et Affichage ---
log_path = os.path.join(model_dir, "evaluation_log.txt")
cm_image_path = os.path.join(model_dir, "confusion_matrix.png")
pr_image_path = os.path.join(model_dir, "pr_curve.png")

summary = (
    f"==================================================\n"
    f" ÉVALUATION MODÈLE : {model_name}\n"
    f"==================================================\n"
    f"Fichiers traités      : {len(gt_files)}\n"
    f"Nombre total de GT    : {total_gt_count}\n"
    f"IoU Moyen (Detections): {np.mean(ious_matched) if ious_matched else 0.0:.4f}\n\n"
    f"--- PERFORMANCES GLOBALES (Average Precision) ---\n"
    f"AP@50 (mAP@50)        : {ap_50:.4f} ({ap_50*100:.2f}%)\n\n"
    f"--- MÉTRIQUES AU SEUIL FIXE (Conf >= {CONF_THRESHOLD}) ---\n"
    f"Vrais Positifs (TP)   : {tp_fixed}\n"
    f"Faux Positifs  (FP)   : {fp_fixed}\n"
    f"Faux Négatifs  (FN)   : {fn_fixed}\n"
    f"Précision             : {precision_fixed:.4f}\n"
    f"Rappel (Recall)       : {recall_fixed:.4f}\n"
    f"Score F1              : {f1_fixed:.4f}\n"
    f"==================================================\n"
)

print(summary)
with open(log_path, "w", encoding="utf-8") as f:
    f.write(summary)

# --- Visualisation 1 : Courbe Précision-Rappel ---
plt.figure(figsize=(8, 6))
plt.plot(mrec, mpre, color="b", lw=2, label=f"Courbe PR (AP@50 = {ap_50:.4f})")
plt.xlabel("Rappel (Recall)")
plt.ylabel("Précision")
plt.title(f"Courbe Précision-Rappel — {model_name}")
plt.grid(True, linestyle="--", alpha=0.6)
plt.legend(loc="lower left")
plt.tight_layout()
plt.savefig(pr_image_path)
print(f"Courbe PR enregistrée dans       : {pr_image_path}")
plt.close()

# --- Visualisation 2 : Matrice de Confusion ---
if y_true_all and y_pred_all:
    unique_classes = sorted(list(set(y_true_all + y_pred_all) - {-1}))
    labels = unique_classes + [-1]
    display_labels = [f"Class {c}" for c in unique_classes] + ["background"]

    cm_normalized = confusion_matrix(y_true_all, y_pred_all, labels=labels, normalize="true")

    plt.figure(figsize=(9, 7))
    sns.heatmap(
        cm_normalized,
        annot=True,
        fmt=".2f",
        cmap="Blues",
        xticklabels=display_labels,
        yticklabels=display_labels,
        vmin=0.0,
        vmax=1.0,
    )
    plt.title(f"Matrice de Confusion Normalisée — {model_name}\n(Conf >= {CONF_THRESHOLD}, IoU >= {IOU_THRESHOLD})")
    plt.xlabel("Prédictions")
    plt.ylabel("Vérité Terrain (GT)")
    plt.tight_layout()
    plt.savefig(cm_image_path)
    print(f"Matrice de confusion enregistrée : {cm_image_path}")
    plt.close()