Update main.py

main.py

@@ -4,7 +4,7 @@ import torch
 import shutil
 from fastapi import FastAPI, UploadFile, File
 from fastapi.middleware.cors import CORSMiddleware
-from transformers import AutoImageProcessor, AutoModelForImageClassification
+from transformers import VideoMAEImageProcessor, VideoMAEForVideoClassification
 from facenet_pytorch import MTCNN
 from PIL import Image
 
@@ -23,32 +23,31 @@ app.add_middleware(
 print("Loading MTCNN Face Detector...")
 mtcnn = MTCNN(keep_all=False, select_largest=True, post_process=False)
 
-print("Loading Hugging Face Deepfake Detector...")
-model_name = "dima806/deepfake_vs_real_image_detection"
-processor = AutoImageProcessor.from_pretrained(model_name)
-model = AutoModelForImageClassification.from_pretrained(model_name)
+print("Loading Hugging Face Temporal Deepfake Detector (VideoMAE)...")
+model_name = "Ammar2k/videomae-base-finetuned-deepfake-subset"
+processor = VideoMAEImageProcessor.from_pretrained(model_name)
+model = VideoMAEForVideoClassification.from_pretrained(model_name)
 
 # Ensure temp directory exists
 os.makedirs("temp", exist_ok=True)
 
-def extract_faces_from_video(video_path, max_frames=8):
-    """Extracts a limited number of frames and crops the face from each."""
+def extract_faces_sequence(video_path, sequence_length=16):
+    """Extracts a sequence of continuous frames and tracks the face temporally."""
     cap = cv2.VideoCapture(video_path)
     frames_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
     
     if frames_count == 0:
         return []
         
-    # Calculate step to get evenly spaced frames
-    step = max(1, frames_count // max_frames)
+    # Try to get frames from the middle of the video
+    start_frame = max(0, (frames_count // 2) - (sequence_length // 2))
+    cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
     
     faces = []
-    current_frame = 0
+    last_box = None
     
-    while cap.isOpened() and len(faces) < max_frames:
-        cap.set(cv2.CAP_PROP_POS_FRAMES, current_frame)
+    for _ in range(sequence_length):
         ret, frame = cap.read()
-        
         if not ret:
             break
             
@@ -58,85 +57,83 @@ def extract_faces_from_video(video_path, max_frames=8):
         
         # Detect and crop face
         boxes, _ = mtcnn.detect(pil_img)
+        
+        # Bounding box smoothing: fallback to last known box if detection fails on a frame
         if boxes is not None and len(boxes) > 0:
             box = boxes[0]
+            last_box = box
+        elif last_box is not None:
+            box = last_box
+        else:
+            continue # Skip if no face found yet
             
-            # Add 10% padding around the face. Reduced from 30% to prevent background noise interference.
-            w = box[2] - box[0]
-            h = box[3] - box[1]
-            pad_w = int(w * 0.1)
-            pad_h = int(h * 0.1)
-            
-            # Ensure box coordinates are within image bounds with padding
-            x1 = max(0, int(box[0]) - pad_w)
-            y1 = max(0, int(box[1]) - pad_h)
-            x2 = min(pil_img.width, int(box[2]) + pad_w)
-            y2 = min(pil_img.height, int(box[3]) + pad_h)
+        # Add 30% padding around the face for better context
+        w = box[2] - box[0]
+        h = box[3] - box[1]
+        pad_w = int(w * 0.3)
+        pad_h = int(h * 0.3)
+        
+        x1 = max(0, int(box[0]) - pad_w)
+        y1 = max(0, int(box[1]) - pad_h)
+        x2 = min(pil_img.width, int(box[2]) + pad_w)
+        y2 = min(pil_img.height, int(box[3]) + pad_h)
+        
+        if x2 > x1 and y2 > y1:
+            face_crop = pil_img.crop((x1, y1, x2, y2))
+            faces.append(face_crop)
             
-            if x2 > x1 and y2 > y1:
-                face_crop = pil_img.crop((x1, y1, x2, y2))
-                faces.append(face_crop)
-                
-        current_frame += step
-        
     cap.release()
+    
+    # VideoMAE requires exactly `sequence_length` frames. Pad by duplicating last frame if short.
+    if len(faces) == 0:
+        return []
+        
+    while len(faces) < sequence_length:
+        faces.append(faces[-1])
+        
     return faces
 
 @app.post("/api/analyze")
 async def analyze_video(file: UploadFile = File(...)):
     print(f"Received file: {file.filename}")
     
-    # Save the uploaded file temporarily
     temp_video_path = os.path.join("temp", file.filename)
     with open(temp_video_path, "wb") as buffer:
         shutil.copyfileobj(file.file, buffer)
         
     try:
-        # Extract faces
-        print("Extracting faces from video...")
-        faces = extract_faces_from_video(temp_video_path, max_frames=15)
+        print("Extracting facial sequence from video...")
+        # Extract 16 consecutive frames
+        faces = extract_faces_sequence(temp_video_path, sequence_length=16)
         
         if not faces:
             return {
                 "isFake": False, 
                 "confidence": 0, 
-                "explanation": "Could not detect any faces in the video. Ensure the subject's face is clearly visible.",
+                "explanation": "Could not detect a clear face in the video sequence.",
                 "details": []
             }
             
-        print(f"Extracted {len(faces)} faces. Running inference...")
+        print(f"Extracted {len(faces)} frame sequence. Running temporal inference...")
         
-        # Prepare for model
-        inputs = processor(images=faces, return_tensors="pt")
+        # Prepare for VideoMAE 3D model
+        inputs = processor(list(faces), return_tensors="pt")
         
         # Run inference
         with torch.no_grad():
             outputs = model(**inputs)
             
-        # Apply softmax to get probabilities
-        probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
-        
-        # For dima806/deepfake_vs_real_image_detection, Fake is index 1 and Real is index 0
-        fake_idx = 1 
-        fake_probs = probabilities[:, fake_idx]
+        # Video classification models return logits for the whole sequence
+        probabilities = torch.nn.functional.softmax(outputs.logits[0], dim=-1)
         
-        # Use Top-K Pooling: Analyze the most "fake" scoring frames
-        # Deepfakes often have glitches in a few frames. Average the worst 3 frames.
-        k = min(3, len(fake_probs))
-        top_k_probs, _ = torch.topk(fake_probs, k)
+        # Get highest probability label
+        predicted_class_idx = probabilities.argmax(-1).item()
+        label = model.config.id2label[predicted_class_idx].lower()
         
-        # Get average probability for 'Fake' class across the worst frames
-        avg_fake_prob = top_k_probs.mean().item()
-        
-        is_fake = avg_fake_prob > 0.75
-        
-        # Calculate confidence score based on the chosen class
-        if is_fake:
-            confidence = round(avg_fake_prob * 100, 2)
-        else:
-            confidence = round((1.0 - avg_fake_prob) * 100, 2)
+        is_fake = 'fake' in label
+        confidence = round(probabilities[predicted_class_idx].item() * 100, 2)
             
-        explanation = "Our AI detected significant spatial artifacts and inconsistencies consistent with synthetic generation or facial manipulation." if is_fake else "No significant manipulation artifacts were detected. The spatial integrity and facial rendering are consistent with genuine media."
+        explanation = "Our 3D Temporal AI analyzed facial motion and detected unnatural movement, micro-expressions, or spatial inconsistencies typical of deepfakes." if is_fake else "Our 3D Temporal AI analyzed the facial movement and found natural micro-expressions and consistent temporal flow."
         
         print(f"Result: isFake={is_fake}, confidence={confidence}%")
         
@@ -145,9 +142,8 @@ async def analyze_video(file: UploadFile = File(...)):
             "confidence": confidence,
             "explanation": explanation,
             "details": [
-                {"title": "Face Detection", "desc": f"Extracted {len(faces)} key frames and isolated the subject's face using MTCNN."},
-                {"title": "Spatial Analysis", "desc": "Evaluated using a Vision Transformer (ViT) deep learning architecture."},
-                {"title": "Temporal Pooling", "desc": f"Applied Top-{k} analysis to identify and flag the highest-risk manipulated frames."}
+                {"title": "Temporal Analysis", "desc": "Analyzed a continuous 16-frame clip using a VideoMAE 3D Transformer."},
+                {"title": "Motion Tracking", "desc": "Tracked facial landmarks across time to detect jitter, blending artifacts, and lip-sync inconsistencies."}
             ]
         }
         
@@ -156,10 +152,9 @@ async def analyze_video(file: UploadFile = File(...)):
         return {"error": str(e)}
         
     finally:
-        # Clean up temp file
         if os.path.exists(temp_video_path):
             os.remove(temp_video_path)
 
 @app.get("/")
 def health_check():
-    return {"status": "Backend is running!"}
+    return {"status": "3D Temporal Backend is running!"}