Update app.py

app.py

@@ -6,18 +6,21 @@ import cv2
 import requests
 from ultralytics import YOLO
 
-# Remove extra CLI arguments (like "--import") from Spaces.
+# Remove extra CLI arguments that Spaces might pass.
 sys.argv = [arg for arg in sys.argv if arg != "--import"]
 
 # Load the YOLO11-pose model (auto-downloads if needed)
 model = YOLO("yolo11n-pose.pt")
 
-def process_input(uploaded_file, youtube_link, image_url):
+def process_input(uploaded_file, youtube_link, image_url, sensitivity):
     """
-    Process an uploaded file, a YouTube link, or an image URL for pose detection.
-    Priority is: YouTube link > Image URL > Uploaded file.
+    Process input from one of the three methods (Upload, YouTube, Image URL).
+    Priority: YouTube link > Image URL > Uploaded file.
+    The sensitivity slider value is passed as the confidence threshold.
+    
     Returns a tuple:
-      (download_file_path, display_file_path, status_message)
+      (download_file_path, display_file_path, status_message, dummy_state)
+    (The dummy_state is used because Gradio requires the same number of outputs.)
     """
     input_path = None
 
@@ -29,90 +32,98 @@ def process_input(uploaded_file, youtube_link, image_url):
             stream = yt.streams.filter(file_extension='mp4', progressive=True)\
                                .order_by("resolution").desc().first()
             if stream is None:
-                return None, None, "No suitable mp4 stream found."
+                return None, None, "No suitable mp4 stream found.", ""
             input_path = stream.download()
         except Exception as e:
-            return None, None, f"Error downloading video: {e}"
+            return None, None, f"Error downloading video: {e}", ""
     # Priority 2: Image URL
     elif image_url and image_url.strip():
         try:
             response = requests.get(image_url, stream=True)
             if response.status_code != 200:
-                return None, None, f"Error downloading image: HTTP {response.status_code}"
+                return None, None, f"Error downloading image: HTTP {response.status_code}", ""
             temp_image_path = os.path.join(tempfile.gettempdir(), "downloaded_image.jpg")
             with open(temp_image_path, "wb") as f:
                 f.write(response.content)
             input_path = temp_image_path
         except Exception as e:
-            return None, None, f"Error downloading image: {e}"
+            return None, None, f"Error downloading image: {e}", ""
     # Priority 3: Uploaded file
     elif uploaded_file is not None:
         input_path = uploaded_file.name
     else:
-        return None, None, "Please provide a YouTube link, image URL, or upload a file."
+        return None, None, "Please provide an input using one of the methods.", ""
 
-    # Run pose detection (with save=True so annotated outputs are written to disk)
     try:
-        results = model.predict(source=input_path, save=True)
+        # Pass the slider value as the confidence threshold.
+        results = model.predict(source=input_path, save=True, conf=sensitivity)
     except Exception as e:
-        return None, None, f"Error running prediction: {e}"
+        return None, None, f"Error running prediction: {e}", ""
 
     output_path = None
     try:
-        # If the result has a save_path attribute, use it.
         if hasattr(results[0], "save_path"):
             output_path = results[0].save_path
         else:
-            # Otherwise, use plot() to get an annotated image and save it.
             annotated = results[0].plot()  # returns a numpy array
             output_path = os.path.join(tempfile.gettempdir(), "annotated.jpg")
             cv2.imwrite(output_path, annotated)
     except Exception as e:
-        return None, None, f"Error processing the file: {e}"
+        return None, None, f"Error processing the file: {e}", ""
 
-    # Clean up temporary file if it was downloaded (from YouTube or Image URL)
-    if (youtube_link or (image_url and image_url.strip())) and input_path and os.path.exists(input_path):
+    # Clean up the temporary input if it was downloaded.
+    if ((youtube_link and youtube_link.strip()) or (image_url and image_url.strip())) and input_path and os.path.exists(input_path):
         os.remove(input_path)
     
-    return output_path, output_path, "Success!"
+    return output_path, output_path, "Success!", ""
 
