Update app.py

app.py

@@ -1,222 +1,196 @@
-import gradio as gr
-import torch
-import torch.nn.functional as F
-from facenet_pytorch import MTCNN, InceptionResnetV1
-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
-from PIL import Image
-import numpy as np
-import warnings
-from typing import Tuple, List, Dict
-import matplotlib.pyplot as plt
-from matplotlib.animation import FuncAnimation
-import io
-
-warnings.filterwarnings("ignore")
-
-# Device configuration
-DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
-
-# Load models
-mtcnn = MTCNN(select_largest=False, post_process=False, device=DEVICE).to(DEVICE).eval()
-model = InceptionResnetV1(pretrained="vggface2", classify=True, num_classes=1, device=DEVICE)
-
-checkpoint = torch.load("df_model.pth", map_location=torch.device('cpu'))
-model.load_state_dict(checkpoint['model_state_dict'])
-model.to(DEVICE)
-model.eval()
-
-def predict_frame(frame: np.ndarray) -> Tuple[str, np.ndarray, Dict[str, float]]:
-    """Predict whether the input frame contains a real or fake face"""
-    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-    frame_pil = Image.fromarray(frame)
-
-    face = mtcnn(frame_pil)
-    if face is None:
-        return None, None, None  # No face detected
-
-    # Preprocess the face
-    face = F.interpolate(face.unsqueeze(0), size=(256, 256), mode='bilinear', align_corners=False)
-    face = face.to(DEVICE, dtype=torch.float32) / 255.0
-
-    # Predict
-    with torch.no_grad():
-        output = torch.sigmoid(model(face).squeeze(0))
-        fake_confidence = output.item()
-        real_confidence = 1 - fake_confidence
-        prediction = "real" if real_confidence > fake_confidence else "fake"
-        
-        confidences = {
-            'real': real_confidence,
-            'fake': fake_confidence
-        }
-
-    # Visualize
-    target_layers = [model.block8.branch1[-1]]
-    cam = GradCAM(model=model, target_layers=target_layers, use_cuda=torch.cuda.is_available())
-    targets = [ClassifierOutputTarget(0)]
-    grayscale_cam = cam(input_tensor=face, targets=targets, eigen_smooth=True)
-    grayscale_cam = grayscale_cam[0, :]
-    face_np = face.squeeze(0).permute(1, 2, 0).cpu().numpy()
-    visualization = show_cam_on_image(face_np, grayscale_cam, use_rgb=True)
-    face_with_mask = cv2.addWeighted((face_np * 255).astype(np.uint8), 1, (visualization * 255).astype(np.uint8), 0.5, 0)
-
-    return prediction, face_with_mask, confidences
-
-def predict_video(input_video: str) -> Tuple[str, np.ndarray, Dict[str, List[float]], List[np.ndarray]]:
-    cap = cv2.VideoCapture(input_video)
-
-    frames = []
-    predictions = []
-    confidences_real = []
-    confidences_fake = []
-    frame_count = 0
-    skip_frames = 20
-
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
-        frame_count += 1
-        if frame_count % skip_frames != 0:
-            continue
-
-        prediction, frame_with_mask, confidence = predict_frame(frame)
-        if prediction is None:
-            continue
-
-        frames.append(frame_with_mask)
-        predictions.append(prediction)
-        confidences_real.append(confidence['real'])
-        confidences_fake.append(confidence['fake'])
-
-    cap.release()
-
-    # Determine the final prediction based on the average confidence
-    avg_real_confidence = sum(confidences_real) / len(confidences_real)
-    avg_fake_confidence = sum(confidences_fake) / len(confidences_fake)
-    final_prediction = 'real' if avg_real_confidence > avg_fake_confidence else 'fake'
-    final_confidence = max(avg_real_confidence, avg_fake_confidence)
-
-    # Create an animated summary plot
-    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 10))
-    
-    def animate(i):
-        ax1.clear()
-        ax2.clear()
-        
-        # Confidence over time
-        ax1.plot(confidences_real[:i+1], label='Real', color='green')
-        ax1.plot(confidences_fake[:i+1], label='Fake', color='red')
-        ax1.set_title('Confidence Scores Over Time')
-        ax1.set_xlabel('Frame')
-        ax1.set_ylabel('Confidence')
-        ax1.legend()
-        ax1.grid(True)
-        ax1.set_ylim(0, 1)
-
-        # Prediction distribution
-        labels, counts = np.unique(predictions[:i+1], return_counts=True)
-        ax2.bar(labels, counts, color=['green', 'red'])
-        ax2.set_title('Distribution of Predictions')
-        ax2.set_xlabel('Prediction')
-        ax2.set_ylabel('Count')
-        ax2.set_ylim(0, len(predictions))
-        
-        plt.tight_layout()
-
-    anim = FuncAnimation(fig, animate, frames=len(confidences_real), repeat=False)
-    
-    # Save the animation as a gif
-    buf = io.BytesIO()
-    anim.save(buf, writer='pillow', fps=5)
-    buf.seek(0)
-    summary_plot = Image.open(buf)
-
-    return final_prediction, final_confidence, summary_plot, {
-        'real': confidences_real,
-        'fake': confidences_fake
-    }, frames
-
-# Custom CSS for a more appealing interface
-custom_css = """
-#output-container {
-    display: flex;
-    justify-content: center;
-    align-items: center;
-    flex-direction: column;
-}
-#confidence-label {
-    font-size: 24px;
-    font-weight: bold;
-    margin-bottom: 10px;
-}
-#confidence-bar {
-    width: 100%;
-    height: 30px;
-    background-color: #f0f0f0;
-    border-radius: 15px;
-    overflow: hidden;
-}
-#confidence-fill {
-    height: 100%;
-    background-color: #4CAF50;
-    transition: width 0.5s ease-in-out;
-}
-"""
-
-# Gradio Interface
-with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🕵️‍♂️ DeepFake Video Detective 🎭")
-    gr.Markdown("Upload a video to determine if it's real or a deepfake. Our AI will analyze it frame by frame!")
-    
-    with gr.Row():
-        input_video = gr.Video(label="📹 Upload Your Video")
-    
-    with gr.Row():
-        submit_btn = gr.Button("🔍 Analyze Video", variant="primary")
-    
-    with gr.Row():
-        with gr.Column():
-            output_label = gr.Label(label="🏷️ Prediction")
-            confidence_output = gr.HTML(
-                """
-                <div id="output-container">
-                    <div id="confidence-label">Confidence: 0%</div>
-                    <div id="confidence-bar">
-                        <div id="confidence-fill" style="width: 0%;"></div>
-                    </div>
-                </div>
-                """
-            )
-        summary_plot = gr.Image(label="📊 Analysis Summary")
-    
-    with gr.Row():
-        output_video = gr.Video(label="🎞️ Processed Video")
-    
-    def update_confidence(prediction, confidence):
-        color = "#4CAF50" if prediction == "real" else "#FF5722"
-        return f"""
-        <div id="output-container">
-            <div id="confidence-label">Confidence: {confidence:.2%}</div>
-            <div id="confidence-bar">
-                <div id="confidence-fill" style="width: {confidence:.2%}; background-color: {color};"></div>
-            </div>
-        </div>
-        """
-    
-    def process_video(video):
-        prediction, confidence, summary, _, frames = predict_video(video)
-        processed_video = np.stack(frames, axis=0)
-        confidence_html = update_confidence(prediction, confidence)
-        return {output_label: prediction, confidence_output: confidence_html, summary_plot: summary, output_video: processed_video}
-    
-    submit_btn.click(
-        process_video,
-        inputs=[input_video],
-        outputs=[output_label, confidence_output, summary_plot, output_video]
-    )
-
+import gradio as gr
+import torch
+import torch.nn.functional as F
+from facenet_pytorch import MTCNN, InceptionResnetV1
+import cv2
+from PIL import Image
+import numpy as np
+import warnings
+from typing import Tuple, Dict
+import matplotlib.pyplot as plt
+import io
+
+warnings.filterwarnings("ignore")
+
+# Device configuration
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+# Load models
+mtcnn = MTCNN(select_largest=False, post_process=False, device=DEVICE).to(DEVICE).eval()
+model = InceptionResnetV1(pretrained="vggface2", classify=True, num_classes=1, device=DEVICE)
+
+checkpoint = torch.load("df_model.pth", map_location=torch.device('cpu'))
+model.load_state_dict(checkpoint['model_state_dict'])
+model.to(DEVICE)
+model.eval()
+
+def predict_frame(frame: np.ndarray) -> Tuple[str, Dict[str, float]]:
+    """Predict whether the input frame contains a real or fake face"""
+    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    frame_pil = Image.fromarray(frame)
+
+    face = mtcnn(frame_pil)
+    if face is None:
+        return None, None  # No face detected
+
+    # Preprocess the face
+    face = F.interpolate(face.unsqueeze(0), size=(256, 256), mode='bilinear', align_corners=False)
+    face = face.to(DEVICE, dtype=torch.float32) / 255.0
+
+    # Predict
+    with torch.no_grad():
+        output = torch.sigmoid(model(face).squeeze(0))
+        fake_confidence = output.item()
+        real_confidence = 1 - fake_confidence
+        prediction = "real" if real_confidence > fake_confidence else "fake"
+        
+        confidences = {
+            'real': real_confidence,
+            'fake': fake_confidence
+        }
+
+    return prediction, confidences
+
+def predict_video(input_video: str) -> Tuple[str, float, np.ndarray]:
+    cap = cv2.VideoCapture(input_video)
+
+    predictions = []
+    confidences_real = []
+    confidences_fake = []
+    frame_count = 0
+    skip_frames = 5  # Analyze every 5th frame for faster processing
+
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frame_count += 1
+        if frame_count % skip_frames != 0:
+            continue
+
+        prediction, confidence = predict_frame(frame)
+        if prediction is None:
+            continue
+
+        predictions.append(prediction)
+        confidences_real.append(confidence['real'])
+        confidences_fake.append(confidence['fake'])
+
+    cap.release()
+
+    # Determine the final prediction based on the average confidence
+    avg_real_confidence = sum(confidences_real) / len(confidences_real)
+    avg_fake_confidence = sum(confidences_fake) / len(confidences_fake)
+    final_prediction = 'real' if avg_real_confidence > avg_fake_confidence else 'fake'
+    final_confidence = max(avg_real_confidence, avg_fake_confidence)
+
+    # Create a summary plot
+    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8, 8))
+    
+    # Confidence over time
+    ax1.plot(confidences_real, label='Real', color='green')
+    ax1.plot(confidences_fake, label='Fake', color='red')
+    ax1.set_title('Confidence Scores Over Time')
+    ax1.set_xlabel('Frame')
+    ax1.set_ylabel('Confidence')
+    ax1.legend()
+    ax1.grid(True)
+
+    # Prediction distribution
+    labels, counts = np.unique(predictions, return_counts=True)
+    ax2.bar(labels, counts, color=['green', 'red'])
+    ax2.set_title('Distribution of Predictions')
+    ax2.set_xlabel('Prediction')
+    ax2.set_ylabel('Count')
+    
+    plt.tight_layout()
+    
+    # Save the plot as an image
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    buf.seek(0)
+    summary_plot = Image.open(buf)
+
+    return final_prediction, final_confidence, summary_plot
+
+# Custom CSS for a more appealing interface
+custom_css = """
+.video-container {
+    max-width: 400px;
+    margin: 0 auto;
+}
+#output-container {
+    display: flex;
+    justify-content: center;
+    align-items: center;
+    flex-direction: column;
+}
+#confidence-label {
+    font-size: 24px;
+    font-weight: bold;
+    margin-bottom: 10px;
+}
+#confidence-bar {
+    width: 100%;
+    height: 30px;
+    background-color: #f0f0f0;
+    border-radius: 15px;
+    overflow: hidden;
+}
+#confidence-fill {
+    height: 100%;
+    background-color: #4CAF50;
+    transition: width 0.5s ease-in-out;
+}
+"""
+
+# Gradio Interface
+with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🕵️‍♂️ DeepFake Video Detective 🎭")
+    gr.Markdown("Upload a video to determine if it's real or a deepfake. Our AI will analyze it frame by frame!")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_video = gr.Video(label="📹 Upload Your Video", elem_classes=["video-container"])
+    
+    with gr.Row():
+        submit_btn = gr.Button("🔍 Analyze Video", variant="primary")
+    
+    with gr.Row():
+        with gr.Column():
+            output_label = gr.Label(label="🏷️ Prediction")
+            confidence_output = gr.HTML(
+                """
+                <div id="output-container">
+                    <div id="confidence-label">Confidence: 0%</div>
+                    <div id="confidence-bar">
+                        <div id="confidence-fill" style="width: 0%;"></div>
+                    </div>
+                </div>
+                """
+            )
+        summary_plot = gr.Image(label="📊 Analysis Summary")
+    
+    def update_confidence(prediction, confidence):
+        color = "#4CAF50" if prediction == "real" else "#FF5722"
+        return f"""
+        <div id="output-container">
+            <div id="confidence-label">Confidence: {confidence:.2%}</div>
+            <div id="confidence-bar">
+                <div id="confidence-fill" style="width: {confidence:.2%}; background-color: {color};"></div>
+            </div>
+        </div>
+        """
+    
+    def process_video(video):
+        prediction, confidence, summary = predict_video(video)
+        confidence_html = update_confidence(prediction, confidence)
+        return {output_label: prediction, confidence_output: confidence_html, summary_plot: summary}
+    
+    submit_btn.click(
+        process_video,
+        inputs=[input_video],
+        outputs=[output_label, confidence_output, summary_plot]
+    )
+
 demo.launch()