rslearn 0.0.24__py3-none-any.whl → 0.0.25__py3-none-any.whl

rslearn/train/dataset.py

@@ -445,6 +445,7 @@ class SplitConfig:
         overlap_ratio: float | None = None,
         load_all_patches: bool | None = None,
         skip_targets: bool | None = None,
+        output_layer_name_skip_inference_if_exists: str | None = None,
     ) -> None:
         """Initialize a new SplitConfig.
 
@@ -467,6 +468,10 @@ class SplitConfig:
                 for each window, read all patches as separate sequential items in the
                 dataset.
             skip_targets: whether to skip targets when loading inputs
+            output_layer_name_skip_inference_if_exists: optional name of the output layer used during prediction.
+                If set, windows that already
+                have this layer completed will be skipped (useful for resuming
+                partial inference runs).
         """
         self.groups = groups
         self.names = names
@@ -477,6 +482,9 @@ class SplitConfig:
         self.sampler = sampler
         self.patch_size = patch_size
         self.skip_targets = skip_targets
+        self.output_layer_name_skip_inference_if_exists = (
+            output_layer_name_skip_inference_if_exists
+        )
 
         # Note that load_all_patches are handled by the RslearnDataModule rather than
         # the ModelDataset.
@@ -504,6 +512,7 @@ class SplitConfig:
             overlap_ratio=self.overlap_ratio,
             load_all_patches=self.load_all_patches,
             skip_targets=self.skip_targets,
+            output_layer_name_skip_inference_if_exists=self.output_layer_name_skip_inference_if_exists,
         )
         if other.groups:
             result.groups = other.groups
@@ -527,6 +536,10 @@ class SplitConfig:
             result.load_all_patches = other.load_all_patches
         if other.skip_targets is not None:
             result.skip_targets = other.skip_targets
+        if other.output_layer_name_skip_inference_if_exists is not None:
+            result.output_layer_name_skip_inference_if_exists = (
+                other.output_layer_name_skip_inference_if_exists
+            )
         return result
 
     def get_patch_size(self) -> tuple[int, int] | None:
@@ -549,16 +562,26 @@ class SplitConfig:
         """Returns whether skip_targets is enabled (default False)."""
         return True if self.skip_targets is True else False
 
+    def get_output_layer_name_skip_inference_if_exists(self) -> str | None:
+        """Returns output layer to use for resume checks (default None)."""
+        return self.output_layer_name_skip_inference_if_exists
+
 
-def check_window(inputs: dict[str, DataInput], window: Window) -> Window | None:
+def check_window(
+    inputs: dict[str, DataInput],
+    window: Window,
+    output_layer_name_skip_inference_if_exists: str | None = None,
+) -> Window | None:
     """Verify that the window has the required layers based on the specified inputs.
 
     Args:
         inputs: the inputs to the dataset.
         window: the window to check.
+        output_layer_name_skip_inference_if_exists: optional name of the output layer to check for existence.
 
     Returns:
-        the window if it has all the required inputs or None otherwise
+        the window if it has all the required inputs and does not need to be skipped
+        due to an existing output layer; or None otherwise
     """
 
     # Make sure window has all the needed layers.
@@ -588,6 +611,16 @@ def check_window(inputs: dict[str, DataInput], window: Window) -> Window | None:
             )
             return None
 
+    # Optionally skip windows that already have the specified output layer completed.
+    if output_layer_name_skip_inference_if_exists is not None:
+        if window.is_layer_completed(output_layer_name_skip_inference_if_exists):
+            logger.debug(
+                "Skipping window %s since output layer '%s' already exists",
+                window.name,
+                output_layer_name_skip_inference_if_exists,
+            )
+            return None
+
     return window
 
 
@@ -648,7 +681,14 @@ class ModelDataset(torch.utils.data.Dataset):
         new_windows = []
         if workers == 0:
             for window in windows:
