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


def download_models(model_id):
    hf_hub_download("kadirnar/yolov10", filename=f"{model_id}", local_dir=f"./")
    return f"./{model_id}"
    
MODEL_PATH = 'yolov10n.pt'
model = YOLOv10(MODEL_PATH)
box_annotator = sv.BoxAnnotator()
model_path = download_models(model_id)

@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()