app.py

"""
Streamlit demo for the deepfake-detection project.

Loads all five trained models (four baselines + the proposed hybrid) from the
Hugging Face Hub, runs each one on an uploaded video, and shows a side-by-side
comparison. Grad-CAM is shown for the hybrid since that's the proposed model.

Run locally:
    streamlit run app.py
"""

import os
import tempfile
from collections import OrderedDict
from pathlib import Path

import numpy as np
import pandas as pd
import streamlit as st
import torch
from huggingface_hub import hf_hub_download

from gradcam_helper import (
    denormalize_frame_for_display,
    generate_gradcam_overlay,
)
from model import (
    CnnLstmClipModel,
    EfficientNetFrameAverager,
    HybridCNNTransformer,
    ViTFrameAverager,
    XceptionFrameAverager,
)
from preprocessing import (
    FRAMES_PER_CLIP,
    IMAGENET_MEAN,
    IMAGENET_STD,
    build_yunet_face_detector,
    process_video_to_clip_tensor,
)


# -------- Constants --------

HUGGINGFACE_REPO_ID = "MUmairAB/deepfake-detection-ff-cn-transformer"

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# The middle frame of the 24-frame clip is the one we show Grad-CAM for.
GRADCAM_FRAME_INDEX = 12

# Model registry. Order here is the order they show up in the table.
# We keep the proposed hybrid last so it reads as the "headline" model.
MODEL_REGISTRY = OrderedDict([
    ("Xception", {
        "constructor": lambda: XceptionFrameAverager(pretrained=False),
        "checkpoint": "xception_baseline/best_model.pt",
        "params_millions": 20.8,
        "is_proposed": False,
    }),
    ("EfficientNet-B3", {
        "constructor": lambda: EfficientNetFrameAverager(pretrained=False),
        "checkpoint": "efficientnet_b3_baseline/best_model.pt",
        "params_millions": 10.7,
        "is_proposed": False,
    }),
    ("CNN + BiLSTM", {
        "constructor": lambda: CnnLstmClipModel(pretrained=False),
        "checkpoint": "cnn_lstm_baseline/best_model.pt",
        "params_millions": 13.7,
        "is_proposed": False,
    }),
    ("ViT-Base/16", {
        "constructor": lambda: ViTFrameAverager(pretrained=False),
        "checkpoint": "vit_base_baseline/best_model.pt",
        "params_millions": 85.8,
        "is_proposed": False,
    }),
    ("Hybrid CNN-Transformer (ours)", {
        "constructor": lambda: HybridCNNTransformer(pretrained=False),
        "checkpoint": "hybrid_cnn_transformer/best_model.pt",
        "params_millions": 74.9,
        "is_proposed": True,
    }),
])

PROPOSED_MODEL_NAME = "Hybrid CNN-Transformer (ours)"


# -------- Page config --------

st.set_page_config(
    page_title="Deepfake Detection - 5 model comparison",
    page_icon="🎭",
    layout="wide",
)


# -------- Cached resource loaders --------

@st.cache_resource(show_spinner=False)
def load_face_detector():
    """Build the YuNet face detector once. The ONNX weights are pulled from
    OpenCV's Hugging Face repo on first run and cached on disk."""
    return build_yunet_face_detector()


@st.cache_resource(show_spinner=False)
def load_all_models():
    """Download all five checkpoints from HF, build each model, load weights.

    Returns an OrderedDict mapping model name -> nn.Module (already in eval
    mode on the right device). Total memory footprint is ~800 MB which fits
    comfortably in the free Space's 16 GB RAM.
    """
    loaded_models = OrderedDict()

    for model_name, model_config in MODEL_REGISTRY.items():
        checkpoint_local_path = hf_hub_download(
            repo_id=HUGGINGFACE_REPO_ID,
            filename=model_config["checkpoint"],
        )
        # weights_only=False is fine here since these are checkpoints we made
        # ourselves and stored on our own HF repo.
        checkpoint = torch.load(checkpoint_local_path, map_location=DEVICE, weights_only=False)

        model_instance = model_config["constructor"]()
        model_instance.load_state_dict(checkpoint["model_state_dict"])
        model_instance.to(DEVICE).eval()
        loaded_models[model_name] = model_instance

    return loaded_models


# -------- Inference --------

