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