Update app.py

app.py

@@ -1,10 +1,12 @@
 import gradio as gr
 import cv2
 import numpy as np
-import pickle
+import tensorflow as tf
 import os
-from PIL import Image
 
+# Load the trained models using Keras
+cnn_model = tf.keras.models.load_model('cnn_model.h5')
+qcnn_model = tf.keras.models.load_model('qcnn_model.h5')
 
 # Directories containing example videos
 examples_dir = 'examples'
@@ -13,16 +15,25 @@ deepfake_roop_dir = os.path.join(examples_dir, 'DeepfakeRoop')
 deepfake_web_dir = os.path.join(examples_dir, 'DeepfakeWeb')
 
 # Function to get video paths from a directory
-def get_video_paths(directory):
-    return [os.path.join(directory, vid) for vid in os.listdir(directory) if vid.endswith('.mp4')]
+def get_video_paths(directory, label):
+    videos = []
+    for vid in os.listdir(directory):
+        if vid.endswith('.mp4'):
+            videos.append({'path': os.path.join(directory, vid), 'label': label})
+    return videos
 
 # Get video paths for each category
-original_videos = get_video_paths(original_dir)
-deepfake_roop_videos = get_video_paths(deepfake_roop_dir)
-deepfake_web_videos = get_video_paths(deepfake_web_dir)
+original_videos = get_video_paths(original_dir, 'Original')
+deepfake_roop_videos = get_video_paths(deepfake_roop_dir, 'DeepfakeRoop')
+deepfake_web_videos = get_video_paths(deepfake_web_dir, 'DeepfakeWeb')
 
-# Function to process video and calculate accuracies
-def process_video(video_path, true_label):
+# Combine all examples
+examples = original_videos + deepfake_roop_videos + deepfake_web_videos
+
+# Map from example video path to label
+example_videos_dict = {example['path']: example['label'] for example in examples}
+
+def process_video(video_path, true_label=None):
     cap = cv2.VideoCapture(video_path)
     fps = cap.get(cv2.CAP_PROP_FPS)
     if fps == 0 or np.isnan(fps):
@@ -46,109 +57,89 @@ def process_video(video_path, true_label):
     
     cnn_correct = 0
     qcnn_correct = 0
+    cnn_class0 = 0
+    cnn_class1 = 0
+    qcnn_class0 = 0
+    qcnn_class1 = 0
     total_frames = len(sampled_frames)
     for frame in sampled_frames:
         cnn_pred = cnn_model.predict(frame)
         cnn_label = np.argmax(cnn_pred)
-        if cnn_label == true_label:
+        if cnn_label == 0:
+            cnn_class0 += 1
+        else:
+            cnn_class1 += 1
+        if true_label is not None and cnn_label == true_label:
             cnn_correct += 1
+        
         qcnn_pred = qcnn_model.predict(frame)
         qcnn_label = np.argmax(qcnn_pred)
-        if qcnn_label == true_label:
+        if qcnn_label == 0:
+            qcnn_class0 += 1
+        else:
+            qcnn_class1 += 1
+        if true_label is not None and qcnn_label == true_label:
             qcnn_correct += 1
     
-    cnn_accuracy = (cnn_correct / total_frames) * 100 if total_frames > 0 else 0
-    qcnn_accuracy = (qcnn_correct / total_frames) * 100 if total_frames > 0 else 0
-    result_label = "Original" if true_label == 0 else "Deepfake"
+    if total_frames > 0:
+        cnn_class0_percent = (cnn_class0 / total_frames) * 100
+        cnn_class1_percent = (cnn_class1 / total_frames) * 100
+        qcnn_class0_percent = (qcnn_class0 / total_frames) * 100
+        qcnn_class1_percent = (qcnn_class1 / total_frames) * 100
+    else:
+        cnn_class0_percent = cnn_class1_percent = qcnn_class0_percent = qcnn_class1_percent = 0
     
-    result = f"Result: {result_label}\n"
-    result += f"CNN Model Accuracy: {cnn_accuracy:.2f}%\n"
-    result += f"QCNN Model Accuracy: {qcnn_accuracy:.2f}%"
+    if true_label is not None:
+        # Calculate accuracy if true_label is provided (example video)
+        cnn_accuracy = (cnn_correct / total_frames) * 100 if total_frames > 0 else 0
+        qcnn_accuracy = (qcnn_correct / total_frames) * 100 if total_frames > 0 else 0
+        result = f"CNN Model Accuracy: {cnn_accuracy:.2f}%\n"
+        result += f"QCNN Model Accuracy: {qcnn_accuracy:.2f}%"
+    else:
+        # Display percent of frames classified from each class
+        result = f"CNN Model Predictions:\nClass 0: {cnn_class0_percent:.2f}%\nClass 1: {cnn_class1_percent:.2f}%\n"
+        result += f"QCNN Model Predictions:\nClass 0: {qcnn_class0_percent:.2f}%\nClass 1: {qcnn_class1_percent:.2f}%"
     return result
 
