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

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.

rslearn/train/data_module.py

@@ -21,6 +21,7 @@ from .all_patches_dataset import (
 )
 from .dataset import (
     DataInput,
+    IndexMode,
     ModelDataset,
     MultiDataset,
     RetryDataset,
@@ -69,6 +70,7 @@ class RslearnDataModule(L.LightningDataModule):
         name: str | None = None,
         retries: int = 0,
         use_in_memory_all_patches_dataset: bool = False,
+        index_mode: IndexMode = IndexMode.OFF,
     ) -> None:
         """Initialize a new RslearnDataModule.
 
@@ -92,6 +94,7 @@ class RslearnDataModule(L.LightningDataModule):
             retries: number of retries to attempt for getitem calls
             use_in_memory_all_patches_dataset: whether to use InMemoryAllPatchesDataset
                 instead of IterableAllPatchesDataset if load_all_patches is set to true.
+            index_mode: controls dataset index caching behavior (default: IndexMode.OFF)
         """
         super().__init__()
         self.inputs = inputs
@@ -103,6 +106,7 @@ class RslearnDataModule(L.LightningDataModule):
         self.name = name
         self.retries = retries
         self.use_in_memory_all_patches_dataset = use_in_memory_all_patches_dataset
+        self.index_mode = index_mode
         self.split_configs = {
             "train": default_config.update(train_config),
             "val": default_config.update(val_config),
@@ -138,6 +142,7 @@ class RslearnDataModule(L.LightningDataModule):
                 workers=self.init_workers,
                 name=self.name,
                 fix_patch_pick=(split != "train"),
+                index_mode=self.index_mode,
             )
             logger.info(f"got {len(dataset)} examples in split {split}")
             if split_config.get_load_all_patches():

rslearn/train/dataset.py

@@ -9,6 +9,7 @@ import tempfile
 import time
 import uuid
 from datetime import datetime
+from enum import StrEnum
 from typing import Any
 
 import torch
@@ -29,6 +30,7 @@ from rslearn.dataset.window import (
     get_layer_and_group_from_dir_name,
 )
 from rslearn.log_utils import get_logger
+from rslearn.train.dataset_index import DatasetIndex
 from rslearn.train.model_context import RasterImage
 from rslearn.utils.feature import Feature
 from rslearn.utils.geometry import PixelBounds, ResolutionFactor
@@ -41,6 +43,19 @@ from .transforms import Sequential
 logger = get_logger(__name__)
 
 
+class IndexMode(StrEnum):
+    """Controls dataset index caching behavior."""
+
+    OFF = "off"
+    """No caching - always load windows from dataset."""
+
+    USE = "use"
+    """Use cached index if available, create if not."""
+
+    REFRESH = "refresh"
+    """Ignore existing cache and rebuild."""
+
+
 def get_torch_dtype(dtype: DType) -> torch.dtype:
     """Convert rslearn DType to torch dtype."""
     if dtype == DType.INT32:
@@ -636,6 +651,7 @@ class ModelDataset(torch.utils.data.Dataset):
         workers: int,
         name: str | None = None,
         fix_patch_pick: bool = False,
+        index_mode: IndexMode = IndexMode.OFF,
     ) -> None:
         """Instantiate a new ModelDataset.
 
@@ -645,9 +661,10 @@ class ModelDataset(torch.utils.data.Dataset):
             inputs: data to read from the dataset for training
             task: the task to train on
             workers: number of workers to use for initializing the dataset
-            name: name of the dataset (default: None)
+            name: name of the dataset
             fix_patch_pick: if True, fix the patch pick to be the same every time
                 for a given window. Useful for testing (default: False)
+            index_mode: controls dataset index caching behavior (default: IndexMode.OFF)
         """
         self.dataset = dataset
         self.split_config = split_config
@@ -668,66 +685,14 @@ class ModelDataset(torch.utils.data.Dataset):
         else:
             self.patch_size = split_config.get_patch_size()
 
