utils.py

import torch
import numpy as np
import os
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import (
    multilabel_confusion_matrix,
    f1_score, accuracy_score, roc_auc_score, auc,
    average_precision_score
)

def evaluate_model(model, dataloader, device, criterion=None, threshold=0.5):
    """
    Evaluates a multi-label classification model on the given dataloader.

    Args:
        model (torch.nn.Module): Trained model.
        dataloader (DataLoader): Validation dataloader.
        device (torch.device): Computation device (CPU or CUDA).
        criterion (loss, optional): Loss function.
        threshold (float): Threshold for converting probabilities to binary predictions.

    Returns:
        dict: Evaluation metrics including loss, F1 score, accuracy, AUC, and mAP.
    """
    model.eval()
    all_labels, all_preds, all_probs = [], [], []
    total_loss, total_samples = 0.0, 0

    with torch.no_grad():
        for images, labels in dataloader:
            images = images.to(device)
            labels = labels.to(device).float()

            outputs = model(images)
            if criterion:
                loss = criterion(outputs, labels)
                total_loss += loss.item() * images.size(0)
                total_samples += images.size(0)

            probs = torch.sigmoid(outputs)
            preds = (probs >= threshold).int()

            all_labels.append(labels.cpu().numpy())
            all_preds.append(preds.cpu().numpy())
            all_probs.append(probs.cpu().numpy())

    labels_np = np.vstack(all_labels)
    preds_np = np.vstack(all_preds)
    probs_np = np.vstack(all_probs)

    metrics = {
        "val_loss": total_loss / total_samples if criterion else 0.0,
        "f1": f1_score(labels_np, preds_np, average='macro'),
        "accuracy": accuracy_score(labels_np, preds_np),
        "auc": roc_auc_score(labels_np, probs_np, average='macro'),
        "mAP": average_precision_score(labels_np, probs_np, average='macro'),
        "labels": labels_np,
        "preds": preds_np,
        "probs": probs_np
    }

    return metrics

def plot_training_metrics(train_stats, val_stats, save_path):
    """
    Plots training and validation metrics over epochs.

    Args:
        train_stats (dict): Training metrics by epoch.
        val_stats (dict): Validation metrics by epoch.
        save_path (str): Path to save the figure.
    """
    metrics = train_stats.keys()
    epochs = range(1, len(next(iter(train_stats.values()))) + 1)

    plt.figure(figsize=(16, 10))
    for i, metric in enumerate(metrics):
        plt.subplot(2, 2, i + 1)
        plt.plot(epochs, train_stats[metric], label="Train", marker='o')
        plt.plot(epochs, val_stats[metric], label="Val", marker='x')
        plt.title(metric)
        plt.xlabel("Epoch")
        plt.grid(True)
        plt.legend()

    plt.tight_layout()
    plt.savefig(save_path)
    plt.close()


def plot_combined_confusion_matrices(train_preds, train_labels, val_preds, val_labels, class_names, save_path):
    """
    Plots multilabel confusion matrices for both training and validation sets.

    Args:
        train_preds (np.ndarray): Training predictions.
        train_labels (np.ndarray): Ground truth training labels.
        val_preds (np.ndarray): Validation predictions.
        val_labels (np.ndarray): Ground truth validation labels.
        class_names (list): Names of the classes.
        save_path (str): Path to save the figure.
    """
    train_cm = multilabel_confusion_matrix(train_labels, train_preds)
    val_cm = multilabel_confusion_matrix(val_labels, val_preds)

    fig, axes = plt.subplots(2, len(class_names), figsize=(5 * len(class_names), 10))
    for i in range(len(class_names)):
        sns.heatmap(train_cm[i], annot=True, fmt='d', cmap='Greens', ax=axes[0, i])
        axes[0, i].set_title(f"Train - {class_names[i]}")
        axes[0, i].set_xlabel("Predicted")
        axes[0, i].set_ylabel("True")

        sns.heatmap(val_cm[i], annot=True, fmt='d', cmap='Blues', ax=axes[1, i])
        axes[1, i].set_title(f"Val - {class_names[i]}")
        axes[1, i].set_xlabel("Predicted")
        axes[1, i].set_ylabel("True")

    plt.tight_layout()
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    plt.savefig(save_path)
    plt.close()


def plot_roc_curves(labels, probs, class_names, save_path):
    """
    Plots ROC curves for each class.

    Args:
        labels (np.ndarray): True binary labels.
        probs (np.ndarray): Predicted probabilities.
        class_names (list): Class labels.
        save_path (str): Output image file.
    """
    plt.figure(figsize=(10, 10))
    for i, cls in enumerate(class_names):
        fpr, tpr, _ = roc_curve(labels[:, i], probs[:, i])
        auc_score = auc(fpr, tpr)
        plt.plot(fpr, tpr, label=f"{cls} (AUC = {auc_score:.2f})")

    plt.plot([0, 1], [0, 1], 'k--')
    plt.xlabel("False Positive Rate")
    plt.ylabel("True Positive Rate")
    plt.title("ROC Curve per Class")
    plt.legend(loc='lower right')
    plt.grid()
    plt.savefig(save_path)
    plt.close()