-# Prediction functions for each step
-def predict_step1(video_file):
-    if video_file is None:
-        return "Please select a video."
-    return process_video(video_file, true_label=0)
-
-def predict_step2(video_file):
-    if video_file is None:
-        return "Please select a video."
-    return process_video(video_file, true_label=1)
-
-def predict_step3(video_file):
-    if video_file is None:
-        return "Please select a video."
-    return process_video(video_file, true_label=1)
+def predict(video_input):
+    if video_input is None:
+        return "Please upload a video or select an example."
+    if isinstance(video_input, dict):
+        video_path = video_input['name']
+    elif isinstance(video_input, str):
+        video_path = video_input
+    else:
+        return "Invalid video input."
+    
+    # Check if video is an example
+    if video_path in example_videos_dict:
+        label = example_videos_dict[video_path]
+        if label == 'Original':
+            true_label = 0
+        else:
+            true_label = 1
+        result = process_video(video_path, true_label=true_label)
+        result = f"Example Video Detected ({label})\n" + result
+    else:
+        result = process_video(video_path)
+    return result
 
-# Create Gradio interface
 with gr.Blocks() as demo:
-    # Centered title
-    gr.HTML("<h1 style='text-align: center;'>Quanvolutional Neural Networks for Deepfake Detection</h1>")
-    gr.HTML("<h2 style='text-align: center;'>Steven Fernandes, Ph.D.</h2>")
-
-    # Step 1: Original Videos
-    gr.Markdown("## Step 1: Select an Original Video")
-    with gr.Row():
-        with gr.Column():
-            original_video = gr.Dropdown(
-                choices=original_videos,
-                label="Select Original Video",
-                interactive=True
-            )
-            predict_button1 = gr.Button("Predict")
-        with gr.Column():
-            output1 = gr.Textbox(label="Result")
-    predict_button1.click(fn=predict_step1, inputs=original_video, outputs=output1)
-
-    # Step 2: Deepfake Videos from DeepfakeWeb
-    gr.Markdown("## Step 2: Select a Deepfake Video Generated Using DeepfakeWeb")
-    with gr.Row():
-        with gr.Column():
-            deepfake_web_video = gr.Dropdown(
-                choices=deepfake_web_videos,
-                label="Select DeepfakeWeb Video",
-                interactive=True
-            )
-            predict_button2 = gr.Button("Predict")
-        with gr.Column():
-            output2 = gr.Textbox(label="Result")
-    predict_button2.click(fn=predict_step2, inputs=deepfake_web_video, outputs=output2)
-
-    # Step 3: Deepfake Videos from Roop Method
-    gr.Markdown("## Step 3: Select a Deepfake Video Generated Using the Roop Method")
+    gr.Markdown("<h1 style='text-align: center;'>Quanvolutional Neural Networks for Deepfake Detection</h1>")
+    gr.Markdown("<h2 style='text-align: center;'>Steven Fernandes, Ph.D.</h2>")
+    
     with gr.Row():
         with gr.Column():
-            deepfake_roop_video = gr.Dropdown(
-                choices=deepfake_roop_videos,
-                label="Select DeepfakeRoop Video",
-                interactive=True
+            video_input = gr.Video(label="Upload Video or Select an Example", type="filepath")
+            gr.Examples(
+                examples=[example['path'] for example in examples],
+                inputs=video_input,
+                label="Examples"
             )
-            predict_button3 = gr.Button("Predict")
+            predict_button = gr.Button("Predict")
         with gr.Column():
-            output3 = gr.Textbox(label="Result")
-    predict_button3.click(fn=predict_step3, inputs=deepfake_roop_video, outputs=output3)
-
-    # Examples section mimicking the design of the provided application
-    gr.Markdown("## Examples")
-    gr.Markdown("### Original Videos")
-    gr.Examples(
-        examples=original_videos,
-        inputs=original_video,
-        label="Original Video Examples"
-    )
-    gr.Markdown("### DeepfakeWeb Videos")
-    gr.Examples(
-        examples=deepfake_web_videos,
-        inputs=deepfake_web_video,
-        label="DeepfakeWeb Video Examples"
-    )
-    gr.Markdown("### DeepfakeRoop Videos")
-    gr.Examples(
-        examples=deepfake_roop_videos,
-        inputs=deepfake_roop_video,
-        label="DeepfakeRoop Video Examples"
-    )
-
-demo.launch()
+            output = gr.Textbox(label="Result")
+    
+    predict_button.click(fn=predict, inputs=video_input, outputs=output)
+    
+    demo.launch()