stage2_explainability.py

"""
Stage 2 Explainability - Grad-CAM++ heatmaps + named artifact detection.
Replaces generic Attention Rollout with decision-tied explanations.
"""

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from typing import List, Dict, Any, Optional, Tuple
import cv2
from pytorch_grad_cam import GradCAMPlusPlus, GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image

from detectors.base import BaseDetector
from detectors.frequency import FrequencyDetector
from stage1_ensemble import EnsembleResult


class _CamModelWrapper(torch.nn.Module):
    """Wraps HuggingFace model so forward returns logits Tensor instead of ModelOutput dict."""
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, x):
        out = self.model(x)
        if hasattr(out, 'logits'):
            return out.logits
        if isinstance(out, dict) and 'logits' in out:
            return out['logits']
        return out


def _vit_reshape_transform(tensor):
    """Reshape ViT token sequences [B, N, C] to 2D spatial maps [B, C, H, W]."""
    if tensor.ndim == 3:
        num_tokens = tensor.size(1)
        side = int(np.sqrt(num_tokens - 1))
        if side * side == num_tokens - 1:
            tokens = tensor[:, 1:, :]
        else:
            side = int(np.sqrt(num_tokens))
            tokens = tensor[:, :side * side, :]
        result = tokens.reshape(tensor.size(0), side, side, tensor.size(2))
        return result.permute(0, 3, 1, 2)
    return tensor


class _CamTarget:
    """Target function for Grad-CAM supporting single-logit binary and multi-class models."""
    def __init__(self, target_class: int):
        self.target_class = target_class

    def __call__(self, model_output):
        s = model_output.squeeze()
        if s.ndim == 0:
            return s if self.target_class == 1 else -s
        if s.ndim == 1:
            if s.shape[0] == 1:
                return s[0] if self.target_class == 1 else -s[0]
            if self.target_class < s.shape[0]:
                return s[self.target_class]
            return s[0]
        return s[0, min(self.target_class, s.shape[-1] - 1)]


def generate_gradcam_heatmap(
    image: Image.Image,
    detector: BaseDetector,
    target_class: int = 1,  # 1 = FAKE class
) -> np.ndarray:
    """
    Generate Grad-CAM / Grad-CAM++ heatmap for a detector's decision.
    Returns normalized 2D heatmap (H, W) in [0, 1].
    """
    target_layer = detector.get_gradcam_target_layer()
    if target_layer is None:
        return np.ones((14, 14)) * 0.5

    # Preprocess image for the detector
    input_tensor = detector._preprocess(image)
    
    # Get the model with gradients enabled
    if hasattr(detector, '_get_model_for_cam'):
        model = detector._get_model_for_cam()
    elif hasattr(detector, '_model'):
        model = detector._model
    elif hasattr(detector, '_clip'):
        model = detector._clip
    elif hasattr(detector, '_siglip'):
        if hasattr(detector, '_get_model_for_cam'):
            model = detector._get_model_for_cam()
        else:
            model = detector._siglip
    else:
        return np.ones((14, 14)) * 0.5
    
    was_training = model.training
    model.train()
    
    for param in model.parameters():
        param.requires_grad = True
    
    wrapped_model = _CamModelWrapper(model)
    targets = [_CamTarget(target_class)]

    try:
        try:
            cam = GradCAM(
                model=wrapped_model,
                target_layers=[target_layer],
                reshape_transform=_vit_reshape_transform,
            )
            grayscale_cam = cam(input_tensor=input_tensor, targets=targets)
        except Exception as cam_err:
            print(f"[Stage 2] GradCAM with reshape_transform failed: {cam_err}, trying direct...")
            cam = GradCAM(
                model=wrapped_model,
                target_layers=[target_layer],
            )
            grayscale_cam = cam(input_tensor=input_tensor, targets=targets)

        if grayscale_cam.ndim == 3:
            grayscale_cam = grayscale_cam[0, :]
            
    finally:
        if not was_training:
            model.eval()
    
    return grayscale_cam


