Add integrated_gradients_xai.py

integrated_gradients_xai.py

@@ -0,0 +1,858 @@
+#!/usr/bin/env python3
+"""
+RetinaSense v3.0 -- Phase 1C: Advanced XAI with Integrated Gradients
+=====================================================================
+Compares Attention Rollout (existing) vs Integrated Gradients (captum)
+on 20 test images (4 per class).
+
+Outputs (all saved to outputs_v3/xai/):
+  - comparison_grid.png       : 20-row x 3-column grid [Original | Rollout | IG]
+  - ig_individual_01..20.png  : Individual IG heatmaps
+  - agreement_heatmap.png     : Spatial correlation matrix between methods
+  - agreement_score.json      : Numerical agreement scores per image
+
+Usage:
+  python integrated_gradients_xai.py
+"""
+
+import os
+import sys
+import json
+import warnings
+import numpy as np
+import cv2
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+from matplotlib.colors import Normalize
+from PIL import Image
+import pandas as pd
+from scipy.stats import pearsonr
+
+warnings.filterwarnings('ignore')
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import transforms
+import timm
+
+from captum.attr import IntegratedGradients
+
+# Maximize CPU parallelism
+torch.set_num_threads(os.cpu_count() or 4)
+
+# ================================================================
+# CONFIGURATION
+# ================================================================
+BASE_DIR   = '/teamspace/studios/this_studio'
+OUTPUT_DIR = os.path.join(BASE_DIR, 'outputs_v3')
+XAI_DIR    = os.path.join(OUTPUT_DIR, 'xai')
+os.makedirs(XAI_DIR, exist_ok=True)
+
+MODEL_PATH       = os.path.join(OUTPUT_DIR, 'best_model.pth')
+TEMPERATURE_PATH = os.path.join(OUTPUT_DIR, 'temperature.json')
+TEST_CSV         = os.path.join(BASE_DIR, 'data', 'test_split.csv')
+NORM_STATS_PATH  = os.path.join(BASE_DIR, 'data', 'fundus_norm_stats.json')
+
+CLASS_NAMES = ['Normal', 'Diabetes/DR', 'Glaucoma', 'Cataract', 'AMD']
+NUM_CLASSES = 5
+IMG_SIZE    = 224
+DROPOUT     = 0.3
+N_PER_CLASS = 4
+
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+print('=' * 65)
+print('   RetinaSense v3.0 -- Phase 1C: Integrated Gradients XAI')
+print('=' * 65)
+print(f'  Device   : {DEVICE}')
+if torch.cuda.is_available():
+    print(f'  GPU      : {torch.cuda.get_device_name(0)}')
+print(f'  Output   : {XAI_DIR}')
+print('=' * 65)
+
+# ================================================================
+# LOAD NORMALISATION STATS
+# ================================================================
+if os.path.exists(NORM_STATS_PATH):
+    with open(NORM_STATS_PATH) as f:
+        norm_stats = json.load(f)
+    NORM_MEAN = norm_stats['mean_rgb']
+    NORM_STD  = norm_stats['std_rgb']
+    print(f'  Fundus norm stats: mean={[round(v,4) for v in NORM_MEAN]}, '
+          f'std={[round(v,4) for v in NORM_STD]}')
+else:
+    NORM_MEAN = [0.485, 0.456, 0.406]
+    NORM_STD  = [0.229, 0.224, 0.225]
+    print('  Using ImageNet normalisation fallback')
+
+with open(TEMPERATURE_PATH) as f:
+    TEMPERATURE = json.load(f)['temperature']
+print(f'  Temperature T = {TEMPERATURE:.4f}')
+
+
+# ================================================================
+# MODEL ARCHITECTURE (mirrors gradcam_v3.py / retinasense_v3.py)
+# ================================================================
+class MultiTaskViT(nn.Module):
+    """ViT-Base-Patch16-224 with disease + severity heads."""
+
+    def __init__(self, n_disease=NUM_CLASSES, n_severity=5, drop=DROPOUT):
+        super().__init__()
+        self.backbone = timm.create_model(
+            'vit_base_patch16_224', pretrained=False, num_classes=0
+        )
+        feat = 768
+
+        self.drop = nn.Dropout(drop)
+
+        self.disease_head = nn.Sequential(
+            nn.Linear(feat, 512), nn.BatchNorm1d(512), nn.ReLU(), nn.Dropout(0.3),
+            nn.Linear(512, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.2),
+            nn.Linear(256, n_disease),
+        )
+        self.severity_head = nn.Sequential(
+            nn.Linear(feat, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.3),
+            nn.Linear(256, n_severity),
+        )
+
+    def forward(self, x):
+        f = self.backbone(x)
+        f = self.drop(f)
+        return self.disease_head(f), self.severity_head(f)
+
+
+# ================================================================
+# DISEASE-LOGITS WRAPPER FOR CAPTUM
+# ================================================================
+class DiseaseLogitModel(nn.Module):
+    """
+    Wraps MultiTaskViT so that forward(x) returns only the disease logits.
+    Captum's IntegratedGradients requires a model whose forward output
+    is either a scalar or a 1-D tensor. We select the target class
+    logit inside the IG call via the `target` parameter, so here we
+    return the full (B, 5) disease logits.
+    """
+
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+
+    def forward(self, x):
+        disease_logits, _ = self.model(x)
+        return disease_logits
+
+
+# ================================================================
+# ATTENTION ROLLOUT (copied from gradcam_v3.py for self-containment)
+# ================================================================
+class ViTAttentionRollout:
+    """
+    Attention Rollout for Vision Transformer.
+    Traces information flow from patches to CLS token across all layers.
+    """
+
+    def __init__(self, model, discard_ratio=0.97):
+        self.model = model
+        self.discard_ratio = discard_ratio
+        self._attention_maps = []
+        self._hooks = []
+
+        # Disable fused attention for explicit weight access
+        for block in model.backbone.blocks:
+            block.attn.fused_attn = False
+
+        # Register forward hooks on all transformer blocks
+        for block in model.backbone.blocks:
+            h = block.attn.register_forward_hook(self._attn_hook)
+            self._hooks.append(h)
+
+    def _attn_hook(self, module, input, output):
+        """Capture softmax attention weights from each block."""
+        x = input[0]
+        B, N, C = x.shape
+        with torch.no_grad():
+            qkv = module.qkv(x).reshape(
+                B, N, 3, module.num_heads, module.head_dim
+            ).permute(2, 0, 3, 1, 4)
+            q, k, _ = qkv.unbind(0)
+            q, k = module.q_norm(q), module.k_norm(k)
+            attn = (q * module.scale @ k.transpose(-2, -1)).softmax(dim=-1)
+            self._attention_maps.append(attn.detach().cpu())
+
+    def generate(self, image_tensor, class_idx=None):
+        """
+        Generate attention rollout heatmap.
+        Returns:
+            heatmap (224, 224) float32 [0, 1], predicted_label, confidence
+        """
+        self.model.eval()
+        self._attention_maps = []
+
+        with torch.no_grad():
+            image_tensor = image_tensor.to(DEVICE)
+            d_out, _ = self.model(image_tensor)
+            probs = torch.softmax(d_out, dim=1)
+            predicted_label = int(probs.argmax(dim=1).item())
+            confidence = float(probs[0, predicted_label].item())
+
+        if class_idx is None:
+            class_idx = predicted_label
+
+        attn_stack = torch.stack(self._attention_maps, dim=0)[:, 0]
+        attn_mean = attn_stack.mean(dim=1)
+
+        if self.discard_ratio > 0:
+            flat = attn_mean.reshape(attn_mean.shape[0], -1)
+            thresh = torch.quantile(flat, self.discard_ratio, dim=1, keepdim=True)
+            thresh = thresh.unsqueeze(-1)
+            attn_mean = torch.where(
+                attn_mean >= thresh, attn_mean, torch.zeros_like(attn_mean)
+            )
+
+        I = torch.eye(attn_mean.shape[-1]).unsqueeze(0)
+        attn_aug = attn_mean + I
+        attn_aug = attn_aug / attn_aug.sum(dim=-1, keepdim=True).clamp(min=1e-8)
+
+        rollout = attn_aug[0]
+        for l in range(1, len(attn_aug)):
+            rollout = rollout @ attn_aug[l]
+
+        cls_attention = rollout[0, 1:]
+        spatial = cls_attention.numpy().reshape(14, 14).astype(np.float32)
+        spatial = cv2.resize(spatial, (IMG_SIZE, IMG_SIZE),
+                             interpolation=cv2.INTER_LINEAR)
+
+        s_min, s_max = spatial.min(), spatial.max()
+        if s_max - s_min > 1e-8:
+            spatial = (spatial - s_min) / (s_max - s_min)
+        else:
+            spatial = np.zeros_like(spatial)
+
+        spatial = np.power(spatial, 0.4)  # gamma stretch
+
+        return spatial.astype(np.float32), predicted_label, confidence
+
+    def remove_hooks(self):
+        for h in self._hooks:
+            h.remove()
+        self._hooks = []
+
+
+# ================================================================
+# IMAGE PREPROCESSING (mirrors gradcam_v3.py)
+# ================================================================
+def ben_graham(path, sz=IMG_SIZE, sigma=10):
+    """Ben Graham high-frequency fundus enhancement (APTOS-style)."""
+    img = cv2.imread(path)
+    if img is None:
+        img = np.array(Image.open(path).convert('RGB'))
+        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    img = cv2.resize(img, (sz, sz))
+    img = cv2.addWeighted(img, 4, cv2.GaussianBlur(img, (0, 0), sigma), -4, 128)
+    mask = np.zeros(img.shape[:2], dtype=np.uint8)
+    cv2.circle(mask, (sz // 2, sz // 2), int(sz * 0.48), 255, -1)
+    return cv2.bitwise_and(img, img, mask=mask)
+
+
+def clahe_preprocess(path, sz=IMG_SIZE):
+    """CLAHE contrast enhancement (ODIR-style)."""
+    img = cv2.imread(path)
+    if img is None:
+        img = np.array(Image.open(path).convert('RGB'))
+        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+    img = cv2.resize(img, (sz, sz))
+    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
+    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+    lab[:, :, 0] = clahe.apply(lab[:, :, 0])
+    img = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
+    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+
+def load_and_preprocess(image_path, dataset='auto'):
+    """
+    Load image with domain-conditional preprocessing.
+    Returns:
+        img_np   : (224, 224, 3) uint8 preprocessed
+        img_orig : (224, 224, 3) uint8 original
+    """
+    if not os.path.isabs(image_path):
+        clean = image_path
+        while clean.startswith('./'):
+            clean = clean[2:]
+        image_path = os.path.join(BASE_DIR, clean)
+
+    if dataset == 'auto':
+        if 'aptos' in image_path.lower() or 'gaussian' in image_path.lower():
+            dataset = 'APTOS'
+        else:
+            dataset = 'ODIR'
+
+    raw = cv2.imread(image_path)
+    if raw is None:
+        raw = np.array(Image.open(image_path).convert('RGB'))
+    else:
+        raw = cv2.cvtColor(raw, cv2.COLOR_BGR2RGB)
+    img_orig = cv2.resize(raw, (IMG_SIZE, IMG_SIZE))
+
+    if dataset == 'APTOS':
+        img_np = ben_graham(image_path, sz=IMG_SIZE)
+    else:
+        img_np = clahe_preprocess(image_path, sz=IMG_SIZE)
+
+    return img_np, img_orig
+
+
+def preprocess_to_tensor(img_np):
+    """Convert preprocessed numpy image to normalised tensor (1, 3, 224, 224)."""
+    transform = transforms.Compose([
+        transforms.ToPILImage(),
+        transforms.ToTensor(),
+        transforms.Normalize(NORM_MEAN, NORM_STD),
+    ])
+    return transform(img_np).unsqueeze(0)
+
+
+# ================================================================
+# CIRCULAR FUNDUS MASK
+# ================================================================
+def create_fundus_mask(h=IMG_SIZE, w=IMG_SIZE):
+    """
+    Create a soft circular mask matching the fundus region.
+    Uses a smooth Gaussian-blurred edge to avoid hard boundaries.
+    Returns float32 mask [0, 1] of shape (h, w).
+    """
+    cy, cx = h // 2, w // 2
+    radius = min(h, w) // 2 - 5
+    mask = np.zeros((h, w), dtype=np.float32)
+    cv2.circle(mask, (cx, cy), radius, 1.0, -1)
+    mask = cv2.GaussianBlur(mask, (21, 21), 0)
+    return mask
+
+
+# ================================================================
+# INTEGRATED GRADIENTS COMPUTATION
+# ================================================================
+def compute_ig_attribution(ig_model, ig_method, img_tensor, target_class,
+                           n_steps=50, internal_batch_size=10, sigma=10):
+    """
+    Compute Integrated Gradients attribution for a single image.
+
+    Uses a Gaussian-blurred baseline (sigma=10) which is more appropriate
+    for fundus images than a black baseline (since the background is already dark).
+
+    Args:
+        ig_model          : DiseaseLogitModel wrapper
+        ig_method         : captum IntegratedGradients instance
+        img_tensor        : (1, 3, 224, 224) normalised tensor on DEVICE
+        target_class      : int, disease class to explain
+        n_steps           : number of interpolation steps
+        internal_batch_size : batch size for internal IG computation
+        sigma             : Gaussian blur sigma for baseline
+
+    Returns:
+        attribution : (224, 224) float32 numpy array, normalised [0, 1]
+    """
+    # Create blurred baseline in pixel space, then normalise
+    # First undo normalisation to get pixel-space tensor
+    mean_t = torch.tensor(NORM_MEAN, device=DEVICE).view(1, 3, 1, 1)
+    std_t  = torch.tensor(NORM_STD, device=DEVICE).view(1, 3, 1, 1)
+
+    # Build the blurred baseline from the input tensor
+    # Denormalise -> blur -> renormalise
+    img_denorm = img_tensor * std_t + mean_t  # approx [0, 1] range
+    img_np_for_blur = (img_denorm[0].permute(1, 2, 0).cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
+    blurred_np = cv2.GaussianBlur(img_np_for_blur, (0, 0), sigma)
+    # Convert blurred back to tensor with normalisation
+    blurred_tensor = transforms.Compose([
+        transforms.ToPILImage(),
+        transforms.ToTensor(),
+        transforms.Normalize(NORM_MEAN, NORM_STD),
+    ])(blurred_np).unsqueeze(0).to(DEVICE)
+
+    # Compute Integrated Gradients
+    img_tensor.requires_grad_(True)
+    attributions = ig_method.attribute(
+        img_tensor,
+        baselines=blurred_tensor,
+        target=target_class,
+        n_steps=n_steps,
+        internal_batch_size=internal_batch_size,
+    )
+
+    # Aggregate across channels: take L2 norm across RGB for spatial map
+    # shape: (1, 3, 224, 224) -> (224, 224)
+    attr_np = attributions[0].detach().cpu().numpy()  # (3, 224, 224)
+    # Use absolute values and sum over channels for a positive attribution map
+    attr_spatial = np.sqrt(np.sum(attr_np ** 2, axis=0))  # (224, 224)
+
+    # Normalise to [0, 1]
+    a_min, a_max = attr_spatial.min(), attr_spatial.max()
+    if a_max - a_min > 1e-8:
+        attr_spatial = (attr_spatial - a_min) / (a_max - a_min)
+    else:
+        attr_spatial = np.zeros_like(attr_spatial)
+
+    return attr_spatial.astype(np.float32)
+
+
+# ================================================================
+# OVERLAY FUNCTION
+# ================================================================
+def overlay_heatmap(original_np, heatmap, alpha=0.6, cmap_name='inferno'):
+    """
+    Blend heatmap onto original image with circular fundus mask.
+
+    Args:
+        original_np : (224, 224, 3) uint8 RGB
+        heatmap     : (224, 224) float32 [0, 1]
+        alpha       : heatmap blending opacity
+        cmap_name   : matplotlib colormap name
+
+    Returns:
+        blended : (224, 224, 3) uint8 RGB
+    """
+    # Apply colormap
+    cmap = plt.get_cmap(cmap_name)
+    colored = cmap(heatmap)[:, :, :3]  # (224, 224, 3) float [0, 1]
+    colored_uint8 = (colored * 255).astype(np.uint8)
+
+    # Get fundus mask
+    mask = create_fundus_mask(heatmap.shape[0], heatmap.shape[1])
+
+    # Blend inside the fundus region only
+    orig = original_np.astype(np.float32)
+    cmap_f = colored_uint8.astype(np.float32)
+    blended = orig.copy()
+    for c in range(3):
+        blended[:, :, c] = (
+            orig[:, :, c] * (1 - alpha * mask)
+            + cmap_f[:, :, c] * (alpha * mask)
+        )
+    return np.clip(blended, 0, 255).astype(np.uint8)
+
+
+# ================================================================
+# SELECT TEST IMAGES (same logic as gradcam_v3.py)
+# ================================================================
+def select_test_images(n_per_class=N_PER_CLASS):
+    """Select n_per_class images per disease class from test split."""
+    df = pd.read_csv(TEST_CSV)
+    samples = []
+    for label in range(NUM_CLASSES):
+        subset = df[df['disease_label'] == label].drop_duplicates(subset='image_path')
+        chosen = subset.head(n_per_class)
+        for _, row in chosen.iterrows():
+            samples.append({
+                'image_path':  row['image_path'],
+                'true_label':  int(row['disease_label']),
+                'dataset':     str(row.get('source', 'auto')),
+            })
+    print(f'  Selected {len(samples)} test images '
+          f'({n_per_class} per class x {NUM_CLASSES} classes)')
+    return samples
+
+
+# ================================================================
+# COMPUTE AGREEMENT METRICS
+# ================================================================
+def compute_agreement(rollout_map, ig_map, fundus_mask):
+    """
+    Compute spatial agreement between Attention Rollout and IG heatmaps.
+
+    Metrics:
+        - Pearson correlation (within fundus mask)
+        - Intersection over Union (IoU) of top-20% activated regions
+        - Cosine similarity of flattened masked vectors
+
+    Returns dict of scores.
+    """
+    # Flatten inside mask
+    mask_bool = fundus_mask > 0.5
+    r_flat = rollout_map[mask_bool]
+    i_flat = ig_map[mask_bool]
+
+    # Pearson correlation
+    if len(r_flat) > 2 and r_flat.std() > 1e-8 and i_flat.std() > 1e-8:
+        pearson_r, pearson_p = pearsonr(r_flat, i_flat)
+    else:
+        pearson_r, pearson_p = 0.0, 1.0
+
+    # IoU of top-20% regions
+    r_thresh = np.percentile(r_flat, 80)
+    i_thresh = np.percentile(i_flat, 80)
+    r_top = rollout_map > r_thresh
+    i_top = ig_map > i_thresh
+    intersection = np.logical_and(r_top, i_top).sum()
+    union = np.logical_or(r_top, i_top).sum()
+    iou = float(intersection / max(union, 1))
+
+    # Cosine similarity
+    r_norm = np.linalg.norm(r_flat)
+    i_norm = np.linalg.norm(i_flat)
+    if r_norm > 1e-8 and i_norm > 1e-8:
+        cosine = float(np.dot(r_flat, i_flat) / (r_norm * i_norm))
+    else:
+        cosine = 0.0
+
+    return {
+        'pearson_r':     round(float(pearson_r), 4),
+        'pearson_p':     round(float(pearson_p), 6),
+        'iou_top20':     round(iou, 4),
+        'cosine_sim':    round(cosine, 4),
+    }
+
+
+# ================================================================
+# MAIN PIPELINE
+# ================================================================
+def main():
+    import time
+    t_start = time.time()
+
+    # ---- 1. Load model ----
+    print('\n[1/6] Loading model...')
+    model = MultiTaskViT().to(DEVICE)
+    ckpt = torch.load(MODEL_PATH, map_location=DEVICE, weights_only=False)
+    model.load_state_dict(ckpt['model_state_dict'])
+    model.eval()
+    print(f'  Loaded: {MODEL_PATH}')
+    print(f'  Checkpoint epoch: {ckpt.get("epoch", "?") + 1}')
+
+    # ---- 2. Select images ----
+    print('\n[2/6] Selecting test images...')
+    samples = select_test_images()
+
+    # ---- 3. Preprocess all images + run Attention Rollout ----
+    # IMPORTANT: Run rollout FIRST then remove hooks BEFORE IG.
+    # This prevents rollout hooks from firing during IG's many
+    # forward passes (50 steps), which would be very slow on CPU.
+    print('\n[3/6] Running Attention Rollout on all images...')
+    rollout = ViTAttentionRollout(model, discard_ratio=0.97)
+    print(f'  Attention Rollout: {len(rollout._hooks)} hooks registered')
+
+    preprocessed = []  # store (img_np, img_orig, img_tensor) per sample
+    rollout_results = []  # store (heatmap, pred_label, confidence) per sample
+
+    for idx, sample in enumerate(samples):
+        img_path   = sample['image_path']
+        true_label = sample['true_label']
+        dataset    = sample['dataset']
+        basename   = os.path.basename(img_path)
+
+        print(f'  [{idx+1:2d}/{len(samples)}] '
+              f'{CLASS_NAMES[true_label]:15s} | {basename}', end='  ')
+
+        try:
+            img_np, img_orig = load_and_preprocess(img_path, dataset=dataset)
+            img_tensor = preprocess_to_tensor(img_np).to(DEVICE)
+            preprocessed.append((img_np, img_orig, img_tensor))
+
+            heatmap, pred_label, pred_conf = rollout.generate(img_tensor)
+            rollout_results.append((heatmap, pred_label, pred_conf))
+            print(f'-> pred={CLASS_NAMES[pred_label]:15s}  conf={pred_conf:.3f}')
+
+        except Exception as e:
+            print(f'FAILED: {e}')
+            preprocessed.append(None)
+            rollout_results.append(None)
+
+    # Remove rollout hooks BEFORE running IG to avoid extra computation
+    rollout.remove_hooks()
+    # Re-enable fused attention for faster forward passes during IG
+    for block in model.backbone.blocks:
+        block.attn.fused_attn = True
+    print('  Rollout hooks removed. fused_attn re-enabled for IG speed.')
+
+    # ---- 4. Run Integrated Gradients (no rollout hooks active) ----
+    print('\n[4/6] Computing Integrated Gradients attributions...')
+    disease_model = DiseaseLogitModel(model)
+    disease_model.eval()
+    ig_method = IntegratedGradients(disease_model)
+    print(f'  Baseline: Gaussian blur (sigma=10), n_steps=50, '
+          f'internal_batch_size=25')
+
+    all_results = []
+    fundus_mask = create_fundus_mask()
+
+    for idx, sample in enumerate(samples):
+        if preprocessed[idx] is None or rollout_results[idx] is None:
+            continue
+
+        img_path   = sample['image_path']
+        true_label = sample['true_label']
+        basename   = os.path.basename(img_path)
+        img_np, img_orig, img_tensor = preprocessed[idx]
+        rollout_heatmap, pred_label, pred_conf = rollout_results[idx]
+
+        print(f'  [{idx+1:2d}/{len(samples)}] '
+              f'{CLASS_NAMES[true_label]:15s} | {basename}', end='  ')
+
+        try:
+            # Fresh tensor (IG needs requires_grad)
+            ig_input = img_tensor.clone().detach().to(DEVICE)
+
+            ig_heatmap = compute_ig_attribution(
+                disease_model, ig_method, ig_input,
+                target_class=pred_label,
+                n_steps=50,
+                internal_batch_size=25,
+                sigma=10,
+            )
+
+            # Agreement
+            agreement = compute_agreement(rollout_heatmap, ig_heatmap,
+                                          fundus_mask)
+
+            print(f'-> pearson={agreement["pearson_r"]:.3f}  '
+                  f'IoU={agreement["iou_top20"]:.3f}')
+
+            all_results.append({
+                'idx':              idx,
+                'image_path':       img_path,
+                'basename':         basename,
+                'true_label':       true_label,
+                'pred_label':       pred_label,
+                'pred_class':       CLASS_NAMES[pred_label],
+                'confidence':       round(pred_conf, 4),
+                'img_orig':         img_orig,
+                'rollout_heatmap':  rollout_heatmap,
+                'ig_heatmap':       ig_heatmap,
+                'agreement':        agreement,
+            })
+
+        except Exception as e:
+            print(f'FAILED: {e}')
+            import traceback
+            traceback.print_exc()
+            continue
+
+    n_success = len(all_results)
+    print(f'\n  Completed: {n_success}/{len(samples)} images')
+
+    if n_success == 0:
+        print('ERROR: No images processed successfully. Exiting.')
+        sys.exit(1)
+
+    # ---- 5. Generate visualisations ----
+    print('\n[5/6] Generating visualisations...')
+
+    # 5a. Individual IG heatmaps
+    print('  Saving individual IG heatmaps...')
+    for r in all_results:
+        fig, axes = plt.subplots(1, 3, figsize=(13, 4.5))
+
+        # Original
+        axes[0].imshow(r['img_orig'])
+        axes[0].set_title(f'Original\nTrue: {CLASS_NAMES[r["true_label"]]}',
+                          fontsize=10, fontweight='bold')
+        axes[0].axis('off')
+
+        # IG heatmap (raw)
+        im = axes[1].imshow(r['ig_heatmap'], cmap='inferno', vmin=0, vmax=1)
+        axes[1].set_title('Integrated Gradients\n(attribution magnitude)',
+                          fontsize=10)
+        axes[1].axis('off')
+
+        # IG overlay on original
+        ig_overlay = overlay_heatmap(r['img_orig'], r['ig_heatmap'],
+                                     alpha=0.6, cmap_name='inferno')
+        axes[2].imshow(ig_overlay)
+        correct = r['pred_label'] == r['true_label']
+        status = 'OK' if correct else 'WRONG'
+        axes[2].set_title(
+            f'IG Overlay\nPred: {r["pred_class"]} ({r["confidence"]:.2f}) [{status}]',
+            fontsize=10,
+            color='green' if correct else 'red',
+            fontweight='bold')
+        axes[2].axis('off')
+
+        plt.tight_layout()
+        save_path = os.path.join(XAI_DIR,
+                                 f'ig_individual_{r["idx"]+1:02d}.png')
+        fig.savefig(save_path, dpi=150, bbox_inches='tight',
+                    facecolor='white')
+        plt.close(fig)
+
+    print(f'    Saved {n_success} individual IG heatmaps')
+
+    # 5b. Comparison grid: 20 rows x 3 columns
+    print('  Generating comparison grid...')
+    n_rows = n_success
+    fig, axes = plt.subplots(n_rows, 3, figsize=(14, 4.2 * n_rows))
+
+    # Handle single-row case
+    if n_rows == 1:
+        axes = axes[np.newaxis, :]
+
+    # Column headers
+    col_titles = ['Original Image', 'Attention Rollout', 'Integrated Gradients']
+
+    for i, r in enumerate(all_results):
+        true_name = CLASS_NAMES[r['true_label']]
+        pred_name = r['pred_class']
+        correct = r['pred_label'] == r['true_label']
+        status = 'OK' if correct else 'WRONG'
+
+        # Column 0: Original
+        axes[i, 0].imshow(r['img_orig'])
+        axes[i, 0].set_ylabel(f'#{r["idx"]+1}\nTrue: {true_name}',
+                               fontsize=9, fontweight='bold', rotation=0,
+                               labelpad=70, va='center')
+        axes[i, 0].set_xticks([])
+        axes[i, 0].set_yticks([])
+        if i == 0:
+            axes[i, 0].set_title(col_titles[0], fontsize=12, fontweight='bold',
+                                 pad=10)
+
+        # Column 1: Attention Rollout overlay
+        rollout_overlay = overlay_heatmap(r['img_orig'], r['rollout_heatmap'],
+                                          alpha=0.6, cmap_name='inferno')
+        axes[i, 1].imshow(rollout_overlay)
+        axes[i, 1].axis('off')
+        if i == 0:
+            axes[i, 1].set_title(col_titles[1], fontsize=12, fontweight='bold',
+                                 pad=10)
+
+        # Column 2: IG overlay
+        ig_overlay = overlay_heatmap(r['img_orig'], r['ig_heatmap'],
+                                     alpha=0.6, cmap_name='inferno')
+        axes[i, 2].imshow(ig_overlay)
+        color = 'green' if correct else 'red'
+        axes[i, 2].set_xlabel(
+            f'Pred: {pred_name} ({r["confidence"]:.2f}) [{status}] | '
+            f'Pearson r={r["agreement"]["pearson_r"]:.2f}',
+            fontsize=8, color=color, fontweight='bold')
+        axes[i, 2].set_xticks([])
+        axes[i, 2].set_yticks([])
+        if i == 0:
+            axes[i, 2].set_title(col_titles[2], fontsize=12, fontweight='bold',
+                                 pad=10)
+
+    plt.suptitle('RetinaSense v3.0 -- Attention Rollout vs Integrated Gradients',
+                 fontsize=14, fontweight='bold', y=1.001)
+    plt.tight_layout()
+    grid_path = os.path.join(XAI_DIR, 'comparison_grid.png')
+    fig.savefig(grid_path, dpi=120, bbox_inches='tight', facecolor='white')
+    plt.close(fig)
+    print(f'    Saved: {grid_path}')
+
+    # 5c. Agreement heatmap (matrix showing per-image spatial correlation)
+    print('  Generating agreement heatmap...')
+
+    # Build per-image metrics matrix
+    image_labels = [
+        f'#{r["idx"]+1} {CLASS_NAMES[r["true_label"]][:6]}'
+        for r in all_results
+    ]
+    metric_names = ['Pearson r', 'IoU (top 20%)', 'Cosine Sim']
+    agreement_matrix = np.zeros((n_success, 3))
+    for i, r in enumerate(all_results):
+        agreement_matrix[i, 0] = r['agreement']['pearson_r']
+        agreement_matrix[i, 1] = r['agreement']['iou_top20']
+        agreement_matrix[i, 2] = r['agreement']['cosine_sim']
+
+    fig, ax = plt.subplots(figsize=(7, max(8, n_success * 0.45)))
+    im = ax.imshow(agreement_matrix, cmap='RdYlGn', aspect='auto',
+                   vmin=-0.2, vmax=1.0)
+
+    ax.set_xticks(range(3))
+    ax.set_xticklabels(metric_names, fontsize=10, fontweight='bold')
+    ax.set_yticks(range(n_success))
+    ax.set_yticklabels(image_labels, fontsize=8)
+
+    # Annotate cells
+    for i in range(n_success):
+        for j in range(3):
+            val = agreement_matrix[i, j]
+            color = 'white' if val < 0.3 else 'black'
+            ax.text(j, i, f'{val:.2f}', ha='center', va='center',
+                    fontsize=8, color=color, fontweight='bold')
+
+    ax.set_title('Rollout vs IG Agreement Scores per Image',
+                 fontsize=12, fontweight='bold', pad=12)
+    plt.colorbar(im, ax=ax, shrink=0.6, label='Score')
+    plt.tight_layout()
+
+    heatmap_path = os.path.join(XAI_DIR, 'agreement_heatmap.png')
+    fig.savefig(heatmap_path, dpi=150, bbox_inches='tight', facecolor='white')
+    plt.close(fig)
+    print(f'    Saved: {heatmap_path}')
+
+    # ---- 6. Save agreement scores JSON ----
+    print('\n[6/6] Saving agreement scores...')
+
+    scores_output = {
+        'summary': {
+            'n_images':          n_success,
+            'mean_pearson_r':    round(float(agreement_matrix[:, 0].mean()), 4),
+            'mean_iou_top20':    round(float(agreement_matrix[:, 1].mean()), 4),
+            'mean_cosine_sim':   round(float(agreement_matrix[:, 2].mean()), 4),
+            'std_pearson_r':     round(float(agreement_matrix[:, 0].std()), 4),
+            'std_iou_top20':     round(float(agreement_matrix[:, 1].std()), 4),
+            'std_cosine_sim':    round(float(agreement_matrix[:, 2].std()), 4),
+        },
+        'per_image': [],
+    }
+    for r in all_results:
+        scores_output['per_image'].append({
+            'image':       r['basename'],
+            'true_label':  r['true_label'],
+            'true_class':  CLASS_NAMES[r['true_label']],
+            'pred_label':  r['pred_label'],
+            'pred_class':  r['pred_class'],
+            'confidence':  r['confidence'],
+            'agreement':   r['agreement'],
+        })
+
+    # Per-class summary
+    per_class = {}
+    for cls_idx in range(NUM_CLASSES):
+        cls_results = [r for r in all_results if r['true_label'] == cls_idx]
+        if cls_results:
+            pearson_vals = [r['agreement']['pearson_r'] for r in cls_results]
+            iou_vals     = [r['agreement']['iou_top20'] for r in cls_results]
+            cosine_vals  = [r['agreement']['cosine_sim'] for r in cls_results]
+            per_class[CLASS_NAMES[cls_idx]] = {
+                'n_images':        len(cls_results),
+                'mean_pearson_r':  round(float(np.mean(pearson_vals)), 4),
+                'mean_iou_top20':  round(float(np.mean(iou_vals)), 4),
+                'mean_cosine_sim': round(float(np.mean(cosine_vals)), 4),
+            }
+    scores_output['per_class'] = per_class
+
+    json_path = os.path.join(XAI_DIR, 'agreement_score.json')
+    with open(json_path, 'w') as f:
+        json.dump(scores_output, f, indent=2)
+    print(f'  Saved: {json_path}')
+
+    # ---- Summary ----
+    elapsed = time.time() - t_start
+    n_correct = sum(1 for r in all_results
+                    if r['pred_label'] == r['true_label'])
+
+    print('\n' + '=' * 65)
+    print('   PHASE 1C COMPLETE: Integrated Gradients XAI')
+    print('=' * 65)
+    print(f'  Images processed  : {n_success}/{len(samples)}')
+    print(f'  Correct preds     : {n_correct}/{n_success} '
+          f'({100*n_correct/max(n_success,1):.1f}%)')
+    print(f'  Mean Pearson r    : {scores_output["summary"]["mean_pearson_r"]:.4f}')
+    print(f'  Mean IoU (top 20%): {scores_output["summary"]["mean_iou_top20"]:.4f}')
+    print(f'  Mean Cosine Sim   : {scores_output["summary"]["mean_cosine_sim"]:.4f}')
+    print(f'  Time elapsed      : {elapsed:.1f}s')
+    print(f'  Outputs in        : {XAI_DIR}')
+    print('=' * 65)
+
+    # List outputs
+    print('\n  Output files:')
+    for fname in sorted(os.listdir(XAI_DIR)):
+        fpath = os.path.join(XAI_DIR, fname)
+        size_kb = os.path.getsize(fpath) / 1024
+        print(f'    {fname:40s}  {size_kb:8.1f} KB')
+
+
+if __name__ == '__main__':
+    main()