import gradio as gr
from ultralytics import YOLO
import PIL

model = YOLO("yolov8n.pt")

def predict_image(img):
    results = model.predict(source=img)
    annotated_img = results[0].plot()
    annotated_img = PIL.Image.fromarray(annotated_img[...,::-1])
    
    return annotated_img

gradio_app = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Object Detection",
    examples=["dog.webp", "zidane.jpg", "huggingface.png"]
)

if __name__ == "__main__":
    gradio_app.launch()