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

rslearn/data_sources/worldpop.py

@@ -80,7 +80,7 @@ class WorldPop(LocalFiles):
         worldpop_upath.mkdir(parents=True, exist_ok=True)
         self.download_worldpop_data(worldpop_upath, timeout)
         super().__init__(
-            src_dir=worldpop_upath,
+            src_dir=worldpop_upath.absolute().as_uri(),
             layer_type=LayerType.RASTER,
             context=context,
         )

rslearn/data_sources/xyz_tiles.py

@@ -19,7 +19,7 @@ from rslearn.config import LayerConfig, QueryConfig
 from rslearn.dataset import Window
 from rslearn.dataset.materialize import RasterMaterializer
 from rslearn.tile_stores import TileStore, TileStoreWithLayer
-from rslearn.utils import PixelBounds, Projection, STGeometry
+from rslearn.utils import PixelBounds, Projection, STGeometry, get_global_raster_bounds
 from rslearn.utils.array import copy_spatial_array
 from rslearn.utils.raster_format import get_transform_from_projection_and_bounds
 
@@ -184,7 +184,7 @@ class XyzTiles(DataSource, TileStore):
             groups.append(cur_groups)
         return groups
 
-    def deserialize_item(self, serialized_item: Any) -> Item:
+    def deserialize_item(self, serialized_item: dict) -> Item:
         """Deserializes an item from JSON-decoded data."""
         return Item.deserialize(serialized_item)
 
