""" explainability/attention_rollout.py ------------------------------------ ViT Attention Rollout — the definitive ViT explainability technique. Algorithm (Abnar & Zuidema, 2020 — "Quantifying Attention Flow in Transformers"): 1. For each of 12 transformer layers, average attention weights across all 12 heads → shape per layer: (seq_len, seq_len) = (197, 197) 2. Add the identity matrix (residual connections preserve information) 3. Re-normalize each row to sum to 1 4. Multiply all 12 matrices together in sequence (matrix chain multiplication) → This propagates attention from input patches all the way to [CLS] 5. Extract the [CLS] row → shape: (197,) = 1 CLS + 196 patches 6. Reshape the 196 patch values to (14, 14) — the spatial patch grid 7. Upsample to (224, 224) — the original image resolution 8. Overlay on the original image as a heatmap Why Attention Rollout > Grad-CAM for ViT: - Grad-CAM was designed for CNNs with spatial feature maps - ViT has no intermediate spatial feature maps — Grad-CAM produces blurry, uninformative results on pure transformers - Attention Rollout correctly accounts for skip connections and is mathematically derived from the transformer's own attention flow Reference: Abnar, S. & Zuidema, W. (2020). Quantifying Attention Flow in Transformers. arXiv:2005.00928 """ from typing import List, Optional, Tuple import numpy as np import torch import torch.nn.functional as F import cv2 import matplotlib.pyplot as plt import matplotlib.cm as cm def compute_attention_rollout( attentions: List[torch.Tensor], head_fusion: str = "mean", discard_ratio: float = 0.9, ) -> np.ndarray: """ Compute Attention Rollout from a list of per-layer attention tensors. Args: attentions: List of 12 tensors, each shape (B, H, N, N). B = batch, H = 12 heads, N = 197 tokens (1 CLS + 196 patches). head_fusion: How to combine attention heads. Options: "mean" — average across heads (standard rollout) "max" — take max across heads (more focused) "min" — take min across heads (conservative) discard_ratio: Fraction of lowest-attention patch weights to zero out. Helps focus the heatmap on the most attended patches. Set to 0.0 to disable. Returns: np.ndarray of shape (14, 14) — attention map over the 14×14 patch grid. Values are normalized to [0, 1]. """ # ── 1. Process each attention layer ─────────────────────────────────────── # Start with identity (represents perfect self-attention at layer 0) batch_size = attentions[0].shape[0] num_tokens = attentions[0].shape[-1] # 197 result = torch.eye(num_tokens, device=attentions[0].device) result = result.unsqueeze(0).expand(batch_size, -1, -1) # (B, 197, 197) for attention in attentions: # attention: (B, H, N, N) # Fuse heads if head_fusion == "mean": attention_fused = attention.mean(dim=1) # (B, N, N) elif head_fusion == "max": attention_fused = attention.max(dim=1).values # (B, N, N) elif head_fusion == "min": attention_fused = attention.min(dim=1).values # (B, N, N) else: raise ValueError(f"Unknown head_fusion: {head_fusion}") # Discard low-attention tokens (remove noise) if discard_ratio > 0.0: flat = attention_fused.view(batch_size, -1) threshold_idx = int(flat.shape[-1] * discard_ratio) # Find the threshold value sorted_flat, _ = flat.sort(dim=-1) threshold = sorted_flat[:, threshold_idx].unsqueeze(-1).unsqueeze(-1) attention_fused = torch.where( attention_fused > threshold, attention_fused, torch.zeros_like(attention_fused), ) # Add residual connection (identity skip) attention_fused = attention_fused + torch.eye( num_tokens, device=attention_fused.device ).unsqueeze(0) # Row-normalize (each token's attention distribution sums to 1) row_sums = attention_fused.sum(dim=-1, keepdim=True) attention_fused = attention_fused / (row_sums + 1e-8) # Chain-multiply: propagate attention through layers result = torch.matmul(attention_fused, result) # (B, N, N) # ── 2. Extract [CLS] → patch attention ──────────────────────────────────── # CLS token is index 0; its row shows which patches it attends to cls_attn = result[:, 0, 1:] # (B, 196) — skip the CLS self-attention # ── 3. Reshape to patch grid ─────────────────────────────────────────────── patch_grid_size = int(cls_attn.shape[-1] ** 0.5) # 14 cls_attn = cls_attn.reshape(batch_size, patch_grid_size, patch_grid_size) # (B, 14, 14) # Normalize to [0, 1] for b in range(batch_size): v_min = cls_attn[b].min() v_max = cls_attn[b].max() cls_attn[b] = (cls_attn[b] - v_min) / (v_max - v_min + 1e-8) # Return as numpy — usually called with batch_size=1 for visualization return cls_attn.detach().cpu().numpy() # (B, 14, 14) def rollout_to_heatmap( rollout_map: np.ndarray, original_image: np.ndarray, colormap: int = cv2.COLORMAP_JET, alpha: float = 0.5, image_size: int = 224, ) -> np.ndarray: """ Upsample the 14×14 rollout map and overlay it on the original image. Args: rollout_map: (14, 14) or (1, 14, 14) float32 array in [0, 1]. original_image: (H, W, 3) uint8 RGB image (before normalization). colormap: OpenCV colormap constant (default: COLORMAP_JET). alpha: Heatmap overlay opacity (0=transparent, 1=opaque). image_size: Target size for both map and image (default: 224). Returns: np.ndarray (H, W, 3) uint8 — heatmap overlay on original image. """ if rollout_map.ndim == 3: rollout_map = rollout_map[0] # (14, 14) # Upsample 14×14 → 224×224 heatmap = cv2.resize(rollout_map, (image_size, image_size), interpolation=cv2.INTER_CUBIC) # Normalize to [0, 255] for colormap heatmap_uint8 = (heatmap * 255).astype(np.uint8) # Apply colormap → (H, W, 3) BGR heatmap_colored = cv2.applyColorMap(heatmap_uint8, colormap) # Convert to RGB for matplotlib/Gradio heatmap_colored_rgb = cv2.cvtColor(heatmap_colored, cv2.COLOR_BGR2RGB) # Ensure original image is 224×224 RGB if original_image.shape[:2] != (image_size, image_size): original_image = cv2.resize(original_image, (image_size, image_size)) if original_image.ndim == 2: original_image = cv2.cvtColor(original_image, cv2.COLOR_GRAY2RGB) # Alpha blend overlay = cv2.addWeighted( original_image.astype(np.float32), 1 - alpha, heatmap_colored_rgb.astype(np.float32), alpha, 0, ).astype(np.uint8) return overlay def visualize_rollout( rollout_map: np.ndarray, original_image: np.ndarray, title: str = "Attention Rollout", disease_scores: Optional[np.ndarray] = None, disease_names: Optional[List[str]] = None, save_path: Optional[str] = None, ) -> plt.Figure: """ Create a rich matplotlib figure with: - Original CLAHE-enhanced X-ray - Attention rollout heatmap - Overlay (blended) - Optional disease probability bar chart Args: rollout_map: (14, 14) attention rollout array. original_image: (H, W, 3) uint8 RGB original image. title: Figure title. disease_scores: Optional (14,) array of sigmoid probabilities. disease_names: Optional list of 14 disease name strings. save_path: If given, saves the figure to this path. Returns: matplotlib Figure object. """ overlay = rollout_to_heatmap(rollout_map, original_image) if disease_scores is not None and disease_names is not None: fig, axes = plt.subplots(1, 4, figsize=(20, 5)) n_panels = 4 else: fig, axes = plt.subplots(1, 3, figsize=(15, 5)) n_panels = 3 fig.suptitle(title, fontsize=14, fontweight="bold", y=1.02) fig.patch.set_facecolor("#1a1a2e") for ax in axes: ax.set_facecolor("#1a1a2e") # Panel 1: Original image axes[0].imshow(original_image, cmap="gray" if original_image.ndim == 2 else None) axes[0].set_title("Original X-Ray (CLAHE)", color="white", fontsize=11) axes[0].axis("off") # Panel 2: Raw rollout map (14×14 upsampled) rollout_upsampled = cv2.resize(rollout_map if rollout_map.ndim == 2 else rollout_map[0], (224, 224), interpolation=cv2.INTER_CUBIC) im = axes[1].imshow(rollout_upsampled, cmap="hot", vmin=0, vmax=1) axes[1].set_title("Attention Rollout Map", color="white", fontsize=11) axes[1].axis("off") plt.colorbar(im, ax=axes[1], fraction=0.046, pad=0.04) # Panel 3: Overlay axes[2].imshow(overlay) axes[2].set_title("Heatmap Overlay", color="white", fontsize=11) axes[2].axis("off") # Panel 4 (optional): Disease scores bar chart if disease_scores is not None and disease_names is not None and n_panels == 4: sorted_idx = np.argsort(disease_scores)[::-1] colors = [ "#ef4444" if disease_scores[i] > 0.5 else "#f97316" if disease_scores[i] > 0.3 else "#3b82f6" for i in sorted_idx ] bars = axes[3].barh( [disease_names[i] for i in sorted_idx], [disease_scores[i] for i in sorted_idx], color=colors, ) axes[3].set_xlim(0, 1) axes[3].axvline(x=0.5, color="white", linestyle="--", alpha=0.5, label="Threshold") axes[3].set_xlabel("Probability", color="white") axes[3].set_title("Disease Predictions", color="white", fontsize=11) axes[3].tick_params(colors="white") axes[3].spines["bottom"].set_color("white") axes[3].spines["left"].set_color("white") axes[3].spines["top"].set_visible(False) axes[3].spines["right"].set_visible(False) for spine in axes[3].spines.values(): spine.set_edgecolor("white") plt.tight_layout() if save_path: plt.savefig(save_path, dpi=150, bbox_inches="tight", facecolor=fig.get_facecolor()) print(f" Visualization saved → {save_path}") return fig @torch.no_grad() def explain_prediction( model, image_tensor: torch.Tensor, original_image: np.ndarray, device: torch.device, disease_names: Optional[List[str]] = None, head_fusion: str = "mean", discard_ratio: float = 0.9, ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: """ Full inference + explainability pipeline for a single image. Args: model: ChestViT model in eval mode. image_tensor: (1, 3, 224, 224) normalized tensor. original_image: (H, W, 3) uint8 original image for overlay. device: torch.device. disease_names: List of 14 disease names. head_fusion: Attention head fusion strategy. discard_ratio: Low-attention token discard ratio. Returns: probs: (14,) numpy array of sigmoid probabilities. rollout: (14, 14) numpy array — attention rollout map. overlay: (224, 224, 3) uint8 numpy array — heatmap overlay. """ model.eval() image_tensor = image_tensor.to(device) logits, attentions = model(image_tensor, output_attentions=True) probs = torch.sigmoid(logits).squeeze().cpu().numpy() # (14,) rollout = compute_attention_rollout( attentions, head_fusion=head_fusion, discard_ratio=discard_ratio ) # (1, 14, 14) overlay = rollout_to_heatmap(rollout[0], original_image) # (224, 224, 3) return probs, rollout[0], overlay