Update app.py

app.py

@@ -14,12 +14,15 @@ model = YOLO("yolo11n-pose.pt")
 
 def process_input(uploaded_file, youtube_link, image_url, sensitivity):
     """
-    Process input from one of the three methods (Upload, YouTube, Image URL).
+    Process input from one of 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, dummy_state)
+    Returns a tuple of 4 items:
+      1. download_file_path (for gr.File)
+      2. image_result (for gr.Image) or None
+      3. video_result (for gr.Video) or None
+      4. status message
     """
     input_path = None
 
@@ -31,9 +34,112 @@ def process_input(uploaded_file, youtube_link, image_url, sensitivity):
             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, None, "No suitable mp4 stream found."
             input_path = stream.download()
         except Exception as e:
+            return None, 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, 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, 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, None, "Please provide an input using one of the methods."
+
+    try:
+        # Run prediction; pass slider value as confidence threshold.
+        results = model.predict(source=input_path, save=True, conf=sensitivity)
+    except Exception as e:
+        return None, None, None, f"Error running prediction: {e}"
+
+    output_path = None
+    try:
+        if hasattr(results[0], "save_path"):
+            output_path = results[0].save_path
+        else:
+            # If no save_path, generate annotated image using plot()
+            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, None, f"Error processing the file: {e}"
+
+    # Clean up 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)
+
+    # Determine if output is video or image based on extension.
+    ext = os.path.splitext(output_path)[1].lower()
+    video_exts = [".mp4", ".mov", ".avi", ".webm"]
+    if ext in video_exts:
+        image_result = None
+        video_result = output_path
+    else:
+        image_result = output_path
+        video_result = None
+
+    return output_path, image_result, video_result, "Success!"
+
+# Build the Gradio interface.
+with gr.Blocks(css="""
+.result_img > img {
+  width: 100%;
+  height: auto;
+  object-fit: contain;
+}
+""") as demo:
+    # Layout: two columns in a row.
+    with gr.Row():
+        # Left column: Header image, title, input method tabs, and shared sensitivity slider.
+        with gr.Column(scale=1):
+            gr.HTML("<div style='text-align:center;'><img src='https://huggingface.co/spaces/tstone87/stance-detection/resolve/main/crowdresult.jpg' style='width:25%;'/></div>")
+            gr.Markdown("## Pose Detection with YOLO11-pose")
+            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: Results displayed at the top.
+        with gr.Column(scale=2):
+            output_image = gr.Image(label="Annotated Output (Image)", elem_classes="result_img")
+            output_video = gr.Video(label="Annotated Output (Video)")
+            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_image, output_video, output_text]
+    )
+    youtube_input.change(
+        fn=process_input,
+        inputs=[gr.State(None), youtube_input, gr.State(""), sensitivity_slider],
+        outputs=[output_file, output_image, output_video, output_text]
+    )
+    image_url_input.change(
+        fn=process_input,
+        inputs=[gr.State(None), gr.State(""), image_url_input, sensitivity_slider],
+        outputs=[output_file, output_image, output_video, output_text]
+    )
+
+if __name__ == "__main__":
+    demo.launch()
             return None, None, f"Error downloading video: {e}", ""
     # Priority 2: Image URL
     elif image_url and image_url.strip():