Update app.py

app.py

@@ -11,6 +11,7 @@ import cv2
 from pytorch_grad_cam import GradCAM
 from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
 from pytorch_grad_cam.utils.image import show_cam_on_image
+import tempfile
 
 with zipfile.ZipFile("examples.zip","r") as zip_ref:
     zip_ref.extractall(".")
@@ -49,56 +50,81 @@ for example_name in examples_names:
 np.random.shuffle(examples) # shuffle
 
 @spaces.GPU
-def predict(input_image:Image.Image, true_label:str):
-    """Predict the label of the input_image"""
-    face = mtcnn(input_image)
+def process_frame(frame, mtcnn, model, cam, targets):
+    face = mtcnn(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)))
     if face is None:
-        raise Exception('No face detected')
-    face = face.unsqueeze(0) # add the batch dimension
+        return frame, None, None
+    
+    face = face.unsqueeze(0)
     face = F.interpolate(face, size=(256, 256), mode='bilinear', align_corners=False)
     
-    # convert the face into a numpy array to be able to plot it
-    prev_face = face.squeeze(0).permute(1, 2, 0).cpu().detach().int().numpy()
-    prev_face = prev_face.astype('uint8')
-
     face = face.to(DEVICE)
     face = face.to(torch.float32)
     face = face / 255.0
-    face_image_to_plot = face.squeeze(0).permute(1, 2, 0).cpu().detach().int().numpy()
-
-    target_layers=[model.block8.branch1[-1]]
-    cam = GradCAM(model=model, target_layers=target_layers)
-    targets = [ClassifierOutputTarget(0)]
+    face_image_to_plot = face.squeeze(0).permute(1, 2, 0).cpu().detach().numpy()
 
     grayscale_cam = cam(input_tensor=face, targets=targets, eigen_smooth=True)
     grayscale_cam = grayscale_cam[0, :]
     visualization = show_cam_on_image(face_image_to_plot, grayscale_cam, use_rgb=True)
-    face_with_mask = cv2.addWeighted(prev_face, 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"
+        confidence = 1 - output.item() if prediction == "real" else output.item()
+    
+    return visualization, prediction, confidence
+
+@spaces.GPU
+def predict_video(input_video: str):
+    """Predict the labels for each frame of the input video"""
+    cap = cv2.VideoCapture(input_video)
+    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))
+    
+    target_layers = [model.block8.branch1[-1]]
+    cam = GradCAM(model=model, target_layers=target_layers)
+    targets = [ClassifierOutputTarget(0)]
+    
+    temp_output = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
+    out = cv2.VideoWriter(temp_output.name, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
+    
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
         
-        real_prediction = 1 - output.item()
-        fake_prediction = output.item()
+        processed_frame, prediction, confidence = process_frame(frame, mtcnn, model, cam, targets)
         
-        confidences = {
-            'real': real_prediction,
-            'fake': fake_prediction
-        }
-    return confidences, true_label, face_with_mask
+        if processed_frame is not None:
+            # Resize the processed frame to match the original video dimensions
+            processed_frame = cv2.resize(processed_frame, (width, height))
+            
+            # Add text with prediction and confidence
+            text = f"{prediction}: {confidence:.2f}"
+            cv2.putText(processed_frame, text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
+            
+            out.write(processed_frame)
+        else:
+            # If no face is detected, write the original frame
+            out.write(frame)
+    
+    cap.release()
+    out.release()
+    
+    return temp_output.name
 
 interface = gr.Interface(
-    fn=predict,
+    fn=predict_video,
     inputs=[
-        gr.components.Image(label="Input Image", type="pil"),  # Updated component import and type
-        gr.components.Text(label="Your Text Input")  # Updated component import
+        gr.Video(label="Input Video")
     ],
     outputs=[
-        gr.components.Label(label="Class"),  # Updated component import
-        gr.components.Text(label="Your Text Output"),  # Updated component import
-        gr.components.Image(label="Face with Explainability", type="numpy")  # Updated component import and type
+        gr.Video(label="Output Video")
     ],
-    examples=[[examples[i]["path"], examples[i]["label"]] for i in range(10)],
-    cache_examples=True  # Adjusted according to the new parameter for caching examples if needed
-).launch()
+    title="Video Deepfake Detection",
+    description="Upload a video to detect deepfakes in each frame."
+)
+
+if __name__ == "__main__":
+    interface.launch()