@torch.no_grad()
def predict_with_all_models(loaded_models, clip_normalized, progress_callback=None):
    """Run every model in `loaded_models` on the clip. Returns a dict
    mapping model name -> fake probability in [0, 1].

    progress_callback(model_name, fraction) is called after each model finishes.
    """
    clip_batched = clip_normalized.unsqueeze(0).to(DEVICE)         # [1, T, 3, H, W]
    fake_probability_by_model = OrderedDict()

    total_models = len(loaded_models)
    for model_index, (model_name, model_instance) in enumerate(loaded_models.items()):
        logit = model_instance(clip_batched).squeeze()
        probability = torch.sigmoid(logit.float()).item()
        fake_probability_by_model[model_name] = probability

        if progress_callback is not None:
            progress_callback(model_name, (model_index + 1) / total_models)

    return fake_probability_by_model


def build_predictions_dataframe(fake_probability_by_model):
    """Turn the dict of per-model probabilities into a tidy DataFrame for display."""
    rows = []
    for model_name, fake_probability in fake_probability_by_model.items():
        config = MODEL_REGISTRY[model_name]
        rows.append({
            "Model": model_name,
            "Params (M)": config["params_millions"],
            "Probability (fake)": fake_probability,
            "Verdict": "FAKE" if fake_probability >= 0.5 else "REAL",
        })
    return pd.DataFrame(rows)


# -------- Sidebar --------

with st.sidebar:
    st.markdown("### About this demo")
    st.markdown(
        "Live comparison of all **five models** from our Deep Machine Learning "
        "course project on video deepfake detection. Four baselines "
        "(Xception, EfficientNet-B3, CNN+BiLSTM, ViT-Base/16) plus our "
        "**proposed Hybrid CNN-Transformer**."
    )
    st.markdown("---")
    st.markdown("### Project links")
    st.markdown(
        f"- [Model weights]({'https://huggingface.co/' + HUGGINGFACE_REPO_ID})\n"
        "- Trained on FaceForensics++ (C23)\n"
        "- Five classes: original, Deepfakes, Face2Face, FaceSwap, NeuralTextures\n"
        "- FaceShifter held out for cross-manipulation evaluation"
    )
    st.markdown("---")
    st.markdown("### Pipeline")
    st.markdown(
        "1. The uploaded video is decoded and 24 frames are sampled evenly.\n"
        "2. YuNet locates the largest face in each frame.\n"
        "3. Each face is cropped (with 30% margin) and resized to 224x224.\n"
        "4. Every model is run on the 24-frame clip and emits one "
        "fake / real probability for the whole clip.\n"
        "5. Grad-CAM is computed on the middle frame of the proposed hybrid model."
    )
    st.markdown("---")
    st.caption(f"Running on **{DEVICE.upper()}**.")


# -------- Main page --------

st.title("Spatiotemporal Deepfake Detection - 5 model comparison")
st.markdown(
    "Upload a short face video. Every one of the five trained models runs on "
    "it and you can compare how they agree (or disagree) on whether the clip "
    "is real or manipulated."
)

# Warm up the cached loaders. First run downloads ~800 MB of checkpoints.
with st.spinner("Loading face detector and all 5 models from Hugging Face (one-time, ~30s)..."):
    face_detector = load_face_detector()
    loaded_models = load_all_models()

st.success(f"All {len(loaded_models)} models loaded. Upload a video to compare them.")

uploaded_video = st.file_uploader(
    "Choose a short face video (MP4 / MOV, ideally under 10 seconds)",
    type=["mp4", "mov", "avi", "mkv"],
)

