xinwei89 commited on
Commit
6825ad3
1 Parent(s): 52981c1

update app

Browse files
Files changed (2) hide show
  1. app.py +41 -21
  2. backend.py +86 -57
app.py CHANGED
@@ -1,26 +1,46 @@
1
- import gradio as gr
2
- from backend import visualize_image
3
 
4
- # gradio inputs
5
- image_input = gr.components.Image(type="pil", label="Input Image")
6
- color_mode_select = gr.components.Radio(choices=["Black/white", "Random", "Segmentation"], label="Color Mode", value="Segmentation")
7
- mode_dropdown = gr.components.Dropdown(choices=["Trees", "Buildings", "Both"], label="Detection Mode", value="Both")
 
 
 
8
 
9
- tree_threshold_slider = gr.components.Slider(minimum=0, maximum=1, step=0.1, value=0.7, label='Set confidence threshold "%" for trees')
10
- building_threshold_slider = gr.components.Slider(minimum=0, maximum=1, step=0.1, value=0.7, label='Set confidence threshold "%" for buildings')
 
 
 
 
 
 
 
 
 
 
 
11
 
12
- # gradio outputs
13
- output_image = gr.components.Image(type="pil", label="Output Image")
14
- title = "Aerial Image Segmentation"
15
- 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"
 
 
 
16
 
17
- # gradio interface
18
- interface = gr.Interface(
19
- fn=visualize_image,
20
- inputs=[image_input, mode_dropdown, tree_threshold_slider, building_threshold_slider, color_mode_select],
21
- outputs=output_image,
22
- title=title,
23
- description=description
24
- )
25
 
26
- interface.launch(debug=True)
 
 
 
 
 
 
 
 
 
 
 
 
1
 
2
+ import os
3
+ os.system("python -m pip install --upgrade pip")
4
+ os.system("pip uninstall -y gradio")
5
+ os.system("pip install gradio==4.1.2")
6
+
7
+ import gradio as gr
8
+ from backend import *
9
 
10
+ with gr.Blocks() as demo:
11
+ gr.Markdown(
12
+ """
13
+ # Aerial Image Segmentation
14
+ 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
15
+ """
16
+ )
17
+ with gr.Row(equal_height=True):
18
+ with gr.Column():
19
+ image_input = gr.components.Image(type="pil", label="Input Image")
20
+ with gr.Column():
21
+ mode_dropdown = gr.Dropdown(choices=["Trees", "Buildings", "Both"], label="Detection Mode", value="Both")
22
+ color_mode_select = gr.components.Radio(choices=["Black/white", "Random", "Segmentation"], label="Color Mode", value="Segmentation")
23
 
24
+ # split tree and building into two rows side by side
25
+ tree_row, building_row = gr.Row(), gr.Row()
26
+ # tree_col, building_col = gr.Column(elem_id="tree_col"), gr.Column(elem_id="building_col")
27
+ with tree_row as tree_options:
28
+ tree_version_dropdown = gr.Dropdown(choices=["treev1", "treev2"], label="Tree Detection Version", value="treev2", visible=True, interactive=True)
29
+ 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)
30
+ 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)
31
 
32
+ with building_row as building_options:
33
+ building_version_dropdown = gr.Dropdown(choices=["buildingv1", "buildingv2"], label="Building Detection Version", value="buildingv1", visible=True, interactive=True)
34
+ 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)
35
+ 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)
 
 
 
 
36
 
37
+ # 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])
38
+ mode_dropdown.change(update_row_visibility, inputs=[mode_dropdown], outputs=[tree_row, building_row])
39
+ tree_version_dropdown.change(update_path_options, inputs=[tree_version_dropdown], outputs=[tree_pth_dropdown])
40
+ building_version_dropdown.change(update_path_options, inputs=[building_version_dropdown], outputs=[building_pth_dropdown])
41
+
42
+ output_image = gr.components.Image(type="pil", label="Output Image")
43
+ run_model = gr.Button("Upload Image and Run Model")
44
+
45
+ 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])
46
+ demo.launch()
backend.py CHANGED
@@ -3,6 +3,8 @@ aerial-segmentation
3
  Proof of concept showing effectiveness of a fine tuned instance segmentation model for detecting trees.
