import os
import gradio as gr
import numpy as np
import supervision as sv
from ultralytics import YOLO


def detect(image, weights):
    model = YOLO(weights)
    result = model(image, verbose=False)[0]
    detections = sv.Detections.from_ultralytics(result)
    box_annotator = sv.BoxAnnotator()
    annotated_image = box_annotator.annotate(image.copy(), detections=detections)

    return annotated_image


if __name__ == '__main__':
    gr.Interface(
        detect, [gr.Image(type="numpy"), gr.Dropdown(choices=["yolov5s", "yolov5s6"])],
        gr.Image(type="numpy"),
        title="Yolov5",
        examples=[["https://media.roboflow.com/notebooks/examples/dog.jpeg", "yolov5s"]],
        description=
        "Gradio based demo for <a href='https://github.com/roboflow/supervision' style='text-decoration: underline' target='_blank'>roboflow/supervision</a>, We write your reusable computer vision tools."
    ).launch()