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

rslearn/data_sources/planetary_computer.py

@@ -3,7 +3,7 @@
 import os
 import tempfile
 import xml.etree.ElementTree as ET
-from datetime import timedelta
+from datetime import datetime, timedelta
 from typing import Any
 
 import affine
@@ -12,6 +12,7 @@ import planetary_computer
 import rasterio
 import requests
 from rasterio.enums import Resampling
+from typing_extensions import override
 from upath import UPath
 
 from rslearn.config import LayerConfig
@@ -24,11 +25,104 @@ from rslearn.tile_stores import TileStore, TileStoreWithLayer
 from rslearn.utils.fsspec import join_upath
 from rslearn.utils.geometry import PixelBounds, Projection, STGeometry
 from rslearn.utils.raster_format import get_raster_projection_and_bounds
+from rslearn.utils.stac import StacClient, StacItem
 
 from .copernicus import get_harmonize_callback
 
 logger = get_logger(__name__)
 
+# Max limit accepted by Planetary Computer API.
+PLANETARY_COMPUTER_LIMIT = 1000
+
+
+class PlanetaryComputerStacClient(StacClient):
+    """A StacClient subclass that handles Planetary Computer's pagination limits.
+
+    Planetary Computer STAC API does not support standard pagination and has a max
+    limit of 1000. If the initial query returns 1000 items, this client paginates
+    by sorting by ID and using gt (greater than) queries to fetch subsequent pages.
+    """
+
+    @override
+    def search(
+        self,
+        collections: list[str] | None = None,
+        bbox: tuple[float, float, float, float] | None = None,
+        intersects: dict[str, Any] | None = None,
+        date_time: datetime | tuple[datetime, datetime] | None = None,
+        ids: list[str] | None = None,
+        limit: int | None = None,
+        query: dict[str, Any] | None = None,
+        sortby: list[dict[str, str]] | None = None,
+    ) -> list[StacItem]:
+        # We will use sortby for pagination, so the caller must not set it.
+        if sortby is not None:
+            raise ValueError("sortby must not be set for PlanetaryComputerStacClient")
+
+        # First, try a simple query with the PC limit to detect if pagination is needed.
+        # We always use PLANETARY_COMPUTER_LIMIT for the request because PC doesn't
+        # support standard pagination, and we need to detect when we hit the limit
+        # to switch to ID-based pagination.
+        # We could just start sorting by ID here and do pagination, but we treate it as
+        # a special case to avoid sorting since that seems to speed up the query.
+        stac_items = super().search(
+            collections=collections,
+            bbox=bbox,
+            intersects=intersects,
+            date_time=date_time,
+            ids=ids,
+            limit=PLANETARY_COMPUTER_LIMIT,
+            query=query,
+        )
+
+        # If we got fewer than the PC limit, we have all the results.
+        if len(stac_items) < PLANETARY_COMPUTER_LIMIT:
+            return stac_items
+
+        # We hit the limit, so we need to paginate by ID.
+        # Re-fetch with sorting by ID to ensure consistent ordering for pagination.
+        logger.debug(
+            "Initial request returned %d items (at limit), switching to ID pagination",
+            len(stac_items),
+        )
+
+        all_items: list[StacItem] = []
+        last_id: str | None = None
+
+        while True:
+            # Build query with id > last_id if we're paginating.
+            combined_query: dict[str, Any] = dict(query) if query else {}
+            if last_id is not None:
+                combined_query["id"] = {"gt": last_id}
+
+            stac_items = super().search(
+                collections=collections,
+                bbox=bbox,
+                intersects=intersects,
+                date_time=date_time,
+                ids=ids,
+                limit=PLANETARY_COMPUTER_LIMIT,
+                query=combined_query if combined_query else None,
+                sortby=[{"field": "id", "direction": "asc"}],
+            )
+
+            all_items.extend(stac_items)
+
+            # If we got fewer than the limit, we've fetched everything.
+            if len(stac_items) < PLANETARY_COMPUTER_LIMIT:
+                break
+
+            # Otherwise, paginate using the last item's ID.
+            last_id = stac_items[-1].id
+            logger.debug(
+                "Got %d items, paginating with id > %s",
+                len(stac_items),
+                last_id,
+            )
+
+        logger.debug("Total items fetched: %d", len(all_items))
+        return all_items
+
 
 class PlanetaryComputer(StacDataSource, TileStore):
     """Modality-agnostic data source for data on Microsoft Planetary Computer.
@@ -100,6 +194,10 @@ class PlanetaryComputer(StacDataSource, TileStore):
             required_assets=required_assets,
             cache_dir=cache_upath,
         )
+
+        # Replace the client with PlanetaryComputerStacClient to handle PC's pagination limits.
+        self.client = PlanetaryComputerStacClient(self.STAC_ENDPOINT)
+
         self.asset_bands = asset_bands
         self.timeout = timeout
         self.skip_items_missing_assets = skip_items_missing_assets

rslearn/data_sources/stac.py

@@ -12,6 +12,7 @@ from rslearn.const import WGS84_PROJECTION
 from rslearn.data_sources.data_source import Item, ItemLookupDataSource
 from rslearn.data_sources.utils import match_candidate_items_to_window
 from rslearn.log_utils import get_logger
+from rslearn.utils.fsspec import open_atomic
 from rslearn.utils.geometry import STGeometry
 from rslearn.utils.stac import StacClient, StacItem
 
@@ -187,7 +188,7 @@ class StacDataSource(ItemLookupDataSource[SourceItem]):
 
         # Finally we cache it if cache_dir is set.
         if cache_fname is not None:
-            with cache_fname.open("w") as f:
+            with open_atomic(cache_fname, "w") as f:
                 json.dump(item.serialize(), f)
 
         return item
@@ -259,7 +260,7 @@ class StacDataSource(ItemLookupDataSource[SourceItem]):
                     cache_fname = self.cache_dir / f"{item.name}.json"
                     if cache_fname.exists():
                         continue
-                    with cache_fname.open("w") as f:
+                    with open_atomic(cache_fname, "w") as f:
                         json.dump(item.serialize(), f)
 
             cur_groups = match_candidate_items_to_window(

rslearn/models/simple_time_series.py

@@ -180,7 +180,7 @@ class SimpleTimeSeries(FeatureExtractor):
         # want to pass 2 timesteps to the model.
         # TODO is probably to make this behaviour clearer but lets leave it like
         # this for now to not break things.
-        num_timesteps = images.shape[1] // image_channels
+        num_timesteps = image_channels // images.shape[1]
         batched_timesteps = images.shape[2] // num_timesteps
         images = rearrange(
             images,

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."""
@@ -59,6 +41,8 @@ class SegmentationTask(BasicTask):
         miou_metric_kwargs: dict[str, Any] = {},
         prob_scales: list[float] | None = None,
         other_metrics: dict[str, Metric] = {},