-                if check_window(self.inputs, window) is None:
+                if (
+                    check_window(
+                        self.inputs,
+                        window,
+                        output_layer_name_skip_inference_if_exists=self.split_config.get_output_layer_name_skip_inference_if_exists(),
+                    )
+                    is None
+                ):
                     continue
                 new_windows.append(window)
         else:
@@ -660,6 +700,7 @@ class ModelDataset(torch.utils.data.Dataset):
                     dict(
                         inputs=self.inputs,
                         window=window,
+                        output_layer_name_skip_inference_if_exists=self.split_config.get_output_layer_name_skip_inference_if_exists(),
                     )
                     for window in windows
                 ],

rslearn/train/tasks/detection.py

@@ -14,27 +14,10 @@ from torchmetrics import Metric, MetricCollection
 
 from rslearn.train.model_context import RasterImage, SampleMetadata
 from rslearn.utils import Feature, STGeometry
+from rslearn.utils.colors import DEFAULT_COLORS
 
 from .task import BasicTask
 
-DEFAULT_COLORS = [
-    (255, 0, 0),
-    (0, 255, 0),
-    (0, 0, 255),
-    (255, 255, 0),
-    (0, 255, 255),
-    (255, 0, 255),
-    (0, 128, 0),
-    (255, 160, 122),
-    (139, 69, 19),
-    (128, 128, 128),
-    (255, 255, 255),
-    (143, 188, 143),
-    (95, 158, 160),
-    (255, 200, 0),
-    (128, 0, 0),
-]
-
 
 class DetectionTask(BasicTask):
     """A point or bounding box detection task."""

rslearn/train/tasks/segmentation.py

@@ -17,28 +17,10 @@ from rslearn.train.model_context import (
     SampleMetadata,
 )
 from rslearn.utils import Feature
+from rslearn.utils.colors import DEFAULT_COLORS
 
 from .task import BasicTask
 
-# TODO: This is duplicated code fix it
-DEFAULT_COLORS = [
-    (255, 0, 0),
-    (0, 255, 0),
-    (0, 0, 255),
-    (255, 255, 0),
-    (0, 255, 255),
-    (255, 0, 255),
-    (0, 128, 0),
-    (255, 160, 122),
-    (139, 69, 19),
-    (128, 128, 128),
-    (255, 255, 255),
-    (143, 188, 143),
-    (95, 158, 160),
-    (255, 200, 0),
-    (128, 0, 0),
-]
-
 
 class SegmentationTask(BasicTask):
     """A segmentation (per-pixel classification) task."""

rslearn/utils/colors.py

@@ -0,0 +1,20 @@
+"""Default color palette for visualizations."""
+
+DEFAULT_COLORS = [
+    (0, 0, 0),
+    (255, 0, 0),
+    (0, 255, 0),
+    (0, 0, 255),
+    (255, 255, 0),
+    (0, 255, 255),
+    (255, 0, 255),
+    (0, 128, 0),
+    (255, 160, 122),
+    (139, 69, 19),
+    (128, 128, 128),
+    (255, 255, 255),
+    (143, 188, 143),
+    (95, 158, 160),
+    (255, 200, 0),
+    (128, 0, 0),
+]

rslearn/vis/__init__.py

@@ -0,0 +1 @@
+"""Visualization module for rslearn datasets."""

rslearn/vis/normalization.py