4
  """
5
  import os
 
 
6
  import cv2
7
  os.system("pip install 'git+https://github.com/facebookresearch/detectron2.git'")
8
  from transformers import DetrFeatureExtractor, DetrForSegmentation
@@ -28,33 +30,67 @@ from detectron2.utils.visualizer import ColorMode
28
  from detectron2.structures import Instances
29
 
30
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
31
  # Model for trees
32
- tree_cfg = get_cfg()
33
- tree_cfg.merge_from_file("tree_model_weights/treev2_cfg.yml")
34
- tree_cfg.MODEL.DEVICE='cpu'
35
- tree_cfg.MODEL.WEIGHTS = "tree_model_weights/treev2_9999.pth"
36
- tree_cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2
37
- tree_predictor = DefaultPredictor(tree_cfg)
 
 
 
38
 
39
  # Model for buildings
40
- building_cfg = get_cfg()
41
- building_cfg.merge_from_file("building_model_weight/buildings_poc_cfg.yml")
42
- building_cfg.MODEL.DEVICE='cpu'
43
- building_cfg.MODEL.WEIGHTS = "building_model_weight/model_final.pth"
44
- building_cfg.MODEL.ROI_HEADS.NUM_CLASSES = 8
45
- building_predictor = DefaultPredictor(building_cfg)
 
 
 
46
 
47
  # A function that runs the buildings model on an given image and confidence threshold
48
- def segment_building(im, confidence_threshold):
49
- building_cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = confidence_threshold
50
  outputs = building_predictor(im)
51
  building_instances = outputs["instances"].to("cpu")
52
 
53
  return building_instances
54
 
55
  # A function that runs the trees model on an given image and confidence threshold
56
- def segment_tree(im, confidence_threshold):
57
- tree_cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = confidence_threshold
58
  outputs = tree_predictor(im)
59
  tree_instances = outputs["instances"].to("cpu")
60
 
@@ -69,50 +105,43 @@ def map_color_mode(color_mode):
69
  elif color_mode == "Segmentation" or color_mode == None:
70
  return ColorMode.SEGMENTATION
71
 
72
- def visualize_image(im, mode, tree_threshold:float, building_threshold:float, color_mode):
73
- im = np.array(im)
74
- color_mode = map_color_mode(color_mode)
 
 
 
 
 
75
 
76
- if mode == "Trees":
77
- instances = segment_tree(im, tree_threshold)
78
- elif mode == "Buildings":
79
- instances = segment_building(im, building_threshold)
80
- elif mode == "Both" or mode == None:
81
- tree_instances = segment_tree(im, tree_threshold)
82
- building_instances = segment_building(im, building_threshold)
83
- instances = Instances.cat([tree_instances, building_instances])
84
-
85
- metadata = MetadataCatalog.get("urban-trees-fdokv_train")
86
- print("metadata", type(metadata), metadata)
87
- print('metadata.get("thing_classes")', type(metadata.get("thing_classes")), metadata.get("thing_classes"))
88
-
89
- visualizer = Visualizer(im[:, :, ::-1],
90
- metadata=metadata,
91
- scale=0.5,
92
- instance_mode=color_mode)
93
-
94
- dataset_names = MetadataCatalog.list()
95
- print(dataset_names)
96
-
97
- metadata = MetadataCatalog.get("urban-small_train")
98
- category_names = metadata.get("thing_classes")
99
- print(category_names)
100
- with open("building_model_weight/_annotations.coco.json", "r") as f:
101
  coco = json.load(f)
102
- categories = coco["categories"]
103
- print("categories", categories)
104
  metadata.thing_classes = [c["name"] for c in categories]
105
- print("metadata.thing_classes", metadata.thing_classes )
106
- # visualizer = Visualizer(im[:, :, ::-1],
107
- # metadata=metadata,
108
- # scale=0.5,
109
- # instance_mode=color_mode)
110
- # # in the visualizer, add category label names to detected instances
111
- # for instance in instances:
112
- # label = category_names[instance["category_id"]]
113
- # visualizer.draw_text(label, instance["bbox"][:2])
114
 
115
- output_image = visualizer.draw_instance_predictions(instances)
 
 
116
 
117
- return Image.fromarray(output_image.get_image()[:, :, ::-1])
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
118
 
 
 
3
  Proof of concept showing effectiveness of a fine tuned instance segmentation model for detecting trees.
4
  """
5
  import os
6
+ import gradio as gr
7
+
8
  import cv2
9
  os.system("pip install 'git+https://github.com/facebookresearch/detectron2.git'")
10
  from transformers import DetrFeatureExtractor, DetrForSegmentation
 
30
  from detectron2.structures import Instances
31
 
32
 
