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

rslearn/train/transforms/normalize.py

@@ -1,5 +1,6 @@
 """Normalization transforms."""
 
+import warnings
 from typing import Any
 
 import torch
@@ -35,14 +36,17 @@ class Normalize(Transform):
             bands: optionally restrict the normalization to these band indices. If set,
                 mean and std must either be one value, or have length equal to the
                 number of band indices passed here.
-            num_bands: the number of bands per image, to distinguish different images
-                in a time series. If set, then the bands list is repeated for each
-                image, e.g. if bands=[2] then we apply normalization on images[2],
-                images[2+num_bands], images[2+num_bands*2], etc. Or if the bands list
-                is not set, then we apply the mean and std on each image in the time
-                series.
+            num_bands: deprecated, no longer used. Will be removed after 2026-04-01.
         """
         super().__init__()
+
+        if num_bands is not None:
+            warnings.warn(
+                "num_bands is deprecated and no longer used. "
+                "It will be removed after 2026-04-01.",
+                FutureWarning,
+            )
+
         self.mean = torch.tensor(mean)
         self.std = torch.tensor(std)
 
@@ -55,92 +59,37 @@ class Normalize(Transform):
 
         self.selectors = selectors
         self.bands = torch.tensor(bands) if bands is not None else None
-        self.num_bands = num_bands
 
-    def apply_image(
-        self, image: torch.Tensor | RasterImage
-    ) -> torch.Tensor | RasterImage:
+    def apply_image(self, image: RasterImage) -> RasterImage:
         """Normalize the specified image.
 
         Args:
             image: the image to transform.
         """
-
-        def _repeat_mean_and_std(
-            image_channels: int, num_bands: int | None, is_raster_image: bool
-        ) -> tuple[torch.Tensor, torch.Tensor]:
-            """Get mean and std tensor that are suitable for applying on the image."""
-            # We only need to repeat the tensor if both of these are true:
-            # - The mean/std are not just one scalar.
-            # - self.num_bands is set, otherwise we treat the input as a single image.
-            if len(self.mean.shape) == 0:
-                return self.mean, self.std
-            if num_bands is None:
-                return self.mean, self.std
-            num_images = image_channels // num_bands
-            if is_raster_image:
-                # add an extra T dimension, CTHW
-                return self.mean.repeat(num_images)[
-                    :, None, None, None
-                ], self.std.repeat(num_images)[:, None, None, None]
-            else:
-                # add an extra T dimension, CTHW
-                return self.mean.repeat(num_images)[:, None, None], self.std.repeat(
-                    num_images
-                )[:, None, None]
+        # Get mean/std with singleton dims for broadcasting over CTHW.
+        if len(self.mean.shape) == 0:
+            # Scalar - broadcasts naturally.
+            mean, std = self.mean, self.std
+        else:
+            # Vector of length C - add singleton dims for T, H, W.
+            mean = self.mean[:, None, None, None]
+            std = self.std[:, None, None, None]
 
         if self.bands is not None:
-            # User has provided band indices to normalize.
-            # If num_bands is set, then we repeat these for each image in the input
-            # image time series.
-            band_indices = self.bands
-            if self.num_bands:
-                num_images = image.shape[0] // self.num_bands
-                band_indices = torch.cat(
-                    [
-                        band_indices + image_idx * self.num_bands
-                        for image_idx in range(num_images)
-                    ],
-                    dim=0,
+            # Normalize only specific band indices.
+            image.image[self.bands] = (image.image[self.bands] - mean) / std
+            if self.valid_min is not None:
+                image.image[self.bands] = torch.clamp(
+                    image.image[self.bands],
+                    min=self.valid_min,
+                    max=self.valid_max,
                 )
-
-            # We use len(self.bands) here because that is how many bands per timestep
-            # we are actually processing with the mean/std.
-            mean, std = _repeat_mean_and_std(
-                image_channels=len(band_indices),
-                num_bands=len(self.bands),
-                is_raster_image=isinstance(image, RasterImage),
-            )
-            if isinstance(image, torch.Tensor):
-                image[band_indices] = (image[band_indices] - mean) / std
-                if self.valid_min is not None:
-                    image[band_indices] = torch.clamp(
-                        image[band_indices], min=self.valid_min, max=self.valid_max
-                    )
-            else:
-                image.image[band_indices] = (image.image[band_indices] - mean) / std
-                if self.valid_min is not None:
-                    image.image[band_indices] = torch.clamp(
-                        image.image[band_indices],
-                        min=self.valid_min,
-                        max=self.valid_max,
-                    )
         else:
-            mean, std = _repeat_mean_and_std(
-                image_channels=image.shape[0],
-                num_bands=self.num_bands,
-                is_raster_image=isinstance(image, RasterImage),
-            )
-            if isinstance(image, torch.Tensor):
-                image = (image - mean) / std
-                if self.valid_min is not None:
-                    image = torch.clamp(image, min=self.valid_min, max=self.valid_max)
-            else:
-                image.image = (image.image - mean) / std
-                if self.valid_min is not None:
-                    image.image = torch.clamp(
-                        image.image, min=self.valid_min, max=self.valid_max
-                    )
+            image.image = (image.image - mean) / std
+            if self.valid_min is not None:
+                image.image = torch.clamp(
+                    image.image, min=self.valid_min, max=self.valid_max
+                )
         return image
 
     def forward(

rslearn/train/transforms/pad.py

@@ -50,9 +50,7 @@ class Pad(Transform):
         """
         return {"size": torch.randint(low=self.size[0], high=self.size[1], size=())}
 
