Spaces:
Build error
Build error
| import os | |
| import torch | |
| import torch.nn.functional as F | |
| import torchvision.transforms as T | |
| from mmdet.apis import init_detector, inference_detector, show_result_pyplot | |
| import mmcv | |
| import gradio as gr | |
| from huggingface_hub import hf_hub_download | |
| # Device on which to run the model | |
| # Set to cuda to load on GPU | |
| device = "cpu" | |
| checkpoint_file = hf_hub_download(repo_id="Andy1621/uniformer", filename="mask_rcnn_3x_ms_hybrid_small.pth") | |
| config_file = './exp/mask_rcnn_3x_ms_hybrid_small/config.py' | |
| # init detector | |
| # build the model from a config file and a checkpoint file | |
| model = init_detector(config_file, checkpoint_file, device='cpu') | |
| def set_example_image(example: list) -> dict: | |
| return gr.Image.update(value=example[0]) | |
| def inference(img): | |
| result = inference_detector(model, img) | |
| res_img = show_result_pyplot(model, img, result) | |
| return res_img | |
| demo = gr.Blocks() | |
| with demo: | |
| gr.Markdown( | |
| """ | |
| # UniFormer-S | |
| Gradio demo for <a href='https://github.com/Sense-X/UniFormer' target='_blank'>UniFormer</a>: To use it, simply upload your image, or click one of the examples to load them. Read more at the links below. | |
| """ | |
| ) | |
| with gr.Box(): | |
| with gr.Row(): | |
| with gr.Column(): | |
| with gr.Row(): | |
| input_image = gr.Image(label='Input Image', type='numpy') | |
| with gr.Row(): | |
| submit_button = gr.Button('Submit') | |
| with gr.Column(): | |
| res_image = gr.Image(type='numpy', label='Detection Resutls') | |
| with gr.Row(): | |
| example_images = gr.Dataset(components=[input_image], samples=[['demo.jpg']]) | |
| gr.Markdown( | |
| """ | |
| <p style='text-align: center'><a href='https://arxiv.org/abs/2201.09450' target='_blank'>UniFormer: Unifying Convolution and Self-attention for Visual Recognition</a> | <a href='https://github.com/Sense-X/UniFormer' target='_blank'>Github Repo</a></p> | |
| """ | |
| ) | |
| submit_button.click(fn=inference, inputs=input_image, outputs=res_image) | |
| example_images.click(fn=set_example_image, inputs=example_images, outputs=example_images.components) | |
| demo.launch(enable_queue=True) |