Update app.py

app.py

@@ -5,16 +5,13 @@ import torch.nn.functional as F
 from facenet_pytorch import MTCNN, InceptionResnetV1
 import os
 import numpy as np
-from PIL import Image as PILImage
+from PIL import Image
 import zipfile
 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
-import matplotlib.pyplot as plt
-from tqdm import tqdm
-import io
 
 with zipfile.ZipFile("examples.zip","r") as zip_ref:
     zip_ref.extractall(".")
@@ -54,7 +51,7 @@ np.random.shuffle(examples) # shuffle
 
 @spaces.GPU
 def process_frame(frame, mtcnn, model, cam, targets):
-    face = mtcnn(PILImage.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)))
+    face = mtcnn(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)))
     if face is None:
         return frame, None, None
     
@@ -78,112 +75,56 @@ def process_frame(frame, mtcnn, model, cam, targets):
     return visualization, prediction, confidence
 
 @spaces.GPU
-def analyze_video(input_video: str):
-    """Analyze the video for deepfake detection"""
+def predict_video(input_video: str):
+    """Predict the labels for each frame of the input video"""
     cap = cv2.VideoCapture(input_video)
-    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    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)]
     
-    frame_confidences = []
-    frame_predictions = []
+    temp_output = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
+    out = cv2.VideoWriter(temp_output.name, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
     
-    for _ in tqdm(range(total_frames), desc="Analyzing video"):
+    while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
             break
         
-        _, prediction, confidence = process_frame(frame, mtcnn, model, cam, targets)
+        processed_frame, prediction, confidence = process_frame(frame, mtcnn, model, cam, targets)
         
-        if prediction is not None and confidence is not None:
-            frame_confidences.append(confidence)
-            frame_predictions.append(1 if prediction == "fake" else 0)
+        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()
     
-    # Calculate metrics
-    fake_percentage = (sum(frame_predictions) / len(frame_predictions)) * 100 if frame_predictions else 0
-    avg_confidence = np.mean(frame_confidences) if frame_confidences else 0
-    
-    # Create graphs
-    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 12))
-    
-    # Confidence over time
-    ax1.plot(frame_confidences)
-    ax1.set_title("Confidence Over Time")
-    ax1.set_xlabel("Frame")
-    ax1.set_ylabel("Confidence")
-    ax1.set_ylim(0, 1)
-    
-    # Prediction distribution
-    ax2.hist(frame_predictions, bins=[0, 0.5, 1], rwidth=0.8)
-    ax2.set_title("Distribution of Predictions")
-    ax2.set_xlabel("Prediction (0: Real, 1: Fake)")
-    ax2.set_ylabel("Count")
-    
-    # Save plot to bytes
-    buf = io.BytesIO()
-    plt.savefig(buf, format='png')
-    buf.seek(0)
-    
-    # Create progress bar image
-    progress_fig, progress_ax = plt.subplots(figsize=(8, 2))
-    progress_ax.barh(["Fake"], [fake_percentage], color='red')
-    progress_ax.barh(["Fake"], [100 - fake_percentage], left=[fake_percentage], color='green')
-    progress_ax.set_xlim(0, 100)
-    progress_ax.set_title("Fake Percentage")
-    progress_ax.set_xlabel("Percentage")
-    progress_ax.text(fake_percentage, 0, f"{fake_percentage:.1f}%", va='center', ha='left')
-    
-    # Save progress bar to bytes
-    progress_buf = io.BytesIO()
-    progress_fig.savefig(progress_buf, format='png')
-    progress_buf.seek(0)
-    
-    return {
-        "fake_percentage": fake_percentage,
-        "avg_confidence": avg_confidence,
-        "analysis_plot": buf,
-        "progress_bar": progress_buf,
-        "total_frames": total_frames,
-        "processed_frames": len(frame_confidences)
-    }
-
-def format_results(results):
-    return f"""
-    Analysis Results:
-    - Fake Percentage: {results['fake_percentage']:.2f}%
-    - Average Confidence: {results['avg_confidence']:.2f}
-    - Total Frames: {results['total_frames']}
-    - Processed Frames: {results['processed_frames']}
-    """
-
-def analyze_and_format(input_video):
-    results = analyze_video(input_video)
-    text_results = format_results(results)
-    
-    # Convert BytesIO to PIL Images
-    analysis_plot = PILImage.open(results['analysis_plot'])
-    progress_bar = PILImage.open(results['progress_bar'])
-    
-    return text_results, analysis_plot, progress_bar
+    return temp_output.name
 
 interface = gr.Interface(
-    fn=analyze_and_format,
+    fn=predict_video,
     inputs=[
         gr.Video(label="Input Video")
     ],
     outputs=[
-        gr.Textbox(label="Analysis Results"),
-        gr.Image(label="Analysis Plots"),
-        gr.Image(label="Fake Percentage")
+        gr.Video(label="Output Video")
     ],
-    title="Video Deepfake Analysis",
-    description="Upload a video to analyze for potential deepfakes.",
-    examples=[]
+    title="Video Deepfake Detection",
+    description="Upload a video to detect deepfakes in each frame."
 )
 
 if __name__ == "__main__":
-    interface.launch(share=True)
+    interface.launch()