import gradio as gr
from PIL import Image
import torch
from ultralyticsplus import YOLO, render_result

available_models = ["YOLOv8n", "YOLOv8n-GhostNet-P5", "YOLOv8n-GhostNet-P6"]
available_models_path = [
    "./models/yolov8n.pt",
    "./models/yolov8n_ghostnet_p5.pt",
    "./models/yolov8n_ghostnet_p6.pt",
]


def launch(
    image: gr.Image = None,
    selectedModel: gr.Dropdown = available_models[0],
    conf_threshold: gr.Slider = 0.4,
    iou_threshold: gr.Slider = 0.50,
):
    selected_model_index = available_models.index(selectedModel)
    image_size = (256,)
    try:
        model = YOLO(available_models_path[selected_model_index])

        # pil_image = Image.fromarray(image)

        results = model.predict(
            image, conf=conf_threshold, iou=iou_threshold, imgsz=image_size
        )
        box = results[0].boxes
        # print(box)

        render = render_result(model=model, image=image, result=results[0])
        return render
    except Exception as e:
        print("error", e)
        return "./download.jpeg"


inputs = [
    gr.Image(type="filepath", label="Input Image"),
    gr.Dropdown(
        info="Choose which model should be used in this task",
        choices=available_models,
        value=available_models[0],
        label="Models",
    ),
    # gr.Slider(minimum=256, maximum=1280, value=640, step=32, label="Image Size"),
    gr.Slider(
        minimum=0.0, maximum=1.0, value=0.4, step=0.1, label="Confidence Threshold"
    ),
    gr.Slider(minimum=0.0, maximum=1.0, value=0.4, step=0.1, label="IOU Threshold"),
]

outputs = gr.Image(type="filepath", label="Output Result")

iface = gr.Interface(fn=launch, inputs=inputs, outputs=outputs)
iface.launch()