Update detector.py

detector.py

@@ -1,5 +1,4 @@
-"""
-Deepfake Detector — Hugging Face ZeroGPU Space
+"""Deepfake Detector — Hugging Face ZeroGPU Space
 V-JEPA 2.1 ViT-Giant + AttentiveClassifier probe
 """
 import glob
@@ -11,6 +10,7 @@ import time
 import uuid
 import json
 import requests
+import shutil
 from datetime import datetime
 
 warnings.filterwarnings("ignore", category=FutureWarning)
@@ -32,7 +32,6 @@ ENCODER_BATCH = 3
 HF_TOKEN = os.environ.get("HF_TOKEN")
 ENCODER_REPO = os.environ.get("ENCODER_REPO", "")
 PROBE_REPO = os.environ.get("PROBE_REPO", ENCODER_REPO)
-
 DATASET_REPO_ID = os.environ.get("DATASET_REPO_ID", "")
 TELEGRAM_BOT_TOKEN = os.environ.get("TELEGRAM_BOT_TOKEN", "")
 TELEGRAM_CHAT_ID = os.environ.get("TELEGRAM_CHAT_ID", "")
@@ -56,13 +55,11 @@ encoder_weights_path = hf_hub_download(
     filename="vjepa2_1_vitg_384.pt",
     token=HF_TOKEN,
 )
-
 ckpt = torch.load(encoder_weights_path, map_location="cpu", weights_only=True)
 ckpt = ckpt.get("ema_encoder", ckpt.get("encoder", ckpt))
 ckpt = {k.replace("module.", "").replace("backbone.", ""): v for k, v in ckpt.items()}
 encoder.load_state_dict(ckpt, strict=False)
 del ckpt
-
 encoder = encoder.to(DEVICE).to(torch.bfloat16)
 encoder.eval()
 print("Encoder ready.")
@@ -73,11 +70,9 @@ probe_path = hf_hub_download(
     filename="attentive_probe_optimized.pt",
     token=HF_TOKEN,
 )
-
 probe_state = torch.load(probe_path, map_location="cpu")
 if "probe_state_dict" in probe_state:
     probe_state = probe_state["probe_state_dict"]
-
 probe = AttentiveClassifier(embed_dim=FEATURE_DIM)
 probe.load_state_dict(probe_state)
 probe = probe.to(DEVICE)
@@ -105,7 +100,6 @@ def extract_platform(url: str) -> str:
 def download_video(url: str) -> tuple[bool, str]:
     tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
     tmp.close()
-
     cmd = [
         "yt-dlp",
         "-f", "best[height<=720]/best",
@@ -117,17 +111,14 @@ def download_video(url: str) -> tuple[bool, str]:
         "-o", tmp.name,
         url,
     ]
-
     try:
         result = subprocess.run(cmd, capture_output=True, text=True, timeout=90)
         candidates = sorted(glob.glob(tmp.name + "*") + [tmp.name])
         actual_file = next((f for f in candidates if os.path.exists(f) and os.path.getsize(f) > 0), None)
-
         if actual_file is None:
             if "does not pass filter" in result.stdout or "does not pass filter" in result.stderr:
                 return False, "Video length exceeds the maximum 2-minute limit."
             return False, f"yt-dlp wrote no output. stderr: {result.stderr[:400]}"
-
         return True, actual_file
     except subprocess.TimeoutExpired:
         return False, "Download timed out (90s limit)."
@@ -166,7 +157,6 @@ def upload_and_notify(video_path: str, meta: dict):
     if not DATASET_REPO_ID or not HF_TOKEN:
         print("Upload skipped: Credentials missing.")
         return
-
     try:
         api = HfApi(token=HF_TOKEN)
         unique_id = str(uuid.uuid4())[:8]
@@ -175,13 +165,13 @@ def upload_and_notify(video_path: str, meta: dict):
         
         repo_video_path = f"data/video_{base_filename}.mp4"
         repo_json_path = f"data/meta_{base_filename}.json"
-
+        
         # Generate the direct download/view asset URL for the video
         direct_video_url = f"https://huggingface.co/datasets/{DATASET_REPO_ID}/resolve/main/{repo_video_path}"
         
         # Store the video direct link inside the metadata dictionary itself
         meta["stored_video_url"] = direct_video_url
-
+        
         # 1. Upload Video File
         api.upload_file(
             path_or_fileobj=video_path,
@@ -197,7 +187,7 @@ def upload_and_notify(video_path: str, meta: dict):
             repo_id=DATASET_REPO_ID,
             repo_type="dataset"
         )
-
+        
         # 3. Send Telegram alert containing the direct video asset link
         if TELEGRAM_BOT_TOKEN and TELEGRAM_CHAT_ID:
             tele_msg = (
@@ -221,11 +211,11 @@ def upload_and_notify(video_path: str, meta: dict):
 # ─────────────────────────────────────────────
 # GRADIO PREDICT
 # ─────────────────────────────────────────────
-def predict(url: str, consent: bool, current_state: dict) -> tuple[str, str, str, dict]:
+def predict(url: str, uploaded_video: str, consent: bool, current_state: dict) -> tuple[str, str, str, dict]:
     # Check mandatory consent gate
     if not consent:
         return "❌ Required: You must check the data collection consent box to process videos.", "", "", current_state
-
+    
     start_time = time.time()
     url = (url or "").strip()
     
@@ -233,16 +223,27 @@ def predict(url: str, consent: bool, current_state: dict) -> tuple[str, str, str
         old_path = current_state["video_path"]
         for f in glob.glob(old_path + "*") + [old_path]:
             if os.path.exists(f):
-                os.remove(f)
-
-    if not url:
-        return "Please enter a URL.", "", "", {}
-
-    ok, path_or_err = download_video(url)
-    if not ok:
-        return f"Download failed: {path_or_err}", "", "", {}
-
-    video_path = path_or_err
+                try:
+                    os.remove(f)
+                except Exception:
+                    pass
+
+    # Source priority: Direct upload over URL
+    if uploaded_video:
+        tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
+        shutil.copy(uploaded_video, tmp.name)
+        video_path = tmp.name
+        platform_name = "Direct Upload"
+        stored_url = "Uploaded File"
+    elif url:
+        ok, path_or_err = download_video(url)
+        if not ok:
+            return f"Download failed: {path_or_err}", "", "", {}
+        video_path = path_or_err
+        platform_name = extract_platform(url)
+        stored_url = url
+    else:
+        return "Please enter a URL or upload a video.", "", "", {}
 
     try:
         windows, total_frames, fps, duration_sec = get_video_windows(video_path)
@@ -250,14 +251,14 @@ def predict(url: str, consent: bool, current_state: dict) -> tuple[str, str, str
         if os.path.exists(video_path):
             os.remove(video_path)
         return f"Video decode failed: {e}", "", "", {}
-    
+        
     scores = encode_and_classify(windows)
-
+    
     if not scores:
         if os.path.exists(video_path):
             os.remove(video_path)
         return "No windows processed.", "", "", {}
-
+        
     final_prob = max(s["prob"] for s in scores)
     label = "⚠️  AI-GENERATED (FAKE)" if final_prob > 0.5 else "✅  REAL"
     elapsed_time = time.time() - start_time
@@ -270,26 +271,29 @@ def predict(url: str, consent: bool, current_state: dict) -> tuple[str, str, str
     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}")
-
+        
     # Build base metadata dict
     meta = {
-        "url": url,
-        "platform": extract_platform(url),
+        "url": stored_url,
+        "platform": platform_name,
         "label": label,
         "confidence": confidence,
         "processing_time_sec": round(elapsed_time, 2),
         "video_duration_sec": round(duration_sec, 2),
         "total_frames": total_frames
     }
-
+    
     # Execute backend background syncing task
     upload_and_notify(video_path, meta)
-
+    
     # Clean up local system file tracking loop paths
     for f in glob.glob(video_path + "*") + [video_path]:
         if os.path.exists(f):
-            os.remove(f)
-
+            try:
+                os.remove(f)
+            except Exception:
+                pass
+                
     return label, confidence, "\n".join(lines), {}
 
 # ─────────────────────────────────────────────
@@ -297,20 +301,27 @@ def predict(url: str, consent: bool, current_state: dict) -> tuple[str, str, str
 # ─────────────────────────────────────────────
 with gr.Blocks(title="Deepfake Detector") as demo:
     session_state = gr.State(value={})
-
+    
     gr.Markdown(
         """
         ## 🎬 Deepfake Video Detector
