""" visualize_gradcam.py Generates Grad-CAM overlays for misclassified examples listed in a CSV (format produced earlier): image_path,true,pred,top1,top2 For each row this script saves a PNG with: - original image - Grad-CAM overlay for the **true** class - Grad-CAM overlay for the **predicted** class - difference overlay (pred - true) Usage: python src/analysis/visualize_gradcam.py \ --checkpoint outputs/swin_mps/best.pth \ --misclassified outputs/analysis/misclassified.csv \ --img-root . \ --model swin --img-size 224 --out-dir outputs/analysis/gradcam_overlays \ --class-names "Comminuted,Greenstick,Healthy,Oblique,Oblique Displaced,Spiral,Transverse,Transverse Displaced" Notes: - Script prefers MPS (Apple Silicon) if available; if Grad-CAM backward on MPS fails it will automatically fall back to CPU for CAM computation. - Requires: torch, timm, torchvision, pillow, numpy, opencv-python """ import os import sys import csv import argparse from pathlib import Path from typing import Optional, List import numpy as np from PIL import Image import cv2 import torch import torch.nn as nn import torchvision.transforms as T import timm import torchvision.models as tvmodels # Add parent directory to path for imports sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..'))) from src.utils import get_device, get_model, get_transforms DEVICE = get_device() print(f"Using device: {DEVICE}") # ----------------------------- Grad-CAM Implementation ----------------------------- class GradCAM: """Hook-based Grad-CAM. Call with a model (in eval mode) and a target conv layer name (optional). If target_layer_name is None, the last nn.Conv2d module is chosen heuristically. """ def __init__(self, model: nn.Module, target_layer_name: Optional[str] = None): self.model = model self.model.eval() self.activations = None self.gradients = None self.handles = [] # pick target layer if target_layer_name is None: target_layer = None for n, m in reversed(list(self.model.named_modules())): if isinstance(m, nn.Conv2d): target_layer_name = n target_layer = m break if target_layer is None: raise RuntimeError('No Conv2d layer found for Grad-CAM') else: target_layer = dict(self.model.named_modules()).get(target_layer_name, None) if target_layer is None: raise RuntimeError(f'layer name {target_layer_name} not found') # register hooks self.handles.append(target_layer.register_forward_hook(self._forward_hook)) # backward hook try: self.handles.append(target_layer.register_backward_hook(self._backward_hook)) except Exception: # fallback for newer pytorch versions: use register_full_backward_hook if available try: self.handles.append(target_layer.register_full_backward_hook(self._backward_hook)) except Exception: # some builds won't allow backward hooks; we'll compute gradients by retaining graph and reading .grad from activations pass def _forward_hook(self, module, inp, out): # out: tensor shape (B,C,H,W) self.activations = out.detach() def _backward_hook(self, module, grad_in, grad_out): # grad_out[0] shape (B,C,H,W) self.gradients = grad_out[0].detach() def clear(self): for h in self.handles: try: h.remove() except Exception: pass self.handles = [] def __call__(self, input_tensor: torch.Tensor, class_idx: Optional[int] = None, device: torch.device = torch.device('cpu')): """Compute CAM for a single input tensor (1,C,H,W). Returns cam resized to input HxW in numpy [0,1].""" self.model.zero_grad() input_tensor = input_tensor.to(device) input_tensor.requires_grad = True outputs = self.model(input_tensor) if class_idx is None: class_idx = int(outputs.argmax(dim=1).item()) loss = outputs[0, class_idx] loss.backward(retain_graph=True) if self.gradients is None or self.activations is None: raise RuntimeError('GradCAM failed to collect gradients/activations (hooks missing)') grads = self.gradients[0] # C,H,W acts = self.activations[0] # C,H,W weights = grads.mean(dim=(1,2)) # C cam = (weights[:, None, None] * acts).sum(dim=0).cpu().numpy() cam = np.maximum(cam, 0) cam = cam - cam.min() if cam.max() > 0: cam = cam / (cam.max() + 1e-8) else: cam = np.zeros_like(cam) # resize to original input spatial size (assume square input) H = input_tensor.shape[-2]; W = input_tensor.shape[-1] cam = cv2.resize(cam, (W, H)) return cam def apply_colormap_on_image(org_img: np.ndarray, activation: np.ndarray, colormap=cv2.COLORMAP_JET, alpha=0.5): """Overlay heatmap on image (org_img: HxW x 3 uint8, activation: HxW float in [0,1])""" if activation is None: raise ValueError('activation is None') # ensure activation is 2D and in [0,1] activation = np.asarray(activation) if activation.ndim == 3: # if somehow a channel dim exists, reduce to single channel activation = activation[..., 0] activation = np.clip(activation, 0.0, 1.0) # Convert activation -> heatmap (BGR) and resize heatmap to match original image heatmap = np.uint8(255 * activation) heatmap = cv2.applyColorMap(heatmap, colormap) # Resize heatmap to original image spatial size before blending h, w = org_img.shape[:2] heatmap = cv2.resize(heatmap, (w, h), interpolation=cv2.INTER_LINEAR) # convert heatmap to RGB to match org_img (which is RGB) heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB) # ensure types match for addWeighted org_uint8 = org_img.astype('uint8') heat_uint8 = heatmap.astype('uint8') overlaid = cv2.addWeighted(org_uint8, 1.0 - alpha, heat_uint8, alpha, 0) return overlaid def pil_to_numpy(img: Image.Image): arr = np.array(img.convert('RGB')) return arr def get_transform(img_size=224): return T.Compose([ T.Resize((img_size, img_size)), T.ToTensor(), T.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]) ]) def main(): parser = argparse.ArgumentParser() parser.add_argument('--checkpoint', required=True) parser.add_argument('--misclassified', required=True) parser.add_argument('--img-root', default='.') parser.add_argument('--model', default='swin') parser.add_argument('--img-size', type=int, default=224) parser.add_argument('--out-dir', default='outputs/analysis/gradcam_overlays') parser.add_argument('--class-names', required=True) parser.add_argument('--target-layer', default=None) parser.add_argument('--max-samples', type=int, default=200, help='max misclassified rows to process') args = parser.parse_args() class_names = [c.strip() for c in args.class_names.split(',')] num_classes = len(class_names) device_pref = detect_device() print('preferred device:', device_pref) model = get_model(args.model, num_classes, pretrained=False) ck = torch.load(args.checkpoint, map_location='cpu') model.load_state_dict(ck['model_state_dict']) # We'll run forward on preferred device, but if backward (for CAM) fails on MPS we'll move to CPU for CAM computation model.to(device_pref) model.eval() transform = get_transform(args.img_size) os.makedirs(args.out_dir, exist_ok=True) rows = [] with open(args.misclassified, 'r') as f: reader = csv.DictReader(f) for r in reader: rows.append(r) rows = rows[:args.max_samples] # initialize GradCAM on device_pref; if backward fails, we will retry on CPU gradcam = None try: gradcam = GradCAM(model, target_layer_name=args.target_layer) cam_device = device_pref except Exception as e: print('GradCAM init failed on preferred device; will try CPU. Error:', e) cam_device = torch.device('cpu') model_cpu = get_model(args.model, num_classes, pretrained=False) model_cpu.load_state_dict(ck['model_state_dict']) model_cpu.to(cam_device) model_cpu.eval() gradcam = GradCAM(model_cpu, target_layer_name=args.target_layer) for i, r in enumerate(rows): img_path = r['image_path'] if os.path.isabs(r['image_path']) else os.path.join(args.img_root, r['image_path']) true_lbl = int(r['true']) pred_lbl = int(r['pred']) try: pil = Image.open(img_path).convert('RGB') except Exception as e: print('failed to open', img_path, e); continue org_np = pil_to_numpy(pil) inp = transform(pil).unsqueeze(0) # forward on preferred device to get outputs and predicted class try: inp_pref = inp.to(device_pref) with torch.no_grad(): out_pref = model(inp_pref) probs = torch.softmax(out_pref, dim=1).cpu().numpy()[0] except Exception as e: print('forward failed on preferred device:', e) # fallback to CPU forward model.cpu(); inp_cpu = inp; model.eval() with torch.no_grad(): out_cpu = model(inp_cpu) probs = torch.softmax(out_cpu, dim=1).numpy()[0] # compute CAMs on gradcam.device (cam_device) cam_true = None; cam_pred = None try: # ensure model used for gradcam is on cam_device cam_model = gradcam.model # move input to cam_device inp_cam = inp.to(cam_device) cam_true = gradcam(inp_cam, class_idx=true_lbl, device=cam_device) cam_pred = gradcam(inp_cam, class_idx=pred_lbl, device=cam_device) except Exception as e: print('Grad-CAM on preferred device failed for', img_path, 'error:', e) # try CPU try: # rebuild cpu model if needed cpu_dev = torch.device('cpu') model_cpu = get_model(args.model, num_classes, pretrained=False) model_cpu.load_state_dict(ck['model_state_dict']) model_cpu.to(cpu_dev); model_cpu.eval() gradcam_cpu = GradCAM(model_cpu, target_layer_name=args.target_layer) cam_true = gradcam_cpu(inp.to(cpu_dev), class_idx=true_lbl, device=cpu_dev) cam_pred = gradcam_cpu(inp.to(cpu_dev), class_idx=pred_lbl, device=cpu_dev) gradcam_cpu.clear() except Exception as e2: print('Grad-CAM CPU retry failed for', img_path, e2) continue # overlay try: over_true = apply_colormap_on_image(org_np, cam_true, alpha=0.5) over_pred = apply_colormap_on_image(org_np, cam_pred, alpha=0.5) diff = cam_pred - cam_true diff = (diff - diff.min()) / (diff.max() - diff.min() + 1e-8) over_diff = apply_colormap_on_image(org_np, diff, alpha=0.6) # concat: original | true | pred | diff h, w, _ = org_np.shape # resize overlays to original size if needed over_true = cv2.resize(over_true, (w, h)) over_pred = cv2.resize(over_pred, (w, h)) over_diff = cv2.resize(over_diff, (w, h)) orig_bgr = cv2.cvtColor(org_np, cv2.COLOR_RGB2BGR) grid = np.vstack([np.hstack([orig_bgr, cv2.cvtColor(over_true, cv2.COLOR_RGB2BGR)]), np.hstack([cv2.cvtColor(over_pred, cv2.COLOR_RGB2BGR), cv2.cvtColor(over_diff, cv2.COLOR_RGB2BGR)])]) out_name = f"{i:04d}_true{true_lbl}_pred{pred_lbl}_{os.path.basename(img_path)}.png" out_path = os.path.join(args.out_dir, out_name) cv2.imwrite(out_path, grid) except Exception as e: print('failed to create overlay for', img_path, e) continue gradcam.clear() print('Saved overlays to', args.out_dir) if __name__ == '__main__': main()