import gradio as gr
from ultralytics import YOLO
import numpy as np
from PIL import Image

# Load YOLO model
model = YOLO("yolo11n-pose.pt")  # Replace with your YOLO model path

# Function to process the input image
def process_image(image):
    # YOLO inference
    results = model(image)
    
    # The result is a list, access the first item
    processed_image = results[0].plot()  # .plot() returns the image with detections
    
    # Return the processed image for display
    return processed_image

# Gradio Interface
iface = gr.Interface(
    fn=process_image,  # Function to call on image input
    inputs=gr.Image(type="pil"),  # Image input (Upload an image file)
    outputs=gr.Image(type="pil"),  # Output processed image
    live=True,  # Enable live feedback (optional)
    title="Object Detection with YOLO",
    description="Upload an image for object detection using YOLO."
)

# Launch the Gradio app
iface.launch()