Update backend.py

backend.py

@@ -82,6 +82,17 @@ def building_model(building_version_dropdown, building_pth_dropdown, building_th
     building_predictor = DefaultPredictor(building_cfg)
     return building_predictor
 
+# Model for LCZs
+def tree_model(lcz_version_dropdown, lcz_pth_dropdown, lcz_threshold, device="cpu"):
+    lcz_cfg = get_cfg()
+    lcz_cfg.merge_from_file(get_version_cfg_yml("lcz_model_weights/lczs_cfg.yaml"))
+    lcz_cfg.MODEL.DEVICE=device
+    lcz_cfg.MODEL.WEIGHTS = f"tree_model_weights/{tree_pth_dropdown}"
+    lcz_cfg.MODEL.ROI_HEADS.NUM_CLASSES = 14  # TODO change this
+    lcz_cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = lcz_threshold
+    lcz_predictor = DefaultPredictor(lcz_cfg)
+    return lcz_predictor
+
 # A function that runs the buildings model on an given image and confidence threshold
 def segment_building(im, building_predictor):
     outputs = building_predictor(im)
@@ -96,6 +107,13 @@ def segment_tree(im, tree_predictor):
 
     return tree_instances
 
+# A function that runs the trees model on an given image and confidence threshold
+def segment_lcz(im, lcz_predictor):
+    outputs = lcz_predictor(im)
+    lcz_instances = outputs["instances"].to("cpu")
+
+    return lcz_instances
+
 # Function to map strings to color mode
 def map_color_mode(color_mode):
     if color_mode == "Black/white":
@@ -122,7 +140,7 @@ def get_metadata(dataset_name, coco_file):
     metadata.thing_classes = [c["name"] for c in categories]
     return metadata
 
-def visualize_image(im, mode, tree_threshold, building_threshold, color_mode, tree_version, tree_pth, building_version, building_pth):
+def visualize_image(im, mode, tree_threshold, building_threshold, color_mode, tree_version, tree_pth, building_version, building_pth, lcz_version, lcz_pth):
     im = np.array(im)
     color_mode = map_color_mode(color_mode)
     
@@ -138,6 +156,11 @@ def visualize_image(im, mode, tree_threshold, building_threshold, color_mode, tr
         building_instances = load_instances(im, building_predictor, segment_building)
         instances = building_instances if mode == "Buildings" else combine_instances(instances, building_instances)
 
+    if mode in {"LCZ", "Both"}:
+        lcz_predictor = load_predictor(lcz_model, lcz_version, lcz_pth, lcz_threshold)
+        lcz_instances = load_instances(im, lcz_predictor, segment_lcz)
+        instances = lcz_instances if mode == "LCZ" else combine_instances(instances, LCZ_instances)   
+
     # Assuming 'urban-small_train' is intended for both Trees and Buildings
     metadata = get_metadata("urban-small_train", "building_model_weight/_annotations.coco.json")
     visualizer = Visualizer(im[:, :, ::-1], metadata=metadata, scale=0.5, instance_mode=color_mode)