rslearn 0.0.13__py3-none-any.whl → 0.0.15__py3-none-any.whl

rslearn/train/dataset.py

@@ -1,7 +1,6 @@
 """Default Dataset for rslearn."""
 
 import hashlib
-import itertools
 import json
 import multiprocessing
 import os
@@ -9,10 +8,8 @@ import random
 import tempfile
 import time
 import uuid
-from collections.abc import Iterable, Iterator
 from typing import Any
 
-import shapely
 import torch
 import tqdm
 from rasterio.warp import Resampling
@@ -29,7 +26,7 @@ from rslearn.dataset.window import Window, get_layer_and_group_from_dir_name
 from rslearn.log_utils import get_logger
 from rslearn.train.tasks import Task
 from rslearn.utils.feature import Feature
-from rslearn.utils.geometry import PixelBounds, STGeometry
+from rslearn.utils.geometry import PixelBounds
 from rslearn.utils.mp import star_imap_unordered
 from rslearn.utils.raster_format import load_raster_format
 from rslearn.utils.vector_format import load_vector_format
@@ -39,70 +36,14 @@ from .transforms import Sequential
 logger = get_logger(__name__)
 
 
-def get_window_patch_options(
-    patch_size: tuple[int, int],
-    overlap_size: tuple[int, int],
-    bounds: PixelBounds,
-) -> list[PixelBounds]:
-    """Get the bounds of each patch within the overall bounds.
-
-    Args:
-        patch_size: the size of the patches to extract.
-        overlap_size: the size of the overlap between patches.
-        bounds: the window bounds to divide up into smaller patches.
-
-    Returns:
-        a list of patch bounds within the overall bounds. The rightmost and
-            bottommost patches may extend beyond the provided bounds.
-    """
-    # We stride the patches by patch_size - overlap_size until the last patch.
-    # We handle the last patch with a special case to ensure it does not exceed the
-    # window bounds. Instead, it may overlap the previous patch.
-    cols = list(
-        range(
-            bounds[0],
-            bounds[2] - patch_size[0],
-            patch_size[0] - overlap_size[0],
-        )
-    ) + [bounds[2] - patch_size[0]]
-    rows = list(
-        range(
-            bounds[1],
-            bounds[3] - patch_size[1],
-            patch_size[1] - overlap_size[1],
-        )
-    ) + [bounds[3] - patch_size[1]]
-
-    patch_bounds: list[PixelBounds] = []
-    for col in cols:
-        for row in rows:
-            patch_bounds.append((col, row, col + patch_size[0], row + patch_size[1]))
-    return patch_bounds
-
-
-def pad_slice_protect(
-    raw_inputs: dict[str, Any],
-    passthrough_inputs: dict[str, Any],
-    patch_size: tuple[int, int],
-) -> tuple[dict[str, Any], dict[str, Any]]:
-    """Pad tensors in-place by patch size to protect slicing near right/bottom edges.
-
-    Args:
-        raw_inputs: the raw inputs to pad.
-        passthrough_inputs: the passthrough inputs to pad.
-        patch_size: the size of the patches to extract.
-
-    Returns:
-        a tuple of (raw_inputs, passthrough_inputs).
-    """
-    for d in [raw_inputs, passthrough_inputs]:
-        for input_name, value in list(d.items()):
-            if not isinstance(value, torch.Tensor):
-                continue
-            d[input_name] = torch.nn.functional.pad(
-                value, pad=(0, patch_size[0], 0, patch_size[1])
-            )
-    return raw_inputs, passthrough_inputs
+def get_torch_dtype(dtype: DType) -> torch.dtype:
+    """Convert rslearn DType to torch dtype."""
+    if dtype == DType.INT32:
+        return torch.int32
+    elif dtype == DType.FLOAT32:
+        return torch.float32
+    else:
+        raise ValueError(f"unable to handle {dtype} as a torch dtype")
 
 
 class SamplerFactory:
