Spaces:
Running
Running
#!/usr/bin/env python | |
from __future__ import annotations | |
import pathlib | |
import gradio as gr | |
from model import Model | |
DESCRIPTION = "# [CBNetV2](https://github.com/VDIGPKU/CBNetV2)" | |
model = Model() | |
with gr.Blocks(css="style.css") as demo: | |
gr.Markdown(DESCRIPTION) | |
with gr.Row(): | |
with gr.Column(): | |
with gr.Row(): | |
input_image = gr.Image(label="Input Image", type="numpy") | |
with gr.Row(): | |
detector_name = gr.Dropdown( | |
label="Detector", choices=list(model.models.keys()), value=model.model_name | |
) | |
with gr.Row(): | |
detect_button = gr.Button("Detect") | |
detection_results = gr.State() | |
with gr.Column(): | |
with gr.Row(): | |
detection_visualization = gr.Image(label="Detection Result", type="numpy") | |
with gr.Row(): | |
visualization_score_threshold = gr.Slider( | |
label="Visualization Score Threshold", minimum=0, maximum=1, step=0.05, value=0.3 | |
) | |
with gr.Row(): | |
redraw_button = gr.Button("Redraw") | |
with gr.Row(): | |
paths = sorted(pathlib.Path("images").rglob("*.jpg")) | |
gr.Examples(examples=[[path.as_posix()] for path in paths], inputs=input_image) | |
detector_name.change(fn=model.set_model_name, inputs=detector_name) | |
detect_button.click( | |
fn=model.detect_and_visualize, | |
inputs=[ | |
input_image, | |
visualization_score_threshold, | |
], | |
outputs=[ | |
detection_results, | |
detection_visualization, | |
], | |
) | |
redraw_button.click( | |
fn=model.visualize_detection_results, | |
inputs=[ | |
input_image, | |
detection_results, | |
visualization_score_threshold, | |
], | |
outputs=detection_visualization, | |
) | |
if __name__ == "__main__": | |
demo.queue(max_size=10).launch() | |