import json
import glob
from collections import Counter

import requests
import gradio as gr
from ultralyticsplus import YOLO, download_from_hub, render_result

hf_model_ids = [
    "chanelcolgate/chamdiemgianhang-vsk",
    "chanelcolgate/chamdiemgianhang-vsk-v2",
    "chanelcolgate/chamdiemgianhang-vsk-v4",
    "chanelcolgate/chamdiemgianhang-vsk-v5",
    "chanelcolgate/chamdiemgianhang-vsk-v6",
]

image_paths = [
    [image_path, "chanelcolgate/chamdiemgianhang-vsk-v2", 640, 0.25, 0.45]
    for image_path in glob.glob("./tmp/*.jpg")
]


def detection_image(
    image=None,
    hf_model_id="chanelcolgate/chamdiemgianhang-vsk-v2",
    image_size=640,
    conf_threshold=0.25,
    iou_threshold=0.45,
):
    model_path = download_from_hub(hf_model_id)
    model = YOLO(model_path)
    results = model(image, imgsz=image_size, conf=conf_threshold, iou=iou_threshold)
    json_result = json.loads(results[0].tojson())
    class_counts = Counter(detection["name"] for detection in json_result)

    render = render_result(model=model, image=image, result=results[0])
    return render, class_counts


def detection_image_link(
    image=None,
    hf_model_id="chanelcolgate/chamdiemgianhang-vsk-v2",
    image_size=640,
    conf_threshold=0.25,
    iou_threshold=0.45,
):
    model_path = download_from_hub(hf_model_id)
    model = YOLO(model_path)
    results = model(image, imgsz=image_size, conf=conf_threshold, iou=iou_threshold)
    json_result = json.loads(results[0].tojson())
    class_counts = Counter(detection["name"] for detection in json_result)

    render = render_result(model=model, image=image, result=results[0])
    return render, class_counts


title = "Cham Diem Gian Hang VSK"

interface = gr.Interface(
    fn=detection_image,
    inputs=[
        gr.Image(type="pil"),
        gr.Dropdown(hf_model_ids),
        gr.Slider(minimum=320, maximum=1280, value=640, step=32, label="Image Size"),
        gr.Slider(
            minimum=0.0,
            maximum=1.0,
            value=0.25,
            step=0.05,
            label="Confidence Threshold",
        ),
        gr.Slider(
            minimum=0.0, maximum=1.0, value=0.45, step=0.05, label="IOU Threshold"
        ),
    ],
    outputs=[gr.Image(type="pil"), gr.Textbox(show_label=False)],
    title=title,
    examples=image_paths,
    cache_examples=True if image_paths else False,
)

interface_link = gr.Interface(
    fn=detection_image,
    inputs=[
        gr.Textbox(label="Image Link"),
        gr.Dropdown(hf_model_ids),
        gr.Slider(minimum=320, maximum=1280, value=640, step=32, label="Image Size"),
        gr.Slider(
            minimum=0.0,
            maximum=1.0,
            value=0.25,
            step=0.05,
            label="Confidence Threshold",
        ),
        gr.Slider(
            minimum=0.0, maximum=1.0, value=0.45, step=0.05, label="IOU Threshold"
        ),
    ],
    outputs=[gr.Image(type="pil"), gr.Textbox(show_label=False)],
    title=title,
)

gr.TabbedInterface(
    [interface, interface_link], tab_names=["Image inference", "Image link inference"]
).queue().launch()