infer_xai.py

"""
ManipShield XAI Inference -- Full Explainability Pipeline
=========================================================

Runs all XAI techniques on a single image and saves annotated PNGs.

Usage:
  python infer_xai.py image.jpg
  python infer_xai.py image.jpg --checkpoint checkpoints/best.pth
  python infer_xai.py image.jpg --build-probes --data-dir data/
  python infer_xai.py image.jpg --build-prototypes --data-dir data/
  python infer_xai.py image.jpg --skip-tcav --skip-prototype
  python infer_xai.py image.jpg --output-dir my_xai_out/

Output folder (default: xai_output/<image_stem>_<timestamp>/):
  annotated_attributes.png  -- per-attribute labeled boxes (paper style)
  gradient_saliency.png     -- gradient x input saliency map
  fft_heatmap.png           -- frequency anomaly map
  counterfactual.png        -- masked region + confidence delta
  forensic_residual.png     -- PRNU noise residual map
  summary.json              -- all scores as JSON
"""

import sys
import argparse
import json
import datetime
from pathlib import Path
from typing import Optional

# ── resolve detection module ──────────────────────────────────────────────────
DETECTION_DIR = Path(__file__).parent / "v1" / "inceptrix" / "ml" / "detection"
sys.path.insert(0, str(DETECTION_DIR))

import torch
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms

from detector import DeepfakeDetector
from xai_techniques import (
    GradientSaliency,
    FrequencyAnalyzer,
    BBoxVisualizer,
    AttributeRegionVisualizer,
    MetadataExtractor,
    CounterfactualExplainer,
    TCavAnalyzer,
    PrototypeLibrary,
    ForensicFingerprint,
    TCAV_CONCEPTS,
)

# ── paths ─────────────────────────────────────────────────────────────────────
ROOT            = Path(__file__).parent
CHECKPOINTS_DIR = ROOT / "checkpoints"
XAI_OUTPUT_DIR  = ROOT / "xai_output"
TCAV_PROBE_PATH = XAI_OUTPUT_DIR / "tcav_probes.pkl"
PROTO_LIB_PATH  = XAI_OUTPUT_DIR / "prototype_lib.pkl"
FINGERPRINT_DIR = XAI_OUTPUT_DIR / "fingerprints"

# ── preprocessing (must match train.py) ──────────────────────────────────────
_PREPROC = transforms.Compose([
    transforms.Resize((768, 768), interpolation=transforms.InterpolationMode.BICUBIC),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])


# ── ANSI colours ──────────────────────────────────────────────────────────────
class C:
    RESET   = "\033[0m"
    BOLD    = "\033[1m"
    RED     = "\033[91m"
    GREEN   = "\033[92m"
    YELLOW  = "\033[93m"
    CYAN    = "\033[96m"
    WHITE   = "\033[97m"
    DIM     = "\033[2m"
    MAGENTA = "\033[95m"


def bar(v: float, w: int = 28, fill: str = "#", empty: str = "-") -> str:
    n = int(round(v * w))
    return fill * n + empty * (w - n)


# ── model loading ──────────────────────────────────────────────────────────────
def load_model(checkpoint_path: Path, device: torch.device) -> DeepfakeDetector:
    print(f"{C.DIM}Loading: {checkpoint_path.name}{C.RESET}")
    ckpt   = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    config = ckpt.get("config", {})
    ep     = ckpt.get("epoch", "?")
    acc    = ckpt.get("val_acc", float("nan"))
    print(f"{C.DIM}  epoch={ep+1 if isinstance(ep, int) else ep}, val_acc={acc:.2f}%{C.RESET}")
    model = DeepfakeDetector(
        model_name=config.get("model_name", "microsoft/Florence-2-base"),
        lora_rank=config.get("lora_rank", 8),
    )
    model.load_state_dict(ckpt["trainable_state_dict"], strict=False)
    model.to(device).eval()
    del ckpt
    torch.cuda.empty_cache()
    return model


# ── terminal helpers ───────────────────────────────────────────────────────────
def print_header(title: str):
    print(f"\n{C.BOLD}{C.WHITE}{'='*62}{C.RESET}")
    print(f"{C.BOLD}{C.WHITE}  {title}{C.RESET}")
    print(f"{C.BOLD}{C.WHITE}{'='*62}{C.RESET}")


