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

os.system("wget https://media.roboflow.com/notebooks/examples/dog.jpeg")
os.system("wget https://media.roboflow.com/notebooks/examples/dog-2.jpeg")


def detect(image, weights, slider_val):
    model = YOLO(weights + '.pt')
    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


inputs_thresh = [
    gr.inputs.Image(type="numpy", label="Input Image"),
    gr.inputs.Radio(label="Detection Methods",
                    choices=[
                        "yolov5s", "yolov8s"
                    ]),
    gr.components.Slider(label="Class Probability Value",
                         value=10, minimum=1, maximum=100, step=1),
]

outputs_thresh = [
    gr.outputs.Image(type="numpy", label="Output Image")
]

detect_tab = gr.Interface(
    detect,
    inputs=inputs_thresh,
    outputs=outputs_thresh,
    title="supervision",
    examples=[["dog.jpeg", "yolov5s"], ["dog-2.jpeg", "yolov8s"]],
    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."
)