xai_engine/methods.py

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