@@ -278,13 +278,9 @@ class XyzTiles(DataSource, TileStore):
         Returns:
             the bounds of the raster in the projection.
         """
-        geom = STGeometry(self.projection, self.shp, None).to_projection(projection)
-        return (
-            int(geom.shp.bounds[0]),
-            int(geom.shp.bounds[1]),
-            int(geom.shp.bounds[2]),
-            int(geom.shp.bounds[3]),
-        )
+        # XyzTiles is a global data source, so we return global raster bounds based on
+        # the projection.
+        return get_global_raster_bounds(projection)
 
     def read_raster(
         self,

rslearn/dataset/materialize.py

@@ -236,7 +236,11 @@ def read_and_stack_raster_windows(
     band_dtype: npt.DTypeLike,
     resampling_method: Resampling = Resampling.bilinear,
 ) -> npt.NDArray[np.generic]:
-    """Create a stack of extent aligned raster windows.
+    """Create a stack of raster images, with one per item in the group.
+
+    We read the portion of each raster item corresponding to the window extent, and
+    stack the resulting images. This is used for the MEAN and MEDIAN compositing
+    methods to it can compute aggregate statistics across the stack.
 
     Args:
         group: Iterable of items (e.g., scene metadata objects) to read data from.

rslearn/models/clay/clay.py

@@ -105,7 +105,7 @@ class Clay(FeatureExtractor):
 
     def _resize_image(self, image: torch.Tensor, original_hw: int) -> torch.Tensor:
         """Resize the image to the input resolution."""
-        new_hw = self.patch_size if original_hw == 1 else DEFAULT_IMAGE_RESOLUTION
+        new_hw = PATCH_SIZE if original_hw == 1 else DEFAULT_IMAGE_RESOLUTION
         return F.interpolate(
             image, size=(new_hw, new_hw), mode="bilinear", align_corners=False
         )
@@ -123,7 +123,8 @@ class Clay(FeatureExtractor):
         device = param.device
 
         chips = torch.stack(
-            [inp[self.modality] for inp in context.inputs], dim=0
+            [inp[self.modality].single_ts_to_chw_tensor() for inp in context.inputs],
+            dim=0,
         )  # (B, C, H, W)
         if self.do_resizing:
             chips = self._resize_image(chips, chips.shape[2])
@@ -203,7 +204,6 @@ class ClayNormalize(Transform):
                 mean=means,
                 std=stds,
                 selectors=[modality],
-                num_bands=len(means),
             )
         self.normalizers = torch.nn.ModuleDict(normalizers)
 

rslearn/models/concatenate_features.py

@@ -3,6 +3,7 @@
 from typing import Any
 
 import torch
+from einops import rearrange
 
 from rslearn.train.model_context import ModelContext
 
@@ -79,7 +80,11 @@ class ConcatenateFeatures(IntermediateComponent):
             )
 
         add_data = torch.stack(
-            [input_data[self.key] for input_data in context.inputs], dim=0
+            [
+                rearrange(input_data[self.key].image, "c t h w -> (c t) h w")
+                for input_data in context.inputs
+            ],
+            dim=0,
         )
         add_features = self.conv_layers(add_data)
 

rslearn/models/detr/detr.py

@@ -468,7 +468,10 @@ class Detr(Predictor):
 
         # Get image sizes.
         image_sizes = torch.tensor(
-            [[inp["image"].shape[2], inp["image"].shape[1]] for inp in context.inputs],
+            [
+                [inp["image"].image.shape[2], inp["image"].image.shape[1]]
+                for inp in context.inputs
+            ],
             dtype=torch.int32,
             device=features.device,
         )

rslearn/models/dinov3.py

@@ -159,7 +159,6 @@ class DinoV3Normalize(Transform):
         self.normalize = Normalize(
             [value * 255 for value in mean],
             [value * 255 for value in std],
-            num_bands=3,
         )
 
     def forward(

rslearn/models/olmoearth_pretrain/model.py

@@ -95,7 +95,9 @@ class OlmoEarth(FeatureExtractor):
         """
         if use_legacy_timestamps:
             warnings.warn(
-                "For new projects, don't use legacy timesteps.", DeprecationWarning
+                "For new projects, don't use legacy timesteps. "
+                "Support will be removed after 2026-04-01.",
+                FutureWarning,
             )
 
         if (

rslearn/models/pooling_decoder.py

@@ -124,6 +124,6 @@ class SegmentationPoolingDecoder(PoolingDecoder):
         """
         output_probs = super().forward(intermediates, context)
         # BC -> BCHW
-        h, w = context.inputs[0][self.image_key].shape[1:3]
+        h, w = context.inputs[0][self.image_key].image.shape[1:3]
         feat_map = output_probs.feature_vector[:, :, None, None].repeat([1, 1, h, w])
         return FeatureMaps([feat_map])

rslearn/models/prithvi.py

@@ -230,7 +230,6 @@ class PrithviNormalize(Transform):
         self.normalizer = Normalize(
             mean=config["mean"],
             std=config["std"],
-            num_bands=len(config["mean"]),
             selectors=[PrithviV2.INPUT_KEY],
         )
 

rslearn/models/simple_time_series.py

@@ -1,5 +1,6 @@
 """SimpleTimeSeries encoder."""
 
+import warnings
 from typing import Any
 
 import torch
@@ -25,13 +26,14 @@ class SimpleTimeSeries(FeatureExtractor):
     def __init__(
         self,
         encoder: FeatureExtractor,
-        image_channels: int | None = None,
+        num_timesteps_per_forward_pass: int = 1,
         op: str = "max",
         groups: list[list[int]] | None = None,
         num_layers: int | None = None,
         image_key: str = "image",
         backbone_channels: list[tuple[int, int]] | None = None,
-        image_keys: dict[str, int] | None = None,
+        image_keys: list[str] | dict[str, int] | None = None,
+        image_channels: int | None = None,
     ) -> None:
         """Create a new SimpleTimeSeries.
 
@@ -39,9 +41,11 @@ class SimpleTimeSeries(FeatureExtractor):
             encoder: the underlying FeatureExtractor. It must provide get_backbone_channels
                 function that returns the output channels, or backbone_channels must be set.
                 It must output a FeatureMaps.
-            image_channels: the number of channels per image of the time series. The
-                input should have multiple images concatenated on the channel axis, so
-                this parameter is used to distinguish the different images.
+            num_timesteps_per_forward_pass: how many timesteps to pass to the encoder
+                in each forward pass. Defaults to 1 (one timestep per forward pass).
+                Set to a higher value to batch multiple timesteps together, e.g. for
+                pre/post change detection where you want 4 pre and 4 post images
+                processed together.
             op: one of max, mean, convrnn, conv3d, or conv1d
             groups: sets of images for which to combine features. Within each set,
                 features are combined using the specified operation; then, across sets,
@@ -51,28 +55,53 @@ class SimpleTimeSeries(FeatureExtractor):
                 combined before features and the combined after features. groups is a
                 list of sets, and each set is a list of image indices.
             num_layers: the number of layers for convrnn, conv3d, and conv1d ops.
-            image_key: the key to access the images.
+            image_key: the key to access the images (used when image_keys is not set).
             backbone_channels: manually specify the backbone channels. Can be set if
                 the encoder does not provide get_backbone_channels function.
-            image_keys: as an alternative to setting image_channels, map from the key
-                in input dict to the number of channels per timestep for that modality.
-                This way SimpleTimeSeries can be used with multimodal inputs. One of
-                image_channels or image_keys must be specified.
+            image_keys: list of keys in input dict to process as multimodal inputs.
+                All keys use the same num_timesteps_per_forward_pass. If not set,
+                only the single image_key is used. Passing a dict[str, int] is
+                deprecated and will be removed on 2026-04-01.
+            image_channels: Deprecated, use num_timesteps_per_forward_pass instead.
+                Will be removed on 2026-04-01.
         """
-        if (image_channels is None and image_keys is None) or (
-            image_channels is not None and image_keys is not None
-        ):
-            raise ValueError(
-                "exactly one of image_channels and image_keys must be specified"
+        # Handle deprecated image_channels parameter
+        if image_channels is not None:
+            warnings.warn(
+                "image_channels is deprecated and will be removed on 2026-04-01. "
+                "Use num_timesteps_per_forward_pass instead. The new parameter directly "
+                "specifies the number of timesteps per forward pass rather than requiring "
+                "image_channels // actual_channels.",
+                FutureWarning,
+                stacklevel=2,
             )
 
+        # Handle deprecated dict form of image_keys
+        deprecated_image_keys_dict: dict[str, int] | None = None
+        if isinstance(image_keys, dict):
+            warnings.warn(
+                "Passing image_keys as a dict is deprecated and will be removed on "
+                "2026-04-01. Use image_keys as a list[str] and set "
+                "num_timesteps_per_forward_pass instead.",
+                FutureWarning,
+                stacklevel=2,
+            )
+            deprecated_image_keys_dict = image_keys
+            image_keys = None  # Will use deprecated path in forward
+
         super().__init__()
         self.encoder = encoder
-        self.image_channels = image_channels
+        self.num_timesteps_per_forward_pass = num_timesteps_per_forward_pass
+        # Store deprecated parameters for runtime conversion
+        self._deprecated_image_channels = image_channels
+        self._deprecated_image_keys_dict = deprecated_image_keys_dict
         self.op = op
         self.groups = groups
-        self.image_key = image_key
-        self.image_keys = image_keys
+        # Normalize image_key to image_keys list form
+        if image_keys is not None:
+            self.image_keys = image_keys
+        else:
+            self.image_keys = [image_key]
 
         if backbone_channels is not None:
             out_channels = backbone_channels
@@ -163,24 +192,25 @@ class SimpleTimeSeries(FeatureExtractor):
         return out_channels
 
     def _get_batched_images(
-        self, input_dicts: list[dict[str, Any]], image_key: str, image_channels: int
+        self, input_dicts: list[dict[str, Any]], image_key: str, num_timesteps: int
     ) -> list[RasterImage]:
         """Collect and reshape images across input dicts.
 
         The BTCHW image time series are reshaped to (B*T)CHW so they can be passed to
         the forward pass of a per-image (unitemporal) model.
+
+        Args:
+            input_dicts: list of input dictionaries containing RasterImage objects.
+            image_key: the key to access the RasterImage in each input dict.
+            num_timesteps: how many timesteps to batch together per forward pass.
         """
         images = torch.stack(
             [input_dict[image_key].image for input_dict in input_dicts], dim=0
         )  # B, C, T, H, W
         timestamps = [input_dict[image_key].timestamps for input_dict in input_dicts]
-        # if image channels is not equal to the actual number of channels, then
-        # then every N images should be batched together. For example, if the
-        # number of input channels c == 2, and image_channels == 4, then we
-        # 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 = image_channels // images.shape[1]
+        # num_timesteps specifies how many timesteps to batch together per forward pass.
+        # For example, if the input has 8 timesteps and num_timesteps=4, we do 2
+        # forward passes, each with 4 timesteps batched together.
         batched_timesteps = images.shape[2] // num_timesteps
         images = rearrange(
             images,
@@ -222,10 +252,22 @@ class SimpleTimeSeries(FeatureExtractor):
         n_batch = len(context.inputs)
         n_images: int | None = None
 
-        if self.image_keys is not None:
-            for image_key, image_channels in self.image_keys.items():
+        if self._deprecated_image_keys_dict is not None:
+            # Deprecated dict form: each key has its own channels_per_timestep.
+            # The channels_per_timestep could be used to group multiple timesteps,
+            # together, so we need to divide by the actual image channel count to get
+            # the number of timesteps to be grouped.
+            for (
+                image_key,
+                channels_per_timestep,
+            ) in self._deprecated_image_keys_dict.items():
+                # For deprecated image_keys dict, the value is channels per timestep,
+                # so we need to compute num_timesteps from the actual image channels
+                sample_image = context.inputs[0][image_key].image
+                actual_channels = sample_image.shape[0]  # C in CTHW
+                num_timesteps = channels_per_timestep // actual_channels
                 batched_images = self._get_batched_images(
-                    context.inputs, image_key, image_channels
+                    context.inputs, image_key, num_timesteps
                 )
 
                 if batched_inputs is None:
@@ -240,12 +282,32 @@ class SimpleTimeSeries(FeatureExtractor):
                     batched_inputs[i][image_key] = image
 
         else:
-            assert self.image_channels is not None
-            batched_images = self._get_batched_images(
-                context.inputs, self.image_key, self.image_channels
-            )
-            batched_inputs = [{self.image_key: image} for image in batched_images]
-            n_images = len(batched_images) // n_batch
+            # Determine num_timesteps - either from deprecated image_channels or
+            # directly from num_timesteps_per_forward_pass
+            if self._deprecated_image_channels is not None:
+                # Backwards compatibility: compute num_timesteps from image_channels
+                # (which should be a multiple of the actual per-timestep channels).
+                sample_image = context.inputs[0][self.image_keys[0]].image
+                actual_channels = sample_image.shape[0]  # C in CTHW
+                num_timesteps = self._deprecated_image_channels // actual_channels
+            else:
+                num_timesteps = self.num_timesteps_per_forward_pass
+
+            for image_key in self.image_keys:
+                batched_images = self._get_batched_images(
+                    context.inputs, image_key, num_timesteps
+                )
+
+                if batched_inputs is None:
+                    batched_inputs = [{} for _ in batched_images]
+                    n_images = len(batched_images) // n_batch
+                elif n_images != len(batched_images) // n_batch:
+                    raise ValueError(
+                        "expected all modalities to have the same number of timesteps"
+                    )
+
+                for i, image in enumerate(batched_images):
+                    batched_inputs[i][image_key] = image
 
         assert n_images is not None
         # Now we can apply the underlying FeatureExtractor.