if uploaded_video is not None:
    # Write the upload to a temp file because cv2.VideoCapture needs a real path.
    with tempfile.NamedTemporaryFile(suffix=Path(uploaded_video.name).suffix, delete=False) as temp_video_file:
        temp_video_file.write(uploaded_video.read())
        temp_video_path = temp_video_file.name

    # Show the uploaded video and pre-processing progress side by side
    left_column, right_column = st.columns([1, 2])
    with left_column:
        st.markdown("**Uploaded clip**")
        st.video(temp_video_path)
    with right_column:
        st.markdown("**Pre-processing**")
        preprocessing_progress = st.progress(0.0, text="Starting...")

        def update_preprocessing_progress(message, fraction):
            preprocessing_progress.progress(min(fraction, 1.0), text=message)

        preprocessing_result = process_video_to_clip_tensor(
            video_path=temp_video_path,
            face_detector=face_detector,
            progress_callback=update_preprocessing_progress,
        )

    try:
        os.remove(temp_video_path)
    except OSError:
        pass

    if preprocessing_result.get("error"):
        st.error(f"Pre-processing failed: {preprocessing_result['error']}")
        st.stop()

    clip_tensor = preprocessing_result["clip_tensor"]
    clip_normalized = preprocessing_result["clip_normalized"]
    face_crops_rgb = preprocessing_result["face_crops_rgb"]

    # --- Run all 5 models with a single progress bar ---
    st.markdown("---")
    st.markdown("## Running all five models")
    inference_progress = st.progress(0.0, text="Starting inference...")

    def update_inference_progress(model_name, fraction):
        inference_progress.progress(fraction, text=f"Done: {model_name}")

    fake_probability_by_model = predict_with_all_models(
        loaded_models=loaded_models,
        clip_normalized=clip_normalized,
        progress_callback=update_inference_progress,
    )
    inference_progress.empty()

    # --- Consensus / headline verdict ---
    predictions_dataframe = build_predictions_dataframe(fake_probability_by_model)
    number_of_fake_votes = int((predictions_dataframe["Probability (fake)"] >= 0.5).sum())
    number_of_real_votes = len(predictions_dataframe) - number_of_fake_votes
    average_fake_probability = float(predictions_dataframe["Probability (fake)"].mean())
    proposed_model_probability = fake_probability_by_model[PROPOSED_MODEL_NAME]

    st.markdown("## Result")

    headline_column, hybrid_column, consensus_column = st.columns([1, 1, 1])

    with headline_column:
        st.markdown("**Majority vote**")
        if number_of_fake_votes > number_of_real_votes:
            st.error(f"### FAKE ({number_of_fake_votes} / 5)")
        elif number_of_real_votes > number_of_fake_votes:
            st.success(f"### REAL ({number_of_real_votes} / 5)")
        else:
            st.warning("### TIED (split decision)")
        st.caption(f"Average fake probability across models: {average_fake_probability:.2%}")

    with hybrid_column:
        st.markdown("**Proposed hybrid model**")
        if proposed_model_probability >= 0.5:
            st.error(f"### FAKE")
        else:
            st.success(f"### REAL")
        st.caption(f"Hybrid says fake with probability {proposed_model_probability:.2%}.")

    with consensus_column:
        st.markdown("**Agreement**")
        if number_of_fake_votes == 5 or number_of_real_votes == 5:
            st.info("### Unanimous")
            st.caption("All five models agree on the verdict.")
        else:
            st.info(f"### {number_of_fake_votes} : {number_of_real_votes}")
            st.caption(f"{number_of_fake_votes} say FAKE, {number_of_real_votes} say REAL.")

    # --- Per-model breakdown table ---
    st.markdown("### Per-model predictions on this clip")
    st.caption(
        "Verdict uses a 0.5 threshold on the sigmoid output. The probability "
        "column shows how confident each model is that the clip is a deepfake."
    )

    st.dataframe(
        predictions_dataframe,
        column_config={
            "Model": st.column_config.TextColumn("Model", width="large"),
            "Params (M)": st.column_config.NumberColumn("Params (M)", format="%.1f"),
            "Probability (fake)": st.column_config.ProgressColumn(
                "Probability of fake",
                help="Sigmoid output of the model on the full 24-frame clip",
                min_value=0.0,
                max_value=1.0,
                format="%.3f",
            ),
            "Verdict": st.column_config.TextColumn("Verdict"),
        },
        hide_index=True,
        use_container_width=True,
    )

    # --- Face crops preview ---
    st.markdown("### Extracted face crops")
    st.caption(
        f"The {FRAMES_PER_CLIP} face crops every model actually sees. "
        "8 evenly-spaced ones shown here to keep the layout compact."
    )
    indices_to_show = np.linspace(0, FRAMES_PER_CLIP - 1, 8, dtype=int)
    preview_columns = st.columns(8)
    for column_index, frame_index in enumerate(indices_to_show):
        with preview_columns[column_index]:
            st.image(
                face_crops_rgb[frame_index],
                caption=f"frame {frame_index}",
                use_container_width=True,
            )

    # --- Grad-CAM (proposed hybrid only) ---
    st.markdown("### Grad-CAM heatmap (proposed hybrid, middle frame)")
    st.caption(
        "Where the EfficientNet-B3 backbone of the proposed hybrid is "
        "looking when scoring this clip. Warmer colors mean stronger contribution."
    )
    with st.spinner("Generating Grad-CAM..."):
        hybrid_model = loaded_models[PROPOSED_MODEL_NAME]
        middle_frame_normalized = clip_normalized[GRADCAM_FRAME_INDEX]
        middle_frame_rgb = denormalize_frame_for_display(
            middle_frame_normalized, IMAGENET_MEAN, IMAGENET_STD
        )
        gradcam_overlay = generate_gradcam_overlay(
            hybrid_model=hybrid_model,
            frame_normalized=middle_frame_normalized,
            frame_rgb_zero_to_one=middle_frame_rgb,
            device=DEVICE,
        )

    original_column, heatmap_column = st.columns(2)
    with original_column:
        st.image(middle_frame_rgb, caption="Original middle frame", use_container_width=True)
    with heatmap_column:
        st.image(gradcam_overlay, caption="Grad-CAM overlay", use_container_width=True)