def print_section(title: str):
    print(f"\n  {C.BOLD}{C.CYAN}>> {title}{C.RESET}")
    print(f"  {C.DIM}{'-'*56}{C.RESET}")


# ── label maps ────────────────────────────────────────────────────────────────
MANIP_LABELS = {
    "face_swap":              "Face Swap",
    "face_reenactment":       "Face Reenactment",
    "entire_synthesis":       "Entire Synthesis (GAN/Diffusion)",
    "attribute_manipulation": "Attribute Manipulation",
    "inpainting":             "Inpainting",
    "unknown":                "Unknown",
}

ATTR_LABELS = {
    "lighting":               "Lighting",
    "texture":                "Texture",
    "shape":                  "Shape",
    "color_consistency":      "Color Consistency",
    "boundary_artifacts":     "Boundary Artifacts",
    "noise_pattern":          "Noise Pattern",
    "compression":            "Compression",
    "reflection":             "Reflection",
    "shadow":                 "Shadow",
    "hair_detail":            "Hair Detail",
    "skin_texture":           "Skin Texture",
    "background_consistency": "Background Consistency",
}

CONCEPT_LABELS = {
    "unnatural_lighting":   "Unnatural Lighting",
    "texture_smoothing":    "Texture Over-smoothing",
    "geometric_distortion": "Geometric Distortion",
    "color_mismatch":       "Color Mismatch",
    "shadow_inconsistency": "Shadow Inconsistency",
    "frequency_anomaly":    "Frequency Anomaly",
}


# ── gradient saliency helper (runs before drawing so boxes use saliency) ──────
def _compute_saliency(
    model: DeepfakeDetector,
    pixel_values: torch.Tensor,
    device: torch.device,
) -> Optional[object]:
    """Compute gradient x input saliency map. Returns (HxW ndarray, score) or None."""
    import numpy as np
    try:
        pv = pixel_values.clone().detach().requires_grad_(True).to(device)
        with torch.enable_grad():
            image_features = model.florence2._encode_image(pv)
            encoder = model.florence2.language_model.model.model.encoder
            enc_out = encoder(inputs_embeds=image_features, output_hidden_states=True)
            hidden_states = list(enc_out.hidden_states[1:])
            features = model.lds(hidden_states)
            fake_prob = model.detection_head(features)["fake_prob"].mean()
            fake_prob.backward()
        grad = pv.grad.detach().cpu()
        sal = (grad * pv.detach().cpu()).abs()[0].max(dim=0).values.numpy()
        sal = (sal - sal.min()) / (sal.max() - sal.min() + 1e-8)
        return sal, float(sal.mean())
    except Exception:
        return None, 0.0


