| |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import pathlib |
|
|
| import gradio as gr |
|
|
| from model import Model |
|
|
| DESCRIPTION = '''# CBNetV2 |
| |
| This is an unofficial demo for [https://github.com/VDIGPKU/CBNetV2](https://github.com/VDIGPKU/CBNetV2).''' |
| FOOTER = '<img id="visitor-badge" alt="visitor badge" src="https://visitor-badge.glitch.me/badge?page_id=hysts.cbnetv2" />' |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--device', type=str, default='cpu') |
| parser.add_argument('--theme', type=str) |
| parser.add_argument('--share', action='store_true') |
| parser.add_argument('--port', type=int) |
| parser.add_argument('--disable-queue', |
| dest='enable_queue', |
| action='store_false') |
| return parser.parse_args() |
|
|
|
|
| def set_example_image(example: list) -> dict: |
| return gr.Image.update(value=example[0]) |
|
|
|
|
| def main(): |
| args = parse_args() |
| model = Model(args.device) |
|
|
| with gr.Blocks(theme=args.theme, 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(list(model.models.keys()), |
| value=model.model_name, |
| label='Detector') |
| with gr.Row(): |
| detect_button = gr.Button(value='Detect') |
| detection_results = gr.Variable() |
| with gr.Column(): |
| with gr.Row(): |
| detection_visualization = gr.Image( |
| label='Detection Result', type='numpy') |
| with gr.Row(): |
| visualization_score_threshold = gr.Slider( |
| 0, |
| 1, |
| step=0.05, |
| value=0.3, |
| label='Visualization Score Threshold') |
| with gr.Row(): |
| redraw_button = gr.Button(value='Redraw') |
|
|
| with gr.Row(): |
| paths = sorted(pathlib.Path('images').rglob('*.jpg')) |
| example_images = gr.Dataset(components=[input_image], |
| samples=[[path.as_posix()] |
| for path in paths]) |
|
|
| gr.Markdown(FOOTER) |
|
|
| detector_name.change(fn=model.set_model_name, |
| inputs=[detector_name], |
| outputs=None) |
| 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]) |
| example_images.click(fn=set_example_image, |
| inputs=[example_images], |
| outputs=[input_image]) |
|
|
| demo.launch( |
| enable_queue=args.enable_queue, |
| server_port=args.port, |
| share=args.share, |
| ) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
