def write_vector_outputs(self, labels: SemanticSegmentationLabels) -> None:
"""Write vectorized outputs for all configs in self.vector_outputs."""
import mask_to_polygons.vectorification as vectorification
import mask_to_polygons.processing.denoise as denoise
log.info('Writing vector output to disk.')
label_arr = self._labels_to_full_label_arr(labels)
with click.progressbar(self.vector_outputs) as bar:
for i, vo in enumerate(bar):
if vo.uri is None:
log.info(f'Skipping VectorOutputConfig at index {i} '
'due to missing uri.')
continue
uri = get_local_path(vo.uri, self.tmp_dir)
denoise_radius = vo.denoise
mode = vo.get_mode()
class_mask = (label_arr == vo.class_id).astype(np.uint8)
def transform(x, y):
return self.crs_transformer.pixel_to_map((x, y))
if denoise_radius > 0:
class_mask = denoise.denoise(class_mask, denoise_radius)
if mode == 'buildings':
geojson = vectorification.geojson_from_mask(
mask=class_mask,
transform=transform,
mode=mode,
min_aspect_ratio=vo.min_aspect_ratio,
min_area=vo.min_area,
width_factor=vo.element_width_factor,
thickness=vo.element_thickness)
elif mode == 'polygons':
geojson = vectorification.geojson_from_mask(
mask=class_mask, transform=transform, mode=mode)
str_to_file(geojson, uri)
upload_or_copy(uri, vo.uri)
def save(self, labels):
"""Save.
Args:
labels - (SemanticSegmentationLabels) labels to be saved
"""
local_path = get_local_path(self.uri, self.tmp_dir)
make_dir(local_path, use_dirname=True)
transform = self.crs_transformer.get_affine_transform()
crs = self.crs_transformer.get_image_crs()
band_count = 1
dtype = np.uint8
if self.class_trans:
band_count = 3
mask = (np.zeros((self.extent.ymax, self.extent.xmax), dtype=np.uint8)
if self.vector_output else None)
# https://github.com/mapbox/rasterio/blob/master/docs/quickstart.rst
# https://rasterio.readthedocs.io/en/latest/topics/windowed-rw.html
with rasterio.open(local_path,
'w',
driver='GTiff',
height=self.extent.ymax,
width=self.extent.xmax,
count=band_count,
dtype=dtype,
transform=transform,
crs=crs) as dataset:
for window in labels.get_windows():
label_arr = labels.get_label_arr(window)
window = window.intersection(self.extent)
label_arr = label_arr[0:window.get_height(),
0:window.get_width()]
if mask is not None:
mask[window.ymin:window.ymax,
window.xmin:window.xmax] = label_arr
window = window.rasterio_format()
if self.class_trans:
rgb_labels = self.class_trans.class_to_rgb(label_arr)
for chan in range(3):
dataset.write_band(chan + 1,
rgb_labels[:, :, chan],
window=window)
else:
img = label_arr.astype(dtype)
dataset.write_band(1, img, window=window)
upload_or_copy(local_path, self.uri)
if self.vector_output:
import mask_to_polygons.vectorification as vectorification
import mask_to_polygons.processing.denoise as denoise
for vo in self.vector_output:
denoise_radius = vo.denoise
uri = vo.uri
mode = vo.get_mode()
class_id = vo.class_id
class_mask = np.array(mask == class_id, dtype=np.uint8)
def transform(x, y):
return self.crs_transformer.pixel_to_map((x, y))
if denoise_radius > 0:
class_mask = denoise.denoise(class_mask, denoise_radius)
if uri and mode == 'buildings':
geojson = vectorification.geojson_from_mask(
mask=class_mask,
transform=transform,
mode=mode,
min_aspect_ratio=vo.min_aspect_ratio,
min_area=vo.min_area,
width_factor=vo.element_width_factor,
thickness=vo.element_thickness)
elif uri and mode == 'polygons':
geojson = vectorification.geojson_from_mask(
mask=class_mask, transform=transform, mode=mode)
str_to_file(geojson, uri)
def save(self, labels):
"""Save.
Args:
labels - (SemanticSegmentationLabels) labels to be saved
"""
local_path = get_local_path(self.uri, self.tmp_dir)
make_dir(local_path, use_dirname=True)
transform = self.crs_transformer.get_affine_transform()
crs = self.crs_transformer.get_image_crs()
band_count = 1
dtype = np.uint8
if self.class_trans:
band_count = 3
if self.vector_output:
# We need to store the whole output mask to run feature extraction.
# If the raster is large, this will result in running out of memory, so
# more work will be needed to get this to work in a scalable way. But this
# is complicated because of the need to merge features that are split
# across windows.
mask = np.zeros((self.extent.ymax, self.extent.xmax),
dtype=np.uint8)
else:
mask = None
# https://github.com/mapbox/rasterio/blob/master/docs/quickstart.rst
# https://rasterio.readthedocs.io/en/latest/topics/windowed-rw.html
with rasterio.open(local_path,
'w',
driver='GTiff',
height=self.extent.ymax,
width=self.extent.xmax,
count=band_count,
dtype=dtype,
transform=transform,
crs=crs) as dataset:
for window in labels.get_windows():
class_labels = labels.get_label_arr(window,
clip_extent=self.extent)
clipped_window = ((window.ymin,
window.ymin + class_labels.shape[0]),
(window.xmin,
window.xmin + class_labels.shape[1]))
if mask is not None:
mask[clipped_window[0][0]:clipped_window[0][1],
clipped_window[1][0]:clipped_window[1]
[1]] = class_labels
if self.class_trans:
rgb_labels = self.class_trans.class_to_rgb(class_labels)
for chan in range(3):
dataset.write_band(chan + 1,
rgb_labels[:, :, chan],
window=clipped_window)
else:
img = class_labels.astype(dtype)
dataset.write_band(1, img, window=clipped_window)
upload_or_copy(local_path, self.uri)
if self.vector_output:
import mask_to_polygons.vectorification as vectorification
import mask_to_polygons.processing.denoise as denoise
for vo in self.vector_output:
denoise_radius = vo['denoise']
uri = vo['uri']
mode = vo['mode']
class_id = vo['class_id']
class_mask = np.array(mask == class_id, dtype=np.uint8)
local_geojson_path = get_local_path(uri, self.tmp_dir)
def transform(x, y):
return self.crs_transformer.pixel_to_map((x, y))
if denoise_radius > 0:
class_mask = denoise.denoise(class_mask, denoise_radius)
if uri and mode == 'buildings':
options = vo['building_options']
geojson = vectorification.geojson_from_mask(
mask=class_mask,
transform=transform,
mode=mode,
min_aspect_ratio=options['min_aspect_ratio'],
min_area=options['min_area'],
width_factor=options['element_width_factor'],
thickness=options['element_thickness'])
elif uri and mode == 'polygons':
geojson = vectorification.geojson_from_mask(
mask=class_mask, transform=transform, mode=mode)
if local_geojson_path:
with open(local_geojson_path, 'w') as file_out:
file_out.write(geojson)
upload_or_copy(local_geojson_path, uri)