"""Medical imaging attention/CAM helpers. Provides: - `compute_attention_rollout(model, pixel_values)`: ViT attention rollout (Abnar & Zuidema 2020) — works for any HuggingFace `ViTForImageClassification` whose forward accepts `output_attentions=True`. - `cam_to_jet_rgb(cam)`: colorize a [0, 1] heatmap as a jet RGB image. - `composite_heatmap_on_rgb(cam, rgb, alpha, gamma)`: blend a jet heatmap on top of an RGB image and return a base64 PNG data-URL. Decision: for the ViT-based skin / brain classifiers we use **attention rollout** rather than Grad-CAM, since these models do not expose the kind of intermediate convolutional feature map Grad-CAM expects, but they do expose per-head attention weights through HF's `output_attentions=True`. """ from __future__ import annotations import base64 import io import math import numpy as np import torch import torch.nn.functional as F from PIL import Image # ───────────────────────── attention rollout ───────────────────────── def compute_attention_rollout( model, pixel_values: torch.Tensor, discard_ratio: float = 0.6, head_fusion: str = "mean", ) -> np.ndarray: """Run a ViT and return a 2-D attention map (H, W) in [0, 1]. Args: model: a HuggingFace ViTForImageClassification (or compatible) whose forward accepts `output_attentions=True`. pixel_values: pre-processed input tensor of shape (1, 3, H, W) — the image processor's `pixel_values`. discard_ratio: per-layer fraction of lowest-attention tokens to zero out before rolling up. 0.6–0.9 typical; higher = sharper map. head_fusion: how to combine multi-head attention. One of "mean", "min", "max". "mean" is the original rollout formulation. Returns the rolled-up attention from CLS token to image patches, reshaped to a square (n_patches × patch_size, n_patches × patch_size) and normalized to [0, 1]. Caller should resize to the original image size. """ if head_fusion not in {"mean", "min", "max"}: raise ValueError(f"head_fusion must be one of mean/min/max, got {head_fusion!r}") model.eval() with torch.inference_mode(): outputs = model(pixel_values=pixel_values, output_attentions=True) attentions = outputs.attentions # tuple of (1, heads, seq, seq) per layer if not attentions: raise RuntimeError("Model did not return attentions; ensure it's a ViT-style model.") # Roll-up per Abnar & Zuidema (2020): per layer, fuse heads, drop low # attention, mix with identity (residual), normalize, then matmul. seq_len = attentions[0].shape[-1] rollup = torch.eye(seq_len, device=attentions[0].device) for attn in attentions: a = attn[0] # (heads, seq, seq) if head_fusion == "mean": fused = a.mean(dim=0) elif head_fusion == "min": fused = a.min(dim=0).values else: # max fused = a.max(dim=0).values # Discard the lowest `discard_ratio` of attention values per row, # but keep CLS-CLS interaction (column 0) intact. if discard_ratio > 0: flat = fused.view(-1) n_drop = int(flat.numel() * discard_ratio) if n_drop > 0: threshold = torch.kthvalue(flat, n_drop).values mask = fused >= threshold # Always keep CLS column to avoid zeroing the entire CLS row mask[:, 0] = True fused = fused * mask.float() # Add identity for residual connection, then row-normalize fused = fused + torch.eye(seq_len, device=fused.device) fused = fused / fused.sum(dim=-1, keepdim=True).clamp(min=1e-8) rollup = fused @ rollup # CLS token's attention to image patches (drop CLS-CLS) cls_attn = rollup[0, 1:] # (n_patches,) n_patches = cls_attn.shape[0] grid = int(round(math.sqrt(n_patches))) if grid * grid != n_patches: # Some ViTs include extra special tokens; truncate to nearest square. grid = int(math.sqrt(n_patches)) cls_attn = cls_attn[: grid * grid] cam = cls_attn.reshape(grid, grid).cpu().numpy().astype(np.float32) if cam.max() > 0: cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8) return cam # ───────────────────────── colorize + composite ───────────────────────── def cam_to_jet_rgb(cam: np.ndarray) -> np.ndarray: """Map a (H, W) [0, 1] heatmap to a (H, W, 3) uint8 jet RGB array. Mirrors the matplotlib jet colormap: 0.00 → dark blue, 0.25 → cyan, 0.50 → green, 0.75 → yellow, 1.00 → red. """ cam = np.clip(cam, 0.0, 1.0) r = np.clip(1.5 - np.abs(4.0 * cam - 3.0), 0, 1) g = np.clip(1.5 - np.abs(4.0 * cam - 2.0), 0, 1) b = np.clip(1.5 - np.abs(4.0 * cam - 1.0), 0, 1) rgb = np.zeros((*cam.shape, 3), dtype=np.uint8) rgb[..., 0] = (r * 255).astype(np.uint8) rgb[..., 1] = (g * 255).astype(np.uint8) rgb[..., 2] = (b * 255).astype(np.uint8) return rgb def resize_cam(cam: np.ndarray, size: tuple[int, int]) -> np.ndarray: """Bilinear-resize a (H, W) cam to (target_h, target_w).""" target_w, target_h = size t = torch.from_numpy(cam).float().unsqueeze(0).unsqueeze(0) t = F.interpolate(t, size=(target_h, target_w), mode="bilinear", align_corners=False) out = t.squeeze().numpy() if out.max() > 0: out = (out - out.min()) / (out.max() - out.min() + 1e-8) return out def composite_heatmap_on_rgb( cam: np.ndarray, rgb: np.ndarray, alpha: float = 0.45, gamma: float = 0.7, ) -> str: """Blend a jet heatmap onto an RGB image and return a base64 PNG data-URL. Args: cam: (H, W) float in [0, 1] — must already match `rgb` dimensions. rgb: (H, W, 3) uint8 — original RGB image. alpha: weight of the heatmap in the blend (0 = only image, 1 = only heatmap). gamma: gamma on the heatmap (<1 brightens mid-range attention). """ if cam.shape != rgb.shape[:2]: raise ValueError( f"cam shape {cam.shape} does not match rgb spatial shape {rgb.shape[:2]}" ) cam = np.clip(cam, 0.0, 1.0) if cam.max() > 0: cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8) cam = np.power(cam, gamma) heatmap_rgb = cam_to_jet_rgb(cam).astype(np.float32) base = rgb.astype(np.float32) composite = alpha * heatmap_rgb + (1.0 - alpha) * base composite = np.clip(composite, 0, 255).astype(np.uint8) img = Image.fromarray(composite, mode="RGB") buf = io.BytesIO() img.save(buf, format="PNG", optimize=True) b64 = base64.b64encode(buf.getvalue()).decode("ascii") return f"data:image/png;base64,{b64}"