code/generate_all_gradcam.py

"""Generate Grad-CAM / attention-rollout figures for all 6 thesis models.

Usage:
    python comparison_experiment/generate_all_gradcam.py \
        --weights-dir final_experiments/weights \
        --manifest holdout_split.json \
        --output-dir gradcam_outputs_final

For each (model, class) pair, picks one representative image from the
held-out test set and overlays the saliency heatmap. CNN models use
GradCAM on their last conv layer. CLIP uses attention rollout on its
visual transformer's last block.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
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


IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.229]


def imagenet_eval_tf(size=224):
    return transforms.Compose([
        transforms.Resize((size, size)),
        transforms.ToTensor(),
        transforms.Normalize(IMAGENET_MEAN, [0.229, 0.224, 0.225]),
    ])


def load_image_for_cam(path, size=224):
    img = Image.open(path).convert("RGB").resize((size, size))
    arr = np.array(img).astype(np.float32) / 255.0
    return img, arr


def make_cnn(name, num_classes):
    if name == "vgg19":
        m = models.vgg19(weights=None)
        m.classifier[-1] = nn.Linear(m.classifier[-1].in_features, num_classes)
        target = m.features[-1]
        return m, target, 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 make_clip(num_classes):
    import open_clip
    clip_model, _, _ = open_clip.create_model_and_transforms("ViT-B-16", pretrained="openai")

    class Wrapper(nn.Module):
        def __init__(self):
            super().__init__()
            self.backbone = clip_model
            with torch.no_grad():
                feat = self.backbone.encode_image(torch.zeros(1, 3, 224, 224)).shape[-1]
            self.head = nn.Linear(feat, num_classes)

        def forward(self, x):
            return self.head(self.backbone.encode_image(x).float())

    return Wrapper(), 224


def clip_attention_rollout(model, image_tensor):
    """Simple attention rollout on the last transformer block of CLIP visual encoder."""
    visual = model.backbone.visual
    attentions = []

    def hook(module, inputs, output):
        # MultiheadAttention returns (attn_output, attn_weights) when need_weights=True
        if isinstance(output, tuple) and len(output) > 1 and output[1] is not None:
            attentions.append(output[1].detach())

    handles = []
    # Hook every attention block in the transformer
    for block in visual.transformer.resblocks:
        h = block.attn.register_forward_hook(hook)
        handles.append(h)

    # Patch the attention layers to return weights
    original_need_weights = {}
    for block in visual.transformer.resblocks:
        original_need_weights[id(block.attn)] = block.attn.batch_first
    # OpenCLIP uses scaled_dot_product_attention; rollout via grad-based attribution instead
    with torch.no_grad():
        _ = model(image_tensor)
    for h in handles:
        h.remove()

    if not attentions:
        # Fall back to gradient saliency
        return None
    # Average heads, multiply across layers
    result = torch.eye(attentions[0].shape[-1], device=attentions[0].device)
    for a in attentions:
        a = a.mean(dim=1)[0]  # heads avg, batch 0
        a = a + torch.eye(a.shape[-1], device=a.device)
        a = a / a.sum(dim=-1, keepdim=True)
        result = a @ result
    mask = result[0, 1:]  # CLS token attention to patches
    grid = int(np.sqrt(mask.shape[0]))
    mask = mask.reshape(grid, grid).cpu().numpy()
    mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-8)
    return mask


def clip_grad_saliency(model, image_tensor, target_idx, image_size):
    """Use input-gradient saliency as a robust CLIP heatmap."""
    image_tensor = image_tensor.clone().requires_grad_(True)
    logits = model(image_tensor)
    score = logits[0, target_idx]
    score.backward()
    sal = image_tensor.grad.detach().abs().max(dim=1)[0][0].cpu().numpy()
    sal = (sal - sal.min()) / (sal.max() - sal.min() + 1e-8)
    return sal


def pick_test_image_per_class(manifest, data_root):
    """Return {class_name: relative_path} for the test split."""
    classes = manifest["classes"]
    test = manifest["splits"]["test"]
    per_class = {}
    for rel_path, cls in test:
        if cls not in per_class:
            per_class[cls] = rel_path
    return per_class, classes


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--weights-dir", default="final_experiments/weights")
    parser.add_argument("--manifest", default="holdout_split.json")
    parser.add_argument("--output-dir", default="gradcam_outputs_final")
    parser.add_argument("--models", nargs="+", default=[
        "vgg19", "resnet50", "resnet101", "densenet121", "inception_v3", "clip_openai",
    ])
    args = parser.parse_args()

    manifest = json.loads(Path(args.manifest).read_text())
    data_root = Path(manifest["data_dir"])
    per_class, classes = pick_test_image_per_class(manifest, data_root)
    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    for model_name in args.models:
        weights_path = Path(args.weights_dir) / f"{model_name}_final.pth"
        if not weights_path.exists():
            print(f"[skip] {model_name}: weights not found at {weights_path}")
            continue
        print(f"[gradcam] {model_name}")
        if model_name == "clip_openai":
            model, image_size = make_clip(len(classes))
            target_layer = None
        else:
            model, target_layer, image_size = make_cnn(model_name, len(classes))
        state = torch.load(weights_path, map_location="cpu")
        model.load_state_dict(state)
        model.to(device).eval()

        tf = imagenet_eval_tf(image_size)
        rows = []
        for cls_name in classes:
            rel = per_class.get(cls_name)
            if not rel:
                continue
            img_path = data_root / rel
            pil, raw = load_image_for_cam(img_path, image_size)
            tensor = tf(pil).unsqueeze(0).to(device)

            if model_name == "clip_openai":
                with torch.no_grad():
                    logits = model(tensor)
                pred_idx = int(logits.argmax(1).item())
                heat = clip_grad_saliency(model, tensor, pred_idx, image_size)
                heat = np.kron(heat, np.ones((1, 1)))  # no-op
                heat_rgb = show_cam_on_image(raw, heat, use_rgb=True)
            else:
                cam = GradCAM(model=model, target_layers=[target_layer])
                grayscale = cam(input_tensor=tensor, targets=None)[0]
                heat_rgb = show_cam_on_image(raw, grayscale, use_rgb=True)
                pred_idx = int(model(tensor).argmax(1).item())

            rows.append((cls_name, raw, heat_rgb, classes[pred_idx]))

        # 2-row grid: original / cam, one column per class
        n = len(rows)
        fig, axes = plt.subplots(2, n, figsize=(2.5 * n, 5.5))
        if n == 1:
            axes = axes.reshape(2, 1)
        for col, (cls_name, raw, heat, pred) in enumerate(rows):
            axes[0, col].imshow(raw)
            axes[0, col].set_title(cls_name[:18], fontsize=7)
            axes[0, col].axis("off")
            axes[1, col].imshow(heat)
            axes[1, col].set_title(f"pred: {pred[:18]}", fontsize=7)
            axes[1, col].axis("off")
        fig.suptitle(f"GradCAM / saliency — {model_name}", fontsize=12)
        fig.tight_layout()
        out_path = out_dir / f"gradcam_{model_name}.png"
        fig.savefig(out_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        print(f"  -> {out_path}")

        del model
        torch.cuda.empty_cache()


if __name__ == "__main__":
    main()