-
Notifications
You must be signed in to change notification settings - Fork 0
/
segment_trees.py
213 lines (183 loc) · 10.9 KB
/
segment_trees.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
import os
import shutil
from glob import glob
import fire
from deepent.config import add_deepent_config
from deepent.data.build_tiled_json_dataset import DataTiler, ignore_black_tiles, remove_no_annotations, \
remove_small_segment_coverage, remove_small_bboxes
from deepent.data.untiler import Untiler
from detectron2.config import get_cfg
from tools.predictor import RGBDPredictor
PREDICTORS = {'fuselong90': {'model': 'output/fuse_long/model_0089999.pth'},
'fuselong75': {'model': 'output/fuse_long/model_0074999.pth'},
'fuselong60': {'model': 'output/fuse_long/model_0059999.pth'},
'fuselong45': {'model': 'output/fuse_long/model_0044999.pth'},
'fuselong30': {'model': 'output/fuse_long/model_0029999.pth'},
'fuselong15': {'model': 'output/fuse_long/model_0014999.pth'},
'fuselong5': {'model': 'output/fuse_long/model_0004999.pth'},
'fuse90': {'model': 'output/baseline_fuse_07_02_2020/model_0089999.pth'},
'fuse75': {'model': 'output/baseline_fuse_07_02_2020/model_0074999.pth'},
'fuse60': {'model': 'output/baseline_fuse_07_02_2020/model_0059999.pth'},
'fuse45': {'model': 'output/baseline_fuse_07_02_2020/model_0044999.pth'},
'fuse30': {'model': 'output/baseline_fuse_07_02_2020/model_0029999.pth'},
'fuse15': {'model': 'output/baseline_fuse_07_02_2020/model_0014999.pth'},
'fuse5': {'model': 'output/baseline_fuse_07_02_2020/model_0004999.pth'},
'fuselateral': {'model': 'output/fuse_lateral/model_final.pth'},
'fuselateral60': {'model': 'output/fuse_lateral/model_0059999.pth'},
'fuselateral45': {'model': 'output/fuse_lateral/model_0044999.pth'},
'fuselateral15': {'model': 'output/fuse_lateral/model_0014999.pth'},
'fuselateral5': {'model': 'output/fuse_lateral/model_0004999.pth'},
'fuse24_90': {'model': 'output/fuse_24/model_0089999.pth'},
'fuse24_75': {'model': 'output/fuse_24/model_0074999.pth'},
'fuse24_60': {'model': 'output/fuse_24/model_0059999.pth'},
'fuse24_45': {'model': 'output/fuse_24/model_0044999.pth'},
'fuse24_30': {'model': 'output/fuse_24/model_0029999.pth'},
'fuse24_15': {'model': 'output/fuse_24/model_0014999.pth'},
'fuse24_5': {'model': 'output/fuse_24/model_0004999.pth'},
'fuse35_90': {'model': 'output/fuse35/model_final.pth'},
'fuse35_75': {'model': 'output/fuse35/model_0074999.pth'},
'fuse35_60': {'model': 'output/fuse35/model_0059999.pth'},
'fuse35_45': {'model': 'output/fuse35/model_0044999.pth'},
'fuse35_30': {'model': 'output/fuse35/model_0029999.pth'},
'fuse35_15': {'model': 'output/fuse35/model_0014999.pth'},
'fuse35_5': {'model': 'output/fuse35/model_0004999.pth'},
'rgbd40': {'model': 'output/rgbd/model_0039999.pth'},
'rgbd30': {'model': 'output/rgbd/model_0029999.pth'},
'rgbd15': {'model': 'output/rgbd/model_0014999.pth'},
'rgbd5': {'model': 'output/rgbd/model_0004999.pth'},
'cpts45': {'model': 'output/cpts/model_0044999.pth'},
'cpts30': {'model': 'output/cpts/model_0029999.pth'},
'cpts15': {'model': 'output/cpts/model_0014999.pth'},
'cpts5': {'model': 'output/cpts/model_0004999.pth'},
'rgb55': {'model': 'output/baseline_25_01_2020/model_0054999.pth'},
'rgb45': {'model': 'output/baseline_25_01_2020/model_0044999.pth'},
'rgb30': {'model': 'output/baseline_25_01_2020/model_0029999.pth'},
'rgb15': {'model': 'output/baseline_25_01_2020/model_0014999.pth'},
'rgb5': {'model': 'output/baseline_25_01_2020/model_0004999.pth'}
}
def run_description(predictors):
for i, ensemble in enumerate(predictors):
print(f'Run {i + 1}: {"+".join(ensemble)}')
for k, model in enumerate(ensemble):
try:
data = PREDICTORS[model]
print(
f'Model {k + 1}: Name: {model}'
f' \n\t Weights: {data["model"]}\n')
except KeyError:
print(f'Model \'{model}\' not in Model Zoo!')
def get_predictor(model, confidence):
cfg = get_cfg()
add_deepent_config(cfg)
model_config = os.path.join(os.path.dirname(model), 'config.yaml')
cfg.merge_from_file(os.path.realpath(model_config))
cfg.MODEL.WEIGHTS = os.path.realpath(model)
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = confidence
cfg.freeze()
return RGBDPredictor(cfg)
class ProjectManager:
"""
Creates a manger to handle data pre-processing, predictions, and post-processing
on a dataset of Ortho, CHM pairs.
:param data:
Option 1:
File path to a folder containing multiple datasets of Ortho+CHM pairs.
Option 2:
Direct path to a single ortho and CHM.
Option 3 (Advanced):
Direct path to inference tiles generated using deepent/data/build_tiled_dataset.py.
For this option, the offsets.json file must be located directly at the given location.
Tiles will not be deleted after run.
:param shapefile_location: File path to where the output shapefiles should be written.
Default drone-tree-id/shapefiles
:param predictors: A string of predictor ensembles separated by commas. Predictors are specified by their names
in the MODELZOO, ensembles are denoted with the '+' symbol.
Default 'redwood+sequoia'
Other examples: 'redwood, sequoia, redwood+sequoia', 'sequoia, redwood+sequoia'
:param datasets: If multiple datasets exist in --data, then this flag may be passed as a comma separated list of
regular expressions to indicate which datasets should be kept
Default '*'
Examples 'CPT*, Kelowna, AM*'
:param confidence: A value in [0, 1]. Confidence threshold for keeping a prediction of a single tree.
Default .5
:param duplicate_tol: A value in [0, 1]. Threshold for keeping a new prediction that is similar to an existing one
i.e a new polygon will be written if it is less then (100*duplicate_tol)% similar to all others
Default .85
:param min_area: A float value. Minimum polygon area in m^2 to keep in final prediction.
Default 4
:param use_generated_tiles: Look for tiles have have been previously generated
Default False
:param retain_tiles: Keep generated tiles after a run
Default False
"""
def __init__(self, data: str, shapefile_location: str = None, predictors='rgb45+fuse30', datasets: str = '*',
confidence: float = .5,
duplicate_tol: float = .8,
min_area: float = 4,
use_generated_tiles: bool = False,
retain_tiles: bool = False):
self.is_tile_path = len(glob(os.path.join(data, 'offsets.json'))) == 1
self.path_to_raw_data = os.path.realpath(data)
self.multiple_datasets = len(
glob(os.path.join(self.path_to_raw_data, '*.png' if self.is_tile_path else '*.tif'))) == 0
assert not (self.multiple_datasets and self.is_tile_path), 'A custom tile path with multiple ' \
'datasets is not currently supported'
self.confidence, self.duplicate_tol, self.min_area = confidence, duplicate_tol, min_area
if isinstance(predictors, tuple):
predictors = ','.join(predictors)
combos = [[y.strip() for y in x.split('+')] for x in predictors.split(',')]
run_description(combos)
self.predictor_combos = [
('-'.join(combo), [get_predictor(**PREDICTORS[model], confidence=self.confidence) for model in combo]) for
combo in combos]
if shapefile_location is None:
self.output = os.path.realpath('shapefiles')
else:
self.output = os.path.realpath(shapefile_location)
if not self.is_tile_path:
self.data_tiler = DataTiler(self.path_to_raw_data, os.path.join(self.path_to_raw_data, 'tmp'),
vertical_overlay=320,
horizontal_overlay=320, dataset_regex=datasets.strip().split(','),
cleanup_on_init=not use_generated_tiles, create_inference_tiles=True)
if not use_generated_tiles:
self.prepare_inference_set()
self.run_predictions()
if not self.is_tile_path:
if use_generated_tiles:
retain_tiles = True
if not retain_tiles:
self.clean_tiles()
def prepare_inference_set(self):
self.data_tiler.tile_dataset(
tile_filtering_function=ignore_black_tiles(thresh=.99),
annotation_filtering_function=lambda an: remove_no_annotations(an) and remove_small_segment_coverage()(an),
bbox_filtering_function=remove_small_bboxes(1000))
def run_predictions(self):
# Case 1, user has passed a path directly to tiles
if self.is_tile_path:
dataset_name = os.path.basename(self.path_to_raw_data)
for predictor_name, predictors in self.predictor_combos:
Untiler(predictors).predict_and_untile(self.path_to_raw_data,
os.path.join(self.output, dataset_name,
f'{dataset_name}_{predictor_name}'),
duplicate_tol=self.duplicate_tol, min_area=self.min_area)
# Case 2, user has passed a path to untiled data
for dataset_name, dataset_path in zip(self.data_tiler.dataset_names, self.data_tiler.dataset_input_paths):
if self.multiple_datasets:
# If a full data collection is passed the tile directory will be in the given path
dataset_path = self.path_to_raw_data
for predictor_name, predictors in self.predictor_combos:
Untiler(predictors).predict_and_untile(os.path.join(dataset_path, 'tmp', 'tiles', dataset_name),
os.path.join(self.output, dataset_name,
f'{dataset_name}_{predictor_name}'),
duplicate_tol=self.duplicate_tol, min_area=self.min_area)
def clean_tiles(self):
if self.multiple_datasets:
shutil.rmtree(os.path.join(self.path_to_raw_data, 'tmp'))
else:
for dp in self.data_tiler.dataset_input_paths:
shutil.rmtree(os.path.join(dp, 'tmp'))
def __str__(self):
return f'Prediction complete. Output written to {self.output}'
if __name__ == '__main__':
fire.Fire(ProjectManager)