import gradio as gr
from ultralytics import YOLOv10 
import supervision as sv

MODEL_PATH = 'yolov10n.pt'
model = YOLOv10(MODEL_PATH)
box_annotator = sv.BoxAnnotator()

def detect(image):
    results = model(source=image, conf=0.25, verbose=False)[0]
    detections = sv.Detections.from_ultralytics(results)
    
    labels = [
        f"{model.model.names[class_id]} {confidence:.2f}" 
        for class_id, confidence in zip(detections.class_id, detections.confidence)
    ]
    annotated_image = box_annotator.annotate(image, detections=detections, labels=labels)

    return annotated_image

gr.Interface(
    fn=detect, 
    inputs=gr.inputs.Image(type="numpy"),
    outputs=gr.outputs.Image(type="numpy", label="Annotated Image"),
    title='YOLOv10 Object Detection',
    description='Detect objects in images using the YOLOv10 model'
).launch()