hlydecker commited on
Commit
4378f9c
1 Parent(s): e06d371

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -0
app.py CHANGED
@@ -21,11 +21,14 @@ from detectron2.engine import DefaultPredictor
21
  from detectron2.config import get_cfg
22
  from detectron2.utils.visualizer import Visualizer
23
  from detectron2.data import MetadataCatalog, DatasetCatalog
 
24
 
25
  cfg = get_cfg()
 
26
  cfg.MODEL.DEVICE='cpu'
27
  cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
28
  cfg.MODEL.WEIGHTS = "model_weights/chatswood_buildings_poc.pth"
 
29
 
30
  def segment_buildings(input_image):
31
 
 
21
  from detectron2.config import get_cfg
22
  from detectron2.utils.visualizer import Visualizer
23
  from detectron2.data import MetadataCatalog, DatasetCatalog
24
+ from detectron2.checkpoint import DetectionCheckpointer
25
 
26
  cfg = get_cfg()
27
+ cfg.merge_from_file(model_zoo.get_config_file('COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml'))
28
  cfg.MODEL.DEVICE='cpu'
29
  cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
30
  cfg.MODEL.WEIGHTS = "model_weights/chatswood_buildings_poc.pth"
31
+ cfg.MODEL.ROI_HEADS.NUM_CLASSES = 8
32
 
33
  def segment_buildings(input_image):
34