Update app.py

app.py

@@ -57,38 +57,19 @@ def process_video(video_path, target, progress=gr.Progress()):
         return None, f"Error: Unable to open video file at {video_path}"
 
     frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-    original_fps = int(cap.get(cv2.CAP_PROP_FPS))
-    output_fps = 3
-
-    output_path = "output_video.mp4"
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(output_path, fourcc, output_fps, (int(cap.get(3)), int(cap.get(4))))
-
-    batch_size = 8
-    frames = []
+    processed_frames = []
 
     for frame in progress.tqdm(range(frame_count)):
         ret, img = cap.read()
         if not ret:
             break
 
-        if frame % (original_fps // output_fps) != 0:
-            continue
-
         pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
-        frames.append(pil_img)
-
-        if len(frames) == batch_size or frame == frame_count - 1:
-            annotated_frames = [detect_objects_in_frame(frame, target) for frame in frames]
-            for annotated_img in annotated_frames:
-                annotated_frame = cv2.cvtColor(np.array(annotated_img), cv2.COLOR_RGB2BGR)
-                out.write(annotated_frame)
-            frames = []
+        annotated_img = detect_objects_in_frame(pil_img, target)
+        processed_frames.append(np.array(annotated_img))
 
     cap.release()
-    out.release()
-
-    return output_path, None
+    return processed_frames, None
 
 def load_sample_frame(video_path):
     cap = cv2.VideoCapture(video_path)
@@ -107,31 +88,47 @@ def gradio_app():
 
         video_input = gr.Video(label="Upload Video")
         target_input = gr.Textbox(label="Target Object", value="Elephant")
-        output_video = gr.Video(label="Output Video")
+        frame_slider = gr.Slider(minimum=0, maximum=100, step=1, label="Frame", value=0)
+        output_image = gr.Image(label="Processed Frame")
         error_output = gr.Textbox(label="Error Messages", visible=False)
         sample_video_frame = gr.Image(value=load_sample_frame("Drone Video of African Wildlife Wild Botswan.mp4"), label="Sample Video Frame")
         use_sample_button = gr.Button("Use Sample Video")
         progress_bar = gr.Progress()
 
+        processed_frames = gr.State([])
+
         def process_and_update(video, target):
-            output_video_path, error = process_video(video, target, progress_bar)
-            return output_video_path, error
+            frames, error = process_video(video, target, progress_bar)
+            if frames is not None:
+                frame_slider.maximum = len(frames) - 1
+                frame_slider.value = 0
+                return frames, frames[0], error, gr.Slider.update(maximum=len(frames) - 1, value=0)
+            return None, None, error, gr.Slider.update(maximum=100, value=0)
+
+        def update_frame(frame_index, frames):
+            if frames and 0 <= frame_index < len(frames):
+                return frames[frame_index]
+            return None
 
         video_input.upload(process_and_update, 
                            inputs=[video_input, target_input], 
-                           outputs=[output_video, error_output])
+                           outputs=[processed_frames, output_image, error_output, frame_slider])
+
+        frame_slider.change(update_frame, 
+                            inputs=[frame_slider, processed_frames], 
+                            outputs=[output_image])
 
         def use_sample_video():
             sample_video_path = "Drone Video of African Wildlife Wild Botswan.mp4"
-            output_video_path, error = process_and_update(sample_video_path, "Elephant")
-            return output_video_path, error
+            frames, output_image, error, slider_update = process_and_update(sample_video_path, "Elephant")
+            return frames, output_image, error, slider_update
 
         use_sample_button.click(use_sample_video, 
                                 inputs=None, 
-                                outputs=[output_video, error_output])
+                                outputs=[processed_frames, output_image, error_output, frame_slider])
 
     return app
 
 if __name__ == "__main__":
     app = gradio_app()
-    app.launch(share=True)
+    app.launch(share=True)