File size: 1,479 Bytes
b94fa0f |
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 |
import os
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_18e_opendet_voc.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
iface = gr.Interface(
inference,
[
"image",
gr.inputs.Radio(
["FR-CNN", "OpenDet-R50", "OpenDet-SwinT"], default='OpenDet-R50'),
],
"image")
iface.launch()
|