Update app.py

app.py

@@ -12,7 +12,7 @@ def check_ffmpeg():
     if shutil.which("ffmpeg") is None:
         raise EnvironmentError("ffmpeg is not installed or not found in the system PATH. Please install ffmpeg to proceed.")
 
-def process_video(video_path):
+def process_video(video_path, output_option):
     # Check if ffmpeg is installed
     check_ffmpeg()
 
@@ -46,17 +46,26 @@ def process_video(video_path):
     if not matched_mp4_files:
         raise Exception(f"No .mp4 video files found in directory {output_dir} after conversion for {video_name}.")
 
-    # Return the first video file found (should correspond to the latest processed video)
-    return matched_mp4_files[0]
+    # Count the number of objects detected
+    object_count = len(results.pandas().xyxy[0])
+
+    # Determine the output based on user's choice
+    if output_option == "Count":
+        return object_count
+    elif output_option == "Video":
+        return matched_mp4_files[0]
+    elif output_option == "Both":
+        return object_count, matched_mp4_files[0]
 
 # Define Gradio inputs
 video_input = gr.Video()
+output_option_input = gr.Radio(choices=["Count", "Video", "Both"], label="Select output type")
 
 video_interface = gr.Interface(
     fn=process_video,
-    inputs=video_input,
-    outputs="video",  # Display the video output
-    description="YOLO Video Tracking"
+    inputs=[video_input, output_option_input],
+    outputs=[gr.Textbox(label="Object Count"), "video"],  # Update outputs to support multiple types
+    description="YOLO Video Tracking with Output Options"
 )
 
 # Deploy the interface with share enabled