File size: 2,051 Bytes
b94fa0f e70c067 b94fa0f e70c067 b94fa0f 519f2af b94fa0f e70c067 989f17e e70c067 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
import os
from pydoc import describe
os.system('pip install torch==1.9 torchvision')
os.system('pip install detectron2==0.5 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu102/torch1.9/index.html')
os.system('pip install timm opencv-python-headless')
import gradio as gr
from demo.predictor import VisualizationDemo
from detectron2.config import get_cfg
from opendet2 import add_opendet_config
model_cfgs = {
"FR-CNN": ["configs/faster_rcnn_R_50_FPN_3x_baseline.yaml", "frcnn_r50.pth"],
"OpenDet-R50": ["configs/faster_rcnn_R_50_FPN_3x_opendet.yaml", "opendet2_r50.pth"],
"OpenDet-SwinT": ["configs/faster_rcnn_Swin_T_FPN_3x_opendet.yaml", "opendet2_swint.pth"],
}
def setup_cfg(model):
cfg = get_cfg()
add_opendet_config(cfg)
model_cfg = model_cfgs[model]
cfg.merge_from_file(model_cfg[0])
cfg.MODEL.WEIGHTS = model_cfg[1]
cfg.MODEL.DEVICE = "cpu"
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
cfg.MODEL.ROI_HEADS.VIS_IOU_THRESH = 0.8
cfg.freeze()
return cfg
def inference(input, model):
cfg = setup_cfg(model)
demo = VisualizationDemo(cfg)
# use PIL, to be consistent with evaluation
predictions, visualized_output = demo.run_on_image(input)
output = visualized_output.get_image()[:, :, ::-1]
return output
examples = [
["demo/000000002157.jpg", "OpenDet-R50"],
["demo/000000020059.jpg", "OpenDet-R50"]]
iface = gr.Interface(
inference,
[
"image",
gr.inputs.Radio(
["FR-CNN", "OpenDet-R50", "OpenDet-SwinT"], default='OpenDet-R50'),
],
"image",
examples=examples,
title="OpenDet",
article="<p style='text-align: center'><a href='https://github.com/csuhan/opendet2' target='_blank'>Github Repo</a> | <a href='https://csuhan.com' target='_blank'>Author Page</a></p>",
description="Online demo for: Expanding Low-Density Latent Regions for Open-Set Object Detection. Please upload your image, or click one of the examples to load them",
)
iface.launch(enable_queue=True, cache_examples=True)
|