def heatmap_to_overlay(
    cam: np.ndarray,
    original_image: Image.Image,
    alpha: float = 0.5,
    colormap: int = cv2.COLORMAP_JET,
) -> bytes:
    """Overlay heatmap on original image, return as PNG bytes."""
    if cam.ndim == 3:
        cam = cam.squeeze(0)
    # Resize CAM to image size
    cam_resized = cv2.resize(cam, original_image.size, interpolation=cv2.INTER_LINEAR)
    
    # Normalize to [0, 1]
    if cam_resized.max() > cam_resized.min():
        cam_resized = (cam_resized - cam_resized.min()) / (cam_resized.max() - cam_resized.min())
    else:
        cam_resized = np.ones_like(cam_resized) * 0.5
    
    # Apply colormap
    cam_uint8 = (cam_resized * 255).astype(np.uint8)
    heatmap_colored = cv2.applyColorMap(cam_uint8, colormap)
    heatmap_colored = cv2.cvtColor(heatmap_colored, cv2.COLOR_BGR2RGB)
    
    # Blend with original
    original_np = np.array(original_image.convert("RGB"))
    blended = cv2.addWeighted(original_np, 1 - alpha, heatmap_colored, alpha, 0)
    
    # Encode to PNG
    success, buffer = cv2.imencode('.png', cv2.cvtColor(blended, cv2.COLOR_RGB2BGR))
    if not success:
        raise RuntimeError("Failed to encode heatmap overlay")
    
    return buffer.tobytes()


def detect_blending_boundary(
    image: Image.Image,
    heatmap: np.ndarray,
    face_region: Optional[Tuple[int, int, int, int]] = None,
) -> Optional[Dict[str, Any]]:
    """
    Detect face-swap blending boundary artifacts.
    Looks for sharp transitions in heatmap at face edges.
    """
    if heatmap.ndim == 3:
        heatmap = heatmap.squeeze(0)
    # Convert heatmap to edges
    heatmap_uint8 = (heatmap * 255).astype(np.uint8)
    edges = cv2.Canny(heatmap_uint8, 50, 150)
    
    # Find contours
    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    if not contours:
        return None
    
    # Filter contours near face region if provided
    suspicious_contours = []
    h, w = heatmap.shape
    
    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area < 50:  # Too small
            continue
        
        x, y, cw, ch = cv2.boundingRect(cnt)
        
        # Check if contour aligns with face boundary
        if face_region:
            fx, fy, fw, fh = face_region
            # Check if contour overlaps face boundary
            if (x < fx + fw and x + cw > fx and 
                y < fy + fh and y + ch > fy):
                suspicious_contours.append({
                    "bbox": [int(x), int(y), int(cw), int(ch)],
                    "area": float(area),
                    "type": "blending_boundary",
                })
        else:
            # General suspicious edge
            suspicious_contours.append({
                "bbox": [int(x), int(y), int(cw), int(ch)],
                "area": float(area),
                "type": "suspicious_edge",
            })
    
    if not suspicious_contours:
        return None
    
    return {
        "type": "blending_boundary",
        "regions": suspicious_contours,
        "severity": min(len(suspicious_contours) / 5.0, 1.0),
    }


