Source code for

from typing import TYPE_CHECKING, Any, Optional

import numpy as np

from import Box
from import ClassConfig
from import SemanticSegmentationLabels
from import LabelSource
from import RasterSource
from import pad_to_window_size

    from import CRSTransformer

[docs]class SemanticSegmentationLabelSource(LabelSource): """A read-only label source for semantic segmentation."""
[docs] def __init__(self, raster_source: RasterSource, class_config: ClassConfig, bbox: Optional[Box] = None): """Constructor. Args: raster_source (RasterSource): A raster source that returns a single channel raster with class_ids as values. null_class_id (int): the null class id used as fill values for when windows go over the edge of the label array. This can be retrieved using class_config.null_class_id. bbox (Optional[Box], optional): User-specified crop of the extent. If None, the full extent available in the source file is used. """ self.raster_source = raster_source self.class_config = class_config if bbox is not None: self.set_bbox(bbox)
[docs] def get_labels(self, window: Optional[Box] = None) -> SemanticSegmentationLabels: """Get labels for a window. Args: window (Optional[Box], optional): Window to get labels for. If None, returns labels covering the full extent of the scene. Returns: SemanticSegmentationLabels: The labels. """ if window is None: window = self.extent label_arr = self.get_label_arr(window) labels = SemanticSegmentationLabels.make_empty( extent=self.extent, num_classes=len(self.class_config), smooth=False) labels[window] = label_arr return labels
[docs] def get_label_arr(self, window: Optional[Box] = None) -> np.ndarray: """Get labels for a window. The returned array will be the same size as the input window. If window overflows the extent, the overflowing region will be filled with the ID of the null class as defined by the class_config. Args: window (Optional[Box], optional): Window (in pixel coords) to get labels for. If None, returns a label array covering the full extent of the scene. Returns: np.ndarray: Label array. """ if window is None: window = self.extent label_arr = self.raster_source.get_chip(window) if label_arr.ndim == 3: label_arr = np.squeeze(label_arr, axis=2) h, w = label_arr.shape if h < window.height or w < window.width: label_arr = pad_to_window_size(label_arr, window, self.extent, self.class_config.null_class_id) return label_arr
@property def bbox(self) -> Box: return self.raster_source.bbox @property def crs_transformer(self) -> 'CRSTransformer': return self.raster_source.crs_transformer
[docs] def set_bbox(self, bbox: 'Box') -> None: self.raster_source.set_bbox(bbox)
[docs] def __getitem__(self, key: Any) -> Any: if isinstance(key, Box): return self.get_label_arr(key) else: return super().__getitem__(key)