Spaces:
Running
Running
| import sys | |
| import gradio as gr | |
| import os | |
| import tempfile | |
| import cv2 | |
| import requests | |
| from ultralytics import YOLO | |
| # 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, sensitivity): | |
| """ | |
| 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 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 | |
| # Priority 1: YouTube link | |
| if youtube_link and youtube_link.strip(): | |
| try: | |
| from pytube import YouTube | |
| yt = YouTube(youtube_link) | |
| stream = yt.streams.filter(file_extension='mp4', progressive=True)\ | |
| .order_by("resolution").desc().first() | |
| if stream is None: | |
| 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(): | |
| try: | |
| response = requests.get(image_url, stream=True) | |
| if response.status_code != 200: | |
| 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}", "" | |
| # Priority 3: Uploaded file | |
| elif uploaded_file is not None: | |
| input_path = uploaded_file.name | |
| else: | |
| return None, None, "Please provide an input using one of the methods.", "" | |
| try: | |
| results = model.predict(source=input_path, save=True, conf=sensitivity) | |
| except Exception as e: | |
| return None, None, f"Error running prediction: {e}", "" | |
| output_path = None | |
| try: | |
| if hasattr(results[0], "save_path"): | |
| output_path = results[0].save_path | |
| else: | |
| 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}", "" | |
| # 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) | |
| return output_path, output_path, "Success!", "" | |
| # 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: | |
| with gr.Row(): | |
| # Left Column: Header image, title, input tabs and 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.0, maximum=1.0, step=0.01, value=0.5, | |
| label="Sensitivity (Confidence Threshold)") | |
| # Right Column: Results display at the top. | |
| 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()] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |