from glob import glob
import os
from timeit import default_timer as timer

import gradio as gr
import numpy as np
import onnxruntime as ort
from PIL import ImageOps

from utils import draw_overlay


sessions = {
    "rtdetr_18vd": ort.InferenceSession("models/rtdetr_r18vd.onnx"),
    "rtdetr_r50vd": ort.InferenceSession("models/rtdetr_r50vd.onnx"),
    "rtdetr_r101vd": ort.InferenceSession("models/rtdetr_r101vd.onnx"),
}


def run(image, threshold=0.5, model="rtdetr_18vd"):
    im = ImageOps.pad(image, (640, 640), color="black")
    input = np.asarray(im)
    input = np.transpose(input, (2, 0, 1))
    input = np.expand_dims(input, axis=0)
    input = input.astype(np.float32) / 255

    inference_start = timer()
    labels, boxes, scores = sessions[model].run(
        None,
        {"images": input, "orig_target_sizes": np.array([[640, 640]])},
    )
    inference_seconds = timer() - inference_start
    draw_overlay(im, threshold, labels, boxes, scores)

    return im, labels[0], boxes[0], scores[0], inference_seconds * 1000


demo = gr.Interface(
    run,
    inputs=[
        gr.Image(type="pil", label="image"),
        gr.Slider(0, 1, 0.5, label="score threshold"),
        gr.Dropdown(
            ["rtdetr_18vd", "rtdetr_r50vd", "rtdetr_r101vd"],
            value="rtdetr_18vd",
            label="model",
        ),
    ],
    outputs=[
        gr.Image(label="result"),
        gr.Textbox(label="labels"),
        gr.Textbox(label="boxes"),
        gr.Textbox(label="scores"),
        gr.Textbox(label="inference ms"),
    ],
    examples=[glob(os.path.join("samples", "*"))],
    title="RT-DETR",
    description="ONNX exported from official pretrained models on COCO + Objects365",
)

demo.launch()