Update app.py

app.py

@@ -1,12 +1,14 @@
-import spaces
 import os
-from flask import Flask, request, jsonify
-from werkzeug.utils import secure_filename
 import cv2
+import numpy as np
 import torch
+import torch.nn as nn
 import torch.nn.functional as F
 from facenet_pytorch import MTCNN, InceptionResnetV1
-import numpy as np
+from collections import deque
+from flask import Flask, request, jsonify
+from werkzeug.utils import secure_filename
+import spaces
 
 app = Flask(__name__)
 
@@ -21,20 +23,50 @@ os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 # Device configuration
 DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'
 
+# Model initialization
 mtcnn = MTCNN(select_largest=False, post_process=False, device=DEVICE).to(DEVICE).eval()
 
-model = InceptionResnetV1(pretrained="vggface2", classify=True, num_classes=1, device=DEVICE)
-checkpoint = torch.load("resnetinceptionv1_epoch_32.pth", map_location=torch.device('cpu'))
+@spaces.GPU
+class EnsembleModel(nn.Module):
+    def __init__(self, num_models=3):
+        super(EnsembleModel, self).__init__()
+        self.models = nn.ModuleList([
+            InceptionResnetV1(pretrained="vggface2", classify=True, num_classes=1)
+            for _ in range(num_models)
+        ])
+    
+    def forward(self, x):
+        outputs = [torch.sigmoid(model(x)) for model in self.models]
+        return torch.mean(torch.stack(outputs), dim=0)
+
+model = EnsembleModel().to(DEVICE)
+checkpoint = torch.load("ensemble_model.pth", map_location=torch.device('cpu'))
 model.load_state_dict(checkpoint['model_state_dict'])
-model.to(DEVICE)
 model.eval()
 
 def allowed_file(filename):
     return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
 
+@spaces.GPU
+def augment_frame(frame):
+    # Random horizontal flip
+    if np.random.rand() > 0.5:
+        frame = cv2.flip(frame, 1)
+    
+    # Random brightness adjustment
+    brightness = np.random.uniform(0.8, 1.2)
+    frame = cv2.convertScaleAbs(frame, alpha=brightness, beta=0)
+    
+    # Random contrast adjustment
+    contrast = np.random.uniform(0.8, 1.2)
+    frame = cv2.addWeighted(frame, contrast, frame, 0, 0)
+    
+    return frame
+
 @spaces.GPU
 def process_frame(frame):
-    face = mtcnn(frame)
+    augmented_frame = augment_frame(frame)
+    face = mtcnn(augmented_frame)
     if face is None:
         return None, None
     
@@ -46,7 +78,7 @@ def process_frame(frame):
     face = face / 255.0
     
     with torch.no_grad():
-        output = torch.sigmoid(model(face).squeeze(0))
+        output = model(face).squeeze(0)
         prediction = "fake" if output.item() >= 0.5 else "real"
     
     return prediction, output.item()
@@ -76,14 +108,10 @@ def preprocess_video(video_path, output_path):
     out.release()
 
 @spaces.GPU
-def analyze_video(video_path):
-    preprocessed_path = os.path.join(app.config['UPLOAD_FOLDER'], 'preprocessed.mp4')
-    preprocess_video(video_path, preprocessed_path)
-    
-    cap = cv2.VideoCapture(preprocessed_path)
-    frame_count = 0
-    fake_count = 0
-    total_processed = 0
+def analyze_temporal(video_path, window_size=5):
+    cap = cv2.VideoCapture(video_path)
+    predictions = deque(maxlen=window_size)
+    frame_predictions = []
     
     while cap.isOpened():
         ret, frame = cap.read()
@@ -94,17 +122,26 @@ def analyze_video(video_path):
         prediction, confidence = process_frame(rgb_frame)
         
         if prediction is not None:
-            total_processed += 1
-            if prediction == "fake":
-                fake_count += 1
+            predictions.append(1 if prediction == "fake" else 0)
         
-        frame_count += 1
+        if len(predictions) == window_size:
+            avg_prediction = sum(predictions) / window_size
+            frame_predictions.append(avg_prediction)
     
     cap.release()
+    return frame_predictions
+
+@spaces.GPU
+def analyze_video(video_path):
+    preprocessed_path = os.path.join(app.config['UPLOAD_FOLDER'], 'preprocessed.mp4')
+    preprocess_video(video_path, preprocessed_path)
+    
+    frame_predictions = analyze_temporal(preprocessed_path)
+    
     os.remove(preprocessed_path)
     
-    if total_processed > 0:
-        fake_percentage = (fake_count / total_processed) * 100
+    if frame_predictions:
+        fake_percentage = (sum(pred > 0.5 for pred in frame_predictions) / len(frame_predictions)) * 100
         return fake_percentage
     else:
         return 0