-# Define the Gradio interface.
-with gr.Blocks() as demo:
-    # Header image scaled down to 25% using HTML.
-    gr.HTML("<div style='text-align: center;'><img src='crowdresult.jpg' style='width:25%;'></div>")
+# Build the Gradio interface with custom CSS for the result image.
+with gr.Blocks(css="""
+.result_img > img {
+  width: 100%;
+  height: auto;
+  object-fit: contain;
+}
+""") as demo:
+    # Header with scaled image (25% width) and title.
+    gr.HTML("<div style='text-align:center;'><img src='crowdresult.jpg' style='width:25%;'/></div>")
     gr.Markdown("## Pose Detection with YOLO11-pose")
-    gr.Markdown("Choose one of the input methods below. The pose detection will run automatically when you provide an input, and the annotated result will be displayed below.")
-
-    # Prepare output components (they will display the annotated result and a download link)
-    output_file = gr.File(label="Download Annotated Output")
-    output_display = gr.Image(label="Annotated Output")
-    output_text = gr.Textbox(label="Status", interactive=False)
-
-    # Create three tabs for the different input methods.
-    with gr.Tabs():
-        with gr.TabItem("Upload File"):
-            file_input = gr.File(label="Upload Image/Video")
-            # Automatically run detection when a file is uploaded.
-            file_input.change(
-                fn=process_input, 
-                inputs=[file_input, gr.State(None), gr.State(None)],
-                outputs=[output_file, output_display, output_text]
-            )
-        with gr.TabItem("YouTube Link"):
-            youtube_input = gr.Textbox(label="YouTube Link", placeholder="https://...")
-            youtube_input.change(
-                fn=process_input, 
-                inputs=[gr.State(None), youtube_input, gr.State(None)],
-                outputs=[output_file, output_display, output_text]
-            )
-        with gr.TabItem("Image URL"):
-            image_url_input = gr.Textbox(label="Image URL", placeholder="https://...")
-            image_url_input.change(
-                fn=process_input, 
-                inputs=[gr.State(None), gr.State(None), image_url_input],
-                outputs=[output_file, output_display, output_text]
-            )
+    
+    # Create two columns.
+    with gr.Row():
+        # Left column: Input tabs and sensitivity slider.
+        with gr.Column(scale=1):
+            with gr.Tabs():
+                with gr.TabItem("Upload File"):
+                    file_input = gr.File(label="Upload Image/Video")
+                with gr.TabItem("YouTube Link"):
+                    youtube_input = gr.Textbox(label="YouTube Link", placeholder="https://...")
+                with gr.TabItem("Image URL"):
+                    image_url_input = gr.Textbox(label="Image URL", placeholder="https://...")
+            sensitivity_slider = gr.Slider(minimum=0.1, maximum=1.0, step=0.05, value=0.5,
+                                           label="Sensitivity (Confidence Threshold)")
+        # Right column: Display result.
+        with gr.Column(scale=2):
+            output_display = gr.Image(label="Annotated Output", elem_classes="result_img")
+            output_file = gr.File(label="Download Annotated Output")
+            output_text = gr.Textbox(label="Status", interactive=False)
+    
+    # Set up automatic triggers for each input type.
+    file_input.change(
+        fn=process_input,
+        inputs=[file_input, gr.State(""), gr.State(""), sensitivity_slider],
+        outputs=[output_file, output_display, output_text, gr.State()]
+    )
+    youtube_input.change(
+        fn=process_input,
+        inputs=[gr.State(None), youtube_input, gr.State(""), sensitivity_slider],
+        outputs=[output_file, output_display, output_text, gr.State()]
+    )
+    image_url_input.change(
+        fn=process_input,
+        inputs=[gr.State(None), gr.State(""), image_url_input, sensitivity_slider],
+        outputs=[output_file, output_display, output_text, gr.State()]
+    )
 
-# Only launch the interface if the script is executed directly.
 if __name__ == "__main__":
     demo.launch()