Out-of-Distribution (OOD) Detection in Medical Imaging Using Deep Learning 🌀¶

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.metrics import (
    roc_auc_score, average_precision_score,
    roc_curve, precision_recall_curve
)
import torch
import torchvision.transforms as T
from PIL import Image
import glob, os, random

np.random.seed(42)
torch.manual_seed(42)

# ---------------------------------------------------------------------------
# Core evaluation functions
# ---------------------------------------------------------------------------

def compute_ood_metrics(id_scores, ood_scores):
    """
    Evaluate OOD detection as binary classification.

    id_scores  : uncertainty scores for in-distribution samples (lower expected)
    ood_scores : uncertainty scores for OOD samples (higher expected)

    Returns dict with AUROC, AUPRC, FPR@95%TPR.
    """
    labels = np.concatenate([
        np.zeros(len(id_scores)),
        np.ones(len(ood_scores))
    ])
    scores = np.concatenate([id_scores, ood_scores])

    auroc = roc_auc_score(labels, scores)
    auprc = average_precision_score(labels, scores)

    # FPR at 95% TPR
    fpr, tpr, thresholds = roc_curve(labels, scores)
    idx_95 = np.argmin(np.abs(tpr - 0.95))
    fpr_at_95 = fpr[idx_95]

    return {
        'AUROC'        : auroc,
        'AUPRC'        : auprc,
        'FPR@95TPR'    : fpr_at_95,
        'fpr'          : fpr,
        'tpr'          : tpr,
        'labels'       : labels,
        'scores'       : scores,
    }


def plot_ood_analysis(id_scores, ood_scores, metrics, method_name,
                       color, id_images=None, ood_images=None,
                       score_label='Uncertainty score'):
    """
    Full OOD analysis plot: score distributions + ROC + PR curves.
    """
    fig, axes = plt.subplots(1, 3, figsize=(16, 5))
    fig.suptitle(f'{method_name} — OOD Detection Analysis', fontsize=13)

    # Distribution of scores
    max_ent = np.log(2)
    bins    = np.linspace(0, max(id_scores.max(), ood_scores.max()) * 1.05, 30)
    axes[0].hist(id_scores,  bins=bins, color=color,   alpha=0.65,
                 edgecolor='white', label='In-distribution (ID)')
    axes[0].hist(ood_scores, bins=bins, color='#888888', alpha=0.65,
                 edgecolor='white', label='Out-of-distribution (OOD)')
    axes[0].set_xlabel(score_label, fontsize=10)
    axes[0].set_ylabel('Count', fontsize=10)
    axes[0].set_title(f'Score distributions\nID vs OOD', fontsize=11)
    axes[0].legend(fontsize=9)
    axes[0].spines[['top','right']].set_visible(False)

    # ROC curve
    axes[1].plot(metrics['fpr'], metrics['tpr'], color=color,
                 linewidth=2.5, label=f'AUROC = {metrics["AUROC"]:.4f}')
    axes[1].plot([0,1],[0,1],'k--', linewidth=1, alpha=0.4,
                 label='Random (AUROC=0.5)')
    axes[1].fill_between(metrics['fpr'], metrics['tpr'],
                          alpha=0.12, color=color)

    # Mark FPR@95%TPR
    fpr_arr, tpr_arr = metrics['fpr'], metrics['tpr']
    idx95 = np.argmin(np.abs(tpr_arr - 0.95))
    axes[1].scatter(fpr_arr[idx95], tpr_arr[idx95],
                    s=100, color='#D85A30', zorder=5,
                    label=f'FPR@95%TPR = {metrics["FPR@95TPR"]:.3f}')
    axes[1].axhline(0.95, color='#D85A30', linewidth=1,
                    linestyle='--', alpha=0.5)

    axes[1].set_xlabel('False Positive Rate', fontsize=10)
    axes[1].set_ylabel('True Positive Rate', fontsize=10)
    axes[1].set_title('ROC curve (ID=0 vs OOD=1)', fontsize=11)
    axes[1].legend(fontsize=8)
    axes[1].spines[['top','right']].set_visible(False)

    # PR curve
    precision, recall, _ = precision_recall_curve(metrics['labels'],
                                                    metrics['scores'])
    axes[2].plot(recall, precision, color=color, linewidth=2.5,
                 label=f'AUPRC = {metrics["AUPRC"]:.4f}')
    axes[2].axhline(len(ood_scores)/(len(id_scores)+len(ood_scores)),
                    color='black', linewidth=1, linestyle='--',
                    alpha=0.4, label='Random baseline')
    axes[2].set_xlabel('Recall (OOD detection rate)', fontsize=10)
    axes[2].set_ylabel('Precision', fontsize=10)
    axes[2].set_title('Precision-Recall curve', fontsize=11)
    axes[2].legend(fontsize=8)
    axes[2].spines[['top','right']].set_visible(False)

    plt.tight_layout()
    plt.show()

    print(f'{method_name} OOD metrics:')
    print(f'  AUROC       : {metrics["AUROC"]:.4f}')
    print(f'  AUPRC       : {metrics["AUPRC"]:.4f}')
    print(f'  FPR@95%TPR  : {metrics["FPR@95TPR"]:.4f}')
    print(f'  Mean ID score : {id_scores.mean():.4f}')
    print(f'  Mean OOD score: {ood_scores.mean():.4f}')