# ── main XAI pipeline ─────────────────────────────────────────────────────────
def run_xai(
    args,
    model: DeepfakeDetector,
    img_pil: Image.Image,
    image_path: Path,
    device: torch.device,
    out_dir: Path,
) -> dict:
    import numpy as np
    out_dir.mkdir(parents=True, exist_ok=True)
    summary = {
        "image":     str(image_path),
        "timestamp": datetime.datetime.now().isoformat(),
    }

    pixel_values = _PREPROC(img_pil).unsqueeze(0).to(device)

    # =========================================================================
    # 1. BASE DETECTION
    # =========================================================================
    print_section("Base Detection")
    with torch.no_grad():
        out = model(pixel_values)

    real_p   = out["real_prob"][0].item()
    fake_p   = out["fake_prob"][0].item()
    is_fake  = fake_p >= args.threshold
    vcol     = C.RED if is_fake else C.GREEN
    verdict  = "FAKE" if is_fake else "REAL"

    print(f"  Verdict : {C.BOLD}{vcol}[ {verdict} ]{C.RESET}  (threshold={args.threshold:.2f})")
    print(f"  {C.GREEN}Real{C.RESET} [{C.GREEN}{bar(real_p)}{C.RESET}] {real_p*100:5.1f}%")
    print(f"  {C.RED}Fake{C.RESET} [{C.RED}{bar(fake_p)}{C.RESET}] {fake_p*100:5.1f}%")

    summary.update({
        "verdict":   verdict,
        "real_prob": round(real_p, 4),
        "fake_prob": round(fake_p, 4),
        "threshold": args.threshold,
    })

    manip_probs = out.get("manipulation_probs", {})
    if manip_probs:
        top_type = max(manip_probs, key=manip_probs.get)
        print(f"\n  Manipulation Type : {C.CYAN}{MANIP_LABELS.get(top_type, top_type)}{C.RESET}"
              f"  ({manip_probs[top_type]*100:.1f}%)")
        for k, v in sorted(manip_probs.items(), key=lambda x: -x[1]):
            col = C.CYAN if k == top_type else C.DIM
            print(f"    {MANIP_LABELS.get(k, k):<38} {col}{bar(v, 12)}{C.RESET} {v*100:5.1f}%")
        summary["manipulation_type"]  = top_type
        summary["manipulation_probs"] = {k: round(v, 4) for k, v in manip_probs.items()}

    attrs = out.get("attributes", {})
    attr_probs: dict = {}
    if attrs:
        attr_probs = {a: F.softmax(l, dim=-1)[0, 1].item() for a, l in attrs.items()}
        flagged    = [a for a, p in attr_probs.items() if p > 0.5]
        print(f"\n  Manipulation Cues : {len(flagged)}/12 triggered")
        for attr, prob in sorted(attr_probs.items(), key=lambda x: -x[1]):
            col    = C.RED if prob > 0.5 else (C.YELLOW if prob > 0.3 else C.DIM)
            marker = " [!]" if prob > 0.5 else ""
            print(f"    {ATTR_LABELS.get(attr, attr):<30} {col}{bar(prob, 14)}{C.RESET}"
                  f" {prob*100:5.1f}%{marker}")
        summary["attribute_cues"] = {k: round(v, 4) for k, v in attr_probs.items()}

    boxes_tensor = out["bounding_boxes"][0]  # (num_boxes, 5)

    # =========================================================================
    # 2. IMAGE METADATA + C2PA PROVENANCE
    # =========================================================================
    print_section("Image Metadata & C2PA Provenance")
    meta_extractor = MetadataExtractor()
    meta = meta_extractor.extract(image_path)
    for line in meta_extractor.format_terminal(meta):
        if "MANIFEST DETECTED" in line:
            print(f"  {C.RED}{C.BOLD}{line.strip()}{C.RESET}")
        elif "AI Markers" in line and "None" not in line:
            print(f"  {C.YELLOW}{line.strip()}{C.RESET}")
        elif "No manifest" in line or "None" in line:
            print(f"  {C.DIM}{line.strip()}{C.RESET}")
        else:
            print(f"  {line.strip()}")
    summary["metadata"] = {
        k: v for k, v in meta.items()
        if k not in ("exif",)   # skip raw exif blob in summary
    }
    summary["exif"] = meta.get("exif", {})

    # =========================================================================
    # 3. GRADIENT SALIENCY  (computed first so boxes can use it)
    # =========================================================================
    saliency_map: Optional[object] = None
    sal_score = 0.0

    if not args.skip_gradient:
        print_section("Gradient Saliency (input x gradient)")
        saliency_map, sal_score = _compute_saliency(model, pixel_values, device)
        if saliency_map is not None:
            # Save overlay
            from xai_techniques import _to_np_rgb, _overlay_heatmap, _save_png
            img_rgb = _to_np_rgb(img_pil)
            overlay = _overlay_heatmap(img_rgb, saliency_map, alpha=0.5)
            _save_png(overlay, out_dir / "gradient_saliency.png")
            print(f"  Mean saliency score : {sal_score:.4f}")
            print(f"  {C.DIM}-> saved gradient_saliency.png{C.RESET}")
            summary["gradient_saliency_score"] = round(sal_score, 4)
        else:
            print(f"  {C.YELLOW}Gradient saliency unavailable{C.RESET}")

    # =========================================================================
    # 4. ATTRIBUTE REGION BOXES  (ManipShield paper style)
    # =========================================================================
    print_section("Attribute Region Localization (ManipShield paper style)")
    attr_vis = AttributeRegionVisualizer()

    if attr_probs:
        _, attr_annotations = attr_vis.draw(
            img_pil,
            attr_probs,
            boxes_tensor,
            out_dir / "annotated_attributes.png",
            saliency_map=saliency_map,
            attr_threshold=0.50,
        )
        if attr_annotations:
            print(f"  {len(attr_annotations)} attribute regions drawn:")
            for ann in attr_annotations:
                r, g, b = ann["color_rgb"]
                bx = ann["box_pixels"]
                print(f"    [{ann['display']:<22}]  {ann['confidence']*100:5.1f}%"
                      f"  box=[{bx[0]},{bx[1]}->{bx[2]},{bx[3]}]")
            print(f"  {C.DIM}-> saved annotated_attributes.png{C.RESET}")
        else:
            print(f"  {C.DIM}No attributes above 0.5 threshold{C.RESET}")
        summary["attribute_annotations"] = attr_annotations
    else:
        print(f"  {C.DIM}No attribute outputs from model{C.RESET}")
        summary["attribute_annotations"] = []

    # Also save the raw localization boxes (generic, no attribute labels)
    bbox_vis = BBoxVisualizer()
    _, drawn_boxes = bbox_vis.draw(
        img_pil, boxes_tensor,
        out_dir / "localization_boxes.png",
        min_conf=0.3,
    )
    summary["localization_boxes"] = drawn_boxes

    # =========================================================================
    # 5. FFT FREQUENCY ANALYSIS
    # =========================================================================
    print_section("Frequency Analysis (FFT / 1-f Spectral Anomaly)")
    freq_analyzer = FrequencyAnalyzer()
    _, spectral_score = freq_analyzer.analyze(img_pil, out_dir / "fft_heatmap.png")
    desc = freq_analyzer.get_spectral_description(spectral_score)
    col  = C.RED if spectral_score > 0.5 else (C.YELLOW if spectral_score > 0.25 else C.GREEN)
    print(f"  Spectral anomaly : {col}{spectral_score*100:.1f}%{C.RESET}")
    print(f"  {desc}")
    print(f"  {C.DIM}-> saved fft_heatmap.png{C.RESET}")
    summary["spectral_anomaly_score"] = round(spectral_score, 4)
    summary["spectral_description"]   = desc

    # =========================================================================
    # 6. COUNTERFACTUAL EXPLANATION
    # =========================================================================
    print_section("Counterfactual Explanation (causal region masking)")
    cf = CounterfactualExplainer(fill_strategy="mean")
    cf_result = cf.explain(
        model, img_pil, boxes_tensor, pixel_values,
        preprocess_fn=_PREPROC,
        out_path=out_dir / "counterfactual.png",
        min_conf=0.3,
    )
    delta  = cf_result["delta"]
    causal = cf_result["causal_score"]
    col    = C.RED if causal > 0.3 else (C.YELLOW if causal > 0.1 else C.GREEN)
    print(f"  {cf_result['message']}")
    if cf_result.get("top_box"):
        print(f"  Causal score : {col}{causal*100:.1f}%{C.RESET}   "
              f"delta-fake-prob = {delta*100:+.1f}pp")
        print(f"  {C.DIM}-> saved counterfactual.png{C.RESET}")
    summary["counterfactual"] = cf_result

    # =========================================================================
    # 7. TCAV  —  Concept Activation Vectors
    # =========================================================================
    if not args.skip_tcav:
        print_section("TCAV -- Concept Activation Vectors")
        data_dir = Path(args.data_dir) if args.data_dir else ROOT / "data"
        tcav = TCavAnalyzer(TCAV_PROBE_PATH, data_dir)

        if args.build_probes or not TCAV_PROBE_PATH.exists():
            tcav.build_probes(model, _PREPROC, device, max_images=args.tcav_images)

        tcav_scores = tcav.score(model, pixel_values)
        for concept, score in sorted(tcav_scores.items(), key=lambda x: -x[1]):
            col   = C.RED if score > 0.65 else (C.YELLOW if score > 0.45 else C.DIM)
            label = CONCEPT_LABELS.get(concept, concept)
            print(f"  {label:<30} {col}{bar(score, 14)}{C.RESET} {score*100:5.1f}%")
        summary["tcav_scores"] = {k: round(v, 4) for k, v in tcav_scores.items()}

    # =========================================================================
    # 8. PROTOTYPE MATCHING
    # =========================================================================
    if not args.skip_prototype:
        print_section("Prototype Matching (contrastive embedding k-NN)")
        data_dir = Path(args.data_dir) if args.data_dir else ROOT / "data"
        proto_lib = PrototypeLibrary(PROTO_LIB_PATH)

        if args.build_prototypes or not PROTO_LIB_PATH.exists():
            proto_lib.build(model, _PREPROC, device, data_dir, max_per_class=100)

        query_emb = out["contrastive_emb"][0].cpu().numpy()
        neighbors = proto_lib.find_nearest(query_emb, k=3)
        if neighbors:
            for nn in neighbors:
                col = C.RED if nn["label"] == "fake" else C.GREEN
                print(f"  sim={nn['similarity']:.3f}  {col}{nn['label']:<5}{C.RESET}"
                      f"  type={nn['subtype']:<25}  ({nn['source']})")
        else:
            print(f"  {C.DIM}No prototype library yet -- run with --build-prototypes{C.RESET}")
        summary["prototype_matches"] = neighbors

    # =========================================================================
    # 9. FORENSIC FINGERPRINT  (PRNU noise residual)
    # =========================================================================
    print_section("Forensic Fingerprint (PRNU Noise Residual)")
    forensic = ForensicFingerprint(FINGERPRINT_DIR)
    similarities, best_match = forensic.analyze(img_pil, out_dir / "forensic_residual.png")
    desc_f = forensic.get_residual_description(similarities, best_match)
    print(f"  {desc_f}")
    if similarities:
        for gen, sim in sorted(similarities.items(), key=lambda x: -x[1]):
            col = C.RED if gen != "real" and sim > 0.5 else C.DIM
            print(f"    {gen:<22} {col}{bar(sim, 14)}{C.RESET} {sim*100:5.1f}%")
    else:
        print(f"  {C.DIM}No reference fingerprints -- run forensic.add_reference() to build{C.RESET}")
    print(f"  {C.DIM}-> saved forensic_residual.png{C.RESET}")
    summary["forensic_fingerprint"] = {"similarities": similarities, "best_match": best_match}

    # =========================================================================
    # Save summary JSON
    # =========================================================================
    summary_path = out_dir / "summary.json"
    with open(summary_path, "w") as f:
        json.dump(summary, f, indent=2, default=str)

    return summary


