Update app.py

app.py

@@ -1,14 +1,15 @@
+import spaces
+from flask import Flask, request, jsonify
 import os
+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
-from collections import deque
-from flask import Flask, request, jsonify
-from werkzeug.utils import secure_filename
-import spaces
+import numpy as np
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
+import os
 
 app = Flask(__name__)
 
@@ -23,154 +24,92 @@ 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()
 
-@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 = InceptionResnetV1(pretrained="vggface2", classify=True, num_classes=1, device=DEVICE)
+# Model Credits: https://huggingface.co/spaces/dhairyashah/deepfake-alpha-version/blob/main/CREDITS.md
+checkpoint = torch.load("resnetinceptionv1_epoch_32.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):
-    augmented_frame = augment_frame(frame)
-    face = mtcnn(augmented_frame)
+    face = mtcnn(frame)
     if face is None:
         return None, None
-    
+
     face = face.unsqueeze(0)
     face = F.interpolate(face, size=(256, 256), mode='bilinear', align_corners=False)
-    
+
     face = face.to(DEVICE)
     face = face.to(torch.float32)
     face = face / 255.0
-    
+
     with torch.no_grad():
-        output = model(face).squeeze(0)
+        output = torch.sigmoid(model(face).squeeze(0))
         prediction = "fake" if output.item() >= 0.5 else "real"
-    
+
     return prediction, output.item()
 
 @spaces.GPU
-def preprocess_video(video_path, output_path):
+def analyze_video(video_path, sample_rate=30):
     cap = cv2.VideoCapture(video_path)
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
-    
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
-        
-        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        face = mtcnn.detect(rgb_frame)
-        
-        if face[0] is not None:
-            out.write(frame)
-    
-    cap.release()
-    out.release()
+    frame_count = 0
+    fake_count = 0
+    total_processed = 0
 
-@spaces.GPU
-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()
         if not ret:
             break
-        
-        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        prediction, confidence = process_frame(rgb_frame)
-        
-        if prediction is not None:
-            predictions.append(1 if prediction == "fake" else 0)
-        
-        if len(predictions) == window_size:
-            avg_prediction = sum(predictions) / window_size
-            frame_predictions.append(avg_prediction)
-    
+
+        if frame_count % sample_rate == 0:
+            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            prediction, confidence = process_frame(rgb_frame)
+
+            if prediction is not None:
+                total_processed += 1
+                if prediction == "fake":
+                    fake_count += 1
+
+        frame_count += 1
+
     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 frame_predictions:
-        fake_percentage = (sum(pred > 0.5 for pred in frame_predictions) / len(frame_predictions)) * 100
+    if total_processed > 0:
+        fake_percentage = (fake_count / total_processed) * 100
         return fake_percentage
     else:
         return 0
 
-@spaces.GPU
 @app.route('/analyze', methods=['POST'])
 def analyze_video_api():
     if 'video' not in request.files:
         return jsonify({'error': 'No video file provided'}), 400
-    
+
     file = request.files['video']
-    
+
     if file.filename == '':
         return jsonify({'error': 'No selected file'}), 400
-    
+
     if file and allowed_file(file.filename):
         filename = secure_filename(file.filename)
         filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
         file.save(filepath)
-        
+
         try:
             fake_percentage = analyze_video(filepath)
             os.remove(filepath)  # Remove the file after analysis
-            
+
             result = {
                 'fake_percentage': round(fake_percentage, 2),
                 'is_likely_deepfake': fake_percentage >= 60
             }
-            
+
             return jsonify(result), 200
         except Exception as e:
             os.remove(filepath)  # Remove the file if an error occurs