-        windows = self._get_initial_windows(split_config, workers)
-
         # If targets are not needed, remove them from the inputs.
         if split_config.get_skip_targets():
             for k in list(self.inputs.keys()):
                 if self.inputs[k].is_target:
                     del self.inputs[k]
 
-        # Eliminate windows that are missing either a requisite input layer, or missing
-        # all target layers.
-        new_windows = []
-        if workers == 0:
-            for window in windows:
-                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:
-            p = multiprocessing.Pool(workers)
-            outputs = star_imap_unordered(
-                p,
-                check_window,
-                [
-                    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
-                ],
-            )
-            for window in tqdm.tqdm(
-                outputs, total=len(windows), desc="Checking available layers in windows"
-            ):
-                if window is None:
-                    continue
-                new_windows.append(window)
-            p.close()
-        windows = new_windows
-
-        # Sort the windows to ensure that the dataset is consistent across GPUs.
-        # Inconsistent ordering can lead to a subset of windows being processed during
-        # "model test" / "model predict" when using multiple GPUs.
-        # We use a hash so that functionality like num_samples limit gets a random
-        # subset of windows (with respect to the hash function choice).
-        windows.sort(
-            key=lambda window: hashlib.sha256(window.name.encode()).hexdigest()
-        )
-
-        # Limit windows to num_samples if requested.
-        if split_config.num_samples:
-            # The windows are sorted by hash of window name so this distribution should
-            # be representative of the population.
-            windows = windows[0 : split_config.num_samples]
+        # Load windows (from index if available, otherwise from dataset)
+        windows = self._load_windows(split_config, workers, index_mode)
 
         # Write dataset_examples to a file so that we can load it lazily in the worker
         # processes. Otherwise it takes a long time to transmit it when spawning each
@@ -796,6 +761,137 @@ class ModelDataset(torch.utils.data.Dataset):
 
         return windows
 
+    def _load_windows(
+        self,
+        split_config: SplitConfig,
+        workers: int,
+        index_mode: IndexMode,
+    ) -> list[Window]:
+        """Load windows, using index if available.
+
+        This method handles:
+        1. Loading from index if index_mode is USE and index exists
+        2. Otherwise, loading from dataset, filtering, sorting, limiting
+        3. Saving to index if index_mode is USE or REFRESH
+
+        Args:
+            split_config: the split configuration.
+            workers: number of worker processes.
+            index_mode: controls caching behavior.
+
+        Returns:
+            list of processed windows ready for training.
+        """
+        # Try to load from index
+        index: DatasetIndex | None = None
+
+        if index_mode != IndexMode.OFF:
+            logger.info(f"Checking index for dataset {self.dataset.path}")
+            index = DatasetIndex(
+                storage=self.dataset.storage,
+                dataset_path=self.dataset.path,
+                groups=split_config.groups,
+                names=split_config.names,
+                tags=split_config.tags,
+                num_samples=split_config.num_samples,
+                skip_targets=split_config.get_skip_targets(),
+                inputs=self.inputs,
+            )
+            refresh = index_mode == IndexMode.REFRESH
+            indexed_windows = index.load_windows(refresh)
+
+            if indexed_windows is not None:
+                logger.info(f"Loaded {len(indexed_windows)} windows from index")
+                return indexed_windows
+
+        # No index available, load and process windows from dataset
+        logger.debug("Loading windows from dataset...")
+        windows = self._get_initial_windows(split_config, workers)
+        windows = self._filter_windows_by_layers(windows, workers)
+        windows = self._sort_and_limit_windows(windows, split_config)
+
+        # Save to index if enabled
+        if index is not None:
+            index.save_windows(windows)
+
+        return windows
+
+    def _filter_windows_by_layers(
+        self, windows: list[Window], workers: int
+    ) -> list[Window]:
+        """Filter windows to only include those with required layers.
+
+        Args:
+            windows: list of windows to filter.
+            workers: number of worker processes for parallel filtering.
+
+        Returns:
+            list of windows that have all required input layers.
+        """
+        output_layer_skip = (
+            self.split_config.get_output_layer_name_skip_inference_if_exists()
+        )
+
+        if workers == 0:
+            return [
+                w
+                for w in windows
+                if check_window(
+                    self.inputs,
+                    w,
+                    output_layer_name_skip_inference_if_exists=output_layer_skip,
+                )
+                is not None
+            ]
+
+        p = multiprocessing.Pool(workers)
+        outputs = star_imap_unordered(
+            p,
+            check_window,
+            [
+                dict(
+                    inputs=self.inputs,
+                    window=window,
+                    output_layer_name_skip_inference_if_exists=output_layer_skip,
+                )
+                for window in windows
+            ],
+        )
+        filtered = []
+        for window in tqdm.tqdm(
+            outputs,
+            total=len(windows),
+            desc="Checking available layers in windows",
+        ):
+            if window is not None:
+                filtered.append(window)
+        p.close()
+        return filtered
+
+    def _sort_and_limit_windows(
+        self, windows: list[Window], split_config: SplitConfig
+    ) -> list[Window]:
+        """Sort windows by hash and apply num_samples limit.
+
+        Sorting ensures consistent ordering across GPUs. Using hash gives a
+        pseudo-random but deterministic order for sampling.
+
+        Args:
+            windows: list of windows to sort and limit.
+            split_config: the split configuration with num_samples.
+
+        Returns:
+            sorted and optionally limited list of windows.
+        """
+        windows.sort(
+            key=lambda window: hashlib.sha256(window.name.encode()).hexdigest()
+        )
+
+        if split_config.num_samples:
+            windows = windows[: split_config.num_samples]
+
+        return windows
+
     def _serialize_item(self, example: Window) -> dict[str, Any]:
         return example.get_metadata()
 