# -------- Static results from the report --------

st.markdown("---")
st.markdown("## Results from the report")
st.caption(
    "Test-set results we reported in the project write-up. Every model used "
    "the exact same training setup (AdamW, lr=1e-4, cosine schedule over 25 "
    "epochs, BCEWithLogitsLoss, balanced sampling) so the differences below "
    "reflect the architecture, not the tuning."
)

in_domain_results = pd.DataFrame(
    [
        {"Model": "Xception",                       "Params (M)": 20.8, "AUC": 0.9944, "Accuracy": 0.9744, "Precision": 0.9823, "Recall": 0.9858, "F1": 0.9840},
        {"Model": "EfficientNet-B3",                "Params (M)": 10.7, "AUC": 0.9976, "Accuracy": 0.9829, "Precision": 0.9825, "Recall": 0.9964, "F1": 0.9894},
        {"Model": "CNN + BiLSTM",                   "Params (M)": 13.7, "AUC": 0.9805, "Accuracy": 0.9744, "Precision": 0.9823, "Recall": 0.9858, "F1": 0.9840},
        {"Model": "ViT-Base/16",                    "Params (M)": 85.8, "AUC": 0.7066, "Accuracy": 0.6695, "Precision": 0.8910, "Recall": 0.6690, "F1": 0.7642},
        {"Model": "Hybrid CNN-Transformer (ours)",  "Params (M)": 74.9, "AUC": 0.9458, "Accuracy": 0.8462, "Precision": 0.9710, "Recall": 0.8327, "F1": 0.8966},
    ]
)

cross_manipulation_results = pd.DataFrame(
    [
        {"Model": "Xception",                       "In-domain AUC": 0.9944, "FaceShifter AUC": 0.7015, "AUC drop": 0.2929, "FaceShifter Acc": 0.3000, "FaceShifter F1": 0.3160},
        {"Model": "EfficientNet-B3",                "In-domain AUC": 0.9976, "FaceShifter AUC": 0.7343, "AUC drop": 0.2633, "FaceShifter Acc": 0.4064, "FaceShifter F1": 0.4746},
        {"Model": "CNN + BiLSTM",                   "In-domain AUC": 0.9805, "FaceShifter AUC": 0.7591, "AUC drop": 0.2214, "FaceShifter Acc": 0.4191, "FaceShifter F1": 0.4916},
        {"Model": "ViT-Base/16",                    "In-domain AUC": 0.7066, "FaceShifter AUC": 0.6880, "AUC drop": 0.0186, "FaceShifter Acc": 0.6170, "FaceShifter F1": 0.7297},
        {"Model": "Hybrid CNN-Transformer (ours)",  "In-domain AUC": 0.9458, "FaceShifter AUC": 0.6273, "AUC drop": 0.3185, "FaceShifter Acc": 0.3404, "FaceShifter F1": 0.3825},
    ]
)

table_tab, cross_tab = st.tabs(["In-domain test set", "Cross-manipulation (FaceShifter)"])

with table_tab:
    st.markdown(
        "Test split of FaceForensics++ C23, drawn from the same five "
        "manipulation classes the models were trained on."
    )
    st.dataframe(in_domain_results, use_container_width=True, hide_index=True)

with cross_tab:
    st.markdown(
        "FaceShifter was **never seen during training**. This measures how well "
        "each model generalizes to a new manipulation method. Temporal models "
        "(CNN + BiLSTM, ViT) hold up better here than pure spatial CNNs do."
    )
    st.dataframe(cross_manipulation_results, use_container_width=True, hide_index=True)