Update app.py

app.py

@@ -3,29 +3,28 @@ import gradio as gr
 import os
 import tempfile
 import cv2
+import requests
 from ultralytics import YOLO
 
-# Optionally remove extra CLI arguments that Spaces might pass
+# Remove extra CLI arguments (like "--import") from Spaces.
 sys.argv = [arg for arg in sys.argv if arg != "--import"]
 
-# Load the YOLO11-pose model (it will auto-download if not present)
+# Load the YOLO11-pose model (will auto-download if needed)
 model = YOLO("yolo11n-pose.pt")
 
-def process_input(uploaded_file, youtube_link):
+def process_input(uploaded_file, youtube_link, image_url):
     """
-    Process an uploaded file or a YouTube link to perform pose detection.
-    Returns a tuple: (annotated_file_path, status_message). 
-    If an error occurs, annotated_file_path is None and status_message describes the error.
+    Process an uploaded file, a YouTube link, or an image URL for pose detection.
+    Returns a tuple: (annotated_file_path, status_message).  
+    Priority is given to the YouTube link first, then image URL, then the uploaded file.
     """
-    error_message = ""
     input_path = None
 
-    # Check for input: either a YouTube link or an uploaded file.
+    # Priority 1: YouTube link
     if youtube_link and youtube_link.strip():
         try:
             from pytube import YouTube
             yt = YouTube(youtube_link)
-            # Get the highest resolution progressive mp4 stream
             stream = yt.streams.filter(file_extension='mp4', progressive=True)\
                                .order_by("resolution").desc().first()
             if stream is None:
@@ -33,54 +32,70 @@ def process_input(uploaded_file, youtube_link):
             input_path = stream.download()
         except Exception as e:
             return 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, f"Error downloading image: HTTP {response.status_code}"
+            # Save the downloaded image to a temporary file.
+            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, f"Error downloading image: {e}"
+    # Priority 3: Uploaded file
     elif uploaded_file is not None:
         input_path = uploaded_file.name
     else:
-        return None, "Please provide an uploaded file or a YouTube link."
+        return None, "Please provide a YouTube link, image URL, or upload a file."
 
-    # Run pose detection (with save=True so that outputs are written to disk)
+    # Run pose detection (with save=True so annotated results are written to disk)
     try:
         results = model.predict(source=input_path, save=True)
     except Exception as e:
         return None, f"Error running prediction: {e}"
 
-    # Try to get the annotated output file:
     output_path = None
     try:
-        # Some YOLO versions may offer a 'save_path' attribute.
+        # If the results object has a save_path attribute, use it.
         if hasattr(results[0], "save_path"):
             output_path = results[0].save_path
         else:
-            # Fallback: generate the annotated image using result.plot()
-            annotated = results[0].plot()  # returns a numpy array with annotations
-            # Save the annotated image to a temporary file.
+            # Otherwise, generate an annotated image using plot() and save it manually.
+            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, f"Error processing the file: {e}"
 
-    # Clean up the downloaded video file if it came from YouTube.
-    if youtube_link and input_path and os.path.exists(input_path):
+    # If the input came from YouTube or image URL, remove the temporary file.
+    if (youtube_link or (image_url and image_url.strip())) and input_path and os.path.exists(input_path):
         os.remove(input_path)
-    
+
     return output_path, "Success!"
 
-# Define the Gradio interface with two outputs: one for the file and one for status text.
+# Define the Gradio interface.
 with gr.Blocks() as demo:
+    # Display the default image (crowdresult.jpg) at the top.
     gr.Markdown("# Pose Detection with YOLO11-pose")
-    gr.Markdown("Upload an image/video or provide a YouTube link to detect human poses.")
-    
+    gr.Image(value="crowdresult.jpg", label="Crowd Result", interactive=False)
+    gr.Markdown("Upload an image/video, provide an image URL, or supply a YouTube link to detect human poses.")
+
     with gr.Row():
         file_input = gr.File(label="Upload Image/Video")
-        youtube_input = gr.Textbox(label="Or enter a YouTube link", placeholder="https://...")
-    
+    with gr.Row():
+        youtube_input = gr.Textbox(label="YouTube Link", placeholder="https://...")
+        image_url_input = gr.Textbox(label="Image URL", placeholder="https://...")
+
     output_file = gr.File(label="Download Annotated Output")
     output_text = gr.Textbox(label="Status", interactive=False)
     run_button = gr.Button("Run Pose Detection")
-    
-    run_button.click(process_input, inputs=[file_input, youtube_input],
+
+    run_button.click(process_input, inputs=[file_input, youtube_input, image_url_input],
                      outputs=[output_file, output_text])
 
-# Only launch the app if the script is executed directly.
+# Only launch the interface if executed directly.
 if __name__ == "__main__":
     demo.launch()