"""Generate GradCAM / attention saliency for all 9 v2 models. Loads fresh v2 weights from weights_v2/ and weights_v3/ and writes per-class grid PNGs into gradcam_v2/. CNN-class models use GradCAM on the final conv block. CLIP uses input-gradient saliency. Vision-Transformer foundation models (Swin-B, DINOv2-L, RETFound) use pytorch-grad-cam with the appropriate reshape_transform. """ from __future__ import annotations import json, sys from pathlib import Path import numpy as np import torch import torch.nn as nn import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt from PIL import Image from torchvision import models, transforms from pytorch_grad_cam import GradCAM from pytorch_grad_cam.utils.image import show_cam_on_image ROOT = Path("/home/bytical/fundus_project") sys.path.insert(0, str(ROOT / "comparison_experiment")) MANIFEST = ROOT / "holdout_split_augmented.json" OUT_DIR = ROOT / "gradcam_v2" OUT_DIR.mkdir(exist_ok=True) W_V2 = ROOT / "weights_v2" W_V3 = ROOT / "weights_v3" DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") IMAGENET_MEAN = [0.485, 0.456, 0.406] IMAGENET_STD = [0.229, 0.224, 0.225] CLASS_FULL = ["CSC", "DR", "Disc Edema", "Glaucoma", "Healthy", "Macular Scar", "Myopia", "Pterygium", "Retinal Det.", "Retinitis Pig."] def tf_for(size): return transforms.Compose([ transforms.Resize((size, size)), transforms.ToTensor(), transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD), ]) # ---------- CNN builders matching v2 training ---------- def build_cnn(name, num_classes=10): if name == "vgg19": m = models.vgg19(weights=None) m.classifier[-1] = nn.Linear(m.classifier[-1].in_features, num_classes) return m, m.features[-1], 224 if name == "resnet50": m = models.resnet50(weights=None); m.fc = nn.Linear(m.fc.in_features, num_classes) return m, m.layer4[-1], 224 if name == "resnet101": m = models.resnet101(weights=None); m.fc = nn.Linear(m.fc.in_features, num_classes) return m, m.layer4[-1], 224 if name == "densenet121": m = models.densenet121(weights=None) m.classifier = nn.Linear(m.classifier.in_features, num_classes) return m, m.features.norm5, 224 if name == "inception_v3": m = models.inception_v3(weights=None, aux_logits=True) m.fc = nn.Linear(m.fc.in_features, num_classes) m.AuxLogits.fc = nn.Linear(m.AuxLogits.fc.in_features, num_classes) return m, m.Mixed_7c, 299 raise ValueError(name) def build_clip(num_classes=10): import open_clip base, _, _ = open_clip.create_model_and_transforms("ViT-B-16", pretrained="openai") class Wrap(nn.Module): def __init__(self): super().__init__() self.backbone = base.visual # match run_v2_experiments.CLIPClf d = self.backbone.output_dim if hasattr(self.backbone, "output_dim") else 512 self.head = nn.Linear(d, num_classes) def forward(self, x): return self.head(self.backbone(x).float()) return Wrap(), 224 # ---------- Foundation builders (mirror run_foundation_models.py) ---------- def build_swin(num_classes=10): import timm m = timm.create_model("swin_base_patch4_window7_224", pretrained=False, num_classes=num_classes) return m, 224 def build_dinov2(num_classes=10): import timm backbone = torch.hub.load("facebookresearch/dinov2", "dinov2_vitl14", source="github") class Wrap(nn.Module): def __init__(self): super().__init__() self.backbone = backbone self.head = nn.Linear(1024, num_classes) def forward(self, x): f = self.backbone(x) return self.head(f) return Wrap(), 224 def build_retfound(num_classes=10): import timm m = timm.create_model("vit_large_patch16_224", pretrained=False, num_classes=num_classes, global_pool="avg") return m, 224 # ---------- Reshape transforms for ViT-style ---------- def vit_reshape(t, h=14, w=14): # t: [B, tokens, dim] -> [B, dim, h, w]; skip CLS if present if t.shape[1] == h*w + 1: t = t[:, 1:, :] elif t.shape[1] != h*w: # Try infer n = t.shape[1] s = int(n ** 0.5) if s*s == n: h = w = s else: return t # give up t = t # keep as-is return t.reshape(t.shape[0], h, w, t.shape[-1]).permute(0, 3, 1, 2) def swin_reshape(t): # Swin outputs [B, H, W, C] from last stage if t.dim() == 4 and t.shape[-1] > t.shape[1]: return t.permute(0, 3, 1, 2) return t # ---------- Image picker: one image per class from test split ---------- M = json.load(open(MANIFEST)) test_items = M["splits"]["test"] per_class = {} for rel, lbl in test_items: if lbl not in per_class: per_class[lbl] = rel classes_sorted = [per_class[i] for i in range(10) if i in per_class] print(f"Found one test image for each of {len(classes_sorted)} classes") # ---------- Render utility ---------- def render_grid(model_name, rows, out_path): n = len(rows) fig, axes = plt.subplots(2, n, figsize=(2.4 * n, 5.4)) if n == 1: axes = axes.reshape(2, 1) for col, (cls_name, raw, heat, pred_name) in enumerate(rows): axes[0, col].imshow(raw); axes[0, col].axis("off") axes[0, col].set_title(cls_name, fontsize=8) axes[1, col].imshow(heat); axes[1, col].axis("off") axes[1, col].set_title(f"pred: {pred_name}", fontsize=7) fig.suptitle(f"Saliency — {model_name} (v2 weights)", fontsize=13) fig.tight_layout() fig.savefig(out_path, dpi=130, bbox_inches="tight") plt.close(fig) print(" ->", out_path) # ---------- Run per model ---------- def gen_for_model(name, weight_path, builder, target_layer_fn, image_size, reshape_transform=None, mode="cam"): print(f"\n[{name}] {weight_path.name}") if not weight_path.exists(): print(f" SKIP missing {weight_path}"); return if name == "clip_openai": model, image_size = builder() elif name in ("swin_b", "retfound"): model, image_size = builder() elif name == "dinov2_l": model, image_size = builder() else: model, _, image_size = builder(name) state = torch.load(weight_path, map_location="cpu", weights_only=False) if isinstance(state, dict) and "state_dict" in state: state = state["state_dict"] try: model.load_state_dict(state, strict=False) except Exception as e: print(f" load_state_dict warning: {e}") model.to(DEVICE).eval() tf = tf_for(image_size) target_layer = target_layer_fn(model) if target_layer_fn else None rows = [] for cls_idx, cls_name in enumerate(CLASS_FULL): rel = per_class.get(cls_idx) if not rel: continue img_path = ROOT / rel if not img_path.exists(): print(f" missing image {img_path}"); continue pil = Image.open(img_path).convert("RGB").resize((image_size, image_size)) raw = np.array(pil).astype(np.float32) / 255.0 x = tf(pil).unsqueeze(0).to(DEVICE) if mode == "saliency": x = x.clone().detach().requires_grad_(True) logits = model(x) if isinstance(logits, tuple): logits = logits[0] pred = int(logits.argmax(1).item()) score = logits[0, pred] model.zero_grad(); score.backward() sal = x.grad.detach().abs().max(dim=1)[0][0].cpu().numpy() sal = (sal - sal.min()) / (sal.max() - sal.min() + 1e-8) heat_rgb = show_cam_on_image(raw, sal, use_rgb=True) else: cam = GradCAM(model=model, target_layers=[target_layer], reshape_transform=reshape_transform) gray = cam(input_tensor=x, targets=None)[0] heat_rgb = show_cam_on_image(raw, gray, use_rgb=True) with torch.no_grad(): out = model(x) if isinstance(out, tuple): out = out[0] pred = int(out.argmax(1).item()) rows.append((cls_name, raw, heat_rgb, CLASS_FULL[pred][:14])) render_grid(name, rows, OUT_DIR / f"gradcam_{name}.png") del model torch.cuda.empty_cache() def _tgt_factory(name): if name == "vgg19": return lambda m: m.features[-1] if name == "resnet50": return lambda m: m.layer4[-1] if name == "resnet101": return lambda m: m.layer4[-1] if name == "densenet121": return lambda m: m.features.norm5 if name == "inception_v3": return lambda m: m.Mixed_7c raise ValueError(name) # ---------- CNN models ---------- for name in ["vgg19", "resnet50", "resnet101", "densenet121", "inception_v3"]: w = W_V2 / f"{name}_v2_final.pth" gen_for_model(name, w, build_cnn, _tgt_factory(name), 224, mode="cam") # CLIP — saliency gen_for_model("clip_openai", W_V2 / "clip_openai_v2_final.pth", build_clip, None, 224, mode="saliency") # Swin-B — use input-gradient saliency for robustness (CAM on hierarchical Swin is fragile) gen_for_model("swin_b", W_V3 / "swin_b_v2.pth", build_swin, None, 224, mode="saliency") # DINOv2-L — input-gradient saliency on the wrapper output (avoids hub-build target-layer issues) gen_for_model("dinov2_l", W_V3 / "dinov2_l_v2.pth", build_dinov2, None, 224, mode="saliency") # RETFound — input-gradient saliency gen_for_model("retfound", W_V3 / "retfound_v2.pth", build_retfound, None, 224, mode="saliency") print("\nALL DONE. PNGs in", OUT_DIR)