Update detector.py

detector.py

@@ -1,18 +1,11 @@
 """
 MintVid Deepfake Detector — Hugging Face ZeroGPU Space
 V-JEPA 2.1 ViT-Giant + AttentiveClassifier probe
-
-ZEROGPU RULES (strictly followed here):
-  - No torch.cuda.is_available() in main process
-  - No map_location="cuda" in main process
-  - No .to("cuda") that actually executes in main process
-    (ZeroGPU *intercepts* .to("cuda") calls and postpones them —
-     models stay on CPU until @spaces.GPU runs)
-  - Only @spaces.GPU functions touch real GPU operations
 """
 
 import os
 import glob
+import math
 import subprocess
 import tempfile
 import warnings
@@ -24,19 +17,23 @@ import torch
 import spaces
 import gradio as gr
 from huggingface_hub import hf_hub_download
+from decord import VideoReader, cpu as decord_cpu
 
 from inference_core import AttentiveClassifier, get_video_windows
 
 # ─────────────────────────────────────────────
 # CONFIG
 # ─────────────────────────────────────────────
-DEVICE          = "cuda"
-FEATURE_DIM     = 1408
-ENCODER_BATCH   = 3
+DEVICE            = "cuda"
+FEATURE_DIM       = 1408
+ENCODER_BATCH     = 3
+FRAMES_PER_CLIP   = 64
+CLIPS_PER_WINDOW  = 3
+FRAMES_PER_WINDOW = FRAMES_PER_CLIP * CLIPS_PER_WINDOW   # 192
 
-HF_TOKEN      = os.environ.get("HF_TOKEN")
-ENCODER_REPO  = os.environ.get("ENCODER_REPO", "")
-PROBE_REPO    = os.environ.get("PROBE_REPO", ENCODER_REPO)
+HF_TOKEN     = os.environ.get("HF_TOKEN")
+ENCODER_REPO = os.environ.get("ENCODER_REPO", "")
+PROBE_REPO   = os.environ.get("PROBE_REPO", ENCODER_REPO)
 
 # ─────────────────────────────────────────────
 # MODEL LOADING  (CPU — ZeroGPU intercepts CUDA calls)
@@ -88,17 +85,30 @@ print("=" * 50)
 
 
 # ─────────────────────────────────────────────
-# VIDEO DOWNLOAD  (pure CPU — no GPU quota used)
+# HELPERS
 # ─────────────────────────────────────────────
+
+def cleanup_file(path: str):
+    """Delete temp video file and any yt-dlp extension variants."""
+    if not path:
+        return
+    for f in glob.glob(path + "*") + [path]:
+        try:
+            if os.path.exists(f):
+                os.remove(f)
+        except Exception:
+            pass
+
+
 def download_video(url: str) -> tuple[bool, str]:
-    """Download via yt-dlp to a temp file. Returns (success, path_or_error)."""
+    """Download via yt-dlp. Returns (success, path_or_error)."""
     tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
     tmp.close()
 
     cmd = [
         "yt-dlp",
         "-f", "best[height<=720]/best",
-        "--max-filesize", "100M",
+        "--max-filesize", "200M",
         "--no-playlist",
         "--no-warnings",
         "--force-overwrites",
@@ -106,53 +116,128 @@ def download_video(url: str) -> tuple[bool, str]:
         url,
     ]
     try:
-        result = subprocess.run(
-            cmd, capture_output=True, text=True, timeout=90
-        )
-
-        # Always print yt-dlp output so it appears in the Space Logs tab
+        result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
         print(f"[yt-dlp] returncode : {result.returncode}")
         print(f"[yt-dlp] stdout     : {result.stdout[:800]}")
         print(f"[yt-dlp] stderr     : {result.stderr[:800]}")
 
         if result.returncode != 0:
