gradcam v2: generator code

code/generate_gradcam_v2.py

@@ -0,0 +1,238 @@
+"""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)