# ── main ──────────────────────────────────────────────────────────────────────
def main():
    parser = argparse.ArgumentParser(
        description="ManipShield XAI Inference -- full explainability pipeline",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument("image",         nargs="?", help="Path to input image")
    parser.add_argument("--checkpoint",  "-c",
                        default=str(CHECKPOINTS_DIR / "best.pth"),
                        help="Checkpoint .pth file (default: checkpoints/best.pth)")
    parser.add_argument("--threshold",   "-t", type=float, default=0.5)
    parser.add_argument("--output-dir",  "-o", default=None,
                        help="Override output dir (default: xai_output/<stem>_<ts>/)")
    parser.add_argument("--data-dir",    default=None,
                        help="Data dir for TCAV probes / prototypes (default: data/)")
    parser.add_argument("--build-probes",      action="store_true",
                        help="Force rebuild TCAV probes from data-dir")
    parser.add_argument("--build-prototypes",  action="store_true",
                        help="Force rebuild prototype library from data-dir")
    parser.add_argument("--tcav-images", type=int, default=200,
                        help="Max images when building TCAV probes (default 200)")
    parser.add_argument("--skip-tcav",       action="store_true")
    parser.add_argument("--skip-prototype",  action="store_true")
    parser.add_argument("--skip-gradient",   action="store_true")
    parser.add_argument("--cpu",             action="store_true", help="Force CPU inference")
    args = parser.parse_args()

    if args.image is None:
        parser.print_help()
        sys.exit(1)

    image_path = Path(args.image)
    if not image_path.exists():
        print(f"{C.RED}Error: image not found: {image_path}{C.RESET}")
        sys.exit(1)

    checkpoint_path = Path(args.checkpoint)
    if not checkpoint_path.exists():
        print(f"{C.RED}Error: checkpoint not found: {checkpoint_path}{C.RESET}")
        sys.exit(1)

    device = torch.device("cpu" if args.cpu or not torch.cuda.is_available() else "cuda")

    ts      = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    out_dir = (Path(args.output_dir) if args.output_dir
               else XAI_OUTPUT_DIR / f"{image_path.stem}_{ts}")

    print_header("ManipShield XAI Inference")
    print(f"  Image      : {image_path}")
    print(f"  Checkpoint : {checkpoint_path.name}")
    print(f"  Device     : {device}")
    print(f"  Output dir : {out_dir}")

    img_pil = Image.open(image_path).convert("RGB")
    model   = load_model(checkpoint_path, device)

    run_xai(args, model, img_pil, image_path, device, out_dir)

    print_header("XAI Output Files")
    for f in sorted(out_dir.glob("*")):
        print(f"  {f.name}")
    print(f"\n  Full path: {out_dir.resolve()}")
    print()


if __name__ == "__main__":
    main()