-            return False, (
-                f"yt-dlp exited {result.returncode}\n"
-                f"stderr: {result.stderr[:600]}"
-            )
+            return False, f"yt-dlp exited {result.returncode}\nstderr: {result.stderr[:600]}"
 
-        # yt-dlp sometimes writes to a different extension than the template
-        # (e.g. tmp.mp4 → tmp.mp4.webm when remux isn't needed).
-        # Search for any file that starts with tmp.name and has content.
         candidates = sorted(glob.glob(tmp.name + "*") + [tmp.name])
-        actual_file = next(
+        actual = next(
             (f for f in candidates if os.path.exists(f) and os.path.getsize(f) > 0),
             None,
         )
-
-        if actual_file is None:
+        if actual is None:
             return False, (
-                f"yt-dlp wrote no output (empty or missing file).\n"
-                f"stdout: {result.stdout[:400]}\n"
-                f"stderr: {result.stderr[:400]}"
+                f"yt-dlp wrote no output.\n"
+                f"stdout: {result.stdout[:400]}\nstderr: {result.stderr[:400]}"
             )
 
-        print(f"[yt-dlp] actual file: {actual_file} ({os.path.getsize(actual_file)} bytes)")
-        return True, actual_file
+        print(f"[yt-dlp] actual file : {actual} ({os.path.getsize(actual)} bytes)")
+        return True, actual
 
     except subprocess.TimeoutExpired:
-        return False, "Download timed out (90s limit)."
+        return False, "Download timed out (120s limit)."
     except FileNotFoundError:
         return False, "yt-dlp not installed. Check requirements.txt."
 
 
+def get_video_info(path: str) -> dict:
+    """Peek at video metadata without decoding all frames."""
+    vr = VideoReader(path, ctx=decord_cpu(0), num_threads=1)
+    total_frames = len(vr)
+    fps          = vr.get_avg_fps()
+    duration     = total_frames / fps if fps > 0 else 0.0
+    return {"path": path, "duration": duration, "fps": fps, "total_frames": total_frames}
+
+
+def make_timeline_html(duration: float, start: float, end: float, fps: float) -> str:
+    """Render a video-trimmer style timeline bar showing the selected segment."""
+    if duration <= 0:
+        return ""
+
+    start = max(0.0, min(float(start), duration))
+    end   = max(start, min(float(end),   duration))
+
+    pct_start = start / duration * 100
+    pct_end   = end   / duration * 100
+    pct_width = pct_end - pct_start
+    mid_pct   = (pct_start + pct_end) / 2
+
+    frames_in_range = int((end - start) * fps)
+    num_windows     = max(1, math.ceil(frames_in_range / FRAMES_PER_WINDOW)) if frames_in_range > 0 else 0
+    selected_sec    = end - start
+    window_sec      = FRAMES_PER_WINDOW / fps
+
+    # One orange tick per window boundary inside the selection
+    ticks = ""
+    if num_windows > 1:
+        for i in range(1, num_windows):
+            tick_pct = (start + i * window_sec) / duration * 100
+            ticks += (
+                f'<div style="position:absolute;left:{tick_pct:.2f}%;top:0;bottom:0;'
+                f'width:1px;background:rgba(249,115,22,0.55);z-index:3;"></div>'
+            )
+
+    return f"""
+<div style="margin:4px 0 8px 0;font-family:system-ui,sans-serif;user-select:none;">
+
+  <div style="position:relative;height:48px;background:#111;border-radius:8px;
+            overflow:hidden;border:1px solid #2a2a2a;">
+
+    <div style="position:absolute;inset:0;pointer-events:none;
+                background:repeating-linear-gradient(90deg,transparent,transparent 19px,
+                rgba(255,255,255,0.025) 19px,rgba(255,255,255,0.025) 20px);"></div>
+
+    <div style="position:absolute;left:0;width:{pct_start:.2f}%;
+                top:0;bottom:0;background:rgba(0,0,0,0.62);z-index:1;"></div>
+
+    <div style="position:absolute;left:{pct_end:.2f}%;right:0;
+                top:0;bottom:0;background:rgba(0,0,0,0.62);z-index:1;"></div>
+
+    <div style="position:absolute;left:{pct_start:.2f}%;width:{pct_width:.2f}%;
+                top:0;bottom:0;border-left:2px solid #f97316;border-right:2px solid #f97316;
+                box-sizing:border-box;z-index:2;"></div>
+
+    {ticks}
+
+    <div style="position:absolute;left:6px;top:50%;transform:translateY(-50%);
+                font-size:10px;color:#444;z-index:4;">0 s</div>
+
+    <div style="position:absolute;right:6px;top:50%;transform:translateY(-50%);
+                font-size:10px;color:#444;z-index:4;">{duration:.1f} s</div>
+
+    <div style="position:absolute;left:{mid_pct:.2f}%;top:50%;
+                transform:translate(-50%,-50%);font-size:13px;font-weight:600;
+                color:#f97316;z-index:5;white-space:nowrap;
+                text-shadow:0 1px 6px rgba(0,0,0,0.95);">{selected_sec:.1f} s</div>
+  </div>
+
+  <div style="display:flex;justify-content:space-between;align-items:center;
+              margin-top:8px;font-size:12px;">
+    <span style="color:#888;">{start:.1f} s</span>
+    <span style="background:rgba(249,115,22,0.1);color:#f97316;
+                 padding:3px 14px;border-radius:20px;font-weight:500;
+                 border:1px solid rgba(249,115,22,0.22);">
+      {num_windows} window{'s' if num_windows != 1 else ''}
+      &nbsp;·&nbsp; ~{num_windows * window_sec:.1f} s
+      &nbsp;·&nbsp; {frames_in_range} frames
+    </span>
+    <span style="color:#888;">{end:.1f} s</span>
+  </div>
+
+</div>
+"""
+
+
 # ─────────────────────────────────────────────
 # GPU INFERENCE  (inside ZeroGPU allocation)
 # ─────────────────────────────────────────────
 @spaces.GPU(duration=120)
 def encode_and_classify(windows: list) -> list:
     scores = []
-
     for window in windows:
         clips = window["clips"]
         batch = torch.stack(clips).to(DEVICE)
@@ -162,100 +247,184 @@ def encode_and_classify(windows: list) -> list:
                 chunks = []
                 for i in range(0, batch.shape[0], ENCODER_BATCH):
                     chunks.append(encoder(batch[i : i + ENCODER_BATCH]))
-                all_feats = torch.cat(chunks, dim=0)
-
+                all_feats   = torch.cat(chunks, dim=0)
                 video_feats = all_feats.view(1, -1, FEATURE_DIM)
-
-                logit, _ = probe(video_feats)
-                prob = torch.sigmoid(logit).item()
+                logit, _    = probe(video_feats)
+                prob        = torch.sigmoid(logit).item()
 
         scores.append({
             "prob":      prob,
             "start_sec": window["start_sec"],
             "end_sec":   window["end_sec"],
         })
-
     return scores
 
 
 # ─────────────────────────────────────────────
-# GRADIO PREDICT
+# GRADIO CALLBACKS
 # ─────────────────────────────────────────────
-def predict(url: str) -> tuple[str, str, str]:
+
+def load_video_cb(url: str, state: dict):
+    """
+    Phase 1 — download the video and expose the segment selector.
+    Deletes any previously loaded temp file before downloading.
+    """
     url = (url or "").strip()
     if not url:
-        return "Please enter a URL.", "", ""
+        return state, "Please enter a URL.", gr.update(visible=False), gr.update(), gr.update(), ""
+
+    # Clean up the last loaded file (if any)
+    if state and state.get("path"):
+        cleanup_file(state["path"])
 
     ok, path_or_err = download_video(url)
     if not ok:
-        return f"Download failed: {path_or_err}", "", ""
+        return None, f"❌  {path_or_err}", gr.update(visible=False), gr.update(), gr.update(), ""
+
+    info     = get_video_info(path_or_err)
+    duration = info["duration"]
+    fps      = info["fps"]
+    step     = round(max(0.1, 1.0 / fps), 3)
+
+    timeline = make_timeline_html(duration, 0.0, duration, fps)
+
+    status = (
+        f"✅  {duration:.1f} s  ·  {fps:.0f} fps  ·  {info['total_frames']} frames  "
+        f"— drag the sliders to choose which segment to analyze"
+    )
+
+    return (
+        info,
+        status,
+        gr.update(visible=True),                                              # segment_panel
+        gr.update(minimum=0.0, maximum=duration, value=0.0,      step=step), # start_slider
+        gr.update(minimum=0.0, maximum=duration, value=duration,  step=step), # end_slider
+        timeline,
+    )
+
+
+def update_timeline_cb(state: dict, start: float, end: float) -> str:
+    """Redraw the timeline whenever a slider moves."""
+    if not state:
+        return ""
+    return make_timeline_html(state["duration"], start, end, state["fps"])
+
+
+def analyze_cb(state: dict, start: float, end: float):
+    """
+    Phase 2 — run inference on the selected [start, end] segment.
+    The video file is kept so the user can re-analyze a different segment
+    without re-downloading. It will be cleaned up on the next Load.
+    """
+    if not state or not state.get("path"):
+        return "No video loaded.", "", "", gr.update(visible=False)
+
+    start = max(0.0, float(start))
+    end   = min(float(end), state["duration"])
+
+    if end <= start:
+        return "Start must be before end.", "", "", gr.update(visible=False)
+
+    video_path = state["path"]
 
-    video_path = path_or_err
     try:
-        try:
-            windows, total_frames, fps, duration_sec = get_video_windows(video_path)
-        except Exception as e:
-            return f"Video decode failed: {e}", "", ""
+        windows, total_frames, fps, duration_sec = get_video_windows(
+            video_path, start_sec=start, end_sec=end
+        )
+    except Exception as e:
+        return f"Video decode failed: {e}", "", "", gr.update(visible=False)
 
-        scores = encode_and_classify(windows)
+    if not windows:
+        return "No frames in selected range.", "", "", gr.update(visible=False)
 
-    finally:
-        # Clean up all files that share the temp base name
-        for f in glob.glob(video_path + "*") + [video_path]:
-            if os.path.exists(f):
-                os.remove(f)
+    scores = encode_and_classify(windows)
 
     if not scores:
-        return "No windows processed.", "", ""
+        return "No windows processed.", "", "", gr.update(visible=False)
 
     final_prob = max(s["prob"] for s in scores)
     label      = "⚠️  AI-GENERATED (FAKE)" if final_prob > 0.5 else "✅  REAL"
     confidence = f"{final_prob * 100:.1f}%"
 
-    lines = [f"Duration: {duration_sec:.1f}s  |  Frames: {total_frames}  |  Windows: {len(scores)}\n"]
+    lines = [f"Segment: {start:.1f} s – {end:.1f} s  |  Frames: {total_frames}  |  Windows: {len(scores)}\n"]
     for s in scores:
         flag = "  ← peak" if (s["prob"] == final_prob and final_prob > 0.5) else ""
         lines.append(
-            f"  [{s['start_sec']:.1f}s – {s['end_sec']:.1f}s]  →  {s['prob']*100:.1f}% fake{flag}"
+            f"  [{s['start_sec']:.1f} s – {s['end_sec']:.1f} s]  →  {s['prob']*100:.1f}% fake{flag}"
         )
 
-    return label, confidence, "\n".join(lines)
+    return label, confidence, "\n".join(lines), gr.update(visible=True)
 
 
 # ─────────────────────────────────────────────
 # GRADIO UI
 # ─────────────────────────────────────────────
 with gr.Blocks(title="MintVid Deepfake Detector") as demo:
+
+    video_state = gr.State(None)
+
     gr.Markdown(
-        """
-        ## 🎬 MintVid — Deepfake Video Detector
-        Paste a public Instagram, TikTok, or direct `.mp4` URL.
-        Uses **V-JEPA 2.1 ViT-Giant** + **AttentiveClassifier** probe.
-        """
+        "## 🎬 MintVid — Deepfake Video Detector\n"
+        "Paste a URL, trim to the segment you want to test, then hit **Analyze**."
     )
 
-    url_box = gr.Textbox(
-        label="Video URL",
-        placeholder="https://www.instagram.com/reels/...",
-        autofocus=True,
+    # ── Phase 1: URL input ──────────────────────────────────────────
+    with gr.Row():
+        url_box  = gr.Textbox(
+            label="Video URL",
+            placeholder="https://www.instagram.com/reels/...   or   https://www.tiktok.com/@user/video/...",
+            scale=5,
+        )
+        load_btn = gr.Button("Load", variant="primary", scale=1, min_width=80)
+
+    status_box = gr.Textbox(
+        label="", interactive=False, show_label=False, container=False,
+        placeholder="Status will appear here after loading."
     )
-    run_btn = gr.Button("Analyze", variant="primary")
 
-    with gr.Row():
-        label_out = gr.Textbox(label="Prediction",  scale=2)
-        conf_out  = gr.Textbox(label="Confidence",  scale=1)
+    # ── Phase 2: Trim / segment selector ───────────────────────────
+    with gr.Column(visible=False) as segment_panel:
+        gr.Markdown("### Select segment to analyze")
 
-    breakdown_out = gr.Textbox(label="Window Breakdown", lines=6, interactive=False)
+        timeline_html = gr.HTML()
 
-    run_btn.click(
-        fn=predict,
-        inputs=url_box,
-        outputs=[label_out, conf_out, breakdown_out],
+        with gr.Row():
+            start_slider = gr.Slider(minimum=0, maximum=100, value=0,   step=0.1, label="Start (s)")
+            end_slider   = gr.Slider(minimum=0, maximum=100, value=100, step=0.1, label="End (s)")
+
+        analyze_btn = gr.Button("Analyze Selection", variant="primary", size="lg")
+
+    # ── Phase 3: Results ────────────────────────────────────────────
+    with gr.Column(visible=False) as results_panel:
+        gr.Markdown("### Result")
+        with gr.Row():
+            label_out = gr.Textbox(label="Prediction", scale=2)
+            conf_out  = gr.Textbox(label="Confidence", scale=1)
+        breakdown_out = gr.Textbox(label="Window Breakdown", lines=6, interactive=False)
+
+    # ── Event wiring ────────────────────────────────────────────────
+
+    _load_inputs  = [url_box, video_state]
+    _load_outputs = [video_state, status_box, segment_panel, start_slider, end_slider, timeline_html]
+
+    load_btn.click(fn=load_video_cb, inputs=_load_inputs, outputs=_load_outputs)
+    url_box.submit(fn=load_video_cb, inputs=_load_inputs, outputs=_load_outputs)
+
+    start_slider.change(
+        fn=update_timeline_cb,
+        inputs=[video_state, start_slider, end_slider],
+        outputs=[timeline_html],
     )
-    url_box.submit(
-        fn=predict,
-        inputs=url_box,
-        outputs=[label_out, conf_out, breakdown_out],
+    end_slider.change(
+        fn=update_timeline_cb,
+        inputs=[video_state, start_slider, end_slider],
+        outputs=[timeline_html],
+    )
+
+    analyze_btn.click(
+        fn=analyze_cb,
+        inputs=[video_state, start_slider, end_slider],
+        outputs=[label_out, conf_out, breakdown_out, results_panel],
     )
 
-demo.launch(share=True)
+demo.launch()