@@ -0,0 +1,127 @@
+"""Normalization functions for raster data visualization."""
+
+from collections.abc import Callable
+from enum import StrEnum
+
+import numpy as np
+
+from rslearn.log_utils import get_logger
+
+logger = get_logger(__name__)
+
+
+class NormalizationMethod(StrEnum):
+    """Normalization methods for raster data visualization."""
+
+    SENTINEL2_RGB = "sentinel2_rgb"
+    """Divide by 10 and clip (for Sentinel-2 B04/B03/B02 bands)."""
+
+    PERCENTILE = "percentile"
+    """Use 2-98 percentile clipping."""
+
+    MINMAX = "minmax"
+    """Use min-max stretch."""
+
+
+def _normalize_sentinel2_rgb(band: np.ndarray) -> np.ndarray:
+    """Normalize band using Sentinel-2 RGB method (divide by 10 and clip).
+
+    Args:
+        band: Input band data
+
+    Returns:
+        Normalized band as uint8 array
+    """
+    band = band / 10.0
+    band = np.clip(band, 0, 255).astype(np.uint8)
+    return band
+
+
+def _normalize_percentile(band: np.ndarray) -> np.ndarray:
+    """Normalize band using 2-98 percentile clipping.
+
+    Args:
+        band: Input band data
+
+    Returns:
+        Normalized band as uint8 array
+    """
+    valid_pixels = band[~np.isnan(band)]
+    if len(valid_pixels) == 0:
+        return np.zeros_like(band, dtype=np.uint8)
+    vmin, vmax = np.nanpercentile(valid_pixels, (2, 98))
+    if vmax == vmin:
+        return np.zeros_like(band, dtype=np.uint8)
+    band = np.clip(band, vmin, vmax)
+    band = ((band - vmin) / (vmax - vmin) * 255).astype(np.uint8)
+    return band
+
+
+def _normalize_minmax(band: np.ndarray) -> np.ndarray:
+    """Normalize band using min-max stretch.
+
+    Args:
+        band: Input band data
+
+    Returns:
+        Normalized band as uint8 array
+    """
+    vmin, vmax = np.nanmin(band), np.nanmax(band)
+    if vmax == vmin:
+        return np.zeros_like(band, dtype=np.uint8)
+    band = np.clip(band, vmin, vmax)
+    band = ((band - vmin) / (vmax - vmin) * 255).astype(np.uint8)
+    return band
+
+
+_NORMALIZATION_FUNCTIONS: dict[
+    NormalizationMethod, Callable[[np.ndarray], np.ndarray]
+] = {
+    NormalizationMethod.SENTINEL2_RGB: _normalize_sentinel2_rgb,
+    NormalizationMethod.PERCENTILE: _normalize_percentile,
+    NormalizationMethod.MINMAX: _normalize_minmax,
+}
+
+
+def normalize_band(
+    band: np.ndarray, method: str | NormalizationMethod = "sentinel2_rgb"
+) -> np.ndarray:
+    """Normalize band to 0-255 range.
+
+    Args:
+        band: Input band data
+        method: Normalization method (string or NormalizationMethod enum)
+            - 'sentinel2_rgb': Divide by 10 and clip (for B04/B03/B02)
+            - 'percentile': Use 2-98 percentile clipping
+            - 'minmax': Use min-max stretch
+
+    Returns:
+        Normalized band as uint8 array
+    """
+    method_enum = NormalizationMethod(method) if isinstance(method, str) else method
+    normalize_func = _NORMALIZATION_FUNCTIONS.get(method_enum)
+    if normalize_func is None:
+        raise ValueError(f"Unknown normalization method: {method_enum}")
+    return normalize_func(band)
+
+
+def normalize_array(
+    array: np.ndarray, method: str | NormalizationMethod = "sentinel2_rgb"
+) -> np.ndarray:
+    """Normalize a multi-band array to 0-255 range.
+
+    Args:
+        array: Input array with shape (channels, height, width) from RasterFormat.decode_raster
+        method: Normalization method (applied per-band, string or NormalizationMethod enum)
+
+    Returns:
+        Normalized array as uint8 with shape (height, width, channels)
+    """
+    if array.ndim == 3:
+        array = np.moveaxis(array, 0, -1)
+
+    normalized = np.zeros_like(array, dtype=np.uint8)
+    for i in range(array.shape[-1]):
+        normalized[..., i] = normalize_band(array[..., i], method)
+
+    return normalized

rslearn/vis/render_raster_label.py

