import time from dataclasses import dataclass import numpy as np import torch from PIL import Image from captum.attr import IntegratedGradients, Occlusion, DeepLift from pytorch_grad_cam import GradCAM from pytorch_grad_cam.utils.model_targets import BinaryClassifierOutputTarget from .chefer import chefer_gradient_rollout from .viz import create_overlay, render_heatmap_raw, to_b64_png, normalize_heatmap, fast_resize, to_b64_jpeg @dataclass class AttributionResult: name: str description: str heatmap: np.ndarray # 2D or 3D numpy array heatmap_norm: np.ndarray # Normalized [0, 1] 2D heatmap (384, 384) overlay_b64: str # Base64 PNG data URL raw_heatmap_b64: str # Base64 PNG data URL metadata: dict compute_time_ms: float def attention_rollout(wrapper, input_tensor) -> np.ndarray: """Attention Rollout for ViT across 12 transformer layers.""" model = wrapper.model model.eval() with torch.no_grad(): outputs = model(input_tensor, output_attentions=True) attentions = outputs.attentions device = attentions[0].device seq_len = attentions[0].size(-1) # 577 result = torch.eye(seq_len, device=device) for attn in attentions: attn_heads = attn.mean(dim=1) # Average over 6 heads -> (B, 577, 577) attn_heads = attn_heads + torch.eye(attn_heads.size(-1), device=attn_heads.device) attn_heads = attn_heads / attn_heads.sum(dim=-1, keepdim=True) result = torch.matmul(attn_heads[0], result) patch_attention = result[0, 1:].reshape(24, 24).cpu().numpy() return patch_attention def integrated_gradients(wrapper, input_tensor, steps=20) -> tuple[np.ndarray, float]: """Integrated Gradients using Captum with memory-safe mini-batching.""" ig = IntegratedGradients(wrapper) baseline = torch.zeros_like(input_tensor) attributions, delta = ig.attribute( input_tensor, baselines=baseline, target=0, n_steps=steps, internal_batch_size=2, # Prevents memory spike on CPU/RAM return_convergence_delta=True ) heatmap = attributions[0].sum(dim=0).detach().cpu().numpy() delta_val = float(delta[0].item()) if hasattr(delta, "__getitem__") else float(delta.item()) return heatmap, delta_val def get_vit_encoder_layer(wrapper, layer_idx: int): """Safely traverse module hierarchy to find ViT encoder layer list regardless of HF transformers version.""" obj = wrapper if hasattr(obj, "model"): obj = obj.model if hasattr(obj, "vit"): obj = obj.vit if hasattr(obj, "encoder"): obj = obj.encoder for attr in ["layer", "layers", "block", "blocks"]: if hasattr(obj, attr): return getattr(obj, attr)[layer_idx] raise AttributeError(f"Could not locate transformer layers list in object of type {type(obj)}") def vit_grad_cam(wrapper, input_tensor, target_layer_idx=11) -> np.ndarray: """GradCAM for ViT using pytorch_grad_cam with custom reshape transform.""" layer = get_vit_encoder_layer(wrapper, target_layer_idx) target_layer = [layer.output if hasattr(layer, "output") else layer] def reshape_transform(tensor): # tensor shape: (B, 577, 384) -> drop CLS token -> (B, 576, 384) -> (B, 24, 24, 384) -> (B, 384, 24, 24) result = tensor[:, 1:, :].reshape(tensor.size(0), 24, 24, tensor.size(2)) result = result.permute(0, 3, 1, 2) return result cam = GradCAM( model=wrapper, target_layers=target_layer, reshape_transform=reshape_transform, ) targets = [BinaryClassifierOutputTarget(0)] grayscale_cam = cam(input_tensor=input_tensor, targets=targets) return grayscale_cam[0] # (24, 24) or (384, 384) def occlusion_sensitivity(wrapper, input_tensor, patch_size=32, stride=16) -> np.ndarray: """Occlusion Sensitivity attribution using Captum with memory-safe mini-batching.""" occ = Occlusion(wrapper) attributions = occ.attribute( input_tensor, target=0, strides=(3, stride, stride), sliding_window_shapes=(3, patch_size, patch_size), baselines=0.0, perturbations_per_eval=4, # Limits memory allocation per evaluation batch ) heatmap = attributions[0].sum(dim=0).detach().cpu().numpy() return heatmap def _cleanup_captum_hooks(model): """Remove all forward/backward hooks Captum registers on model modules. Captum's DeepLift (and LayerGradCam, etc.) attach hooks via register_forward_hook / register_full_backward_hook. If attribute() throws mid-execution these hooks survive on the singleton model and corrupt every subsequent attribution call on the same process. """ for module in model.modules(): module._forward_hooks.clear() module._forward_pre_hooks.clear() module._backward_hooks.clear() def deeplift(wrapper, input_tensor) -> np.ndarray: """DeepLIFT attribution using Captum with hook cleanup to prevent cross-method contamination.""" try: dl = DeepLift(wrapper) baseline = torch.zeros_like(input_tensor) attributions = dl.attribute(input_tensor, baselines=baseline, target=0) heatmap = attributions[0].sum(dim=0).detach().cpu().numpy() return heatmap except Exception as e: print(f"[deeplift] Warning: Captum DeepLIFT failed ({e}), falling back to Input x Gradient.") input_tensor_copy = input_tensor.clone().detach().requires_grad_(True) logits = wrapper(input_tensor_copy) logits[0, 0].backward() heatmap = (input_tensor_copy * input_tensor_copy.grad).sum(dim=1)[0].detach().cpu().numpy() return heatmap finally: _cleanup_captum_hooks(wrapper) def run_per_head_attention(wrapper, image: Image.Image, layer_idx: int = 11, cmap: str = "jet") -> dict: """Extract and render attention maps for all 6 heads at a specific layer.""" input_tensor = wrapper.preprocess(image) device = next(wrapper.model.parameters()).device input_tensor = input_tensor.to(device) with torch.no_grad(): outputs = wrapper.model(input_tensor, output_attentions=True) attns = outputs.attentions[layer_idx][0] # (6, 577, 577) head_results = [] for h in range(6): head_map = attns[h, 0, 1:].reshape(24, 24).cpu().numpy() overlay = create_overlay(image, head_map, alpha=0.6, cmap=cmap) head_results.append({ "head": h + 1, "overlay_b64": to_b64_png(overlay), "raw_b64": to_b64_png(render_heatmap_raw(head_map, cmap=cmap)), "min": round(float(head_map.min()), 4), "max": round(float(head_map.max()), 4), }) return { "layer": layer_idx + 1, "heads": head_results, } import threading _method_lock = threading.Lock() def run_single_method( wrapper, image: Image.Image, method_name: str, alpha: float = 0.6, cmap: str = "jet", steps: int = 20, patch_size: int = 32, stride: int = 16, target_layer_idx: int = 11, ) -> AttributionResult: """Execute a single specified XAI method with thread safety on model state.""" with _method_lock: input_tensor = wrapper.preprocess(image) device = next(wrapper.model.parameters()).device input_tensor = input_tensor.to(device) t0 = time.perf_counter() metadata = {} if method_name == "rollout": disp_name = "Attention Rollout" desc = "Global multi-layer attention flow across all 12 transformer blocks." heatmap = attention_rollout(wrapper, input_tensor) elif method_name == "integrated": disp_name = "Integrated Gradients" desc = "Path-integral gradient attribution relative to a black baseline." heatmap, delta = integrated_gradients(wrapper, input_tensor, steps=steps) metadata["delta"] = delta metadata["steps"] = steps elif method_name == "gradcam": disp_name = "GradCAM" desc = f"Gradient-weighted activation map targeting layer {target_layer_idx + 1}." heatmap = vit_grad_cam(wrapper, input_tensor, target_layer_idx=target_layer_idx) metadata["target_layer"] = target_layer_idx + 1 elif method_name == "chefer": disp_name = "Gradient Attention Rollout (Chefer)" desc = "Gradient-weighted attention rollout combining gradients and self-attention (CVPR 2021)." heatmap = chefer_gradient_rollout(wrapper, input_tensor) elif method_name == "occlusion": disp_name = "Occlusion Sensitivity" desc = f"Sliding window perturbation ({patch_size}x{patch_size} patch, stride {stride})." heatmap = occlusion_sensitivity(wrapper, input_tensor, patch_size=patch_size, stride=stride) metadata["patch_size"] = patch_size metadata["stride"] = stride elif method_name == "deeplift": disp_name = "DeepLIFT" desc = "Non-linear feature attribution backpropagating differences from baseline." heatmap = deeplift(wrapper, input_tensor) else: raise ValueError(f"Unknown XAI method: {method_name}") elapsed_ms = (time.perf_counter() - t0) * 1000 # Resize heatmap to 384x384 if needed if heatmap.shape != (384, 384): heatmap_resized = fast_resize(heatmap, (384, 384)) else: heatmap_resized = heatmap norm_h = normalize_heatmap(heatmap_resized) overlay = create_overlay(image, heatmap, alpha=alpha, cmap=cmap) raw_render = render_heatmap_raw(heatmap, cmap=cmap) metadata.update({ "min": round(float(heatmap.min()), 5), "max": round(float(heatmap.max()), 5), "mean": round(float(heatmap.mean()), 5), "std": round(float(heatmap.std()), 5), }) return AttributionResult( name=disp_name, description=desc, heatmap=heatmap, heatmap_norm=norm_h, overlay_b64=to_b64_jpeg(overlay), raw_heatmap_b64=to_b64_jpeg(raw_render), metadata=metadata, compute_time_ms=round(elapsed_ms, 1), ) def run_all_methods( wrapper, image: Image.Image, alpha: float = 0.6, cmap: str = "jet", ) -> list[AttributionResult]: """Run all 6 primary spatial attribution methods sequentially with per-method exception handling.""" methods = ["rollout", "chefer", "gradcam", "integrated", "deeplift", "occlusion"] results = [] for m in methods: try: res = run_single_method(wrapper, image, m, alpha=alpha, cmap=cmap) results.append(res) except Exception as err: print(f"[xai_engine] Error running method {m}: {err}") return results