SmitaMahajan commited on
Commit
df12bee
·
verified ·
1 Parent(s): f78be86

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +28 -0
app.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import cv2
3
+ from detectron2.engine import DefaultPredictor
4
+ from detectron2.config import get_cfg
5
+ from detectron2 import model_zoo
6
+ from detectron2.utils.visualizer import Visualizer
7
+ from detectron2.data import MetadataCatalog
8
+
9
+ # Setup Detectron2
10
+ cfg = get_cfg()
11
+ cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
12
+ cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
13
+ cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
14
+ predictor = DefaultPredictor(cfg)
15
+
16
+ def detect_objects(image):
17
+ image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
18
+ outputs = predictor(image_rgb)
19
+ v = Visualizer(image_rgb, MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.0)
20
+ out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
21
+ return out.get_image()
22
+
23
+ demo = gr.Interface(fn=detect_objects,
24
+ inputs=gr.Image(type="numpy"),
25
+ outputs="image",
26
+ title="Detectron2 Object Detection")
27
+
28
+ demo.launch()