evaluate.py

import os
import numpy as np
import torch
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import StratifiedKFold
from sklearn.metrics import confusion_matrix, classification_report, f1_score
from torch.utils.data import DataLoader
from tqdm import tqdm

from config import (DEVICE, SAVE_DIR, TRAIN_DIR, BATCH_SIZE,
                    K_FOLDS, NUM_WORKERS, NUM_GROUPS)
from dataset import WBCDataset, get_transforms


def plot_confusion_matrix(all_labels, all_preds, class_names, title="Matrice de Confusion"):
    cm      = confusion_matrix(all_labels, all_preds)
    cm_perc = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]

    plt.figure(figsize=(14, 10))
    sns.heatmap(cm_perc, annot=True, fmt='.2f', cmap='magma',
                xticklabels=class_names, yticklabels=class_names)
    plt.title(title)
    plt.xlabel('Prédictions')
    plt.ylabel('Vrai label')
    plt.tight_layout()
    plt.show()


def evaluate_oof(df, le, model_class, tag="model", use_clahe=True):
    """
    evalue les modeles sauvegardes en out-of-fold et affiche la matrice de
    confusion globale + le rapport par classe."""
    _, val_tf = get_transforms()
    skf = StratifiedKFold(n_splits=K_FOLDS, shuffle=True, random_state=42)

    all_preds, all_labels = [], []

    for fold, (train_idx, val_idx) in enumerate(skf.split(df, df['label_encoded'])):
        model_path = f"{SAVE_DIR}/best_{tag}_fold{fold+1}.pth"

        model = model_class(NUM_GROUPS, len(le.classes_)).to(DEVICE)
        model.load_state_dict(torch.load(model_path, map_location=DEVICE))
        model.eval()

        val_df     = df.iloc[val_idx]
        val_loader = DataLoader(
            WBCDataset(val_df, TRAIN_DIR, val_tf, use_clahe=use_clahe),
            batch_size=BATCH_SIZE, num_workers=NUM_WORKERS)

        with torch.no_grad():
            for imgs, _, fines in tqdm(val_loader, desc=f"Eval Fold {fold+1}"):
                _, out_f = model(imgs.to(DEVICE))
                all_preds.extend(out_f.argmax(1).cpu().numpy())
                all_labels.extend(fines.numpy())

    plot_confusion_matrix(all_labels, all_preds, le.classes_,
                          title="Matrice de Confusion")
    print(classification_report(all_labels, all_preds, target_names=le.classes_))

    macro_f1 = f1_score(all_labels, all_preds, average='macro')
    print(f"F1 macro global : {macro_f1:.4f}")
    return macro_f1