Added error handling for video loading, added output flagging

app.py

@@ -10,6 +10,8 @@ theme = gr.themes.Default(
     font=[gr.themes.GoogleFont("IBM Plex Mono"), "system-ui"]
 )
 
+callback = gr.CSVLogger()
+
 with gr.Blocks(theme=theme) as demo:
     # DEFINE COMPONENTS
 
@@ -49,11 +51,20 @@ with gr.Blocks(theme=theme) as demo:
         visible=False
     )
 
+    # Button for flagging the output
+    flagBtn = gr.Button(value="Flag Output", visible=False)
+
     # DEFINE FUNCTIONS
     # Load video from URL, display sample frames, and enable prediction button
     loadVideoBtn.click(fn=load_video_from_url, inputs=[urlInput], outputs=[videoTitle, sampleFrames, predVideoBtn, predOutput])
 
     # Generate video prediction
-    predVideoBtn.click(fn=detect_deepfake, outputs=[predOutput])
+    predVideoBtn.click(fn=detect_deepfake, outputs=[predOutput, flagBtn])
+
+    # Define flag callback
+    callback.setup([urlInput], "flagged_data_points")
+
+    # Flag output
+    flagBtn.click(fn=lambda *args: callback.flag(args), inputs=[urlInput], outputs=None)
 
 demo.launch()

helpers.py

@@ -9,40 +9,47 @@ import pickle
 
 
 def load_video_from_url(youtube_url):
-    # DOWNLOAD THE VIDEO USING THE GIVEN URL
-    yt = YouTube(youtube_url)
-    yt_stream = yt.streams.filter(file_extension='mp4').first()
-    title = yt_stream.title
-    src = yt_stream.download()
-    capture = cv2.VideoCapture(src)
-
-    # SAMPLE FRAMES FROM VIDEO FILE
-    sampled_frames = sample_frames_from_video_file(capture)
-
-    # PICK EXAMPLE FRAME FROM THE MIDDLE OF THE SAMPLED FRAMES
-    example_frames = [
-        sampled_frames[len(sampled_frames) // 4],
-        sampled_frames[len(sampled_frames) // 2],
-        sampled_frames[3 * len(sampled_frames) // 4],
-    ]
-
-    # DELETE VIDEO FILE
-    if os.path.exists(src):
-        os.remove(src)
-
-    # CONVERT SAMPLED FRAMES TO TENSOR
-    frames_tensor = tf.expand_dims(tf.convert_to_tensor(sampled_frames, dtype=tf.float32), axis=0)
-
-    # SAVE TENSOR TO FILE
-    pickle.dump(frames_tensor, open("frames_tf.pkl", "wb"))
+    visible = True
+    try:
+        # DOWNLOAD THE VIDEO USING THE GIVEN URL
+        yt = YouTube(youtube_url)
+        yt_stream = yt.streams.filter(file_extension='mp4').first()
+        title = yt_stream.title
+        src = yt_stream.download()
+        capture = cv2.VideoCapture(src)
+
+        # SAMPLE FRAMES FROM VIDEO FILE
+        sampled_frames = sample_frames_from_video_file(capture)
+
+        # PICK EXAMPLE FRAME FROM THE MIDDLE OF THE SAMPLED FRAMES
+        example_frames = [
+            sampled_frames[len(sampled_frames) // 4],
+            sampled_frames[len(sampled_frames) // 2],
+            sampled_frames[3 * len(sampled_frames) // 4],
+        ]
+
+        # DELETE VIDEO FILE
+        if os.path.exists(src):
+            os.remove(src)
+
+        # CONVERT SAMPLED FRAMES TO TENSOR
+        frames_tensor = tf.expand_dims(tf.convert_to_tensor(sampled_frames, dtype=tf.float32), axis=0)
+
+        # SAVE TENSOR TO FILE
+        pickle.dump(frames_tensor, open("frames_tf.pkl", "wb"))
+
+    except Exception as e:
+        title = "Error while loading video: " + str(e)
+        visible = False
+        example_frames = [np.zeros((256, 256, 3)) for _ in range(3)]
 
     # Define visible prediction components to show upon video loaded
-    predVideoBtn = gr.Button(value="Classify Video", visible=True)
+    predVideoBtn = gr.Button(value="Classify Video", visible=visible)
 
     predOutput = gr.Label(
         label="DETECTED LABEL (AND CONFIDENCE LEVEL)",
         num_top_classes=2,
-        visible=True
+        visible=visible
     )
 
     return title, example_frames, predVideoBtn, predOutput
@@ -66,8 +73,11 @@ def detect_deepfake():
     fake_confidence = 1 - real_confidence
     confidence_dict = {"FAKE": fake_confidence, "REAL": real_confidence}
 
+    # MAKE FLAG BUTTON VISIBLE
+    flagBtn = gr.Button(value="Flag Output", visible=True)
+
     # RETURN THE OUTPUT LABEL AND EXAMPLE FRAMES
-    return confidence_dict
+    return confidence_dict, flagBtn
 
 
 def sample_frames_from_video_file(capture, sample_count=10, frames_per_sample=10, frame_step=10,