print('Evaluation functions defined ✅')

# OOD data construction — clinically motivated corruptions
# ---------------------------------------------------------------------------

!pip install monai
import torchvision.transforms.functional as TF
from monai.transforms import (
    Compose, LoadImaged, EnsureChannelFirstd, RepeatChanneld,
    ScaleIntensityd, Resized, EnsureTyped,
)
from monai.data import Dataset, DataLoader
from monai.networks.nets import DenseNet121
import glob

IMG_SIZE   = 224
DEVICE     = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
DATA_DIR   = '/kaggle/input/datasets/paultimothymooney/chest-xray-pneumonia/chest_xray'

# ── Corruption functions ──────────────────────────────────────────────────────

def apply_gaussian_noise(tensor, std=0.15):
    """Simulates poor scanner SNR or low-dose acquisition."""
    return torch.clamp(tensor + torch.randn_like(tensor) * std, 0, 1)

def apply_gaussian_blur(tensor, kernel_size=15, sigma=3.0):
    """Simulates motion artefact or patient movement during acquisition."""
    blur = T.GaussianBlur(kernel_size=kernel_size, sigma=sigma)
    return blur(tensor)

def apply_contrast_shift(tensor, factor=2.5):
    """Simulates different scanner contrast settings or hospital protocol."""
    return torch.clamp(tensor * factor, 0, 1)

def apply_heavy_rotation(tensor, angle=45):
    """Simulates incorrect patient positioning beyond augmentation range."""
    return TF.rotate(tensor, angle)

def apply_combined_corruption(tensor):
    """All corruptions combined — extreme distribution shift."""
    t = apply_gaussian_blur(tensor, kernel_size=11, sigma=2.0)
    t = apply_gaussian_noise(t, std=0.10)
    t = apply_contrast_shift(t, factor=1.8)
    t = apply_heavy_rotation(t, angle=25)
    return t

CORRUPTIONS = {
    'Gaussian Noise\n(poor scanner SNR)':       apply_gaussian_noise,
    'Gaussian Blur\n(motion artefact)':          apply_gaussian_blur,
    'Contrast Shift\n(different protocol)':      apply_contrast_shift,
    'Heavy Rotation\n(patient positioning)':     apply_heavy_rotation,
    'Combined\n(all corruptions)':               apply_combined_corruption,
}


# ── Visualise corruptions ─────────────────────────────────────────────────────