def detect_ela_inconsistency(image: Image.Image, block_size: int = 8) -> Optional[Dict[str, Any]]:
    """
    Error Level Analysis - detects JPEG recompression inconsistencies.
    Classic splicing indicator: different compression levels in different regions.
    """
    # Save as JPEG at known quality, reload, compute difference
    import io
    
    # Original
    original = np.array(image.convert("RGB"))
    
    # Re-compress at quality 90
    buf = io.BytesIO()
    image.save(buf, format='JPEG', quality=90)
    buf.seek(0)
    recompressed = np.array(Image.open(buf).convert("RGB"))
    
    # Difference
    diff = np.abs(original.astype(np.float32) - recompressed.astype(np.float32))
    diff = diff.mean(axis=2)  # Average across channels
    
    # Block-wise analysis
    h, w = diff.shape
    block_means = []
    
    for y in range(0, h - block_size, block_size):
        for x in range(0, w - block_size, block_size):
            block = diff[y:y+block_size, x:x+block_size]
            block_means.append(block.mean())
    
    block_means = np.array(block_means)
    
    if len(block_means) == 0:
        return None
    
    # Look for outliers (blocks with significantly different error levels)
    median_err = np.median(block_means)
    mad = np.median(np.abs(block_means - median_err))
    
    # Modified z-score
    z_scores = 0.6745 * (block_means - median_err) / (mad + 1e-8)
    outliers = np.where(np.abs(z_scores) > 3.5)[0]
    
    if len(outliers) == 0:
        return None
    
    # Convert outlier indices to coordinates
    suspicious_blocks = []
    for idx in outliers[:10]:  # Limit to top 10
        by = idx // ((w - block_size) // block_size)
        bx = idx % ((w - block_size) // block_size)
        suspicious_blocks.append({
            "x": int(bx * block_size),
            "y": int(by * block_size),
            "size": block_size,
            "z_score": float(z_scores[idx]),
        })
    
    return {
        "type": "ela_inconsistency",
        "suspicious_blocks": suspicious_blocks,
        "severity": min(len(outliers) / 20.0, 1.0),
        "description": f"Found {len(outliers)} blocks with inconsistent JPEG compression levels",
    }


def detect_frequency_artifacts(
    image: Image.Image,
    frequency_detector: FrequencyDetector,
) -> List[Dict[str, Any]]:
    """Get artifacts from frequency domain analysis."""
    return frequency_detector.get_artifacts(image)


def run_stage2_explainability(
    image: Image.Image,
    ensemble_result: EnsembleResult,
    detectors: List[BaseDetector],
    mode: str,
) -> Dict[str, Any]:
    """
    Run full Stage 2 explainability pipeline.
    Returns heatmap bytes and named artifact findings.
    """
    # Use primary detector (highest weight) for Grad-CAM++
    primary_detector = max(detectors, key=lambda d: d.model_name in ensemble_result.primary_model)
    
    # Target class: 1 for FAKE (assuming binary with FAKE=1)
    target_class = 1 if ensemble_result.ensemble_label == "FAKE" else 0
    
    # Generate Grad-CAM++ heatmap
    try:
        cam = generate_gradcam_heatmap(image, primary_detector, target_class)
    except Exception as e:
        print(f"[Stage 2] Grad-CAM++ failed: {e}, using fallback")
        cam = np.ones((14, 14)) * 0.5
    
    # Create overlay
    heatmap_bytes = heatmap_to_overlay(cam, image)
    
    # Detect named artifacts
    artifacts = []
    
    # 1. Frequency artifacts (always run)
    freq_detector = next((d for d in detectors if isinstance(d, FrequencyDetector)), None)
    if freq_detector:
        freq_artifacts = detect_frequency_artifacts(image, freq_detector)
        artifacts.extend(freq_artifacts)
    
    # 2. Face-swap blending boundary (deepfake mode)
    if mode == "deepfake":
        # Try to get face region (simplified - would use face detector in production)
        blend_result = detect_blending_boundary(image, cam)
        if blend_result:
            artifacts.append(blend_result)
    
    # 3. ELA inconsistency (both modes)
    ela_result = detect_ela_inconsistency(image)
    if ela_result:
        artifacts.append(ela_result)
    
    # 4. High activation regions from heatmap
    high_activation = np.percentile(cam, 90)
    high_regions = np.argwhere(cam >= high_activation)
    if len(high_regions) > 0:
        artifacts.append({
            "type": "high_model_attention",
            "description": f"Model focused on {len(high_regions)} regions with top 10% activation",
            "severity": float(np.mean(cam[cam >= high_activation])),
            "locations": high_regions.tolist()[:20],  # Limit
        })
    
    return {
        "heatmap_bytes": heatmap_bytes,
        "heatmap_array": cam,
        "artifacts": artifacts,
        "primary_detector": primary_detector.model_name,
    }