""" GradCAM analysis for Phase 5 — logit_mlp_head checkpointed run. Produces (all in figures/explainability/gradcam/): mean_cam_normal.png — average heatmap across all normal test eyes mean_cam_glaucoma.png — average heatmap across all glaucoma test eyes mean_cam_comparison.png — side-by-side normal vs glaucoma mean CAMs overlay_grid_normal.png — grid of individual overlays (normal eyes) overlay_grid_glaucoma.png — grid of individual overlays (glaucoma eyes) Checkpoints loaded from: v3/results/phase5/logit_mlp_head_ckpt/rep00/binary/ensemble/fold{0..4}/best_single.pt Usage: python -m v3.scripts.output_analysis.explainability.gradcam_phase5 """ from __future__ import annotations import sys from pathlib import Path import matplotlib matplotlib.use("Agg") import matplotlib.cm as cm import matplotlib.pyplot as plt import numpy as np import torch import torch.nn.functional as F from PIL import Image from tqdm import tqdm REPO_ROOT = Path(__file__).resolve().parents[4] CKPT_RUN = REPO_ROOT / "v3" / "results" / "phase5" / "logit_mlp_head_ckpt" FIGURES_ROOT = REPO_ROOT / "v3" / "figures" / "explainability" / "gradcam" CLINICAL_DIR = REPO_ROOT / "Papila" / "ClinicalData" IMAGE_DIR = REPO_ROOT / "Papila" / "FundusImages" CONTOUR_DIR = REPO_ROOT / "Papila" / "ExpertsSegmentations" / "Contours" DISC_SPAN = 5 # patch side = DISC_SPAN × disc diameter PATCH_SIZE = 96 # output thumbnail pixels # Model hyperparameters (inferred from checkpoint weight shapes) BACKBONE = "resnet50" NUM_CLASSES = 2 CD_HIDDEN = 128 FUSION_DIM = 256 LABEL_NAMES = {0: "Normal", 1: "Glaucoma"} # ── GradCAM ────────────────────────────────────────────────────────────────── class GradCAM: """Minimal GradCAM via forward/backward hooks.""" def __init__(self, target_layer: torch.nn.Module) -> None: self._acts = None self._grads = None self._h1 = target_layer.register_forward_hook(self._save_acts) self._h2 = target_layer.register_full_backward_hook(self._save_grads) def _save_acts(self, _m, _i, output): self._acts = output.detach() def _save_grads(self, _m, _gi, grad_output): self._grads = grad_output[0].detach() def compute(self, img: torch.Tensor, meta: torch.Tensor, model: torch.nn.Module, target_class: int | None = None) -> tuple[np.ndarray, int]: """Return (cam [H,W] normalised 0-1, predicted_class).""" model.eval() with torch.enable_grad(): out = model(img, meta) pred = int(out.argmax(1).item()) tc = pred if target_class is None else target_class model.zero_grad() out[0, tc].backward() weights = self._grads.mean(dim=(2, 3), keepdim=True) cam = F.relu((weights * self._acts).sum(dim=1, keepdim=True)) cam = F.interpolate(cam, img.shape[-2:], mode="bilinear", align_corners=False) cam_np = cam.squeeze().cpu().numpy() lo, hi = cam_np.min(), cam_np.max() return (cam_np - lo) / (hi - lo + 1e-8), pred def remove(self) -> None: self._h1.remove(); self._h2.remove() def overlay_gradcam(pil: Image.Image, cam: np.ndarray, alpha: float = 0.45) -> Image.Image: cam_u8 = (cam * 255).astype(np.uint8) cam_r = np.array(Image.fromarray(cam_u8).resize(pil.size, Image.BILINEAR)) / 255.0 colored = (cm.jet(cam_r)[:, :, :3] * 255).astype(np.uint8) return Image.blend(pil.convert("RGB"), Image.fromarray(colored), alpha) # ── Disc-centred attention helpers ──────────────────────────────────────────── def _disc_contour_path(pid: int, eye: str, expert: int = 1) -> Path: return CONTOUR_DIR / f"RET{pid:03d}{eye}_disc_exp{expert}.txt" def _load_disc_mask(pid: int, eye: str, cam_h: int, cam_w: int) -> np.ndarray | None: """Load expert disc contour, polygon-fill, resize to (cam_h, cam_w).""" from PIL import ImageDraw as _ID p = _disc_contour_path(pid, eye) if not p.exists(): return None try: arr = np.loadtxt(str(p), dtype=np.float32) except Exception: return None if arr.ndim == 1: arr = arr.reshape(-1, 2) if arr.shape[0] < 3: return None # Get original image size img_path = get_image_path(pid, eye) try: with Image.open(img_path) as im: orig_w, orig_h = im.size except Exception: return None canvas = Image.new("L", (orig_w, orig_h), 0) _ID.Draw(canvas).polygon([tuple(pt) for pt in arr[:, :2]], fill=1) return np.array(canvas.resize((cam_w, cam_h), Image.NEAREST), dtype=bool) def _disc_centred_patch(cam: np.ndarray, disc_mask: np.ndarray, span: int = DISC_SPAN, out: int = PATCH_SIZE ) -> tuple[np.ndarray | None, float | None]: """Translate+scale cam so disc centroid is centred; return (patch, disc_r_out).""" if disc_mask is None or disc_mask.sum() == 0: return None, None ys, xs = np.where(disc_mask) cy, cx = ys.mean(), xs.mean() disc_r = float(np.sqrt(disc_mask.sum() / np.pi)) half = max(1, int(round(span * disc_r / 2))) h, w = cam.shape y0, y1 = int(round(cy)) - half, int(round(cy)) + half x0, x1 = int(round(cx)) - half, int(round(cx)) + half pt = max(0, -y0); pb = max(0, y1 - h) pl = max(0, -x0); pr = max(0, x1 - w) cam_pad = np.pad(cam, ((pt, pb), (pl, pr)), constant_values=0.0) patch = cam_pad[y0 + pt: y1 + pt, x0 + pl: x1 + pl] patch_out = np.array( Image.fromarray((np.clip(patch, 0, 1) * 255).astype(np.uint8)) .resize((out, out), Image.BILINEAR) ) / 255.0 disc_r_out = out * disc_r / (2 * half) return patch_out.astype(np.float32), disc_r_out def make_disc_attention_detail( mean_patches: dict, stats_rows: list[dict], out_path: Path, ) -> None: # noqa: C901 """ 2-row (Normal / Glaucoma) × 3-col (correct cam | incorrect cam | disc_frac strip). mean_patches: {(cls_name, split): (mean_patch_array, mean_disc_r, count)} stats_rows: list of {true_name, correct, disc_frac} dicts (floats only, no arrays) """ import pandas as pd from matplotlib.patches import Circle classes = ["Normal", "Glaucoma"] splits = ["correct", "incorrect"] corr_colors = {"correct": "steelblue", "incorrect": "tomato"} stats = pd.DataFrame(stats_rows) # 2 rows (Normal / Glaucoma) × 3 cols (correct cam | incorrect cam | disc_frac strip) fig, axes = plt.subplots(2, 3, figsize=(13, 8), gridspec_kw={"width_ratios": [1, 1, 0.75]}) fig.patch.set_facecolor("#f4f4f4") rng = np.random.default_rng(42) for ri, cls in enumerate(classes): # Col 0 & 1: correct / incorrect mean CAMs for ci, split in enumerate(splits): ax = axes[ri, ci] ax.set_facecolor("#222") key = (cls, split) if key in mean_patches: mp, disc_r_out, count = mean_patches[key] ax.imshow(mp, cmap="jet", vmin=0, vmax=1, origin="upper", extent=[0, PATCH_SIZE, PATCH_SIZE, 0]) cx = cy = PATCH_SIZE / 2 ax.add_patch(Circle((cx, cy), disc_r_out, fill=False, edgecolor="white", linewidth=2, linestyle="--")) ax.set_title(f"{split.capitalize()} (N={count})", fontsize=9) else: ax.text(0.5, 0.5, "no data", ha="center", va="center", transform=ax.transAxes, fontsize=9, color="grey") ax.set_title(split.capitalize(), fontsize=9) ax.axis("off") # Row label on leftmost column axes[ri, 0].set_ylabel(cls, fontsize=11, fontweight="bold", labelpad=8) # Col 2: disc_frac strip ax = axes[ri, 2] ax.set_facecolor("#f4f4f4") sub = stats[stats["true_name"] == cls].dropna(subset=["disc_frac"]) for xi, split in enumerate(splits): pts = sub[sub["correct"] == (split == "correct")]["disc_frac"].values if len(pts) == 0: continue color = corr_colors[split] jitter = rng.uniform(-0.18, 0.18, size=len(pts)) ax.scatter(xi + jitter, pts, color=color, alpha=0.7, s=28, edgecolors="none") ax.hlines(pts.mean(), xi - 0.28, xi + 0.28, colors=color, linewidth=2.5, zorder=5) ax.set_xticks([0, 1]) ax.set_xticklabels(["Correct", "Incorrect"], fontsize=9) ax.set_xlim(-0.55, 1.55) ax.set_ylim(0, 1) ax.set_title("Disc fraction", fontsize=9) ax.grid(axis="y", linestyle="--", alpha=0.3) if ri == 0: ax.set_ylabel("Attention mass inside GT disc", fontsize=9) fig.suptitle( "Disc-centred GradCAM attention | dashed circle = GT disc boundary\n" "Phase 5, logit_mlp_head, fold 0–4", fontsize=11, fontweight="bold", ) fig.tight_layout() out_path.parent.mkdir(parents=True, exist_ok=True) fig.savefig(out_path, dpi=150, bbox_inches="tight") plt.close(fig) print(f"Saved: {out_path}") # ── Model loading ───────────────────────────────────────────────────────────── def build_model(ckpt_path: Path, device: torch.device): """Reconstruct SingleEyeHT from checkpoint and load weights.""" from types import SimpleNamespace from v3.classes.hypertower_models import SingleEyeHT sd = torch.load(ckpt_path, map_location="cpu") # ClinicalTower only reads clinical_data.feature_dim at init time cd_in = sd["cd_tower.block0.0.weight"].shape[1] clinical_shim = SimpleNamespace(feature_dim=cd_in) model = SingleEyeHT( backbone=BACKBONE, freeze_ratio=0.0, augment=False, clinical_data=clinical_shim, num_classes=NUM_CLASSES, cd_hidden_dim=CD_HIDDEN, fusion_dim=FUSION_DIM, ) model.load_state_dict(sd) model.to(device).eval() return model # ── Data helpers ────────────────────────────────────────────────────────────── def build_data_bundle(): """Build the PAPILA DataBundle matching the checkpointed run's feature config.""" from v3.classes.papila_builders import build_papila_data import torch as _t # Auto-detect cd_in from the majority of checkpoints (excludes stale reps). import collections as _col all_ckpts = list(CKPT_RUN.glob("rep*/binary/ensemble/fold*/best_single.pt")) if all_ckpts: counts = _col.Counter( _t.load(c, map_location="cpu")["cd_tower.block0.0.weight"].shape[1] for c in all_ckpts ) cd_in = counts.most_common(1)[0][0] else: cd_in = 25 drop_raw = cd_in <= 21 excl = ["Axial_Length"] if cd_in in (21, 23) else [] return build_papila_data( image_dir=str(IMAGE_DIR), clinical_dir=str(CLINICAL_DIR), label_col="Diagnosis", cat_cols=["Gender", "Phakic/Pseudophakic"], iop_corr_method="ratio", iop_drop_raw=drop_raw, exclude_cols=excl, ) def get_image_path(pid: int, eye: str) -> Path: return IMAGE_DIR / f"RET{pid:03d}{eye}.jpg" def build_meta_vector(row, data) -> torch.Tensor: """Build the training-compatible feature vector via DataBundle.vectorize_row.""" vec = data.vectorize_row(row) return torch.tensor(vec, dtype=torch.float32).unsqueeze(0) # ── Eval transform ──────────────────────────────────────────────────────────── def get_eval_transform(): from torchvision import transforms return transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), ]) # ── Main loop ───────────────────────────────────────────────────────────────── def _discover_checkpoints(ckpt_run: Path) -> list[tuple[int, int, Path]]: """ Scan ckpt_run for all available best_single.pt files. Skips checkpoints whose cd_in doesn't match the majority (to exclude stale reps). Returns sorted list of (rep_idx, fold_idx, ckpt_path). """ import collections candidates = [] for rep_dir in sorted(ckpt_run.glob("rep*")): try: rep_idx = int(rep_dir.name.replace("rep", "")) except ValueError: continue for fold_dir in sorted((rep_dir / "binary" / "ensemble").glob("fold[0-9]")): ckpt = fold_dir / "best_single.pt" if ckpt.exists(): fold_idx = int(fold_dir.name.replace("fold", "")) cd_in = torch.load(ckpt, map_location="cpu")[ "cd_tower.block0.0.weight" ].shape[1] candidates.append((rep_idx, fold_idx, ckpt, cd_in)) if not candidates: return [] # Use the majority cd_in so stale reps are automatically excluded counts = collections.Counter(c[3] for c in candidates) target_cd_in = counts.most_common(1)[0][0] skipped = sum(1 for c in candidates if c[3] != target_cd_in) if skipped: print(f" [discover] skipping {skipped} checkpoint(s) with cd_in≠{target_cd_in}") return [(rep, fold, ckpt) for rep, fold, ckpt, cd in candidates if cd == target_cd_in] def run(n_grid: int = 16, alpha: float = 0.45, target_class: int | None = None): """ Loop over all available checkpoints in the run directory (all reps × folds). Aggregate CAMs per class, collect overlay grids. """ import pandas as pd from v3.scripts.output_analysis.explainability.fold_patient_ids import ( get_test_patient_ids, ) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Device: {device}") print("Building DataBundle ...") data = build_data_bundle() clinical = data.df print(f" feature_dim={data.feature_dim} rows={len(clinical)}") transform = get_eval_transform() checkpoints = _discover_checkpoints(CKPT_RUN) print(f"Found {len(checkpoints)} checkpoint(s) across " f"{len(set(r for r,f,_ in checkpoints))} rep(s)") if not checkpoints: print("No checkpoints found — run with --save-checkpoints first.") return # ── Incremental accumulators (no full-res arrays kept after each eye) ──────── # Mean CAM per class: running sum cam_sum = {0: None, 1: None} cam_count = {0: 0, 1: 0} # Overlay grid: keep at most n_grid PIL images per class (capped) overlay_items = {0: [], 1: []} # Disc-attention detail: running sum of disc patches (not list of arrays) disc_patch_sum = {} # (cls_name, split) → np.ndarray sum disc_patch_count = {} # (cls_name, split) → int disc_radius_sum = {} # (cls_name, split) → float sum disc_stats_rows = [] # floats only — no arrays ckpt_bar = tqdm(checkpoints, desc="Folds", unit="fold") for rep_idx, fold_idx, ckpt_path in ckpt_bar: ckpt_bar.set_postfix(rep=rep_idx, fold=fold_idx) model = build_model(ckpt_path, device) # GradCAM target: last ResNet block target_layer = model.img_tower.backbone.layer4[-1] gcam = GradCAM(target_layer) pids = get_test_patient_ids(rep_idx, fold_idx, clinical_dir=CLINICAL_DIR) for pid in tqdm(pids, desc=f" rep{rep_idx:02d}/fold{fold_idx}", leave=False, unit="pt"): for eye in ("OD", "OS"): img_path = get_image_path(pid, eye) if not img_path.exists(): continue row = clinical[ (clinical["Patient ID"] == pid) & (clinical["eyeID"] == eye) ] if len(row) == 0: continue row = row.iloc[0] label = int(row["Diagnosis"]) pil_orig = Image.open(img_path).convert("RGB") img_t = transform(pil_orig).unsqueeze(0).to(device) meta_t = build_meta_vector(row, data).to(device) cam_np, pred = gcam.compute(img_t, meta_t, model, target_class=target_class) # Running mean CAM if cam_sum[label] is None: cam_sum[label] = cam_np.copy() else: cam_sum[label] += cam_np cam_count[label] += 1 # Overlay grid — only keep up to n_grid per class if len(overlay_items[label]) < n_grid: ov = overlay_gradcam(pil_orig, cam_np, alpha=alpha) overlay_items[label].append((ov, pid, eye, pred)) # Disc-attention: extract patch now, accumulate into running sum h, w = cam_np.shape disc_mask = _load_disc_mask(pid, eye, h, w) disc_frac = None if disc_mask is not None and disc_mask.sum() > 0: disc_frac = float(cam_np[disc_mask].sum() / (cam_np.sum() + 1e-8)) cls_name = LABEL_NAMES[label] split = "correct" if (pred == label) else "incorrect" key = (cls_name, split) patch, disc_r_out = _disc_centred_patch(cam_np, disc_mask) if patch is not None: if key not in disc_patch_sum: disc_patch_sum[key] = patch.copy() disc_patch_count[key] = 1 disc_radius_sum[key] = disc_r_out else: disc_patch_sum[key] += patch disc_patch_count[key] += 1 disc_radius_sum[key] += disc_r_out disc_stats_rows.append({ "true_name": cls_name, "correct": pred == label, "disc_frac": disc_frac, }) # Release per-eye tensors immediately del img_t, meta_t, cam_np if disc_mask is not None: del disc_mask gcam.remove() del model torch.cuda.empty_cache() if torch.cuda.is_available() else None # Build mean_patches dict for disc detail plot mean_patches = { key: ( disc_patch_sum[key] / disc_patch_count[key], disc_radius_sum[key] / disc_patch_count[key], disc_patch_count[key], ) for key in disc_patch_sum } # ── Save outputs ───────────────────────────────────────────────────────── FIGURES_ROOT.mkdir(parents=True, exist_ok=True) for cls in [0, 1]: if cam_count[cls] == 0: continue mean_cam = cam_sum[cls] / cam_count[cls] lo, hi = mean_cam.min(), mean_cam.max() mean_cam = (mean_cam - lo) / (hi - lo + 1e-8) fig, ax = plt.subplots(figsize=(5, 5)) ax.imshow(mean_cam, cmap="jet", vmin=0, vmax=1) ax.axis("off") ax.set_title(f"Mean GradCAM — {LABEL_NAMES[cls]}\n(n={cam_count[cls]} eyes, fold 0–4)", fontsize=11, fontweight="bold") plt.colorbar(ax.images[0], ax=ax, fraction=0.046, pad=0.04) out = FIGURES_ROOT / f"mean_cam_{LABEL_NAMES[cls].lower()}.png" fig.savefig(out, dpi=180, bbox_inches="tight") plt.close(fig) print(f"Saved: {out}") # Side-by-side comparison if cam_count[0] > 0 and cam_count[1] > 0: fig, axes = plt.subplots(1, 2, figsize=(10, 5)) fig.suptitle("Mean GradCAM — Normal vs Glaucoma (Phase 5, fold 0–4)", fontsize=12, fontweight="bold") for ax, cls in zip(axes, [0, 1]): mean_cam = cam_sum[cls] / cam_count[cls] lo, hi = mean_cam.min(), mean_cam.max() mean_cam = (mean_cam - lo) / (hi - lo + 1e-8) im = ax.imshow(mean_cam, cmap="jet", vmin=0, vmax=1) ax.axis("off") ax.set_title(f"{LABEL_NAMES[cls]} (n={cam_count[cls]})", fontsize=11) plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) out = FIGURES_ROOT / "mean_cam_comparison.png" fig.savefig(out, dpi=180, bbox_inches="tight") plt.close(fig) print(f"Saved: {out}") # Overlay grids for cls in [0, 1]: items = overlay_items[cls] if not items: continue # Sort: misclassified first (more interesting) items.sort(key=lambda x: x[3] == cls) # wrong preds first items = items[:n_grid] ncols = 4 nrows = int(np.ceil(len(items) / ncols)) fig, axes = plt.subplots(nrows, ncols, figsize=(ncols * 3.2, nrows * 3.2)) axes = np.array(axes).reshape(-1) fig.suptitle(f"GradCAM Overlays — {LABEL_NAMES[cls]} (Phase 5)", fontsize=12, fontweight="bold") for i, ax in enumerate(axes): if i < len(items): ov, pid, eye, pred = items[i] ax.imshow(ov) correct = pred == cls col = "#2e7d32" if correct else "#c62828" ax.set_title(f"RET{pid:03d}{eye}\n→ {LABEL_NAMES[pred]}", fontsize=7.5, color=col) ax.axis("off") out = FIGURES_ROOT / f"overlay_grid_{LABEL_NAMES[cls].lower()}.png" fig.tight_layout() fig.savefig(out, dpi=150, bbox_inches="tight") plt.close(fig) print(f"Saved: {out}") # Disc-centred detail plot if disc_patch_sum: make_disc_attention_detail(mean_patches, disc_stats_rows, FIGURES_ROOT / "disc_attention_detail.png") if __name__ == "__main__": import argparse ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--n-grid", type=int, default=16, help="Max overlays per class in grid (default 16)") ap.add_argument("--alpha", type=float, default=0.45, help="GradCAM overlay opacity (default 0.45)") ap.add_argument("--target-class", type=int, default=None, help="GradCAM target class (default: predicted class)") args = ap.parse_args() run(n_grid=args.n_grid, alpha=args.alpha, target_class=args.target_class)