33
+ def list_pth_files_in_directory(directory, version="v1"):
34
+ files = os.listdir(directory)
35
+ version = version.split("v")[1]
36
+ # return files that contains substring version and end with .pth
37
+ pth_files = [f for f in files if version in f and f.endswith(".pth")]
38
+ return pth_files
39
+
40
+ def get_version_cfg_yml(path):
41
+ directory = path.split("/")[0]
42
+ version = path.split("/")[1]
43
+ files = os.listdir(directory)
44
+ cfg_file = [f for f in files if (f.endswith(".yml") or f.endswith(".yaml")) and version in f]
45
+ return directory + "/" + cfg_file[0]
46
+
47
+ def update_row_visibility(mode):
48
+ visibility = {
49
+ "tree": mode in ["Trees", "Both"],
50
+ "building": mode in ["Buildings", "Both"]
51
+ }
52
+ tree_row, building_row = gr.Row(visible=visibility["tree"]), gr.Row(visible=visibility["building"])
53
+
54
+ return tree_row, building_row
55
+
56
+ def update_path_options(version):
57
+ if "tree" in version:
58
+ directory = "tree_model_weights"
59
+ else:
60
+ directory = "building_model_weight"
61
+ 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)
62
+
63
  # Model for trees
64
+ def tree_model(tree_version_dropdown, tree_pth_dropdown, tree_threshold, device="cpu"):
65
+ tree_cfg = get_cfg()
66
+ tree_cfg.merge_from_file(get_version_cfg_yml(f"tree_model_weights/{tree_version_dropdown}"))
67
+ tree_cfg.MODEL.DEVICE=device
68
+ tree_cfg.MODEL.WEIGHTS = f"tree_model_weights/{tree_pth_dropdown}"
69
+ tree_cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2 # TODO change this
70
+ tree_cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = tree_threshold
71
+ tree_predictor = DefaultPredictor(tree_cfg)
72
+ return tree_predictor
73
 
74
  # Model for buildings
75
+ def building_model(building_version_dropdown, building_pth_dropdown, building_threshold, device="cpu"):
76
+ building_cfg = get_cfg()
77
+ building_cfg.merge_from_file(get_version_cfg_yml(f"building_model_weight/{building_version_dropdown}"))
78
+ building_cfg.MODEL.DEVICE=device
79
+ building_cfg.MODEL.WEIGHTS = f"building_model_weight/{building_pth_dropdown}"
80
+ building_cfg.MODEL.ROI_HEADS.NUM_CLASSES = 8 # TODO change this
81
+ building_cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = building_threshold
82
+ building_predictor = DefaultPredictor(building_cfg)
83
+ return building_predictor
84
 
85
  # A function that runs the buildings model on an given image and confidence threshold
86
+ def segment_building(im, building_predictor):
 
87
  outputs = building_predictor(im)
88
  building_instances = outputs["instances"].to("cpu")
89
 
90
  return building_instances
91
 
92
  # A function that runs the trees model on an given image and confidence threshold
93
+ def segment_tree(im, tree_predictor):
 
94
  outputs = tree_predictor(im)
95
  tree_instances = outputs["instances"].to("cpu")
96
 
 
105
  elif color_mode == "Segmentation" or color_mode == None:
106
  return ColorMode.SEGMENTATION
107
 
108
+ def load_predictor(model, version, pth, threshold):
109
+ return model(version, pth, threshold)
110
+
111
+ def load_instances(image, predictor, segment_function):
112
+ return segment_function(image, predictor)
113
+
114
+ def combine_instances(tree_instances, building_instances):
115
+ return Instances.cat([tree_instances, building_instances])
116
 
117
+ def get_metadata(dataset_name, coco_file):
118
+ metadata = MetadataCatalog.get(dataset_name)
119
+ with open(coco_file, "r") as f:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
120
  coco = json.load(f)
121
+ categories = coco["categories"]
 
122
  metadata.thing_classes = [c["name"] for c in categories]
123
+ return metadata
 
 
 
 
 
 
 
 
124
 
125
+ def visualize_image(im, mode, tree_threshold, building_threshold, color_mode, tree_version, tree_pth, building_version, building_pth):
126
+ im = np.array(im)
127
+ color_mode = map_color_mode(color_mode)
128
 
129
+ instances = None
130
+
131
+ if mode in {"Trees", "Both"}:
132
+ tree_predictor = load_predictor(tree_model, tree_version, tree_pth, tree_threshold)
133
+ tree_instances = load_instances(im, tree_predictor, segment_tree)
134
+ instances = tree_instances
135
+
136
+ if mode in {"Buildings", "Both"}:
137
+ building_predictor = load_predictor(building_model, building_version, building_pth, building_threshold)
138
+ building_instances = load_instances(im, building_predictor, segment_building)
139
+ instances = building_instances if mode == "Buildings" else combine_instances(instances, building_instances)
140
+
141
+ # Assuming 'urban-small_train' is intended for both Trees and Buildings
142
+ metadata = get_metadata("urban-small_train", "building_model_weight/_annotations.coco.json")
143
+ visualizer = Visualizer(im[:, :, ::-1], metadata=metadata, scale=0.5, instance_mode=color_mode)
144
+
145
+ output_image = visualizer.draw_instance_predictions(instances)
146
 
147
+ return Image.fromarray(output_image.get_image()[:, :, ::-1])