code/generate_gradcam_v2.py

"""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)