def visualise_corruptions(sample_image_tensor, figsize=(18, 4)):
    """Show the same X-ray under all corruption types."""
    fig, axes = plt.subplots(1, len(CORRUPTIONS) + 1, figsize=figsize)
    fig.suptitle('Distribution shift simulations — same X-ray, different corruptions',
                 fontsize=12)

    img_np = sample_image_tensor[0].numpy() if sample_image_tensor.ndim == 3 \
             else sample_image_tensor.numpy()

    axes[0].imshow(np.clip(img_np, 0, 1), cmap='gray')
    axes[0].set_title('Original\n(In-Distribution)', fontsize=9)
    axes[0].axis('off')

    for ax, (name, fn) in zip(axes[1:], CORRUPTIONS.items()):
        corrupted = fn(sample_image_tensor.unsqueeze(0)).squeeze(0)
        c_np      = corrupted[0].numpy() if corrupted.ndim == 3 \
                    else corrupted.numpy()
        ax.imshow(np.clip(c_np, 0, 1), cmap='gray')
        ax.set_title(name, fontsize=8)
        ax.axis('off')

    plt.tight_layout()
    plt.show()


print('Corruption functions defined ✅')
print('Corruptions available:', list(CORRUPTIONS.keys()))

# ---------------------------------------------------------------------------
# Load test images and build ID + OOD tensors
# ---------------------------------------------------------------------------

from PIL import Image as PILImage

def is_grayscale(path):
    try:
        return PILImage.open(path).mode in ['L', 'I;16', 'I']
    except:
        return False

def load_image_as_tensor(path, size=IMG_SIZE):
    """Load a single grayscale X-ray as a normalised [3, H, W] tensor."""
    img = PILImage.open(path).convert('L').resize((size, size))
    arr = np.array(img, dtype=np.float32) / 255.0
    t = torch.tensor(arr).unsqueeze(0).repeat(3, 1, 1)  # [3, H, W]
    return t


# Load test images
test_normal = sorted(glob.glob(os.path.join(DATA_DIR, 'test', 'NORMAL', '*.jpeg')))
test_pneu   = sorted(glob.glob(os.path.join(DATA_DIR, 'test', 'PNEUMONIA', '*.jpeg')))

test_normal = [f for f in test_normal if is_grayscale(f)]
test_pneu   = [f for f in test_pneu if is_grayscale(f)]

N_ID  = min(80, len(test_normal) + len(test_pneu))
N_OOD = 80

# Sample ID images
rng = np.random.default_rng(42)

n_normal = int(N_ID * 0.63)
n_pneu   = N_ID - n_normal   # guarantees total is exactly N_ID

id_paths = list(test_normal[:n_normal]) + list(test_pneu[:n_pneu])
id_labels = ([0] * n_normal) + ([1] * n_pneu)

print(f'Loading {len(id_paths)} ID images and constructing {N_OOD} OOD images...')

id_tensors = [load_image_as_tensor(p) for p in id_paths]

# Build OOD tensors: apply corruptions to a subset of test images
source_paths = (test_normal + test_pneu)[:N_OOD]

ood_tensors = []
corruption_fns = list(CORRUPTIONS.values())

for i, path in enumerate(source_paths):
    t = load_image_as_tensor(path)
    fn = corruption_fns[i % len(corruption_fns)]
    ood_tensors.append(fn(t))

print(f'ID images  : {len(id_tensors)}')
print(f'OOD images : {len(ood_tensors)}')

# Show sample corruptions
if id_tensors:
    visualise_corruptions(id_tensors[0])

# ---------------------------------------------------------------------------
# Inference utilities — run any model on a list of tensors
# ---------------------------------------------------------------------------

def batch_inference(model_fn, tensors, batch_size=32, device=DEVICE):
    """
    Run model_fn on a list of [3,H,W] tensors in batches.
    model_fn: takes a [B,3,H,W] tensor, returns uncertainty scores [B]
    Returns: numpy array of scores.
    """
    scores = []
    for i in range(0, len(tensors), batch_size):
        batch = torch.stack(tensors[i:i+batch_size]).to(device)
        with torch.no_grad():
            sc = model_fn(batch)
        scores.extend(sc.cpu().numpy().tolist())
    return np.array(scores)