@@ -296,7 +237,7 @@ def read_raster_layer_for_data_input(
 
     image = torch.zeros(
         (len(needed_bands), bounds[3] - bounds[1], bounds[2] - bounds[0]),
-        dtype=data_input.dtype.get_torch_dtype(),
+        dtype=get_torch_dtype(data_input.dtype),
     )
 
     for band_set, src_indexes, dst_indexes in needed_sets_and_indexes:
@@ -893,383 +834,6 @@ class ModelDataset(torch.utils.data.Dataset):
         self.name = name
 
 
-class IterableAllPatchesDataset(torch.utils.data.IterableDataset):
-    """This wraps a ModelDataset to iterate over all patches in that dataset.
-
-    This should be used when SplitConfig.load_all_patches is enabled. The ModelDataset
-    is configured with no patch size (load entire windows), and the dataset is wrapped
-    in an AllPatchesDataset.
-
-    Similar to DistributedSampler, we add extra samples at each rank to ensure
-    consistent number of batches across all ranks.
-    """
-
-    def __init__(
-        self,
-        dataset: ModelDataset,
-        patch_size: tuple[int, int],
-        overlap_ratio: float = 0.0,
-        rank: int = 0,
-        world_size: int = 1,
-    ):
-        """Create a new IterableAllPatchesDataset.
-
-        Args:
-            dataset: the ModelDataset to wrap.
-            patch_size: the size of the patches to extract.
-            overlap_ratio: whether to include overlap between the patches. Note that
-                the right/bottom-most patches may still overlap since we ensure that
-                all patches are contained in the window bounds.
-            rank: the global rank of this train worker process.
-            world_size: the total number of train worker processes.
-        """
-        super().__init__()
-        self.dataset = dataset
-        self.patch_size = patch_size
-        self.overlap_size = (
-            round(self.patch_size[0] * overlap_ratio),
-            round(self.patch_size[1] * overlap_ratio),
-        )
-        self.rank = rank
-        self.world_size = world_size
-        self.windows = self.dataset.get_dataset_examples()
-
-    def set_name(self, name: str) -> None:
-        """Sets dataset name.
-
-        Args:
-            name: dataset name
-        """
-        self.dataset.set_name(name)
-
-    def get_window_num_patches(self, bounds: PixelBounds) -> int:
-        """Get the number of patches for these bounds.
-
-        This corresponds to the length of the list returned by get_patch_options.
-        """
-        num_cols = (
-            len(
-                range(
-                    bounds[0],
-                    bounds[2] - self.patch_size[0],
-                    self.patch_size[0] - self.overlap_size[0],
-                )
-            )
-            + 1
-        )
-        num_rows = (
-            len(
-                range(
-                    bounds[1],
-                    bounds[3] - self.patch_size[1],
-                    self.patch_size[1] - self.overlap_size[1],
-                )
-            )
-            + 1
-        )
-        return num_cols * num_rows
-
-    def _get_worker_iteration_data(self) -> tuple[Iterable[int], int]:
-        """Get the windows we should iterate over.
-
-        This is split both by training worker (self.rank) and data loader worker (via
-        get_worker_info).
-
-        We also compute the total number of samples that each data loader worker should
-        yield. This is important for DDP to ensure that all ranks see the same number
-        of batches.
-
-        Returns:
-            a tuple (window_ids, num_samples_per_worker).
-        """
-        # Figure out the total number of data loader workers and our worker ID.
-        worker_info = torch.utils.data.get_worker_info()
-        if worker_info is None:
-            worker_id = 0
-            num_workers = 1
-        else:
-            worker_id = worker_info.id
-            num_workers = worker_info.num_workers
-        global_worker_id = self.rank * num_workers + worker_id
-        global_num_workers = self.world_size * num_workers
-
-        # Split up the windows evenly among the workers.
-        # We compute this for all workers since we will need to see the maximum number
-        # of samples under this assignment across workers.
-        window_indexes = range(len(self.windows))
-        windows_by_worker = [
-            window_indexes[cur_rank :: self.world_size][cur_worker_id::num_workers]
-            for cur_rank in range(self.world_size)
-            for cur_worker_id in range(num_workers)
-        ]
-
-        # Now compute the maximum number of samples across workers.
-        max_num_patches = 0
-        for worker_windows in windows_by_worker:
-            worker_num_patches = 0
-            for window_id in worker_windows:
-                worker_num_patches += self.get_window_num_patches(
-                    self.windows[window_id].bounds
-                )
-            max_num_patches = max(max_num_patches, worker_num_patches)
-
-        # Each worker needs at least one window, otherwise it won't be able to pad.
-        # Unless there are zero windows total, which is fine.
-        # Previously we would address this by borrowing the windows from another
-        # worker, but this causes issues with RslearnWriter: if we yield the same
-        # window from parallel workers, it may end up writing an empty output for that
-        # window in the end.
-        # So now we raise an error instead, and require the number of workers to be
-        # less than the number of windows.
-        if len(windows_by_worker[global_worker_id]) == 0 and max_num_patches > 0:
-            raise ValueError(
-                f"the number of workers {global_num_workers} must be <= the number of windows {len(self.windows)}"
-            )
-
-        return (windows_by_worker[global_worker_id], max_num_patches)
-
-    def __iter__(
-        self,
-    ) -> Iterator[tuple[dict[str, Any], dict[str, Any], dict[str, Any]]]:
-        """Iterate over all patches in each element of the underlying ModelDataset."""
-        # Iterate over the window IDs until we have returned enough samples.
-        window_ids, num_samples_needed = self._get_worker_iteration_data()
-        num_samples_returned = 0
-
-        for iteration_idx in itertools.count():
-            for window_id in window_ids:
-                raw_inputs, passthrough_inputs, metadata = self.dataset.get_raw_inputs(
-                    window_id
-                )
-                bounds = metadata["bounds"]
-
-                # For simplicity, pad tensors by patch size to ensure that any patch bounds
-                # extending outside the window bounds will not have issues when we slice
-                # the tensors later.
-                pad_slice_protect(raw_inputs, passthrough_inputs, self.patch_size)
-
-                # Now iterate over the patches and extract/yield the crops.
-                # Note that, in case user is leveraging RslearnWriter, it is important that
-                # the patch_idx be increasing (as we iterate) within one window.
-                patches = get_window_patch_options(
-                    self.patch_size, self.overlap_size, bounds
-                )
-                for patch_idx, patch_bounds in enumerate(patches):
-                    cur_geom = STGeometry(
-                        metadata["projection"], shapely.box(*patch_bounds), None
-                    )
-                    start_offset = (
-                        patch_bounds[0] - bounds[0],
-                        patch_bounds[1] - bounds[1],
-                    )
-                    end_offset = (
-                        patch_bounds[2] - bounds[0],
-                        patch_bounds[3] - bounds[1],
-                    )
-
-                    # Define a helper function to handle each input dict.
-                    def crop_input_dict(d: dict[str, Any]) -> dict[str, Any]:
-                        cropped = {}
-                        for input_name, value in d.items():
-                            if isinstance(value, torch.Tensor):
-                                # Crop the CHW tensor.
-                                cropped[input_name] = value[
-                                    :,
-                                    start_offset[1] : end_offset[1],
-                                    start_offset[0] : end_offset[0],
-                                ].clone()
-                            elif isinstance(value, list):
-                                cropped[input_name] = [
-                                    feat
-                                    for feat in value
-                                    if cur_geom.intersects(feat.geometry)
-                                ]
-                            else:
-                                raise ValueError(
-                                    "got input that is neither tensor nor feature list"
-                                )
-                        return cropped
-
-                    cur_raw_inputs = crop_input_dict(raw_inputs)
-                    cur_passthrough_inputs = crop_input_dict(passthrough_inputs)
-
-                    # Adjust the metadata as well.
-                    cur_metadata = metadata.copy()
-                    cur_metadata["bounds"] = patch_bounds
-                    cur_metadata["patch_idx"] = patch_idx
-                    cur_metadata["num_patches"] = len(patches)
-
-                    # Now we can compute input and target dicts via the task.
-                    input_dict, target_dict = self.dataset.task.process_inputs(
-                        cur_raw_inputs,
-                        metadata=cur_metadata,
-                        load_targets=not self.dataset.split_config.get_skip_targets(),
-                    )
-                    input_dict.update(cur_passthrough_inputs)
-                    input_dict, target_dict = self.dataset.transforms(
-                        input_dict, target_dict
-                    )
-                    input_dict["dataset_source"] = self.dataset.name
-
-                    if num_samples_returned < num_samples_needed:
-                        yield input_dict, target_dict, cur_metadata
-                        num_samples_returned += 1
-                    else:
-                        assert iteration_idx > 0
-
-            if num_samples_returned >= num_samples_needed:
-                break
-
-    def get_dataset_examples(self) -> list[Window]:
-        """Returns a list of windows in this dataset."""
-        return self.dataset.get_dataset_examples()
-
-
-class InMemoryAllPatchesDataset(torch.utils.data.Dataset):
-    """This wraps a ModelDataset to iterate over all patches in that dataset.
-
-    This should be used when SplitConfig.load_all_patches is enabled.
-
-    This is a simpler version of IterableAllPatchesDataset that caches all windows in memory.
-    This is useful for small datasets that fit in memory.
-    """
-
-    def __init__(
-        self,
-        dataset: ModelDataset,
-        patch_size: tuple[int, int],
-        overlap_ratio: float = 0.0,
-    ):
-        """Create a new InMemoryAllPatchesDataset.
-
-        Args:
-            dataset: the ModelDataset to wrap.
-            patch_size: the size of the patches to extract.
-            overlap_ratio: whether to include overlap between the patches. Note that
-                the right/bottom-most patches may still overlap since we ensure that
-                all patches are contained in the window bounds.
-        """
-        super().__init__()
-        self.dataset = dataset
-        self.patch_size = patch_size
-        self.overlap_size = (
-            round(self.patch_size[0] * overlap_ratio),
-            round(self.patch_size[1] * overlap_ratio),
-        )
-        self.windows = self.dataset.get_dataset_examples()
-        self.window_cache: dict[
-            int, tuple[dict[str, Any], dict[str, Any], dict[str, Any]]
-        ] = {}
-
-        # Precompute the batch boundaries for each window
-        self.patches = []
-        for window_id, window in enumerate(self.windows):
-            patch_bounds = get_window_patch_options(
-                self.patch_size, self.overlap_size, window.bounds
-            )
-            for i, patch_bound in enumerate(patch_bounds):
-                self.patches.append((window_id, patch_bound, (i, len(patch_bounds))))
-
-    def get_raw_inputs(
-        self, index: int
-    ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
-        """Get the raw inputs for a single patch. Retrieve from cache if possible.
-
-        Also crops/pads the tensors by patch size to protect slicing near right/bottom edges.
-
-        Args:
-            index: the index of the patch.
-
-        Returns:
-            a tuple of (raw_inputs, passthrough_inputs, metadata).
-        """
-        if index in self.window_cache:
-            return self.window_cache[index]
-
-        raw_inputs, passthrough_inputs, metadata = self.dataset.get_raw_inputs(index)
-        pad_slice_protect(raw_inputs, passthrough_inputs, self.patch_size)
-
-        self.window_cache[index] = (raw_inputs, passthrough_inputs, metadata)
-        return self.window_cache[index]
-
-    @staticmethod
-    def _crop_input_dict(
-        d: dict[str, Any],
-        start_offset: tuple[int, int],
-        end_offset: tuple[int, int],
-        cur_geom: STGeometry,
-    ) -> dict[str, Any]:
-        """Crop a dictionary of inputs to the given bounds."""
-        cropped = {}
-        for input_name, value in d.items():
-            if isinstance(value, torch.Tensor):
-                cropped[input_name] = value[
-                    :,
-                    start_offset[1] : end_offset[1],
-                    start_offset[0] : end_offset[0],
-                ].clone()
-            elif isinstance(value, list):
-                cropped[input_name] = [
-                    feat for feat in value if cur_geom.intersects(feat.geometry)
-                ]
-            else:
-                raise ValueError("got input that is neither tensor nor feature list")
-        return cropped
-
-    def __len__(self) -> int:
-        """Return the total number of patches in the dataset."""
-        return len(self.patches)
-
-    def __getitem__(
-        self, index: int
-    ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
-        """Return (input_dict, target_dict, metadata) for a single flattened patch."""
-        (window_id, patch_bounds, (patch_idx, num_patches)) = self.patches[index]
-        raw_inputs, passthrough_inputs, metadata = self.get_raw_inputs(window_id)
-        bounds = metadata["bounds"]
-
-        cur_geom = STGeometry(metadata["projection"], shapely.box(*patch_bounds), None)
-        start_offset = (patch_bounds[0] - bounds[0], patch_bounds[1] - bounds[1])
-        end_offset = (patch_bounds[2] - bounds[0], patch_bounds[3] - bounds[1])
-
-        cur_raw_inputs = self._crop_input_dict(
-            raw_inputs, start_offset, end_offset, cur_geom
-        )
-        cur_passthrough_inputs = self._crop_input_dict(
-            passthrough_inputs, start_offset, end_offset, cur_geom
-        )
-
-        # Adjust the metadata as well.
-        cur_metadata = metadata.copy()
-        cur_metadata["bounds"] = patch_bounds
-        cur_metadata["patch_idx"] = patch_idx
-        cur_metadata["num_patches"] = num_patches
-
-        # Now we can compute input and target dicts via the task.
-        input_dict, target_dict = self.dataset.task.process_inputs(
-            cur_raw_inputs,
-            metadata=cur_metadata,
-            load_targets=not self.dataset.split_config.get_skip_targets(),
-        )
-        input_dict.update(cur_passthrough_inputs)
-        input_dict, target_dict = self.dataset.transforms(input_dict, target_dict)
-        input_dict["dataset_source"] = self.dataset.name
-
-        return input_dict, target_dict, cur_metadata
-
-    def get_dataset_examples(self) -> list[Window]:
-        """Returns a list of windows in this dataset."""
-        return self.dataset.get_dataset_examples()
-
-    def set_name(self, name: str) -> None:
-        """Sets dataset name.
-
-        Args:
-            name: dataset name
-        """
-        self.dataset.set_name(name)
-
-
 class RetryDataset(torch.utils.data.Dataset):
     """A dataset wrapper that retries getitem upon encountering error."""
 

rslearn/utils/array.py

@@ -1,14 +1,16 @@
 """Array util functions."""
 
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy.typing as npt
-import torch
+
+if TYPE_CHECKING:
+    import torch
 
 
 def copy_spatial_array(
-    src: torch.Tensor | npt.NDArray[Any],
-    dst: torch.Tensor | npt.NDArray[Any],
+    src: "torch.Tensor | npt.NDArray[Any]",
+    dst: "torch.Tensor | npt.NDArray[Any]",
     src_offset: tuple[int, int],
     dst_offset: tuple[int, int],
 ) -> None:

{rslearn-0.0.13.dist-info → rslearn-0.0.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rslearn
-Version: 0.0.13
+Version: 0.0.15
 Summary: A library for developing remote sensing datasets and models
 Author: OlmoEarth Team
 License:                                  Apache License

{rslearn-0.0.13.dist-info → rslearn-0.0.15.dist-info}/RECORD

@@ -1,12 +1,13 @@
 rslearn/__init__.py,sha256=fFmAen3vxZyosEfPbG0W46IttujYGVxzrGkJ0YutmmY,73
 rslearn/arg_parser.py,sha256=GNlJncO6Ck_dCNrcg7z_SSG61I-2gKn3Ix2tAxIk9CI,1428
 rslearn/const.py,sha256=FUCfsvFAs-QarEDJ0grdy0C1HjUjLpNFYGo5I2Vpc5Y,449
+rslearn/lightning_cli.py,sha256=io1Agb2fr-fUu9yOODNJhP8-vJp_v9UbJJA2hkLubKA,2435
 rslearn/log_utils.py,sha256=unD9gShiuO7cx5Nnq8qqVQ4qrbOOwFVgcHxN5bXuiAo,941
-rslearn/main.py,sha256=fLYmm2ZsUTCaJBKZvxu3pc4fB2thaf-p2Qv0AifDlXM,31292
+rslearn/main.py,sha256=JMNMhAHqpb9bDUoKzj6kN659Ft_-gZv_rKUieJcJNwI,29087
 rslearn/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rslearn/template_params.py,sha256=Vop0Ha-S44ctCa9lvSZRjrMETznJZlR5y_gJrVIwrPg,791
 rslearn/config/__init__.py,sha256=Bhf2VVncdMYRC8Wfb4GsJJ13OAJYNCO_ODLSNTmBOHM,638
-rslearn/config/dataset.py,sha256=VpXUGKCr45kzE-W27rgF4tPQuyICfwQkJTxb2z9aXQM,21685
+rslearn/config/dataset.py,sha256=lIuFgJG0Hz7nxacFIpbwOyNJqjlkOlaMfWt91Chjb_M,21338
 rslearn/data_sources/__init__.py,sha256=8_7Pi3agKsatNoxXw74-U5G-QAP-rbdfcH8EkZfJbH4,1449
 rslearn/data_sources/aws_landsat.py,sha256=GA9H04KagBDm-N37jFdh_aHCX2ZneVdnqT1SNOyAwTs,20829
 rslearn/data_sources/aws_open_data.py,sha256=nU_D5cqc-wibxq4uyUNb0z-XD0Puf1gZ8v5FMiMAN5w,30258
@@ -39,7 +40,7 @@ rslearn/dataset/add_windows.py,sha256=pwCEvwLE1jQCoqQxw6CJ-sP46ayWppFa2hGYIB6VVk
 rslearn/dataset/dataset.py,sha256=bjf9nI55j-MF0bIQWSNPjNbpfqnLK4jy-96TAcwO0MM,5214
 rslearn/dataset/handler_summaries.py,sha256=wI99RDk5erCWkzl1A7Uc4chatQ9KWIr4F_0Hxr9Co6s,2607
 rslearn/dataset/index.py,sha256=Wni5m6h4gisRB54fPLnCfUrRTEsJ5EvwS0fs9sYc2wg,6025
-rslearn/dataset/manage.py,sha256=owelBiBqvoIQYLhFMDK4ULzcoGBNE27JV8kl68jf3wg,18563
+rslearn/dataset/manage.py,sha256=IURlbCtm9a5f4d52AXfte1yyodlf6MgjfYn3__GdIL4,19062
 rslearn/dataset/materialize.py,sha256=-z47svc_JqGhzkp8kq5Hd9fykWNqFEUCQezo887TWBw,22056
 rslearn/dataset/remap.py,sha256=6MaImsY02GNACpvRM81RvWmjZWRfAHxo_R3Ox6XLF6A,2723
 rslearn/dataset/window.py,sha256=I5RqZ12jlIXhohw4qews1x_I4tSDpml709DZRtLiN24,12546
@@ -47,7 +48,7 @@ rslearn/models/__init__.py,sha256=_vWoF9d2Slah8-6XhYhdU4SRsy_CNxXjCGQTD2yvu3Q,22
 rslearn/models/anysat.py,sha256=3Oh2gWxicVdUzOjevBEZf0PuolmCy0KC5Ad7JY-0Plc,7949
 rslearn/models/clip.py,sha256=u5aqYnVB4Jag7o1h8EzPDAc1t2BAHeALA9FcUwP5tfo,2238
 rslearn/models/conv.py,sha256=fWyByeswIOKKzyPmP3erYUlZaKEV0huWHA4CyKTBbfY,1703
-rslearn/models/croma.py,sha256=cOazTp3l2PNJltKrmPqD5Gy4pi3CI03-X9G4T10cX2k,9529
+rslearn/models/croma.py,sha256=n7yunpT7lo8vWWaOpx4yt8jZSXjgWqfgZcZKFW5zuEQ,10591
 rslearn/models/dinov3.py,sha256=9k9kNlXCorQQwKjLGptooANd48TUBsITQ1e4fUomlM4,6337
 rslearn/models/faster_rcnn.py,sha256=uaxX6-E1f0BibaA9sorEg3be83C7kTdTc39pC5jRqwE,8286
 rslearn/models/feature_center_crop.py,sha256=24eOrvLEGGVWPw7kPHyUes5HtYNAX7GZ_NpqDGMILEY,1553
@@ -62,18 +63,18 @@ rslearn/models/prithvi.py,sha256=AIzcO5xk1ggR0MjbfhIzqPVgUKFN7odxygmgyAelfW8,401
 rslearn/models/registry.py,sha256=yCcrOvLkbn07Xtln1j7hAB_kmGw0MGsiR2TloJq9Bmk,504
 rslearn/models/resize_features.py,sha256=asKXWrLHIBrU6GaAV0Ory9YuK7IK104XjhkB4ljzI3A,1289
 rslearn/models/sam2_enc.py,sha256=gNlPokr7eNxO2KvnzDMXNxYM2WRO0YkQPjR4110n6cw,3508
-rslearn/models/satlaspretrain.py,sha256=YpjXl-uClhTZMDmyhN64Fg3AszzT-ymZgJB0fO9RyoY,2419
+rslearn/models/satlaspretrain.py,sha256=b6FR_il6MnWU4UpB9OxInZSK9n0IS0PcQuLrWH4YD8g,3046
 rslearn/models/simple_time_series.py,sha256=oTg_akabYFBExJu7JCpbuM211-ZgQS4WerG2nEYrIZY,12774
 rslearn/models/singletask.py,sha256=z4vN9Yvzz0I-U4KJdVZxLJK2ZV-MIv9tzwCGcOWoUPY,1604
 rslearn/models/ssl4eo_s12.py,sha256=sOGEHcDo-rNdmEyoLu2AVEqfxRM_cv6zpfAmyn5c6tw,3553
 rslearn/models/swin.py,sha256=bMlGePXMFou4A_YSUZzjHgN9NniGXaCWdGQ31xHDKis,5511
 rslearn/models/task_embedding.py,sha256=Z6sf61BLCtvdrdnvjh8500b-KiFp3GeWbT4mOqpaCKk,9100
-rslearn/models/terramind.py,sha256=kipar8sMaHJJ3b8vIgL0-s4qhHcA0Vb854vmlZ9cWh4,7524
+rslearn/models/terramind.py,sha256=5POVk_y29LlbVswa6ojd9gdB70iO41yB9Y2aqVY4WdQ,8327
 rslearn/models/trunk.py,sha256=H1QPQGAKsmocq3OiF66GW8MQI4LffupTDrgZR4Ta7QM,4708
 rslearn/models/unet.py,sha256=WUgLgvvlgV8l_6MIDBl6aX1HNOkb24DfnVRIyYXHCjo,6865
 rslearn/models/upsample.py,sha256=3kWbyWZIk56JJxj8en9pieitbrk3XnbIsTKlEkiDQQY,938
 rslearn/models/use_croma.py,sha256=OSBqMuLp-pDtqPNWAVBfmX4wckmyYCKtUDdGCjJk_K8,17966
-rslearn/models/clay/clay.py,sha256=5RO5H8EM0tKjCwWMQ4xDkKkUCwKpm2K_Yw1alnhvVhU,7773
+rslearn/models/clay/clay.py,sha256=29CGCOysx9duEX4Y6LUNHXck_sHjCFrlV4w8CP_hKmI,8460
 rslearn/models/clay/configs/metadata.yaml,sha256=rZTFh4Yb9htEfbQNOPl4HTbFogEhzwIRqFzG-1uT01Y,4652
 rslearn/models/detr/__init__.py,sha256=GGAnTIhyuvl34IRrJ_4gXjm_01OlM5rbQQ3c3TGfbK8,84
 rslearn/models/detr/box_ops.py,sha256=ORCF6EwMpMBB_VgQT05SjR47dCR2rN2gPhL_gsuUWJs,3236
@@ -107,8 +108,9 @@ rslearn/tile_stores/__init__.py,sha256=o_tWVKu6UwFzZbO9jn_3cmIDqc_Q3qDd6tA9If0T_
 rslearn/tile_stores/default.py,sha256=PYaDNvBxhJTDKJGw0EjDTSE1OKajR7_iJpMbOjj-mE8,15054
 rslearn/tile_stores/tile_store.py,sha256=9AeYduDYPp_Ia2NMlq6osptpz_AFGIOQcLJrqZ_m-z0,10469
 rslearn/train/__init__.py,sha256=fnJyY4aHs5zQqbDKSfXsJZXY_M9fbTsf7dRYaPwZr2M,30
-rslearn/train/data_module.py,sha256=K-nQgnOZn-KGq_G2pVOQFtWRrlWih0212i_bkXZ2bEE,23515
-rslearn/train/dataset.py,sha256=YiskNlYYcKqZxyw0Xzop1RGLbjMc-oK_rmhrSMVbTQg,51857
+rslearn/train/all_patches_dataset.py,sha256=xFJ96HU3CodrUBzXTsgrmEShosKH79T2SxI0xDVSH3Q,18217
+rslearn/train/data_module.py,sha256=pgut8rEWHIieZ7RR8dUvhtlNqk0egEdznYF3tCvqdHg,23552
+rslearn/train/dataset.py,sha256=8F3bpus25g_NG0-CwMCuznwKxOvBDClNBCOEvDbMyN8,34312
 rslearn/train/lightning_module.py,sha256=ZLBiId3secUlVs2yzkN-mwVv4rMdh5TkdZYl4vv_Cw0,14466
 rslearn/train/optimizer.py,sha256=EKSqkmERalDA0bF32Gey7n6z69KLyaUWKlRsGJfKBmE,927
 rslearn/train/prediction_writer.py,sha256=mDvREwEB5k5_tNuBnYIvAGnxS3sYFWQYvV07V3UEe2k,14106
@@ -138,7 +140,7 @@ rslearn/train/transforms/select_bands.py,sha256=uDfD9G8Z4VTt88QZsjj1FB20QEmzSefh
 rslearn/train/transforms/sentinel1.py,sha256=FrLaYZs2AjqWQCun8DTFtgo1l0xLxqaFKtDNIehtpDg,1913
 rslearn/train/transforms/transform.py,sha256=n1Qzqix2dVvej-Q7iPzHeOQbqH79IBlvqPoymxhNVpE,4446
 rslearn/utils/__init__.py,sha256=GNvdTUmXakiEMnLdje7k1fe5aC7SFVqP757kbpN6Fzw,558
-rslearn/utils/array.py,sha256=JwZi7o0uj-dftREzJmqrRVR2joIwBikm3Er9KeHVIZU,2402
+rslearn/utils/array.py,sha256=RC7ygtPnQwU6Lb9kwORvNxatJcaJ76JPsykQvndAfes,2444
 rslearn/utils/feature.py,sha256=lsg0WThZDJzo1mrbaL04dXYI5G3x-n5FG9aEjj7uUaI,1649
 rslearn/utils/fsspec.py,sha256=9QwN46heBhjUnth3qFeRNE3W6Wlr6dM3twYVswPnS9o,5300
 rslearn/utils/geometry.py,sha256=oZllq1aBFcDewTTDYAMnTeP1xR0EdB5Xz3ILmfASo-8,18455
@@ -152,10 +154,10 @@ rslearn/utils/spatial_index.py,sha256=eomJAUgzmjir8j9HZnSgQoJHwN9H0wGTjmJkMkLLfs
 rslearn/utils/sqlite_index.py,sha256=YGOJi66544e6JNtfSft6YIlHklFdSJO2duxQ4TJ2iu4,2920
 rslearn/utils/time.py,sha256=2ilSLG94_sxLP3y5RSV5L5CG8CoND_dbdzYEHVtN-I8,387
 rslearn/utils/vector_format.py,sha256=EIChYCL6GLOILS2TO2JBkca1TuaWsSubWv6iRS3P2ds,16139
-rslearn-0.0.13.dist-info/licenses/LICENSE,sha256=_99ZWPoLdlUbqZoSC5DF4ihiNwl5rTEmBaq2fACecdg,11352
-rslearn-0.0.13.dist-info/licenses/NOTICE,sha256=wLPr6rwV_jCg-xEknNGwhnkfRfuoOE9MZ-lru2yZyLI,5070
-rslearn-0.0.13.dist-info/METADATA,sha256=44oDmbvkIrjJ0unVNaYeO5OypD6RavmG7l5HUz9Re48,36319
-rslearn-0.0.13.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-rslearn-0.0.13.dist-info/entry_points.txt,sha256=doTBQ57NT7nq-dgYGgTTw6mafcGWb_4PWYtYR4rGm50,46
-rslearn-0.0.13.dist-info/top_level.txt,sha256=XDKo90WBH8P9RQumHxo0giLJsoufT4r9odv-WE6Ahk4,8
-rslearn-0.0.13.dist-info/RECORD,,
+rslearn-0.0.15.dist-info/licenses/LICENSE,sha256=_99ZWPoLdlUbqZoSC5DF4ihiNwl5rTEmBaq2fACecdg,11352
+rslearn-0.0.15.dist-info/licenses/NOTICE,sha256=wLPr6rwV_jCg-xEknNGwhnkfRfuoOE9MZ-lru2yZyLI,5070
+rslearn-0.0.15.dist-info/METADATA,sha256=HRkJjQfvxCEosmCBvLcLd9nZnXjbmfAgPIknMy_ORBo,36319
+rslearn-0.0.15.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+rslearn-0.0.15.dist-info/entry_points.txt,sha256=doTBQ57NT7nq-dgYGgTTw6mafcGWb_4PWYtYR4rGm50,46
+rslearn-0.0.15.dist-info/top_level.txt,sha256=XDKo90WBH8P9RQumHxo0giLJsoufT4r9odv-WE6Ahk4,8
+rslearn-0.0.15.dist-info/RECORD,,