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()