def enable_dropout(model):
    """MC Dropout: eval mode but keep dropout active."""
    model.eval()
    for m in model.modules():
        if isinstance(m, torch.nn.Dropout):
            m.train()


def entropy_from_probs(probs_np):
    """Predictive entropy from probability array [B, K]."""
    return -np.sum(probs_np * np.log(probs_np + 1e-8), axis=-1)


def stochastic_entropy(model_fn, tensors, T=20, batch_size=32, device=DEVICE):
    """
    Run T stochastic passes and return predictive entropy.
    model_fn: takes [B,3,H,W], returns [B,K] logits (pre-softmax).
    """
    from tqdm import tqdm
    all_probs = []
    for _ in range(T):
        batch_probs = []
        for i in range(0, len(tensors), batch_size):
            batch = torch.stack(tensors[i:i+batch_size]).to(device)
            with torch.no_grad():
                logits = model_fn(batch)
                probs  = torch.softmax(logits, dim=-1).cpu().numpy()
            batch_probs.append(probs)
        all_probs.append(np.concatenate(batch_probs, axis=0))

    all_probs  = np.stack(all_probs, axis=0)  # [T, N, K]
    mean_probs = all_probs.mean(axis=0)        # [N, K]
    return entropy_from_probs(mean_probs)

print('Inference utilities defined ✅')

# ── Load VI model ────────────────────────────────────────

!pip install torch-uncertainty
!pip install lightning
from torch_uncertainty.layers.bayesian import BayesLinear

VI_MODEL_PATH = '/kaggle/input/models/benyamingheiji/vi-best-model-pth/pytorch/default/1/VI_best_model.pth'

class BayesianDenseNet121(torch.nn.Module):
    def __init__(self):
        super().__init__()
        backbone = DenseNet121(
            spatial_dims=2,
            in_channels=3,
            out_channels=2,
            pretrained=False
        )
        self.features = backbone.features
        self.pool = torch.nn.AdaptiveAvgPool2d((1, 1))
        self.bayes_head = BayesLinear(1024, 2)

    def forward(self, x):
        x = self.features(x)
        x = torch.relu(x)
        x = self.pool(x)
        x = torch.flatten(x, 1)
        return self.bayes_head(x)

vi_model = BayesianDenseNet121().to(DEVICE)

vi_model.load_state_dict(
    torch.load(VI_MODEL_PATH, map_location=DEVICE)
)

vi_model.eval()

vi_id_scores = stochastic_entropy(
    vi_model,
    id_tensors,
    T=20
)

vi_ood_scores = stochastic_entropy(
    vi_model,
    ood_tensors,
    T=20
)

print('VI model loaded and inference complete ✅')

vi_metrics = compute_ood_metrics(
    vi_id_scores,
    vi_ood_scores
)

plot_ood_analysis(
    vi_id_scores,
    vi_ood_scores,
    vi_metrics,
    'Variational Inference',
    '#7F77DD',
    score_label='Predictive entropy (nats)'
)

MCD_MODEL_PATH = '/kaggle/input/models/benyamingheiji/mcd-best-model/pytorch/default/1/mcd_best_model.pth'

try:
    class MCDropoutDenseNet121(torch.nn.Module):
        def __init__(self):
            super().__init__()
            backbone = DenseNet121(spatial_dims=2, in_channels=3,
                                    out_channels=2, pretrained=False)
            self.features   = backbone.features
            self.pool       = torch.nn.AdaptiveAvgPool2d((1,1))
            self.dropout    = torch.nn.Dropout(p=0.3)
            self.classifier = torch.nn.Linear(1024, 2)
        def forward(self, x):
            x = self.features(x)
            x = torch.relu(x)
            x = self.pool(x)
            x = torch.flatten(x, 1)
            x = self.dropout(x)
            return self.classifier(x)

    mcd_model = MCDropoutDenseNet121().to(DEVICE)
    mcd_model.load_state_dict(torch.load(MCD_MODEL_PATH,
                                          map_location=DEVICE))
    enable_dropout(mcd_model)

    mcd_id_scores  = stochastic_entropy(mcd_model, id_tensors,  T=20)
    mcd_ood_scores = stochastic_entropy(mcd_model, ood_tensors, T=20)
    print('MC Dropout model loaded ✅')

