Update app.py

app.py

@@ -5,7 +5,7 @@ Serves the frontend UI + deepfake detection API from one HF Space.
 Routes:
   GET  /          → Serves index.html (the web UI)
   GET  /health    → JSON health check
-  POST /predict   → Video upload → REAL/FAKE prediction
+  POST /predict   → Video/Photo upload → REAL/FAKE prediction
 """
 
 import os
@@ -15,84 +15,27 @@ import shutil
 import logging
 import tempfile
 from pathlib import Path
-from functools import lru_cache
 
 import cv2
 import torch
-import torch.nn as nn
 import numpy as np
 from PIL import Image, ImageFile
-from facenet_pytorch import MTCNN
 from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse, FileResponse
 from fastapi.staticfiles import StaticFiles
-import torchvision.transforms as T
 
 ImageFile.LOAD_TRUNCATED_IMAGES = True
 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
 logger = logging.getLogger(__name__)
 
-# ─────────────────────────────────────────────
-# Model Definition (self-contained)
-# ─────────────────────────────────────────────
-
-class DINOv2Extractor(nn.Module):
-    def __init__(self, variant: str = "dinov2_vitb14"):
-        super().__init__()
-        logger.info(f"Loading {variant} from torch.hub...")
-        self.backbone = torch.hub.load(
-            "facebookresearch/dinov2", variant, pretrained=True
-        )
-        self.feature_dim = 768
-        for p in self.backbone.parameters():
-            p.requires_grad = False
-        logger.info("DINOv2 backbone loaded (frozen).")
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.backbone(x)
-
-
-class MLPClassifier(nn.Module):
-    def __init__(self, input_dim: int = 1536, num_classes: int = 2, dropout: float = 0.3):
-        super().__init__()
-        self.net = nn.Sequential(
-            nn.Linear(input_dim, 512),
-            nn.LayerNorm(512),
-            nn.GELU(),
-            nn.Dropout(dropout),
-            nn.Linear(512, 256),
-            nn.LayerNorm(256),
-            nn.GELU(),
-            nn.Dropout(dropout / 2),
-            nn.Linear(256, num_classes),
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.net(x)
-
-
-class DeepfakeDetector(nn.Module):
-    def __init__(self, dual_input: bool = True):
-        super().__init__()
-        self.dual_input = dual_input
-        self.extractor = DINOv2Extractor()
-        feat_dim = 1536 if dual_input else 768
-        self.classifier = MLPClassifier(input_dim=feat_dim)
-
-    def forward(self, full_img: torch.Tensor, face_img: torch.Tensor = None) -> torch.Tensor:
-        full_feat = self.extractor(full_img)
-        if self.dual_input and face_img is not None:
-            face_feat = self.extractor(face_img)
-            feats = torch.cat([full_feat, face_feat], dim=1)
-        else:
-            feats = full_feat
-        return self.classifier(feats)
-
-
-# ─────────────────────────────────────────────
-# App Setup
-# ─────────────────────────────────────────────
+# --- START OF IMPORT EXACTLY AS test_real.py ---
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from src.utils.helpers import load_config, get_device
+from src.data.transforms import get_val_transforms
+from src.models.classifier import DeepfakeClassifier
+from src.utils.face_detect import FaceDetector
+# --- END OF IMPORT ---
 
 app = FastAPI(
     title="DeepShield AI",
@@ -108,132 +51,137 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+DEVICE = get_device('cuda' if torch.cuda.is_available() else 'cpu')
 CHECKPOINT_PATH = Path("best_model.pth")
 MAX_FRAMES = 20
 MAX_FILE_MB = 30
 MAX_DURATION_SEC = 60
 
-# MTCNN face detector (initialized once, CPU is fine for detection)
-try:
-    MTCNN_DETECTOR = MTCNN(
-        image_size=224,
-        margin=40,
-        min_face_size=20,
-        thresholds=[0.6, 0.7, 0.9],
-        keep_all=False,
-        device='cpu'
-    )
-    logger.info("MTCNN face detector initialized.")
-except Exception as e:
-    MTCNN_DETECTOR = None
-    logger.warning(f"MTCNN init failed (will use full frame fallback): {e}")
-
-TRANSFORM = T.Compose([
-    T.Resize((224, 224)),
-    T.ToTensor(),
-    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-])
-
-
-def detect_face_crop(img: Image.Image) -> Image.Image:
-    """Detect face with MTCNN and return cropped face, or None if not found."""
-    if MTCNN_DETECTOR is None:
-        return None
-    try:
-        # MTCNN returns the cropped tensor directly
-        face_tensor = MTCNN_DETECTOR(img)
-        if face_tensor is not None:
-            # Convert tensor back to PIL Image
-            face_np = face_tensor.permute(1, 2, 0).numpy()
-            face_np = ((face_np * 128) + 127.5).clip(0, 255).astype(np.uint8)
-            return Image.fromarray(face_np)
-    except Exception:
-        pass
-    return None
-
-
-@lru_cache(maxsize=1)
-def load_model() -> DeepfakeDetector:
-    if not CHECKPOINT_PATH.exists():
-        raise RuntimeError("best_model.pth not found. Upload it to this HF Space.")
-
-    logger.info(f"Loading checkpoint on {DEVICE}...")
-    ckpt = torch.load(CHECKPOINT_PATH, map_location=DEVICE)
-    state = ckpt.get("model_state_dict", ckpt)
-
-    mlp_w = state.get("classifier.net.0.weight", None)
-    dual = (mlp_w.shape[1] == 1536) if mlp_w is not None else True
+# Model caching
+global_model = None
+global_face_detector = None
+global_transform = None
+global_dual_input = True
 
-    model = DeepfakeDetector(dual_input=dual).to(DEVICE)
-    model.load_state_dict(state, strict=False)
-    model.eval()
-    logger.info(f"Model ready. dual_input={dual}, device={DEVICE}")
-    return model
+def load_model_and_detector():
+    global global_model, global_face_detector, global_transform, global_dual_input
+    if global_model is not None:
+        return global_model, global_face_detector, global_transform, global_dual_input
 
+    ckpt_path_to_load = None
+    if not CHECKPOINT_PATH.exists():
+        base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+        fallback_path = os.path.join(base_dir, 'models2/checkpoints/best_model.pth')
+        if not os.path.exists(fallback_path):
+            fallback_path = os.path.join(base_dir, 'models2/checkpoints/best_mlp.pth')
+        
+        if os.path.exists(fallback_path):
+            ckpt_path_to_load = fallback_path
+        else:
+            raise RuntimeError("best_model.pth not found. Upload it to this HF Space.")
+    else:
+        ckpt_path_to_load = str(CHECKPOINT_PATH)
+
+    logger.info(f"Loading checkpoint on {DEVICE} from {ckpt_path_to_load} ...")
+    checkpoint = torch.load(ckpt_path_to_load, map_location=str(DEVICE))
+    
+    is_fast_mlp = ckpt_path_to_load.endswith('best_mlp.pth')
+    dual_input = True
+    if is_fast_mlp and 'feat_dim' in checkpoint:
+        dual_input = (checkpoint['feat_dim'] == 1536)
+
+    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    config_path = os.path.join(base_dir, 'configs/config2.yaml')
+    if os.path.exists(config_path):
+        config = load_config(config_path)
+    else:
+        # Fallback if config is missing in HF Spaces standalone
+        config = {
+            'model': {'dino_variant': 'dinov2_vitb14', 'unfreeze_last_n_blocks': 0, 'dual_input': True},
+            'face_detection': {'margin': 40, 'confidence_threshold': 0.9},
+            'data': {'image_size': 224}
+        }
+    
+    if not is_fast_mlp:
+        dual_input = config['model'].get('dual_input', True)
+
+    face_detector = FaceDetector(
+        margin=config['face_detection']['margin'],
+        confidence_threshold=config['face_detection']['confidence_threshold'],
+        image_size=config['data']['image_size'],
+        device=str(DEVICE)
+    ) if dual_input else None
+
+    model = DeepfakeClassifier(
+        dino_variant=config['model']['dino_variant'],
+        freeze_backbone=not is_fast_mlp,
+        unfreeze_last_n_blocks=config['model']['unfreeze_last_n_blocks'] if not is_fast_mlp else 0, 
+        dual_input=dual_input
+    )
+    
+    if is_fast_mlp:
+        model.classifier.load_state_dict(checkpoint['model_state_dict'])
+    else:
+        model.load_state_dict(checkpoint['model_state_dict'] if 'model_state_dict' in checkpoint else checkpoint)
+        
+    model = model.to(DEVICE).eval()
+    transform = get_val_transforms(config['data']['image_size'])
 
-def extract_frames(video_path: str, output_dir: str, num_frames: int = MAX_FRAMES) -> list:
-    cap = cv2.VideoCapture(video_path)
-    if not cap.isOpened():
-        raise ValueError("Cannot open video file.")
-
-    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-    fps = cap.get(cv2.CAP_PROP_FPS) or 25
-    duration = total_frames / fps if fps > 0 else 0
-
-    if duration > MAX_DURATION_SEC:
-        cap.release()
-        raise ValueError(f"Video too long ({duration:.0f}s). Max: {MAX_DURATION_SEC}s.")
+    global_model = model
+    global_face_detector = face_detector
+    global_transform = transform
+    global_dual_input = dual_input
 
-    if total_frames <= 0:
-        total_frames = int(fps * MAX_DURATION_SEC)
+    logger.info(f"Model ready. dual_input={dual_input}, device={DEVICE}, is_fast_mlp={is_fast_mlp}")
+    return model, face_detector, transform, dual_input
 
-    step = max(1, total_frames // num_frames)
-    target_indices = set(range(0, total_frames, step))
-    saved_paths = []
-    frame_idx = 0
 
-    while len(saved_paths) < num_frames:
+def extract_frames(video_path: str, temp_dir: str, num_frames: int = MAX_FRAMES) -> list:
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened(): return []
+    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    step = max(1, (total if total > 0 else 300) // num_frames)
+    indices = set(range(0, total if total > 0 else 300, step))
+    
+    saved = []
+    for i in range(total if total > 0 else 300):
         ret, frame = cap.read()
-        if not ret:
-            break
-        if frame_idx in target_indices:
+        if not ret: break
+        if i in indices:
             rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-            path = os.path.join(output_dir, f"frame_{len(saved_paths):04d}.jpg")
-            Image.fromarray(rgb).save(path, quality=90)
-            saved_paths.append(path)
-        frame_idx += 1
-
+            path = os.path.join(temp_dir, f"frame_{len(saved):03d}.jpg")
+            Image.fromarray(rgb).save(path)
+            saved.append(path)
+            if len(saved) >= num_frames: break
     cap.release()
-    return saved_paths
+    return saved
 
 
-def run_inference(model: DeepfakeDetector, frame_paths: list) -> dict:
+def run_inference(frame_paths: list) -> dict:
+    model, face_detector, transform, dual_input = load_model_and_detector()
     fake_probs = []
+    
     with torch.no_grad():
-        for fpath in frame_paths:
+        for i, p in enumerate(frame_paths):
             try:
-                img = Image.open(fpath).convert("RGB")
-                t_img = TRANSFORM(img).unsqueeze(0).to(DEVICE)
-
-                # Try MTCNN face detection first (same as test_real.py)
-                t_face = t_img  # default fallback = full frame
-                if model.dual_input:
-                    face_crop = detect_face_crop(img)
-                    if face_crop is not None:
-                        t_face = TRANSFORM(face_crop).unsqueeze(0).to(DEVICE)
-                    # else: fallback to full image (face not detected)
-
-                logits = model(t_img, t_face if model.dual_input else None)
-                prob = torch.softmax(logits, dim=1)[0, 1].item()
+                img = Image.open(p).convert('RGB')
+                t_img = transform(img).unsqueeze(0).to(DEVICE)
+                t_face = t_img
+                
+                if dual_input:
+                    face = face_detector.detect_and_crop(p)
+                    if face is not None:
+                        t_face = transform(face).unsqueeze(0).to(DEVICE)
+                        
+                probs = torch.softmax(model(t_img, t_face if dual_input else None), dim=1)
+                prob = probs[0, 1].item()
                 fake_probs.append(prob)
             except Exception as e:
-                logger.warning(f"Skipping frame {fpath}: {e}")
+                logger.warning(f"Skipping frame {p}: {e}")
 
     if not fake_probs:
         raise ValueError("No frames could be processed.")
 
-    # 1. Simple Aggregation (Mean) to match test_real.py
     video_fake_prob = float(np.mean(fake_probs))
 
     is_fake = video_fake_prob > 0.5
@@ -249,25 +197,26 @@ def run_inference(model: DeepfakeDetector, frame_paths: list) -> dict:
     }
 
 
-# ─────────────────────────────────────────────
-# API Routes (must be defined BEFORE static mount)
-# ─────────────────────────────────────────────
-
 @app.on_event("startup")
 async def startup_event():
     try:
-        load_model()
+        load_model_and_detector()
     except Exception as e:
         logger.error(f"Startup model load failed: {e}")
 
 
 @app.get("/health")
 def health_check():
+    try:
+        model_loaded = CHECKPOINT_PATH.exists() or os.path.exists(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'models2/checkpoints/best_model.pth'))
+    except:
+        model_loaded = False
+        
     return {
         "status": "ok",
         "model": "DINO-G50 Deepfake Detector",
         "device": str(DEVICE),
-        "model_loaded": CHECKPOINT_PATH.exists(),
+        "model_loaded": model_loaded,
     }
 
 
@@ -295,7 +244,6 @@ async def predict(file: UploadFile = File(...)):
             f.write(content)
         del content
 
-        model = load_model()
         logger.info(f"[{job_id}] Processing: {file.filename} ({size_mb:.1f} MB)")
 
         if ext in {".mp4", ".mov", ".avi", ".mkv"}:
@@ -307,7 +255,7 @@ async def predict(file: UploadFile = File(...)):
             shutil.copy(video_path, img_path)
             frame_paths = [str(img_path)]
 
-        result = run_inference(model, frame_paths)
+        result = run_inference(frame_paths)
         result["filename"] = file.filename
         result["file_size_mb"] = round(size_mb, 2)
         result["job_id"] = job_id
@@ -326,8 +274,4 @@ async def predict(file: UploadFile = File(...)):
         shutil.rmtree(temp_dir, ignore_errors=True)
         logger.info(f"[{job_id}] Cleanup done.")
 
-
-# ──────────────────────────────────────���──────
-# Static Frontend (mounted LAST — serves index.html at /)
-# ─────────────────────────────────────────────
 app.mount("/", StaticFiles(directory="static", html=True), name="static")