import gradio as gr
from PIL import Image
from ultralytics import YOLO
from ultralytics.utils.plotting import Annotator

model_path = 'yolo11n_cs2.pt'
model = YOLO(model_path, task="detect")

def predict(image):
    global model
    results = model.predict(image, imgsz=1280, conf=0.2)
    annotator = Annotator(image)
    for r in results:
        boxes = r.boxes
        for box in boxes:
            b = box.xyxy[0]
            c = box.cls
            annotator.box_label(b, model.names[int(c)])
    image = annotator.result()
    return image


gr.Interface(
    title="Detect players",
    fn=predict,
    inputs=[
        gr.Image(type="pil", label="Load image"),
    ],
    outputs=gr.Image(type="pil", label="Detected players"),
).launch()