rslearn/train/dataset_index.py

@@ -0,0 +1,156 @@
+"""Dataset index for caching window lists to speed up ModelDataset initialization."""
+
+import hashlib
+import json
+from datetime import datetime
+from typing import TYPE_CHECKING, Any
+
+from upath import UPath
+
+from rslearn.dataset.window import Window
+from rslearn.log_utils import get_logger
+from rslearn.utils.fsspec import open_atomic
+
+if TYPE_CHECKING:
+    from rslearn.dataset.storage.storage import WindowStorage
+
+logger = get_logger(__name__)
+
+# Increment this when the index format changes to force rebuild
+INDEX_VERSION = 1
+
+# Directory name for storing index files
+INDEX_DIR_NAME = ".rslearn_dataset_index"
+
+
+class DatasetIndex:
+    """Manages indexed window lists for faster ModelDataset initialization.
+
+    Note: The index does NOT automatically detect when windows are added or removed
+    from the dataset. Use refresh=True after modifying dataset windows.
+    """
+
+    def __init__(
+        self,
+        storage: "WindowStorage",
+        dataset_path: UPath,
+        groups: list[str] | None,
+        names: list[str] | None,
+        tags: dict[str, Any] | None,
+        num_samples: int | None,
+        skip_targets: bool,
+        inputs: dict[str, Any],
+    ) -> None:
+        """Initialize DatasetIndex with specific configuration.
+
+        Args:
+            storage: WindowStorage for deserializing windows.
+            dataset_path: Path to the dataset directory.
+            groups: list of window groups to include.
+            names: list of window names to include.
+            tags: tags to filter windows by.
+            num_samples: limit on number of samples.
+            skip_targets: whether targets are skipped.
+            inputs: dict mapping input names to DataInput objects.
+        """
+        self.storage = storage
+        self.dataset_path = dataset_path
+        self.index_dir = dataset_path / INDEX_DIR_NAME
+
+        # Compute index key from configuration
+        inputs_data = {}
+        for name, inp in inputs.items():
+            inputs_data[name] = {
+                "layers": inp.layers,
+                "required": inp.required,
+                "load_all_layers": inp.load_all_layers,
+                "is_target": inp.is_target,
+            }
+
+        key_data = {
+            "groups": groups,
+            "names": names,
+            "tags": tags,
+            "num_samples": num_samples,
+            "skip_targets": skip_targets,
+            "inputs": inputs_data,
+        }
+        self.index_key = hashlib.sha256(
+            json.dumps(key_data, sort_keys=True).encode()
+        ).hexdigest()
+
+    def _get_config_hash(self) -> str:
+        """Get hash of config.json for quick validation.
+
+        Returns:
+            A 16-character hex string hash of the config, or empty string if no config.
+        """
+        config_path = self.dataset_path / "config.json"
+        if config_path.exists():
+            with config_path.open() as f:
+                return hashlib.sha256(f.read().encode()).hexdigest()[:16]
+        return ""
+
+    def load_windows(self, refresh: bool = False) -> list[Window] | None:
+        """Load indexed window list if valid, else return None.
+
+        Args:
+            refresh: If True, ignore existing index and return None.
+
+        Returns:
+            List of Window objects if index is valid, None otherwise.
+        """
+        if refresh:
+            logger.info("refresh=True, rebuilding index")
+            return None
+
+        index_file = self.index_dir / f"{self.index_key}.json"
+        if not index_file.exists():
+            logger.info(f"No index found at {index_file}, will build")
+            return None
+
+        try:
+            with index_file.open() as f:
+                index_data = json.load(f)
+        except (OSError, json.JSONDecodeError):
+            logger.warning(f"Corrupted index file at {index_file}, will rebuild")
+            return None
+
+        # Check index version
+        if index_data.get("version") != INDEX_VERSION:
+            logger.info(
+                f"Index version mismatch (got {index_data.get('version')}, "
+                f"expected {INDEX_VERSION}), will rebuild"
+            )
+            return None
+
+        # Quick validation: check config hash
+        if index_data.get("config_hash") != self._get_config_hash():
+            logger.info("Config hash mismatch, index invalidated")
+            return None
+
+        # Deserialize windows
+        return [Window.from_metadata(self.storage, w) for w in index_data["windows"]]
+
+    def save_windows(self, windows: list[Window]) -> None:
+        """Save processed windows to index with atomic write.
+
+        Args:
+            windows: List of Window objects to index.
+        """
+        self.index_dir.mkdir(parents=True, exist_ok=True)
+        index_file = self.index_dir / f"{self.index_key}.json"
+
+        # Serialize windows
+        serialized_windows = [w.get_metadata() for w in windows]
+
+        index_data = {
+            "version": INDEX_VERSION,
+            "config_hash": self._get_config_hash(),
+            "created_at": datetime.now().isoformat(),
+            "num_windows": len(windows),
+            "windows": serialized_windows,
+        }
+        with open_atomic(index_file, "w") as f:
+            json.dump(index_data, f)
+        logger.info(f"Saved {len(windows)} windows to index at {index_file}")

rslearn/train/model_context.py

@@ -43,6 +43,22 @@ class RasterImage:
             raise ValueError(f"Expected a single timestep, got {self.image.shape[1]}")
         return self.image[:, 0]
 
+    def get_hw_tensor(self) -> torch.Tensor:
+        """Get a 2D HW tensor from a single-channel, single-timestep RasterImage.
+
+        This function checks that C=1 and T=1, then returns the HW tensor.
+        Useful for per-pixel labels like segmentation masks.
+        """
+        if self.image.shape[0] != 1:
+            raise ValueError(
+                f"Expected single channel (C=1), got {self.image.shape[0]}"
+            )
+        if self.image.shape[1] != 1:
+            raise ValueError(
+                f"Expected single timestep (T=1), got {self.image.shape[1]}"
+            )
+        return self.image[0, 0]
+
 
 @dataclass
 class SampleMetadata: