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


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

@spaces.GPU(duration=200)
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

gradio_app = gr.Blocks()
with gradio_app:
    gr.HTML(
        """
    <h1 style='text-align: center'>
    YOLOv10: Real-Time End-to-End Object Detection
    </h1>
    """)
    gr.HTML(
        """
        <h3 style='text-align: center'>
        Follow me for more!
        <a href='https://twitter.com/kadirnar_ai' target='_blank'>Twitter</a> | <a href='https://github.com/kadirnar' target='_blank'>Github</a> | <a href='https://www.linkedin.com/in/kadir-nar/' target='_blank'>Linkedin</a> | <a href='https://www.huggingface.co/kadirnar/' target='_blank'>HuggingFace</a>
        </h3>
        """)

    input_image = gr.Image(type="numpy")
    output_image = gr.Image(type="numpy", label="Annotated Image")
    
    gr.Interface(
        fn=detect,
        inputs=input_image,
        outputs=output_image,
    )

gradio_app.launch()