app.py

"""
DeepGuard — Explainable Forensic Deepfake & AI Image Detection System
Hosted on Hugging Face Spaces (ZeroGPU Accelerated).
"""

import os
import io
import base64
from typing import Optional, List, Dict, Any, Tuple
from PIL import Image
import numpy as np

# ZeroGPU Setup: import spaces before torch
try:
    import spaces
    has_spaces = True
except ImportError:
    has_spaces = False
    class spaces:
        @staticmethod
        def GPU(duration=60):
            def decorator(fn):
                return fn
            return decorator

import torch
import gradio as gr
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

# Internal Forensic Modules
from stage1_ensemble import create_ensemble, WEIGHT_PRESETS
from stage2_explainability import run_stage2_explainability
from metadata_extractor import extract_metadata
from llm_interpreter import generate_forensic_report, chat_with_forensic_context

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"[DeepGuard] Running on device: {DEVICE} (ZeroGPU: {has_spaces})")


# ─────────────────────────────────────────────────────────────────────────────
# Core Prediction Pipeline Function (ZeroGPU Decorated)
# ─────────────────────────────────────────────────────────────────────────────
@spaces.GPU(duration=30)
def analyze_forensics(
    image: Image.Image,
    mode_label: str,
    preset_choice: str,
    w_cf: float,
    w_itsnotai: float,
    w_freq: float,
) -> Tuple[str, Image.Image, List[List[Any]], str, str, Dict[str, Any]]:
    """
    Executes the 4-stage explainable forensic pipeline.
    """
    if image is None:
        return (
            "### ⚠️ No Image Provided\nPlease upload an image to begin forensic inspection.",
            None,
            [],
            "No metadata available.",
            "No narrative generated.",
            {},
        )

    # 1. Normalize image
    image = image.convert("RGB")
    buf = io.BytesIO()
    image.save(buf, format="JPEG", quality=95)
    image_bytes = buf.getvalue()

    # 2. Map mode
    mode = "ai_image" if "Synthetic" in mode_label else "deepfake"

    # 3. Handle Weights & Presets
    preset_map = {
        "Auto-Dynamic (Recommended)": "auto",
        "Modern AI & Art (SDXL, Midjourney, FLUX)": "modern_art",
        "Balanced (50 / 50)": "balanced",
        "Photographic Sensor": "photographic",
        "Custom Sliders": "custom",
    }
    selected_preset = preset_map.get(preset_choice, "auto")

    custom_weights = None
    use_dynamic = True

    if selected_preset == "custom":
        use_dynamic = False
        if mode == "ai_image":
            custom_weights = {
                "Community-Forensics (ViT-Small)": float(w_cf),
                "ItsNotAI v2 (BEiT-Large)": float(w_itsnotai),
                "Frequency Analysis (FFT + LR)": float(w_freq),
            }
        else:
            custom_weights = {
                "Face-Swap CLIP (ViT-L/14)": float(w_cf),
                "Face-Swap Backup (ViT-Base)": float(w_itsnotai),
                "Frequency Analysis (FFT + LR)": float(w_freq),
            }
    elif selected_preset != "auto":
        use_dynamic = False

    # Stage 2.5: Metadata extraction
    metadata = extract_metadata(image_bytes)

    # Stage 1: Ensemble Classification
    ensemble = create_ensemble(mode, DEVICE)
    ensemble_res = ensemble.predict(
        image,
        custom_weights=custom_weights,
        weight_preset=selected_preset if selected_preset != "custom" else "auto",
        dynamic_optimization=use_dynamic,
    )

    # Stage 2: Grad-CAM++ Explainability & Artifacts
    explainability = run_stage2_explainability(image, ensemble_res, ensemble.detectors, mode=mode)

    # Stage 3: Constrained VLM Narrative
    use_llm = bool(os.environ.get("OPENROUTER_API_KEY"))
    heatmap_bytes = explainability.get("heatmap_bytes") or image_bytes
    artifacts = explainability.get("artifacts", [])

    import asyncio
    try:
        forensic = asyncio.run(
            generate_forensic_report(
                original_bytes=image_bytes,
                heatmap_bytes=heatmap_bytes,
                ensemble_result=ensemble_res,
                artifacts=artifacts,
                metadata=metadata,
                mode=mode,
                use_llm=use_llm,
            )
        )
    except Exception as e:
        print(f"[Stage 3] Exception running LLM report: {e}, using template fallback")
        from llm_interpreter import generate_template_narrative
        findings = {
            "ensemble": ensemble_res.__dict__ if hasattr(ensemble_res, '__dict__') else ensemble_res,
            "artifacts": artifacts,
            "metadata": metadata,
            "uncertainty": {
                "ensemble_disagreement": getattr(ensemble_res, 'ensemble_disagreement', 0),
                "oos_flag": getattr(ensemble_res, 'oos_flag', False),
            }
        }
        forensic = generate_template_narrative(findings)

    # ── Format UI Outputs ────────────────────────────────────────────────────
    label = ensemble_res.ensemble_label
    conf = ensemble_res.ensemble_confidence * 100
    disagreement = ensemble_res.ensemble_disagreement * 100

    if label == "FAKE":
        badge_color = "#ef4444"
        badge_text = "AI-GENERATED / DEEPFAKE"
    elif label == "REAL":
        badge_color = "#10b981"
        badge_text = "AUTHENTIC / CAMERA-CAPTURED"
    else:
        badge_color = "#f59e0b"
        badge_text = "UNCERTAIN / INCONCLUSIVE"

    verdict_html = f"""
    <div style="background: rgba(15, 23, 42, 0.7); border: 1px solid rgba(255, 255, 255, 0.1); border-radius: 12px; padding: 20px; margin-bottom: 15px;">
        <div style="display: flex; align-items: center; justify-content: space-between; flex-wrap: wrap; gap: 10px;">
            <div>
                <span style="background-color: {badge_color}; color: white; padding: 6px 14px; border-radius: 20px; font-weight: 700; font-size: 14px; letter-spacing: 0.5px;">
                    {badge_text}
                </span>
                <span style="color: #94a3b8; font-size: 13px; margin-left: 10px;">
                    Mode: <b>{mode_label}</b>
                </span>
            </div>
            <div style="font-size: 24px; font-weight: 800; color: #f8fafc;">
                {conf:.1f}% <span style="font-size: 14px; font-weight: 400; color: #94a3b8;">confidence</span>
            </div>
        </div>
        <div style="margin-top: 14px; display: grid; grid-template-columns: 1fr 1fr; gap: 10px; font-size: 13px;">
            <div style="background: rgba(30, 41, 59, 0.6); padding: 8px 12px; border-radius: 6px; border: 1px solid rgba(255, 255, 255, 0.05);">
                <span style="color: #94a3b8;">Ensemble Disagreement:</span>
                <b style="color: #f1f5f9; float: right;">{disagreement:.1f}%</b>
            </div>
            <div style="background: rgba(30, 41, 59, 0.6); padding: 8px 12px; border-radius: 6px; border: 1px solid rgba(255, 255, 255, 0.05);">
                <span style="color: #94a3b8;">Primary Detector:</span>
                <b style="color: #f1f5f9; float: right;">{ensemble_res.primary_model.split('(')[0].strip()}</b>
            </div>
        </div>
    """

    if getattr(ensemble_res, "optimization_note", None):
        verdict_html += f"""
        <div style="margin-top: 12px; background: rgba(59, 130, 246, 0.15); border: 1px solid rgba(59, 130, 246, 0.3); border-radius: 6px; padding: 10px 14px; color: #60a5fa; font-size: 13px;">
            {ensemble_res.optimization_note}
        </div>
        """
    verdict_html += "</div>"

    # Overlay Image (Original + Heatmap)
    heatmap_bytes = explainability.get("heatmap_bytes")
    if heatmap_bytes:
        try:
            overlay_img = Image.open(io.BytesIO(heatmap_bytes))
        except Exception:
            overlay_img = image
    else:
        overlay_img = image

    # Consensus Table
    consensus_data = []
    for vote in ensemble_res.votes:
        vote_label = vote.get("label", "N/A")
        raw_score = vote.get("raw_score", 0.5)
        weight_val = vote.get("weight", 0.0)
        consensus_data.append([
            vote.get("model_name", "Detector"),
            vote_label,
            f"{raw_score:.3f}",
            f"{weight_val * 100:.1f}%",
        ])

    # Metadata Dossier Markdown
    exif_found = metadata.get("has_exif", False)
    meta_status = "Intact EXIF Metadata Found" if exif_found else "No EXIF Metadata (Stripped/Generated)"
    dossier_md = f"### File Provenance Dossier\n\n**Status:** `{meta_status}`\n\n"
    dossier_md += f"- **Dimensions:** {metadata.get('width', 'Unknown')} × {metadata.get('height', 'Unknown')} px\n"
    dossier_md += f"- **Format:** {metadata.get('format', 'JPEG')}\n"
    dossier_md += f"- **Camera Make/Model:** {metadata.get('camera_make', 'N/A')} {metadata.get('camera_model', '')}\n"
    dossier_md += f"- **Software:** {metadata.get('software', 'None')}\n"

    if metadata.get("flags"):
        dossier_md += "\n**Forensic Flags:**\n"
        for flag in metadata["flags"]:
            dossier_md += f"- ⚠️ {flag}\n"

    # Stage 3 Narrative
    narrative_md = forensic.get("forensic_narrative", "Narrative unavailable.")

    # Context dict for Investigator Chatbot
    chat_context = {
        "label": label,
        "fake_probability": ensemble_res.ensemble_confidence,
        "disagreement": ensemble_res.ensemble_disagreement,
        "key_artifacts": forensic.get("key_artifacts", []),
        "recommendation": forensic.get("recommendation", ""),
        "metadata": metadata,
        "ensemble_votes": ensemble_res.votes,
    }

    return verdict_html, overlay_img, consensus_data, dossier_md, narrative_md, chat_context


