Update app.py

app.py

@@ -20,7 +20,7 @@ mtcnn = MTCNN(
     select_largest=False,
     post_process=False,
     device=DEVICE
-).to(DEVICE).eval()
+).eval()
 
 model = InceptionResnetV1(
     pretrained="vggface2",
@@ -35,8 +35,38 @@ model.to(DEVICE)
 model.eval()
 
 # Model Inference
+def predict_frame(frame):
+    """Predict whether the input frame contains real or fake faces"""
+    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    frame_pil = Image.fromarray(frame)
+
+    face = mtcnn(frame_pil)
+    if face is None:
+        raise Exception('No face detected')
+    face = face.unsqueeze(0)  # add the batch dimension
+    face = F.interpolate(face, size=(256, 256), mode='bilinear', align_corners=False)
+
+    face = face.to(DEVICE, dtype=torch.float32) / 255.0
+
+    target_layers = [model.block8.branch1[-1]]
+    use_cuda = True if torch.cuda.is_available() else False
+    cam = GradCAM(model=model, target_layers=target_layers, use_cuda=use_cuda)
+    targets = [ClassifierOutputTarget(0)]
+
+    grayscale_cam = cam(input_tensor=face, targets=targets, eigen_smooth=True)
+    grayscale_cam = grayscale_cam[0, :]
+    visualization = show_cam_on_image(frame, grayscale_cam, use_rgb=True)
+    face_with_mask = cv2.addWeighted(frame, 1, visualization, 0.5, 0)
+
+    with torch.no_grad():
+        output = torch.sigmoid(model(face).squeeze(0))
+        prediction = "real" if output.item() < 0.5 else "fake"
+
+    return prediction, face_with_mask
+
+
+# Function to process video
 def predict_video(input_video):
-    """Predict whether the input video contains real or fake faces"""
     cap = cv2.VideoCapture(input_video.name)
     frames = []
     confidences = []
@@ -45,44 +75,18 @@ def predict_video(input_video):
         ret, frame = cap.read()
         if not ret:
             break
-        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        frame_pil = Image.fromarray(frame)
 
-        face = mtcnn(frame_pil)
-        if face is None:
-            raise Exception('No face detected')
-        face = face.unsqueeze(0)  # add the batch dimension
-        face = F.interpolate(face, size=(256, 256), mode='bilinear', align_corners=False)
+        prediction, frame_with_mask = predict_frame(frame)
 
-        face = face.to(DEVICE, dtype=torch.float32) / 255.0
-
-        target_layers = [model.block8.branch1[-1]]
-        use_cuda = True if torch.cuda.is_available() else False
-        cam = GradCAM(model=model, target_layers=target_layers, use_cuda=use_cuda)
-        targets = [ClassifierOutputTarget(0)]
-
-        grayscale_cam = cam(input_tensor=face, targets=targets, eigen_smooth=True)
-        grayscale_cam = grayscale_cam[0, :]
-        visualization = show_cam_on_image(frame, grayscale_cam, use_rgb=True)
-        face_with_mask = cv2.addWeighted(frame, 1, visualization, 0.5, 0)
-
-        with torch.no_grad():
-            output = torch.sigmoid(model(face).squeeze(0))
-            prediction = "real" if output.item() < 0.5 else "fake"
-
-            real_prediction = 1 - output.item()
-            fake_prediction = output.item()
-
-            confidences.append({
-                'real': real_prediction,
-                'fake': fake_prediction
-            })
-
-        frames.append(face_with_mask)
+        frames.append(frame_with_mask)
+        confidences.append(prediction)
 
     cap.release()
 
-    return confidences, frames
+    # Determine the final prediction based on the maximum occurrence of predictions
+    final_prediction = max(set(confidences), key=confidences.count)
+
+    return final_prediction, frames
 
 
 # Gradio Interface