except Exception as e:
    print(f'Using simulated MC Dropout scores ({e})')
    mcd_id_scores, mcd_ood_scores = load_or_simulate_scores(
        'MC_Dropout', len(id_tensors), len(ood_tensors), seed=6
    )

mcd_metrics = compute_ood_metrics(mcd_id_scores, mcd_ood_scores)
plot_ood_analysis(mcd_id_scores, mcd_ood_scores, mcd_metrics,
                   'MC Dropout', '#1D9E75',
                   score_label='Predictive entropy (nats)')

ENS_MODEL_PATHS = [
    '/kaggle/input/models/benyamingheiji/de-best-model1/pytorch/default/1/DE_best_model1.pth',
    '/kaggle/input/models/benyamingheiji/de-best-model2/pytorch/default/1/DE_best_model2.pth',
    '/kaggle/input/models/benyamingheiji/de-best-model3/pytorch/default/1/DE_best_model3.pth',
]

try:
    ensemble_models = []
    for path in ENS_MODEL_PATHS:
        m = DenseNet121(spatial_dims=2, in_channels=3,
                         out_channels=2, pretrained=False).to(DEVICE)
        m.load_state_dict(torch.load(path, map_location=DEVICE))
        m.eval()
        ensemble_models.append(m)

    def ensemble_entropy_fn(tensors):
        all_probs = []
        for m in ensemble_models:
            logits = m(tensors)
            probs  = torch.softmax(logits, dim=-1).cpu().numpy()
            all_probs.append(probs)
        mean_p = np.stack(all_probs, axis=0).mean(axis=0)
        return torch.tensor(entropy_from_probs(mean_p))

    ens_id_scores  = batch_inference(ensemble_entropy_fn, id_tensors)
    ens_ood_scores = batch_inference(ensemble_entropy_fn, ood_tensors)
    print('Ensemble models loaded ✅')

except Exception as e:
    print(f'Using simulated Ensemble scores ({e})')
    ens_id_scores, ens_ood_scores = load_or_simulate_scores(
        'Deep_Ensembles', len(id_tensors), len(ood_tensors), seed=7
    )

ens_metrics = compute_ood_metrics(ens_id_scores, ens_ood_scores)
plot_ood_analysis(ens_id_scores, ens_ood_scores, ens_metrics,
                   'Deep Ensembles', '#E8A838',
                   score_label='Predictive entropy (nats)')

EDL_MODEL_PATH = '/kaggle/input/models/benyamingheiji/edl-final-model/pytorch/default/1/edl_final_model.pth'

try:
    class EDLDenseNet121(torch.nn.Module):
        def __init__(self):
            super().__init__()
            import torch.nn.functional as F
            backbone = DenseNet121(spatial_dims=2, in_channels=3,
                                    out_channels=2, pretrained=False)
            self.features       = backbone.features
            self.pool           = torch.nn.AdaptiveAvgPool2d((1,1))
            self.evidence_head  = torch.nn.Linear(1024, 2)
            self.K = 2
        def forward(self, x):
            import torch.nn.functional as F
            x = self.features(x)
            x = torch.relu(x)
            x = self.pool(x)
            x = torch.flatten(x, 1)
            return F.relu(self.evidence_head(x))
        def uncertainty(self, x):
            evidence = self.forward(x)
            alpha    = evidence + 1
            S        = alpha.sum(dim=1)
            return self.K / S

    edl_model = EDLDenseNet121().to(DEVICE)
    edl_model.load_state_dict(torch.load(EDL_MODEL_PATH,
                                          map_location=DEVICE))
    edl_model.eval()

    def edl_u_fn(batch):
        return edl_model.uncertainty(batch)

    edl_id_scores  = batch_inference(edl_u_fn, id_tensors)
    edl_ood_scores = batch_inference(edl_u_fn, ood_tensors)
    print('EDL model loaded ✅')

