|
import gradio as gr |
|
import supervision as sv |
|
import os |
|
from time import perf_counter |
|
|
|
from detr import SimpleDetr, PanopticDetrResenet101 |
|
|
|
ASSETS_DIR = os.path.abspath(os.curdir) + "/data/assets" |
|
|
|
print("Assets:", ASSETS_DIR) |
|
|
|
|
|
def run_inference(image, confidence, model_name, progress=gr.Progress(track_tqdm=True)): |
|
progress(0.1, "loading model..") |
|
t0 = perf_counter() |
|
if model_name == "detr_demo_boxes": |
|
model = SimpleDetr() |
|
else: |
|
model = PanopticDetrResenet101() |
|
t1 = perf_counter() |
|
progress(0.1, "Inference..") |
|
|
|
annotated_img = model.detect(image, confidence) |
|
t2 = perf_counter() |
|
return annotated_img, {"load_model": t1 - t0, "inference": t2 - t1}, None |
|
|
|
|
|
with gr.Blocks() as inference_gradio: |
|
gr.Markdown("# DETR inference") |
|
with gr.Row(): |
|
with gr.Column(): |
|
img_file = gr.Image(type="pil") |
|
|
|
model_name = gr.Dropdown( |
|
label="Model", |
|
scale=3, |
|
choices=["detr_demo_boxes", "detr_resnet101_panoptic"], |
|
value="detr_demo_boxes", |
|
) |
|
|
|
conf = gr.Slider(label="Confidence", minimum=0, maximum=0.99, value=0.5) |
|
|
|
with gr.Row(): |
|
start_btn = gr.Button("Start", variant="primary") |
|
with gr.Column(): |
|
annotated_img = gr.Image(label="Annotated Image") |
|
speed = gr.JSON(label="speed") |
|
examples = gr.Examples( |
|
examples=[ |
|
[path] |
|
for path in sv.list_files_with_extensions( |
|
directory=ASSETS_DIR, extensions=["jpeg", "jpg", "png"] |
|
) |
|
], |
|
inputs=[img_file], |
|
) |
|
start_btn.click( |
|
fn=run_inference, |
|
inputs=[img_file, conf, model_name], |
|
outputs=[annotated_img, speed], |
|
) |
|
|
|
if __name__ == "__main__": |
|
inference_gradio.queue(2).launch( |
|
debug=True, |
|
server_name="0.0.0.0", |
|
server_port=7000, |
|
) |
|
|