-        Paste a public Instagram, TikTok, or direct `.mp4` URL. (Maximum duration: 2 minutes).
+        Paste a public Instagram, TikTok, or direct `.mp4` URL, **OR** upload a video directly. (Maximum duration: 2 minutes).
         Uses **V-JEPA 2.1 ViT-Giant** + **AttentiveClassifier** probe.
         """
     )
     
-    url_box = gr.Textbox(
-        label="Video URL",
-        placeholder="https://www.instagram.com/reels/...",
-        autofocus=True,
-    )
+    with gr.Row():
+        url_box = gr.Textbox(
+            label="Video URL",
+            placeholder="https://www.instagram.com/reels/...",
+            autofocus=True,
+            scale=1
+        )
+        upload_box = gr.Video(
+            label="Or Upload Video directly",
+            sources=["upload"],
+            scale=1
+        )
     
     # Mandatory Consent Checkbox
     consent_box = gr.Checkbox(
@@ -329,12 +340,13 @@ with gr.Blocks(title="Deepfake Detector") as demo:
 
     run_btn.click(
         fn=predict,
-        inputs=[url_box, consent_box, session_state],
+        inputs=[url_box, upload_box, consent_box, session_state],
         outputs=[label_out, conf_out, breakdown_out, session_state],
     )
+    
     url_box.submit(
         fn=predict,
-        inputs=[url_box, consent_box, session_state],
+        inputs=[url_box, upload_box, consent_box, session_state],
         outputs=[label_out, conf_out, breakdown_out, session_state],
     )