dynamic choosing of model versions

app.py

@@ -25,13 +25,13 @@ with gr.Blocks() as demo:
     tree_row, building_row , lcz_row = gr.Row(), gr.Row(), gr.Row()
     # tree_col, building_col = gr.Column(elem_id="tree_col"), gr.Column(elem_id="building_col")
     with tree_row as tree_options:
-        tree_version_dropdown = gr.Dropdown(choices=["treev1", "treev2"], label="Tree Detection Version", value="treev2", visible=True, interactive=True)
+        tree_version_dropdown = gr.Dropdown(choices=list_cfg_file_versions("tree_model_weights"), label="Tree Detection Version", value="treev2", visible=True, interactive=True)
         tree_pth_dropdown = gr.Dropdown(choices=list_pth_files_in_directory("tree_model_weights", "v2"), label="Select a tree model file", visible=True, interactive=True)
         tree_threshold_slider = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.7, label='Set confidence threshold "%" for trees', visible=True, interactive=True)
 
     with building_row as building_options:
-        building_version_dropdown = gr.Dropdown(choices=["buildingv1", "buildingv2"], label="Building Detection Version", value="buildingv1", visible=True, interactive=True)
-        building_pth_dropdown = gr.Dropdown(choices=list_pth_files_in_directory("building_model_weight", "v1"), label="Select a building model file", visible=True, interactive=True)
+        building_version_dropdown = gr.Dropdown(choices=list_cfg_file_versions("building_model_weights"), label="Building Detection Version", value="buildingv1", visible=True, interactive=True)
+        building_pth_dropdown = gr.Dropdown(choices=list_pth_files_in_directory("building_model_weights", "v1"), label="Select a building model file", visible=True, interactive=True)
         building_threshold_slider = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.7, label='Set confidence threshold "%" for buildings', visible=True, interactive=True)
 
     with lcz_row as lcz_options:

backend.py

@@ -29,6 +29,11 @@ from detectron2.checkpoint import DetectionCheckpointer
 from detectron2.utils.visualizer import ColorMode
 from detectron2.structures import Instances
 
+def list_cfg_file_versions(directory):
+    files = os.listdir(directory)
+    # return files that contains substring version and end with .yml
+    cfg_files = [f.split("_")[0] for f in files if (f.endswith(".yml") or f.endswith(".yaml")) and f.startswith(f"{directory.split('_')[0]}v")]
+    return cfg_files
 
 def list_pth_files_in_directory(directory, version="v1"):
     files = os.listdir(directory)
@@ -57,7 +62,7 @@ def update_path_options(version):
     if "tree" in version:
         directory = "tree_model_weights"
     else:
-        directory = "building_model_weight"
+        directory = "building_model_weights"
     return gr.Dropdown(choices=list_pth_files_in_directory(directory, version), label=f"Select a {version.split('v')[0]} model file", visible=True, interactive=True)
 
 # Model for trees
@@ -74,9 +79,9 @@ def tree_model(tree_version_dropdown, tree_pth_dropdown, tree_threshold, device=
 # Model for buildings
 def building_model(building_version_dropdown, building_pth_dropdown, building_threshold, device="cpu"):
     building_cfg = get_cfg()
-    building_cfg.merge_from_file(get_version_cfg_yml(f"building_model_weight/{building_version_dropdown}"))
+    building_cfg.merge_from_file(get_version_cfg_yml(f"building_model_weights/{building_version_dropdown}"))
     building_cfg.MODEL.DEVICE=device
-    building_cfg.MODEL.WEIGHTS = f"building_model_weight/{building_pth_dropdown}"
+    building_cfg.MODEL.WEIGHTS = f"building_model_weights/{building_pth_dropdown}"
     building_cfg.MODEL.ROI_HEADS.NUM_CLASSES = 8  # TODO change this
     building_cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = building_threshold
     building_predictor = DefaultPredictor(building_cfg)
@@ -162,7 +167,7 @@ def visualize_image(im, mode, tree_threshold, building_threshold, color_mode, tr
         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")
+    metadata = get_metadata("urban-small_train", "building_model_weights/_annotations.coco.json")
     visualizer = Visualizer(im[:, :, ::-1], metadata=metadata, scale=0.5, instance_mode=color_mode)
 
     output_image = visualizer.draw_instance_predictions(instances)