update app

app.py

@@ -1,26 +1,46 @@
-import gradio as gr
-from backend import visualize_image
 
-# gradio inputs
-image_input = gr.components.Image(type="pil", label="Input Image")
-color_mode_select = gr.components.Radio(choices=["Black/white", "Random", "Segmentation"], label="Color Mode", value="Segmentation")
-mode_dropdown = gr.components.Dropdown(choices=["Trees", "Buildings", "Both"], label="Detection Mode", value="Both")
+import os 
+os.system("python -m pip install --upgrade pip")
+os.system("pip uninstall -y gradio")
+os.system("pip install gradio==4.1.2")
+
+import gradio as gr
+from backend import *
 
-tree_threshold_slider = gr.components.Slider(minimum=0, maximum=1, step=0.1, value=0.7, label='Set confidence threshold "%" for trees')
-building_threshold_slider = gr.components.Slider(minimum=0, maximum=1, step=0.1, value=0.7, label='Set confidence threshold "%" for buildings')
+with gr.Blocks() as demo:
+    gr.Markdown(
+    """
+    #                 Aerial Image Segmentation
+    An instance segmentation demo for identifying boundaries of buildings and trees in aerial images using DETR (End-to-End Object Detection) model with MaskRCNN-101 backbone
+    """
+    )
+    with gr.Row(equal_height=True):
+        with gr.Column():
+            image_input = gr.components.Image(type="pil", label="Input Image")
+        with gr.Column():
+            mode_dropdown = gr.Dropdown(choices=["Trees", "Buildings", "Both"], label="Detection Mode", value="Both")
+            color_mode_select = gr.components.Radio(choices=["Black/white", "Random", "Segmentation"], label="Color Mode", value="Segmentation")
 
-# gradio outputs
-output_image = gr.components.Image(type="pil", label="Output Image")
-title = "Aerial Image Segmentation"
-description = "An instance segmentation demo for identifying boundaries of buildings and trees in aerial images using DETR (End-to-End Object Detection) model with MaskRCNN-101 backbone"
+    # split tree and building into two rows side by side
+    tree_row, building_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_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)
 
-# gradio interface
-interface = gr.Interface(
-    fn=visualize_image,
-    inputs=[image_input, mode_dropdown, tree_threshold_slider, building_threshold_slider, color_mode_select],
-    outputs=output_image,
-    title=title,
-    description=description
-)
+    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_threshold_slider = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.7, label='Set confidence threshold "%" for buildings', visible=True, interactive=True)
 
-interface.launch(debug=True)
+    # mode_dropdown.change(update_visibility, inputs=[mode_dropdown], outputs=[tree_version_dropdown, tree_pth_dropdown, tree_threshold_slider, building_version_dropdown, building_pth_dropdown, building_threshold_slider])
+    mode_dropdown.change(update_row_visibility, inputs=[mode_dropdown], outputs=[tree_row, building_row])
+    tree_version_dropdown.change(update_path_options, inputs=[tree_version_dropdown], outputs=[tree_pth_dropdown])
+    building_version_dropdown.change(update_path_options, inputs=[building_version_dropdown], outputs=[building_pth_dropdown])
+    
+    output_image = gr.components.Image(type="pil", label="Output Image")
+    run_model = gr.Button("Upload Image and Run Model")
+    
+    run_model.click(visualize_image, inputs=[image_input, mode_dropdown, tree_threshold_slider, building_threshold_slider, color_mode_select, tree_version_dropdown, tree_pth_dropdown, building_version_dropdown, building_pth_dropdown], outputs=[output_image])
+    demo.launch()

backend.py

