Update app.py

app.py

@@ -1,6 +1,6 @@
 """
 building-segmentation
-Proof of concept showing effectiveness of a fine tuned instance segmentation model for deteting buildings.
+Proof of concept showing effectiveness of a fine tuned instance segmentation model for detecting buildings.
 """
 import os
 import cv2
@@ -25,14 +25,16 @@ from detectron2.data import MetadataCatalog, DatasetCatalog
 from detectron2.checkpoint import DetectionCheckpointer
 
 cfg = get_cfg()
-cfg.merge_from_file("model_weights/buildings_poc_cfg.yml")
+cfg.merge_from_file("model_weights/lczs_cfg.yml")
 cfg.MODEL.DEVICE='cpu'
 cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.25 
-cfg.MODEL.WEIGHTS = "model_weights/model_final.pth"  
-cfg.MODEL.ROI_HEADS.NUM_CLASSES = 8
+#cfg.MODEL.WEIGHTS = "model_weights/model_final.pth"  
+#cfg.MODEL.ROI_HEADS.NUM_CLASSES = 8
+cfg.MODEL.WEIGHTS = "model_weights/lczs_v2.pth"  
+cfg.MODEL.ROI_HEADS.NUM_CLASSES = 14
 predictor = DefaultPredictor(cfg)
 
-def segment_buildings(im):
+def segment(im):
 
     im = np.array(im)
     outputs = predictor(im)
@@ -58,7 +60,7 @@ title = "Building Segmentation"
 description = "An instance segmentation demo for identifying boundaries of buildings in aerial images using DETR (End-to-End Object Detection) model with MaskRCNN-101 backbone"
 
 # Create user interface and launch
-gr.Interface(segment_buildings, 
+gr.Interface(segment, 
                 inputs = inputs,
                 outputs = outputs,
                  title = title,