Source code for

from typing import TYPE_CHECKING, Optional, Sequence, Tuple
from os.path import join
import logging

import numpy as np
import rasterio as rio
from import tqdm

from rastervision.pipeline.file_system import (
    get_local_path, json_to_file, make_dir, sync_to_dir, file_exists,
    download_if_needed, NotReadableError, get_tmp_dir)
from import Box
from import (CRSTransformer, ClassConfig)
from import (SemanticSegmentationLabels,
from import LabelStore
from import SemanticSegmentationLabelSource
from import RGBClassTransformer
from import RasterioSource

    from import (VectorOutputConfig,

log = logging.getLogger(__name__)

[docs]class SemanticSegmentationLabelStore(LabelStore): """Storage for semantic segmentation predictions. Can store predicted class ID raster and class scores raster as GeoTIFFs, and can optionally vectorize predictions and store them as GeoJSON files. """
[docs] def __init__( self, uri: str, extent: Box, crs_transformer: CRSTransformer, class_config: ClassConfig, tmp_dir: Optional[str] = None, vector_outputs: Optional[Sequence['VectorOutputConfig']] = None, save_as_rgb: bool = False, discrete_output: bool = True, smooth_output: bool = False, smooth_as_uint8: bool = False, rasterio_block_size: int = 512): """Constructor. Args: uri (str): Path to directory where the predictions are/will be stored. Smooth scores will be saved as "uri/scores.tif", discrete labels will be stored as "uri/labels.tif", and vector outputs will be saved in "uri/vector_outputs/". extent (Box): The extent of the scene. crs_transformer (CRSTransformer): CRS transformer for correctly mapping from pixel coords to map coords. tmp_dir (Optional[str], optional): Temporary directory to use. If None, will be auto-generated. Defaults to None. vector_outputs (Optional[Sequence[VectorOutputConfig]], optional): List of VectorOutputConfig's containing vectorization configuration information. Only classes for which a VectorOutputConfig is specified will be saved as vectors. If None, no vector outputs will be produced. Defaults to None. class_config (ClassConfig): Class config. save_as_rgb (bool, optional): If True, saves labels as an RGB image, using the class-color mapping in the class_config. Defaults to False. discrete_output (bool, optional): If True, saves labels as a raster of class IDs (one band). Defaults to False. smooth_output (bool, optional): If True, saves labels as a raster of class scores (one band for each class). Defaults to False. smooth_as_uint8 (bool, optional): If True, stores smooth class scores as np.uint8 (0-255) values rather than as np.float32 discrete labels, to help save memory/disk space. Defaults to False. rasterio_block_size (int, optional): Value to set blockxsize and blockysize to. Defaults to 512. """ self.root_uri = uri self.label_uri = join(uri, 'labels.tif') self.score_uri = join(uri, 'scores.tif') self.hits_uri = join(uri, 'pixel_hits.npy') self.tmp_dir = tmp_dir if self.tmp_dir is None: self._tmp_dir = get_tmp_dir() self.tmp_dir = self.vector_outputs = vector_outputs self.extent = extent self.crs_transformer = crs_transformer self.class_config = class_config self.discrete_output = discrete_output self.smooth_output = smooth_output self.smooth_as_uint8 = smooth_as_uint8 self.rasterio_block_size = rasterio_block_size self.class_transformer = None if save_as_rgb: self.class_transformer = RGBClassTransformer(class_config) self.label_source = None self.score_source = None if file_exists(self.label_uri): if self.class_transformer is not None: tfs = [self.class_transformer] else: tfs = [] label_raster_source = RasterioSource( self.label_uri, raster_transformers=tfs) self.label_source = SemanticSegmentationLabelSource( label_raster_source, class_config) if self.smooth_output: if file_exists(self.score_uri): raster_source = RasterioSource(self.score_uri) extents_equal = raster_source.extent == self.extent bands_equal = raster_source.num_channels == len(class_config) self_dtype = np.uint8 if self.smooth_as_uint8 else np.float32 dtypes_equal = raster_source.dtype == self_dtype if extents_equal and bands_equal and dtypes_equal: self.score_source = raster_source else: raise FileExistsError(f'{self.score_uri} already exists ' 'and is incompatible.')
[docs] def get_labels(self) -> SemanticSegmentationLabels: """Get all labels. Returns: SemanticSegmentationLabels """ if self.smooth_output: return self.get_scores() else: return self.get_discrete_labels()
[docs] def get_discrete_labels(self) -> 'SemanticSegmentationDiscreteLabels': """Get all labels. Returns: SemanticSegmentationDiscreteLabels """ if self.label_source is None: raise FileNotFoundError( f'Raster source at {self.label_uri} does not exist.') return self.label_source.get_labels()
[docs] def get_scores(self) -> 'SemanticSegmentationSmoothLabels': """Get all scores. Returns: SemanticSegmentationSmoothLabels """ if self.score_source is None: raise Exception( f'Raster source at {self.score_uri} does not exist ' 'or is not consistent with the current params.')'Loading scores...') extent = self.score_source.extent try: hits_uri_local = download_if_needed(self.hits_uri) hits_arr = np.load(hits_uri_local) except NotReadableError: log.warn(f'Pixel hits array not found at {self.hits_uri}.' 'Setting all pixels hits to 1.') hits_arr = np.ones(extent.size, dtype=np.uint8) score_arr = self.score_source.get_chip(extent) # (H, W, C) --> (C, H, W) score_arr = score_arr.transpose(2, 0, 1) # convert to float if score_arr.dtype == np.uint8: score_arr = score_arr.astype(np.float16) score_arr /= 255 labels = self.empty_labels() assert isinstance(labels, SemanticSegmentationSmoothLabels) labels.pixel_scores = score_arr * hits_arr labels.pixel_hits = hits_arr return labels
[docs] def save(self, labels: SemanticSegmentationLabels, profile: Optional[dict] = None) -> None: """Save labels to disk. More info on rasterio IO: - - Args: labels - (SemanticSegmentationLabels) labels to be saved """ local_root = get_local_path(self.root_uri, self.tmp_dir) make_dir(local_root) height, width = self.extent.size out_profile = dict( driver='GTiff', height=height, width=width, transform=self.crs_transformer.transform, crs=self.crs_transformer.image_crs, blockxsize=min(self.rasterio_block_size, width), blockysize=min(self.rasterio_block_size, height)) if profile is not None: out_profile.update(profile) # if old scores exist, combine them with the new ones if self.score_source is not None:'Old scores found. Merging with current scores.') old_labels = self.get_scores() labels += old_labels if self.discrete_output: labels_path = get_local_path(self.label_uri, self.tmp_dir) self.write_discrete_raster_output(out_profile, labels_path, labels) if self.smooth_output: scores_path = get_local_path(self.score_uri, self.tmp_dir) hits_path = get_local_path(self.hits_uri, self.tmp_dir) self.write_smooth_raster_output(out_profile, scores_path, hits_path, labels) if self.vector_outputs is not None: self.write_vector_outputs(labels) sync_to_dir(local_root, self.root_uri)
[docs] def write_smooth_raster_output( self, out_profile: dict, scores_path: str, hits_path: str, labels: SemanticSegmentationSmoothLabels) -> None: num_bands = labels.num_classes dtype = np.uint8 if self.smooth_as_uint8 else np.float32 out_profile.update(dict(count=num_bands, dtype=dtype)) with, 'w', **out_profile) as ds: windows = [Box.from_rasterio(w) for _, w in ds.block_windows(1)] with tqdm(windows, desc='Saving pixel scores') as bar: for window in bar: window, _ = self._clip_to_extent(self.extent, window) score_arr = labels.get_score_arr(window) if dtype == np.uint8: score_arr = self._scores_to_uint8(score_arr) else: score_arr = score_arr.astype(dtype) self._write_array(ds, window, score_arr) # save pixel hits too, labels.pixel_hits)
[docs] def write_discrete_raster_output( self, out_profile: dict, path: str, labels: SemanticSegmentationLabels) -> None: num_bands = 1 if self.class_transformer is None else 3 dtype = np.uint8 out_profile.update(dict(count=num_bands, dtype=dtype)) null_class_id = self.class_config.null_class_id with, 'w', **out_profile) as ds: windows = [Box.from_rasterio(w) for _, w in ds.block_windows(1)] with tqdm(windows, desc='Saving pixel labels') as bar: for window in bar: label_arr = labels.get_label_arr( window, null_class_id).astype(dtype) window, label_arr = self._clip_to_extent( self.extent, window, label_arr) if self.class_transformer is not None: label_arr = self.class_transformer.class_to_rgb( label_arr) label_arr = label_arr.transpose(2, 0, 1) self._write_array(ds, window, label_arr)
[docs] def write_vector_outputs(self, labels: SemanticSegmentationLabels) -> None: """Write vectorized outputs for all configs in self.vector_outputs.""" from import (denoise, geoms_to_geojson, mask_to_building_polygons, mask_to_polygons)'Writing vector output to disk.') label_arr = labels.get_label_arr(self.extent, self.class_config.null_class_id) with tqdm(self.vector_outputs, desc='Vectorizing predictions') as bar: for i, vo in enumerate(bar): bar.set_postfix( dict( class_id=vo.class_id, mode=vo.get_mode(), denoise_radius=vo.denoise)) if vo.uri is None:'Skipping VectorOutputConfig at index {i} ' 'due to missing uri.') continue class_mask = (label_arr == vo.class_id).astype(np.uint8) if vo.denoise > 0: class_mask = denoise(class_mask, radius=vo.denoise) mode = vo.get_mode() if mode == 'polygons': polys = mask_to_polygons(class_mask) elif mode == 'buildings': polys = mask_to_building_polygons( mask=class_mask, min_area=vo.min_area, width_factor=vo.element_width_factor, thickness=vo.element_thickness) else: raise NotImplementedError() polys = [self.crs_transformer.pixel_to_map(p) for p in polys] geojson = geoms_to_geojson(polys) json_to_file(geojson, vo.uri)
[docs] def empty_labels(self, **kwargs) -> SemanticSegmentationLabels: """Returns an empty SemanticSegmentationLabels object.""" args = dict( extent=self.extent, num_classes=len(self.class_config), smooth=self.smooth_output) args.update(**kwargs) labels = SemanticSegmentationLabels.make_empty(**args) return labels
def _write_array(self, dataset: rio.DatasetReader, window: Box, arr: np.ndarray) -> None: """Write array out to a rasterio dataset. Array must be of shape (C, H, W). """ rio_window = window.rasterio_format() if len(arr.shape) == 2: dataset.write_band(1, arr, window=rio_window) else: for i, band in enumerate(arr, start=1): dataset.write_band(i, band, window=rio_window) def _clip_to_extent(self, extent: Box, window: Box, arr: Optional[np.ndarray] = None ) -> Tuple[Box, Optional[np.ndarray]]: clipped_window = window.intersection(extent) if arr is not None: h, w = clipped_window.size arr = arr[:h, :w] return clipped_window, arr def _scores_to_uint8(self, score_arr: np.ndarray) -> np.ndarray: """Quantize scores to uint8 (0-255).""" score_arr *= 255 score_arr = np.around(score_arr, out=score_arr) score_arr = score_arr.astype(np.uint8) return score_arr