| import gradio as gr | |
| import cv2 | |
| from detectron2.engine import DefaultPredictor | |
| from detectron2.config import get_cfg | |
| from detectron2 import model_zoo | |
| from detectron2.utils.visualizer import Visualizer | |
| from detectron2.data import MetadataCatalog | |
| # Setup Detectron2 | |
| cfg = get_cfg() | |
| cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")) | |
| cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 | |
| cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml") | |
| predictor = DefaultPredictor(cfg) | |
| def detect_objects(image): | |
| image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| outputs = predictor(image_rgb) | |
| v = Visualizer(image_rgb, MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.0) | |
| out = v.draw_instance_predictions(outputs["instances"].to("cpu")) | |
| return out.get_image() | |
| demo = gr.Interface(fn=detect_objects, | |
| inputs=gr.Image(type="numpy"), | |
| outputs="image", | |
| title="Detectron2 Object Detection") | |
| demo.launch() | |