src/evaluate.py

import logging
import torch
import numpy as np
from sklearn.metrics import roc_curve, auc, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
from tqdm import tqdm
import torchvision.transforms.functional as TF

from config import Config
from file_io_manager import FileIOManager

logger = logging.getLogger(__name__)


class Evaluator:
    def __init__(self, model, criterion, io: FileIOManager | None = None):
        self.model = model
        self.criterion = criterion
        self._device = Config.get_training_config()['device']
        self._io = io

    def evaluate(self, val_loader, use_tta=False):
        self.model.eval()
        total_loss = 0.0
        all_preds: list[int]   = []
        all_labels: list[int]  = []
        all_probs: list[float] = []

        with torch.no_grad():
            for images_batch_tensor, metadata_batch, labels_batch in tqdm(val_loader, desc="Evaluating"):
                metadata_batch = metadata_batch.to(self._device)
                all_labels.extend(labels_batch.cpu().numpy().flatten())

                for i in range(images_batch_tensor.size(0)):
                    img_tensor   = images_batch_tensor[i]
                    meta_single  = metadata_batch[i:i+1]
                    label_single = labels_batch[i:i+1].unsqueeze(1).float().to(self._device)

                    outputs = self.model(img_tensor.unsqueeze(0).to(self._device), meta_single)
                    loss    = self.criterion(outputs, label_single)
                    total_loss += loss.item()

                    if use_tta:
                        probs_tta = [torch.sigmoid(outputs).item()]
                        hflip_out = self.model(TF.hflip(img_tensor).unsqueeze(0).to(self._device), meta_single)
                        probs_tta.append(torch.sigmoid(hflip_out).item())
                        final_prob = float(np.mean(probs_tta))
                    else:
                        final_prob = torch.sigmoid(outputs).item()

                    all_probs.append(final_prob)
                    all_preds.append(1 if final_prob > 0.5 else 0)

        avg_loss = total_loss / len(all_labels) if all_labels else 0.0
        return avg_loss, all_preds, all_labels, all_probs

    def plot_roc_curve(self, labels, probs):
        fpr, tpr, _ = roc_curve(labels, probs)
        roc_auc = auc(fpr, tpr)
        plt.figure()
        plt.plot(fpr, tpr, color='darkorange', lw=2, label='ROC curve (area = %0.2f)' % roc_auc)
        plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
        plt.xlim([0.0, 1.0])
        plt.ylim([0.0, 1.05])
        plt.xlabel('False Positive Rate')
        plt.ylabel('True Positive Rate')
        plt.title('Receiver Operating Characteristic (ROC) Curve')
        plt.legend(loc="lower right")
        path = self._io.roc_curve_path() if self._io else FileIOManager.for_run("default").roc_curve_path()
        plt.savefig(path)
        plt.close()
        logger.info("ROC curve saved to %s", path)

    def plot_confusion_matrix(self, labels, preds):
        cm = confusion_matrix(labels, preds)
        plt.figure(figsize=(8, 6))
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                    xticklabels=['Benign', 'Malignant'],
                    yticklabels=['Benign', 'Malignant'])
        plt.xlabel('Predicted Label')
        plt.ylabel('True Label')
        plt.title('Confusion Matrix')
        path = self._io.confusion_matrix_path() if self._io else FileIOManager.for_run("default").confusion_matrix_path()
        plt.savefig(path)
        plt.close()
        logger.info("Confusion matrix saved to %s", path)

    def plot_shap(self, val_loader, feature_names: list[str]) -> None:
        """Generate SHAP feature importance plot for the metadata branch.

        Uses KernelExplainer (model-agnostic) to attribute the model's output
        to each of the 14 metadata features while holding image features fixed
        at their validation-set mean.
        """
        import shap

        self.model.eval()

        # Collect metadata + compute mean image features across validation set
        all_metadata: list = []
        all_img_feats: list = []
        with torch.no_grad():
            for img_batch, metadata_batch, _ in val_loader:
                all_metadata.append(metadata_batch)
                img_feats = self.model.cnn_dropout(
                    self.model.image_backbone(img_batch.to(self._device))
                )
                all_img_feats.append(img_feats)
        all_metadata_np = torch.cat(all_metadata, dim=0).cpu().numpy()
        mean_img_feats = torch.cat(all_img_feats, dim=0).mean(dim=0).unsqueeze(0)

        background = shap.kmeans(all_metadata_np, 25)
        test_sample = all_metadata_np[:200]

        explainer = shap.KernelExplainer(
            lambda m: self.model.predict_metadata_proba(m, mean_img_feats),
            background,
        )
        shap_values = explainer.shap_values(test_sample, silent=True)

        plt.figure()
        shap.summary_plot(
            shap_values,
            features=test_sample,
            feature_names=feature_names,
            show=False,
        )
        io = self._io or FileIOManager.for_run("default")
        path = io.shap_plot_path()
        plt.savefig(path, bbox_inches='tight')
        plt.close()
        logger.info("SHAP feature importance plot saved to %s", path)

    def compute_ood_stats(self, val_loader) -> None:
        """Compute mean + inverse covariance of backbone features for Mahalanobis OOD detection.

        Based on: Lee et al., "A Simple Unified Framework for Detecting
        Out-of-Distribution Samples and Adversarial Attacks", NeurIPS 2018.
        https://arxiv.org/abs/1807.03888
        """
        self.model.eval()
        features_list: list[torch.Tensor] = []
        with torch.no_grad():
            for images, _, _ in tqdm(val_loader, desc="Computing OOD stats"):
                images = images.to(self._device)
                feats = self.model.cnn_dropout(self.model.image_backbone(images))
                features_list.append(feats.cpu())

        features = torch.cat(features_list, dim=0).float()
        mean = features.mean(dim=0)
        centered = features - mean
        cov = (centered.T @ centered) / (centered.shape[0] - 1)
        cov += 1e-5 * torch.eye(cov.shape[0])  # regularise for invertibility
        cov_inv = torch.linalg.inv(cov)

        # Compute empirical threshold from validation distances
        dists = (centered @ cov_inv * centered).sum(dim=1)
        threshold = float(dists.mean() + 3 * dists.std())
        logger.info("OOD distances — mean: %.1f, std: %.1f, threshold (μ+3σ): %.1f",
                    dists.mean(), dists.std(), threshold)

        io = self._io or FileIOManager.for_run("default")
        io.save_ood_stats(mean, cov_inv, threshold)