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.
    
    For video files (mp4, mov, avi, webm), we use streaming mode to obtain annotated frames and encode them into a video.
    For images, we use the normal prediction and either use the built‑in save_path or plot() method.
    
    Returns a tuple:
      - download_file_path (for gr.File)
      - image_result (for gr.Image) or None
      - video_result (for gr.Video) or None
      - 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."

    # Determine if input is a video (by extension).
    ext_input = os.path.splitext(input_path)[1].lower()
    video_exts = [".mp4", ".mov", ".avi", ".webm"]

    output_path = None

    if ext_input in video_exts:
        # Process video using streaming mode.
        try:
            # Open video to get properties.
            cap = cv2.VideoCapture(input_path)
            if not cap.isOpened():
                return None, None, None, "Error opening video file."
            fps = cap.get(cv2.CAP_PROP_FPS)
            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            cap.release()

            # Use streaming mode to process each frame.
            frames = []
            for result in model.predict(source=input_path, stream=True, conf=sensitivity):
                # result.plot() returns an annotated frame (numpy array)
                annotated_frame = result.plot()
                frames.append(annotated_frame)
            if not frames:
                return None, None, None, "No detections were returned from video streaming."
            # Write frames to a temporary video file.
            temp_video_path = os.path.join(tempfile.gettempdir(), "annotated_video.mp4")
            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            out = cv2.VideoWriter(temp_video_path, fourcc, fps, (width, height))
            for frame in frames:
                out.write(frame)
            out.release()
            output_path = temp_video_path
        except Exception as e:
            return None, None, None, f"Error processing video: {e}"
    else:
        # Process as an image.
        try:
            results = model.predict(source=input_path, save=True, conf=sensitivity)
        except Exception as e:
            return None, None, None, f"Error running prediction: {e}"

        try:
            if not results or len(results) == 0:
                return None, None, None, "No detections were returned."
            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, None, f"Error processing the file: {e}"

    # Clean up temporary input if 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)

    # Set outputs based on output file extension.
    ext_output = os.path.splitext(output_path)[1].lower()
    if ext_output 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!"

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

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