@@ -3,6 +3,8 @@ aerial-segmentation
 Proof of concept showing effectiveness of a fine tuned instance segmentation model for detecting trees.
 """
 import os
+import gradio as gr
+
 import cv2
 os.system("pip install 'git+https://github.com/facebookresearch/detectron2.git'")
 from transformers import DetrFeatureExtractor, DetrForSegmentation
@@ -28,33 +30,67 @@ from detectron2.utils.visualizer import ColorMode
 from detectron2.structures import Instances
 
 
+def list_pth_files_in_directory(directory, version="v1"):
+    files = os.listdir(directory)
+    version = version.split("v")[1]
+    # return files that contains substring version and end with .pth
+    pth_files = [f for f in files if version in f and f.endswith(".pth")]
+    return pth_files
+
+def get_version_cfg_yml(path):
+    directory = path.split("/")[0]
+    version = path.split("/")[1]
+    files = os.listdir(directory)
+    cfg_file = [f for f in files if (f.endswith(".yml") or f.endswith(".yaml")) and version in f]
+    return directory + "/" + cfg_file[0]
+
+def update_row_visibility(mode):
+    visibility = {
+        "tree": mode in ["Trees", "Both"],
+        "building": mode in ["Buildings", "Both"]
+    }
+    tree_row, building_row = gr.Row(visible=visibility["tree"]), gr.Row(visible=visibility["building"])
+    
+    return tree_row, building_row
+
+def update_path_options(version):
+    if "tree" in version:
+        directory = "tree_model_weights"
+    else:
+        directory = "building_model_weight"
+    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
-tree_cfg = get_cfg()
-tree_cfg.merge_from_file("tree_model_weights/treev2_cfg.yml")
-tree_cfg.MODEL.DEVICE='cpu'
-tree_cfg.MODEL.WEIGHTS = "tree_model_weights/treev2_9999.pth"
-tree_cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2
-tree_predictor = DefaultPredictor(tree_cfg)
+def tree_model(tree_version_dropdown, tree_pth_dropdown, tree_threshold, device="cpu"):
+    tree_cfg = get_cfg()
+    tree_cfg.merge_from_file(get_version_cfg_yml(f"tree_model_weights/{tree_version_dropdown}"))
+    tree_cfg.MODEL.DEVICE=device
+    tree_cfg.MODEL.WEIGHTS = f"tree_model_weights/{tree_pth_dropdown}"
+    tree_cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2  # TODO change this
+    tree_cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = tree_threshold
+    tree_predictor = DefaultPredictor(tree_cfg)
+    return tree_predictor
 
 # Model for buildings
-building_cfg = get_cfg()
-building_cfg.merge_from_file("building_model_weight/buildings_poc_cfg.yml")
-building_cfg.MODEL.DEVICE='cpu'
-building_cfg.MODEL.WEIGHTS = "building_model_weight/model_final.pth"  
-building_cfg.MODEL.ROI_HEADS.NUM_CLASSES = 8
-building_predictor = DefaultPredictor(building_cfg)
+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.MODEL.DEVICE=device
+    building_cfg.MODEL.WEIGHTS = f"building_model_weight/{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)
+    return building_predictor
 
 # A function that runs the buildings model on an given image and confidence threshold
-def segment_building(im, confidence_threshold):
-    building_cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = confidence_threshold
+def segment_building(im, building_predictor):
     outputs = building_predictor(im)
     building_instances = outputs["instances"].to("cpu")
 
     return building_instances
 
 # A function that runs the trees model on an given image and confidence threshold
-def segment_tree(im, confidence_threshold):
-    tree_cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = confidence_threshold
+def segment_tree(im, tree_predictor):
     outputs = tree_predictor(im)
     tree_instances = outputs["instances"].to("cpu")
 
@@ -69,50 +105,43 @@ def map_color_mode(color_mode):
     elif color_mode == "Segmentation" or color_mode == None:
         return ColorMode.SEGMENTATION
 
-def visualize_image(im, mode, tree_threshold:float, building_threshold:float, color_mode):
-    im = np.array(im)
-    color_mode = map_color_mode(color_mode)
+def load_predictor(model, version, pth, threshold):
+    return model(version, pth, threshold)
+
+def load_instances(image, predictor, segment_function):
+    return segment_function(image, predictor)
+
+def combine_instances(tree_instances, building_instances):
+    return Instances.cat([tree_instances, building_instances])
 
-    if mode == "Trees":
-        instances = segment_tree(im, tree_threshold)
-    elif mode == "Buildings":
-        instances = segment_building(im, building_threshold)
-    elif mode == "Both" or mode == None:
-        tree_instances = segment_tree(im, tree_threshold)
-        building_instances = segment_building(im, building_threshold)
-        instances = Instances.cat([tree_instances, building_instances])
-
-    metadata = MetadataCatalog.get("urban-trees-fdokv_train")
-    print("metadata", type(metadata), metadata)
-    print('metadata.get("thing_classes")', type(metadata.get("thing_classes")), metadata.get("thing_classes"))
-
-    visualizer = Visualizer(im[:, :, ::-1],
-                            metadata=metadata,
-                            scale=0.5,
-                            instance_mode=color_mode)
-
-    dataset_names = MetadataCatalog.list()
-    print(dataset_names)
-
-    metadata = MetadataCatalog.get("urban-small_train")
-    category_names = metadata.get("thing_classes")
-    print(category_names)
-    with open("building_model_weight/_annotations.coco.json", "r") as f:
+def get_metadata(dataset_name, coco_file):
+    metadata = MetadataCatalog.get(dataset_name)
+    with open(coco_file, "r") as f:
         coco = json.load(f)
-        categories = coco["categories"]
-    print("categories", categories)
+    categories = coco["categories"]
     metadata.thing_classes = [c["name"] for c in categories]
-    print("metadata.thing_classes", metadata.thing_classes )
-    # visualizer = Visualizer(im[:, :, ::-1],
-    #                         metadata=metadata,
-    #                         scale=0.5,
-    #                         instance_mode=color_mode)
-    # # in the visualizer, add category label names to detected instances
-    # for instance in instances:
-    #     label = category_names[instance["category_id"]]
-    #     visualizer.draw_text(label, instance["bbox"][:2])
+    return metadata
 
-    output_image = visualizer.draw_instance_predictions(instances)
+def visualize_image(im, mode, tree_threshold, building_threshold, color_mode, tree_version, tree_pth, building_version, building_pth):
+    im = np.array(im)
+    color_mode = map_color_mode(color_mode)
     
-    return Image.fromarray(output_image.get_image()[:, :, ::-1])
+    instances = None
+
+    if mode in {"Trees", "Both"}:
+        tree_predictor = load_predictor(tree_model, tree_version, tree_pth, tree_threshold)
+        tree_instances = load_instances(im, tree_predictor, segment_tree)
+        instances = tree_instances
+
+    if mode in {"Buildings", "Both"}:
+        building_predictor = load_predictor(building_model, building_version, building_pth, building_threshold)
+        building_instances = load_instances(im, building_predictor, segment_building)
+        instances = building_instances if mode == "Buildings" else combine_instances(instances, building_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)
+
+    output_image = visualizer.draw_instance_predictions(instances)
 
+    return Image.fromarray(output_image.get_image()[:, :, ::-1])