-    def apply_image(
-        self, image: RasterImage | torch.Tensor, state: dict[str, bool]
-    ) -> RasterImage | torch.Tensor:
+    def apply_image(self, image: RasterImage, state: dict[str, bool]) -> RasterImage:
         """Apply the sampled state on the specified image.
 
         Args:
@@ -105,16 +103,12 @@ class Pad(Transform):
             horizontal_pad = (horizontal_half, horizontal_extra - horizontal_half)
             vertical_pad = (vertical_half, vertical_extra - vertical_half)
 
-        if isinstance(image, RasterImage):
-            image.image = apply_padding(
-                image.image, True, horizontal_pad[0], horizontal_pad[1]
-            )
-            image.image = apply_padding(
-                image.image, False, vertical_pad[0], vertical_pad[1]
-            )
-        else:
-            image = apply_padding(image, True, horizontal_pad[0], horizontal_pad[1])
-            image = apply_padding(image, False, vertical_pad[0], vertical_pad[1])
+        image.image = apply_padding(
+            image.image, True, horizontal_pad[0], horizontal_pad[1]
+        )
+        image.image = apply_padding(
+            image.image, False, vertical_pad[0], vertical_pad[1]
+        )
         return image
 
     def apply_boxes(self, boxes: Any, state: dict[str, bool]) -> torch.Tensor:

rslearn/train/transforms/resize.py

@@ -2,7 +2,6 @@
 
 from typing import Any
 
-import torch
 import torchvision
 from torchvision.transforms import InterpolationMode
 
@@ -40,32 +39,16 @@ class Resize(Transform):
         self.selectors = selectors
         self.interpolation = INTERPOLATION_MODES[interpolation]
 
-    def apply_resize(
-        self, image: torch.Tensor | RasterImage
-    ) -> torch.Tensor | RasterImage:
+    def apply_resize(self, image: RasterImage) -> RasterImage:
         """Apply resizing on the specified image.
 
-        If the image is 2D, it is unsqueezed to 3D and then squeezed
-        back after resizing.
-
         Args:
             image: the image to transform.
         """
-        if isinstance(image, torch.Tensor):
-            if image.dim() == 2:
-                image = image.unsqueeze(0)  # (H, W) -> (1, H, W)
-                result = torchvision.transforms.functional.resize(
-                    image, self.target_size, self.interpolation
-                )
-                return result.squeeze(0)  # (1, H, W) -> (H, W)
-            return torchvision.transforms.functional.resize(
-                image, self.target_size, self.interpolation
-            )
-        else:
-            image.image = torchvision.transforms.functional.resize(
-                image.image, self.target_size, self.interpolation
-            )
-            return image
+        image.image = torchvision.transforms.functional.resize(
+            image.image, self.target_size, self.interpolation
+        )
+        return image
 
     def forward(
         self, input_dict: dict[str, Any], target_dict: dict[str, Any]

rslearn/train/transforms/select_bands.py

@@ -1,9 +1,8 @@
 """The SelectBands transform."""
 
+import warnings
 from typing import Any
 
-from rslearn.train.model_context import RasterImage
-
 from .transform import Transform, read_selector, write_selector
 
 
@@ -17,60 +16,41 @@ class SelectBands(Transform):
         output_selector: str = "image",
         num_bands_per_timestep: int | None = None,
     ):
-        """Initialize a new Concatenate.
+        """Initialize a new SelectBands.
 
         Args:
-            band_indices: the bands to select.
+            band_indices: the bands to select from the channel dimension.
             input_selector: the selector to read the input image.
             output_selector: the output selector under which to save the output image.
-            num_bands_per_timestep: the number of bands per image, to distinguish
-                between stacked images in an image time series. If set, then the
-                band_indices are selected for each image in the time series.
+            num_bands_per_timestep: deprecated, no longer used. Will be removed after
+                2026-04-01.
         """
         super().__init__()
