Update app.py

app.py

@@ -21,11 +21,14 @@ from detectron2.engine import DefaultPredictor
 from detectron2.config import get_cfg
 from detectron2.utils.visualizer import Visualizer
 from detectron2.data import MetadataCatalog, DatasetCatalog
+from detectron2.checkpoint import DetectionCheckpointer
 
 cfg = get_cfg()
+cfg.merge_from_file(model_zoo.get_config_file('COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml'))
 cfg.MODEL.DEVICE='cpu'
 cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 
 cfg.MODEL.WEIGHTS = "model_weights/chatswood_buildings_poc.pth"  
+cfg.MODEL.ROI_HEADS.NUM_CLASSES = 8
 
 def segment_buildings(input_image):