Achouche/detectron2

1import gradio as gr2import cv23import torch4from detectron2.engine import DefaultPredictor5from detectron2.config import get_cfg6from detectron2 import model_zoo7from detectron2.utils.visualizer import Visualizer8 9# Load model10cfg = get_cfg()11cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))12cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")13cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.514predictor = DefaultPredictor(cfg)15 16def detect_objects(image):17    outputs = predictor(image)18    v = Visualizer(image[:, :, ::-1])19    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))20    return out.get_image()[:, :, ::-1]21 22gr.Interface(fn=detect_objects, inputs=gr.Image(type="numpy"), outputs=gr.Image(type="numpy")).launch()23