Update app.py

app.py

@@ -1,58 +1,86 @@
+import sys
 import gradio as gr
 import os
+import tempfile
+import cv2
 from ultralytics import YOLO
 
-# Load the YOLO11-pose model (auto-downloads if not present)
+# Optionally remove extra CLI arguments that Spaces might pass
+sys.argv = [arg for arg in sys.argv if arg != "--import"]
+
+# Load the YOLO11-pose model (it will auto-download if not present)
 model = YOLO("yolo11n-pose.pt")
 
 def process_input(uploaded_file, youtube_link):
     """
     Process an uploaded file or a YouTube link to perform pose detection.
-    Returns the path to the annotated output.
+    Returns a tuple: (annotated_file_path, status_message). 
+    If an error occurs, annotated_file_path is None and status_message describes the error.
     """
+    error_message = ""
+    input_path = None
+
+    # Check for input: either a YouTube link or an uploaded file.
     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:
-                return "No suitable mp4 stream found."
+                return None, "No suitable mp4 stream found."
             input_path = stream.download()
         except Exception as e:
-            return f"Error downloading video: {e}"
+            return None, f"Error downloading video: {e}"
     elif uploaded_file is not None:
         input_path = uploaded_file.name
     else:
-        return "Please provide an uploaded file or a YouTube link."
-    
-    # Run pose detection and save the annotated output.
-    results = model.predict(source=input_path, save=True)
-    
+        return None, "Please provide an uploaded file or a YouTube link."
+
+    # Run pose detection (with save=True so that outputs are written to disk)
     try:
-        output_path = results[0].save_path
+        results = model.predict(source=input_path, save=True)
     except Exception as e:
-        return f"Error processing the file: {e}"
-    
-    # Optionally remove the downloaded video if applicable.
-    if youtube_link and os.path.exists(input_path):
+        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 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.
+            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):
         os.remove(input_path)
     
-    return output_path
-
-# Define the Gradio Blocks interface as a global variable.
-demo = gr.Blocks()
+    return output_path, "Success!"
 
-with demo:
+# Define the Gradio interface with two outputs: one for the file and one for status text.
+with gr.Blocks() as demo:
     gr.Markdown("# Pose Detection with YOLO11-pose")
     gr.Markdown("Upload an image/video or provide 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://...")
+    
     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], outputs=output_file)
+    
+    run_button.click(process_input, inputs=[file_input, youtube_input],
+                     outputs=[output_file, output_text])
 
-# Only launch the interface if this file is executed directly.
+# Only launch the app if the script is executed directly.
 if __name__ == "__main__":
     demo.launch()