# ─────────────────────────────────────────────────────────────────────────────
# Interactive Forensic Investigator Chat (Multi-turn)
# ─────────────────────────────────────────────────────────────────────────────
async def chat_investigator_fn(message: str, history: Optional[List[Any]], context: Dict[str, Any]):
    if not message or not message.strip():
        return history or [], ""

    messages = []
    if history:
        for msg in history:
            if isinstance(msg, dict):
                messages.append({"role": msg.get("role", "user"), "content": msg.get("content", "")})
            elif isinstance(msg, (list, tuple)) and len(msg) >= 2:
                messages.append({"role": "user", "content": str(msg[0])})
                messages.append({"role": "assistant", "content": str(msg[1])})
    messages.append({"role": "user", "content": message})

    result = await chat_with_forensic_context(
        messages=messages,
        forensic_context=context or {},
    )
    reply = result.get("reply", "No response generated.")

    new_history = list(history) if history else []
    new_history.append({"role": "user", "content": message})
    new_history.append({"role": "assistant", "content": reply})

    return new_history, ""


# ─────────────────────────────────────────────────────────────────────────────
# Gradio Blocks UI Definition
# ─────────────────────────────────────────────────────────────────────────────
custom_theme = gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="slate",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui"],
)

with gr.Blocks(title="DeepGuard — Forensic Detection") as demo:
    forensic_context_state = gr.State({})

    gr.HTML("""
    <div style="text-align: center; margin-bottom: 25px; padding-top: 10px;">
        <h1 style="font-size: 2.2rem; font-weight: 800; color: #3b82f6; margin-bottom: 8px;">
            🛡️ DeepGuard Visual Forensics
        </h1>
        <p style="color: #94a3b8; font-size: 1.05rem; max-width: 720px; margin: 0 auto;">
            Court-ready visual forensics with weighted ensemble verdicts, Grad-CAM++ attention maps, file provenance dossier, and AI investigator chat.
        </p>
    </div>
    """)

    with gr.Row():
        # Left Column: Inputs & Calibration
        with gr.Column(scale=5):
            input_image = gr.Image(type="pil", label="Suspect Image", height=320)

            with gr.Row():
                mode_selector = gr.Radio(
                    choices=["Synthetic AI Media", "Face-Swap & Deepfake"],
                    value="Synthetic AI Media",
                    label="Forensic Mode",
                )

            with gr.Row():
                preset_selector = gr.Dropdown(
                    choices=[
                        "Auto-Dynamic (Recommended)",
                        "Modern AI & Art (SDXL, Midjourney, FLUX)",
                        "Balanced (50 / 50)",
                        "Photographic Sensor",
                        "Custom Sliders",
                    ],
                    value="Auto-Dynamic (Recommended)",
                    label="Weight Preset Calibration",
                )

            with gr.Accordion("⚙️ Custom Detector Weights", open=False):
                gr.Markdown("Adjust weights when **Custom Sliders** preset is selected:")
                slider_cf = gr.Slider(minimum=0.0, maximum=1.0, value=0.48, step=0.05, label="Community-Forensics / Face-Swap CLIP")
                slider_itsnotai = gr.Slider(minimum=0.0, maximum=1.0, value=0.47, step=0.05, label="ItsNotAI v2 / Face-Swap Backup")
                slider_freq = gr.Slider(minimum=0.0, maximum=1.0, value=0.05, step=0.05, label="Frequency Analysis (FFT + LR)")

            analyze_btn = gr.Button("🔍 Run Forensic Inspection", variant="primary", size="lg")

            gr.Examples(
                examples=[
                    ["examples/test_ai.jpg", "Synthetic AI Media", "Auto-Dynamic (Recommended)"],
                    ["examples/test_real.jpg", "Synthetic AI Media", "Auto-Dynamic (Recommended)"],
                ],
                inputs=[input_image, mode_selector, preset_selector],
                label="Sample Images",
            )

        # Right Column: Analysis Findings
        with gr.Column(scale=7):
            verdict_box = gr.HTML(label="Verdict")

            with gr.Row():
                gradcam_display = gr.Image(label="Grad-CAM++ Attention Overlay", height=340)

            with gr.Tabs():
                with gr.TabItem("📊 Consensus Matrix"):
                    consensus_table = gr.Dataframe(
                        headers=["Detector", "Vote", "P(Fake)", "Assigned Weight"],
                        datatype=["str", "str", "str", "str"],
                        label="Ensemble Consensus Breakdown",
                    )
                with gr.TabItem("📑 File Provenance"):
                    dossier_display = gr.Markdown()
                with gr.TabItem("📝 Forensic Narrative"):
                    narrative_display = gr.Markdown()

    # Interactive Forensic Investigator Chatbot Section
    gr.Markdown("---")
    gr.Markdown("### 🕵️ Interactive Forensic Investigator Chat")
    gr.Markdown("Inquire deeper into model decision boundaries, specific attention zones, or contradictory detector signals.")

    chatbot = gr.Chatbot(height=320, label="DeepGuard Forensic Investigator")
    with gr.Row():
        chat_input = gr.Textbox(
            placeholder="Ask about this image's verdict, Grad-CAM peaks, or provenance flags...",
            scale=8,
            show_label=False,
        )
        chat_submit = gr.Button("Send", variant="primary", scale=1)

    # ── Wire Handlers ────────────────────────────────────────────────────────
    analyze_btn.click(
        fn=analyze_forensics,
        inputs=[input_image, mode_selector, preset_selector, slider_cf, slider_itsnotai, slider_freq],
        outputs=[verdict_box, gradcam_display, consensus_table, dossier_display, narrative_display, forensic_context_state],
    )

    chat_submit.click(
        fn=chat_investigator_fn,
        inputs=[chat_input, chatbot, forensic_context_state],
        outputs=[chatbot, chat_input],
    )
    chat_input.submit(
        fn=chat_investigator_fn,
        inputs=[chat_input, chatbot, forensic_context_state],
        outputs=[chatbot, chat_input],
    )


# ─────────────────────────────────────────────────────────────────────────────
# Launch ZeroGPU Space
# ─────────────────────────────────────────────────────────────────────────────
demo.queue(default_concurrency_limit=5).launch()