app.py

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()