import gradio as gr
from ultralytics import YOLO
import cv2

# Load YOLO11-Pose model (downloads if not cached)
model = YOLO("yolo11n-pose.pt")  # You can also use yolo11s/m/l/x-pose.pt

def predict_pose(frame):
    # Run YOLO11-Pose inference
    results = model.predict(frame, imgsz=480, conf=0.5, iou=0.5)[0]
    
    # Draw results on frame
    annotated_frame = results.plot(labels=False, boxes=False)  # YOLO11 handles drawing
    
    return annotated_frame

# Set up Gradio interface
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(sources=["webcam"])
        with gr.Column():
            output_img = gr.Image(streaming=True)
        input_img.stream(predict_pose, input_img, output_img,
                                time_limit=30, stream_every=1, concurrency_limit=30)
        
demo.launch(debug=True)