@@ -0,0 +1,96 @@
+"""Functions for rendering raster label masks (e.g., segmentation masks)."""
+
+import numpy as np
+from PIL import Image
+from rasterio.warp import Resampling
+
+from rslearn.config import DType, LayerConfig
+from rslearn.dataset import Window
+from rslearn.log_utils import get_logger
+from rslearn.train.dataset import DataInput, read_raster_layer_for_data_input
+from rslearn.utils.geometry import PixelBounds, ResolutionFactor
+
+logger = get_logger(__name__)
+
+
+def render_raster_label(
+    label_array: np.ndarray,
+    label_colors: dict[str, tuple[int, int, int]],
+    layer_config: LayerConfig,
+) -> np.ndarray:
+    """Render a raster label array as a colored mask numpy array.
+
+    Args:
+        label_array: Raster label array with shape (bands, height, width) - typically single band
+        label_colors: Dictionary mapping label class names to RGB color tuples
+        layer_config: LayerConfig object (to access class_names if available)
+
+    Returns:
+        Array with shape (height, width, 3) as uint8
+    """
+    if label_array.ndim == 3:
+        label_values = label_array[0, :, :]
+    else:
+        label_values = label_array
+
+    height, width = label_values.shape
+    mask_img = np.zeros((height, width, 3), dtype=np.uint8)
+    valid_mask = ~np.isnan(label_values)
+
+    if not layer_config.class_names:
+        raise ValueError(
+            "class_names must be specified in config for raster label layer"
+        )
+
+    label_int = label_values.astype(np.int32)
+    for idx in range(len(layer_config.class_names)):
+        class_name = layer_config.class_names[idx]
+        color = label_colors.get(str(class_name), (0, 0, 0))
+        mask = (label_int == idx) & valid_mask
+        mask_img[mask] = color
+
+    img = Image.fromarray(mask_img, mode="RGB")
+    return np.array(img)
+
+
+def read_raster_layer(
+    window: Window,
+    layer_name: str,
+    layer_config: LayerConfig,
+    band_names: list[str],
+    group_idx: int = 0,
+    bounds: PixelBounds | None = None,
+) -> np.ndarray:
+    """Read a raster layer for visualization.
+
+    This reads bands from potentially multiple band sets to get the requested bands.
+    Uses read_raster_layer_for_data_input from rslearn.train.dataset.
+
+    Args:
+        window: The window to read from
+        layer_name: The layer name
+        layer_config: The layer configuration
+        band_names: List of band names to read (e.g., ["B04", "B03", "B02"])
+        group_idx: The item group index (default 0)
+        bounds: Optional bounds to read. If None, uses window.bounds
+
+    Returns:
+        Array with shape (bands, height, width) as float32
+    """
+    if bounds is None:
+        bounds = window.bounds
+
+    data_input = DataInput(
+        data_type="raster",
+        layers=[layer_name],
+        bands=band_names,
+        dtype=DType.FLOAT32,
+        resolution_factor=ResolutionFactor(),  # Default 1/1, no scaling
+        resampling=Resampling.nearest,
+    )
+
+    image_tensor = read_raster_layer_for_data_input(
+        window, bounds, layer_name, group_idx, layer_config, data_input
+    )
+
+    return image_tensor.numpy().astype(np.float32)

rslearn/vis/render_sensor_image.py

@@ -0,0 +1,27 @@
+"""Functions for rendering raster sensor images (e.g., Sentinel-2, Landsat)."""
+
+import numpy as np
+
+from .normalization import normalize_array
+
+
+def render_sensor_image(
+    array: np.ndarray,
+    normalization_method: str,
+) -> np.ndarray:
+    """Render a raster sensor image array as a numpy array.
+
+    Args:
+        array: Array with shape (channels, height, width) from RasterFormat.decode_raster
+        normalization_method: Normalization method to apply
+
+    Returns:
+        Array with shape (height, width, channels) as uint8
+    """
+    normalized = normalize_array(array, normalization_method)
+
+    # If more than 3 channels, take only the first 3 for RGB
+    if normalized.shape[-1] > 3:
+        normalized = normalized[:, :, :3]
+
+    return normalized