Update app.py

app.py

@@ -12,6 +12,9 @@ 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(".")
@@ -75,59 +78,108 @@ def process_frame(frame, mtcnn, model, cam, targets):
     return visualization, prediction, confidence
 
 @spaces.GPU
-def predict_video(input_video: str):
-    """Predict the labels for each frame of the input video"""
+def analyze_video(input_video: str):
+    """Analyze the video for deepfake detection"""
     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))
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
     
     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))
+    frame_confidences = []
+    frame_predictions = []
     
-    while cap.isOpened():
+    for _ in tqdm(range(total_frames), desc="Analyzing video"):
         ret, frame = cap.read()
         if not ret:
             break
         
-        processed_frame, prediction, confidence = process_frame(frame, mtcnn, model, cam, targets)
+        _, prediction, confidence = process_frame(frame, mtcnn, model, cam, targets)
         
-        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
-            if prediction is not None and confidence is not None:
-                text = f"{prediction}: {confidence:.2f}"
-            else:
-                text = "No prediction available"
-            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)
+        if prediction is not None and confidence is not None:
+            frame_confidences.append(confidence)
+            frame_predictions.append(1 if prediction == "fake" else 0)
     
     cap.release()
-    out.release()
     
-    return temp_output.name
+    # Calculate metrics
+    fake_percentage = (sum(frame_predictions) / len(frame_predictions)) * 100
+    avg_confidence = np.mean(frame_confidences)
+    
+    # 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)
+    return text_results, results['analysis_plot'], results['progress_bar']
 
 interface = gr.Interface(
-    fn=predict_video,
+    fn=analyze_and_format,
     inputs=[
         gr.Video(label="Input Video")
     ],
     outputs=[
-        gr.Video(label="Output Video")
+        gr.Textbox(label="Analysis Results"),
+        gr.Image(label="Analysis Plots"),
+        gr.Image(label="Fake Percentage")
     ],
-    title="Video Deepfake Detection",
-    description="Upload a video to detect deepfakes in each frame."
+    title="Video Deepfake Analysis",
+    description="Upload a video to analyze for potential deepfakes.",
+    examples=[],
+    interpretation="default"
 )
 
 if __name__ == "__main__":
-    interface.launch()
+    interface.launch(share=True)