+        output_probs: bool = False,
+        output_class_idx: int | None = None,
         **kwargs: Any,
     ) -> None:
         """Initialize a new SegmentationTask.
@@ -92,6 +76,10 @@ class SegmentationTask(BasicTask):
                 this is only applied during prediction, not when computing val or test
                 metrics.
             other_metrics: additional metrics to configure on this task.
+            output_probs: if True, output raw softmax probabilities instead of class IDs
+                during prediction.
+            output_class_idx: if set along with output_probs, only output the probability
+                for this specific class index (single-channel output).
             kwargs: additional arguments to pass to BasicTask
         """
         super().__init__(**kwargs)
@@ -116,6 +104,8 @@ class SegmentationTask(BasicTask):
         self.miou_metric_kwargs = miou_metric_kwargs
         self.prob_scales = prob_scales
         self.other_metrics = other_metrics
+        self.output_probs = output_probs
+        self.output_class_idx = output_class_idx
 
     def process_inputs(
         self,
@@ -171,7 +161,9 @@ class SegmentationTask(BasicTask):
             metadata: metadata about the patch being read
 
         Returns:
-            CHW numpy array with one channel, containing the predicted class IDs.
+            CHW numpy array. If output_probs is False, returns one channel with
+            predicted class IDs. If output_probs is True, returns softmax probabilities
+            (num_classes channels, or 1 channel if output_class_idx is set).
         """
         if not isinstance(raw_output, torch.Tensor) or len(raw_output.shape) != 3:
             raise ValueError("the output for SegmentationTask must be a CHW tensor")
@@ -183,6 +175,15 @@ class SegmentationTask(BasicTask):
                     self.prob_scales, device=raw_output.device, dtype=raw_output.dtype
                 )[:, None, None]
             )
+
+        if self.output_probs:
+            # Return raw softmax probabilities
+            probs = raw_output.cpu().numpy()
+            if self.output_class_idx is not None:
+                # Return only the specified class probability
+                return probs[self.output_class_idx : self.output_class_idx + 1, :, :]
+            return probs
+
         classes = raw_output.argmax(dim=0).cpu().numpy()
         return classes[None, :, :]
 

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/utils/raster_format.py

@@ -476,6 +476,7 @@ class GeotiffRasterFormat(RasterFormat):
         bounds: PixelBounds,
         array: npt.NDArray[Any],
         fname: str | None = None,
+        nodata_val: int | float | None = None,
     ) -> None:
         """Encodes raster data.
 
@@ -485,6 +486,7 @@ class GeotiffRasterFormat(RasterFormat):
             bounds: the bounds of the raster data in the projection
             array: the raster data
             fname: override the filename to save as
+            nodata_val: set the nodata value when writing the raster.
         """
         if fname is None:
             fname = self.fname
@@ -520,6 +522,9 @@ class GeotiffRasterFormat(RasterFormat):
             profile["tiled"] = True
             profile["blockxsize"] = self.block_size
             profile["blockysize"] = self.block_size
+        # Set nodata_val if provided.
+        if nodata_val is not None:
+            profile["nodata"] = nodata_val
 
         profile.update(self.geotiff_options)
 
@@ -535,6 +540,7 @@ class GeotiffRasterFormat(RasterFormat):
         bounds: PixelBounds,
         resampling: Resampling = Resampling.bilinear,
         fname: str | None = None,
+        nodata_val: int | float | None = None,
     ) -> npt.NDArray[Any]:
         """Decodes raster data.
 
@@ -544,6 +550,16 @@ class GeotiffRasterFormat(RasterFormat):
             bounds: the bounds to read in the given projection.
             resampling: resampling method to use in case resampling is needed.
             fname: override the filename to read from
+            nodata_val: override the nodata value in the raster when reading. Pixels in
+                bounds that are not present in the source raster will be initialized to
+                this value. Note that, if the raster specifies a nodata value, and
+                some source pixels have that value, they will still be read under their
+                original value; overriding the nodata value is primarily useful if the
+                user wants out of bounds pixels to have a different value from the
+                source pixels, e.g. if the source data has background and foreground
+                classes (with background being nodata) but we want to read it in a
+                different projection and have out of bounds pixels be a third "invalid"
+                value.
 
         Returns:
             the raster data
@@ -561,6 +577,7 @@ class GeotiffRasterFormat(RasterFormat):
                 width=bounds[2] - bounds[0],
                 height=bounds[3] - bounds[1],
                 resampling=resampling,
+                src_nodata=nodata_val,
             ) as vrt:
                 return vrt.read()
 

rslearn/utils/stac.py

@@ -101,6 +101,7 @@ class StacClient:
         ids: list[str] | None = None,
         limit: int | None = None,
         query: dict[str, Any] | None = None,
+        sortby: list[dict[str, str]] | None = None,
     ) -> list[StacItem]:
         """Execute a STAC item search.
 
@@ -117,6 +118,7 @@ class StacClient:
             limit: number of items per page. We will read all the pages.
             query: query dict, if STAC query extension is supported by this API. See
                 https://github.com/stac-api-extensions/query.
+            sortby: list of sort specifications, e.g. [{"field": "id", "direction": "asc"}].
 
         Returns:
             list of matching STAC items.
@@ -142,6 +144,8 @@ class StacClient:
             request_data["limit"] = limit
         if query is not None:
             request_data["query"] = query
+        if sortby is not None:
+            request_data["sortby"] = sortby
 
         # Handle pagination.
         cur_url = self.endpoint + "/search"

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