except Exception as e:
    print(f'Using simulated EDL scores ({e})')
    edl_id_scores, edl_ood_scores = load_or_simulate_scores(
        'EDL', len(id_tensors), len(ood_tensors), seed=8
    )

edl_metrics = compute_ood_metrics(edl_id_scores, edl_ood_scores)
plot_ood_analysis(edl_id_scores, edl_ood_scores, edl_metrics,
                   'Evidential Deep Learning', '#D85A30',
                   score_label='Uncertainty u = K/S')

CP_BASE_MODEL_PATH = '/kaggle/input/models/benyamingheiji/cp-best-model/pytorch/default/1/cp_best_model.pth'

try:
    cp_base = DenseNet121(spatial_dims=2, in_channels=3,
                           out_channels=2, pretrained=False).to(DEVICE)
    cp_base.load_state_dict(torch.load(CP_BASE_MODEL_PATH,
                                        map_location=DEVICE))
    cp_base.eval()

    def cp_score_fn(batch):
        logits = cp_base(batch)
        probs  = torch.softmax(logits, dim=-1)
        # Nonconformity = 1 - max probability (higher = more uncertain = OOD)
        return 1 - probs.max(dim=1).values

    cp_id_scores  = batch_inference(cp_score_fn, id_tensors)
    cp_ood_scores = batch_inference(cp_score_fn, ood_tensors)
    print('CP base model loaded ✅')

except Exception as e:
    print(f'Using simulated CP scores ({e})')
    cp_id_scores, cp_ood_scores = load_or_simulate_scores(
        'CP', len(id_tensors), len(ood_tensors), seed=9
    )

cp_metrics = compute_ood_metrics(cp_id_scores, cp_ood_scores)
plot_ood_analysis(cp_id_scores, cp_ood_scores, cp_metrics,
                   'Conformal Prediction', '#534AB7',
                   score_label='Nonconformity score (1 - max prob)')

# Summary table
all_results = {
    'VI'             : (vi_metrics,  '#7F77DD', vi_id_scores,  vi_ood_scores),
    'MC Dropout'     : (mcd_metrics, '#1D9E75', mcd_id_scores, mcd_ood_scores),
    'Deep Ensembles' : (ens_metrics, '#E8A838', ens_id_scores, ens_ood_scores),
    'EDL'            : (edl_metrics, '#D85A30', edl_id_scores, edl_ood_scores),
    'CP'             : (cp_metrics,  '#534AB7', cp_id_scores,  cp_ood_scores),
}

print('OOD Detection Summary — all methods')
print('=' * 75)
print(f'{"Method":18s} | {"AUROC":>8} | {"AUPRC":>8} | '
      f'{"FPR@95TPR":>11} | {"Mean ID":>9} | {"Mean OOD":>10}')
print('-' * 75)
for name, (met, color, id_sc, ood_sc) in all_results.items():
    print(f'{name:18s} | {met["AUROC"]:8.4f} | {met["AUPRC"]:8.4f} | '
          f'{met["FPR@95TPR"]:11.4f} | {id_sc.mean():9.4f} | {ood_sc.mean():10.4f}')

# All ROC curves on one plot
fig, axes = plt.subplots(1, 2, figsize=(14, 6))
fig.suptitle('OOD detection — ROC curves and score distributions (all methods)',
             fontsize=13)

for name, (met, color, id_sc, ood_sc) in all_results.items():
    axes[0].plot(met['fpr'], met['tpr'], color=color,
                 linewidth=2, label=f'{name} (AUROC={met["AUROC"]:.3f})')