+
+        if num_bands_per_timestep is not None:
+            warnings.warn(
+                "num_bands_per_timestep is deprecated and no longer used. "
+                "It will be removed after 2026-04-01.",
+                FutureWarning,
+            )
+
         self.input_selector = input_selector
         self.output_selector = output_selector
         self.band_indices = band_indices
-        self.num_bands_per_timestep = num_bands_per_timestep
 
     def forward(
         self, input_dict: dict[str, Any], target_dict: dict[str, Any]
     ) -> tuple[dict[str, Any], dict[str, Any]]:
-        """Apply concatenation over the inputs and targets.
+        """Apply band selection over the inputs and targets.
 
         Args:
             input_dict: the input
             target_dict: the target
 
         Returns:
-            normalized (input_dicts, target_dicts) tuple
+            (input_dicts, target_dicts) tuple with selected bands
         """
         image = read_selector(input_dict, target_dict, self.input_selector)
-        num_bands_per_timestep = (
-            self.num_bands_per_timestep
-            if self.num_bands_per_timestep is not None
-            else image.shape[0]
-        )
-        if isinstance(image, RasterImage):
-            assert num_bands_per_timestep == image.shape[0], (
-                "Expect a seperate dimension for timesteps in RasterImages."
-            )
-
-        if image.shape[0] % num_bands_per_timestep != 0:
-            raise ValueError(
-                f"channel dimension {image.shape[0]} is not multiple of bands per timestep {num_bands_per_timestep}"
-            )
-
-        # Copy the band indices for each timestep in the input.
-        wanted_bands: list[int] = []
-        for start_channel_idx in range(0, image.shape[0], num_bands_per_timestep):
-            wanted_bands.extend(
-                [(start_channel_idx + band_idx) for band_idx in self.band_indices]
-            )
-
-        if isinstance(image, RasterImage):
-            image.image = image.image[wanted_bands]
-        else:
-            image = image[wanted_bands]
+        image.image = image.image[self.band_indices]
         write_selector(input_dict, target_dict, self.output_selector, image)
         return input_dict, target_dict

rslearn/train/transforms/sentinel1.py

@@ -33,31 +33,19 @@ class Sentinel1ToDecibels(Transform):
         self.from_decibels = from_decibels
         self.epsilon = epsilon
 
-    def apply_image(
-        self, image: torch.Tensor | RasterImage
-    ) -> torch.Tensor | RasterImage:
+    def apply_image(self, image: RasterImage) -> RasterImage:
         """Normalize the specified image.
 
         Args:
             image: the image to transform.
         """
-        if isinstance(image, torch.Tensor):
-            image_to_process = image
-        else:
-            image_to_process = image.image
         if self.from_decibels:
             # Decibels to linear scale.
-            image_to_process = torch.pow(10.0, image_to_process / 10.0)
+            image.image = torch.pow(10.0, image.image / 10.0)
         else:
             # Linear scale to decibels.
-            image_to_process = 10 * torch.log10(
-                torch.clamp(image_to_process, min=self.epsilon)
-            )
-        if isinstance(image, torch.Tensor):
-            return image_to_process
-        else:
-            image.image = image_to_process
-            return image
+            image.image = 10 * torch.log10(torch.clamp(image.image, min=self.epsilon))
+        return image
 
     def forward(
         self, input_dict: dict[str, Any], target_dict: dict[str, Any]

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