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

os.system("wget https://datasets-server.huggingface.co/cached-assets/imagenet-1k/--/default/test/27/image/image1.jpg")
os.system("wget https://datasets-server.huggingface.co/cached-assets/imagenet-1k/--/default/test/22/image/image2.jpg")

response = requests.get("https://git.io/JJkYN")
labels = response.text.split("\n")


def clasify(image, radio_choice, slider_val):
    print(radio_choice)
    model = YOLO(radio_choice + '.pt')
    result = model(image, verbose=False)[0]
    classifications = sv.Classifications.from_ultralytics(result)
    out_dic = {}
    cls_out = classifications.get_top_k(4)
    for idx in range(4):
        cls_id = cls_out[0][idx]
        cls_prob = cls_out[1][idx]
        out_dic[labels[int(cls_id)]] = float(cls_prob)
    return out_dic


inputs_thresh = [
    gr.inputs.Image(type="filepath", label="Input Image"),
    gr.inputs.Radio(label="Classification Methods",
                    choices=[
                        "yolov8n-cls", "yolov8s-cls"
                    ]),
    gr.components.Slider(label="Class Probability Value",
                         value=10, minimum=1, maximum=100, step=1),
]

classify_tab = gr.Interface(
    clasify,
    inputs=inputs_thresh,
    outputs=gr.outputs.Label(num_top_classes=4),
    title="supervision",
    examples=[["image1.jpg", "yolov8s-cls"], ["image2.jpg", "yolov8n-cls"]],
    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."
)