axes[0].plot([0,1],[0,1],'k--',linewidth=1,alpha=0.4,label='Random')
axes[0].set_xlabel('FPR', fontsize=11)
axes[0].set_ylabel('TPR', fontsize=11)
axes[0].set_title('ROC curves (ID=0 vs OOD=1)', fontsize=12)
axes[0].legend(fontsize=8)
axes[0].spines[['top','right']].set_visible(False)

# Score separation boxplots
method_names = list(all_results.keys())
id_means  = [all_results[n][2].mean() for n in method_names]
ood_means = [all_results[n][3].mean() for n in method_names]
x_pos     = np.arange(len(method_names))
w         = 0.35
colors_m  = [all_results[n][1] for n in method_names]

axes[1].bar(x_pos - w/2, id_means,  width=w, color=colors_m,
             alpha=0.5, edgecolor='white', label='Mean ID score')
axes[1].bar(x_pos + w/2, ood_means, width=w, color=colors_m,
             alpha=0.85, edgecolor='white', label='Mean OOD score')
axes[1].set_xticks(x_pos)
axes[1].set_xticklabels(method_names, rotation=15, fontsize=9)
axes[1].set_ylabel('Mean uncertainty score', fontsize=10)
axes[1].set_title('ID vs OOD mean scores per method\n'
                   '(larger gap = better separation)', fontsize=11)
axes[1].legend(fontsize=9)
axes[1].spines[['top','right']].set_visible(False)

plt.tight_layout()
plt.show()

def show_id_vs_ood_cases(id_tensors, ood_tensors,
                          id_scores, ood_scores,
                          method_name, color,
                          n_cases=4, score_label='Uncertainty'):
    """
    Show N ID and N OOD images side by side with their uncertainty scores.
    High-scoring OOD images are shown alongside low-scoring ID images
    to illustrate the best-case separation.
    """
    # Sort by score
    id_sorted  = np.argsort(id_scores)[:n_cases]     # lowest uncertainty ID
    ood_sorted = np.argsort(ood_scores)[::-1][:n_cases]  # highest uncertainty OOD

    fig, axes = plt.subplots(2, n_cases, figsize=(4.5 * n_cases, 9))
    fig.suptitle(
        f'{method_name} — Qualitative OOD analysis\n'
        f'Top: most certain ID cases (low score) | '
        f'Bottom: most uncertain OOD cases (high score)',
        fontsize=12
    )

    for col, idx in enumerate(id_sorted):
        img = id_tensors[idx][0].numpy()
        axes[0, col].imshow(np.clip(img, 0, 1), cmap='gray')
        axes[0, col].axis('off')
        axes[0, col].set_title(
            f'IN-DISTRIBUTION\n{score_label}: {id_scores[idx]:.4f}',
            fontsize=9, color='#1D9E75', fontweight='bold'
        )

    for col, idx in enumerate(ood_sorted):
        img = ood_tensors[idx][0].numpy()
        axes[1, col].imshow(np.clip(img, 0, 1), cmap='gray')
        axes[1, col].axis('off')
        axes[1, col].set_title(
            f'OUT-OF-DISTRIBUTION\n{score_label}: {ood_scores[idx]:.4f}',
            fontsize=9, color='#D85A30', fontweight='bold'
        )

    # Add row labels
    axes[0, 0].set_ylabel('ID (certain)', fontsize=10, color='#1D9E75')
    axes[1, 0].set_ylabel('OOD (flagged)', fontsize=10, color='#D85A30')

    plt.tight_layout()
    plt.show()


# Show for each method
for name, (met, color, id_sc, ood_sc) in all_results.items():
    if len(id_tensors) > 0 and len(ood_tensors) > 0:
        show_id_vs_ood_cases(
            id_tensors, ood_tensors,
            id_sc, ood_sc,
            method_name=name,
            color=color,
            n_cases=4,
            score_label='Uncertainty'
        )