rslearn 0.0.18__py3-none-any.whl → 0.0.19__py3-none-any.whl

rslearn/models/simple_time_series.py

@@ -229,7 +229,13 @@ class SimpleTimeSeries(FeatureExtractor):
 
         # Now we can apply the underlying FeatureExtractor.
         # Its output must be a FeatureMaps.
-        encoder_output = self.encoder(batched_inputs)
+        assert batched_inputs is not None
+        encoder_output = self.encoder(
+            ModelContext(
+                inputs=batched_inputs,
+                metadatas=context.metadatas,
+            )
+        )
         if not isinstance(encoder_output, FeatureMaps):
             raise ValueError(
                 "output of underlying FeatureExtractor in SimpleTimeSeries must be a FeatureMaps"

rslearn/train/all_patches_dataset.py

@@ -9,7 +9,7 @@ import shapely
 import torch
 
 from rslearn.dataset import Window
-from rslearn.train.dataset import ModelDataset
+from rslearn.train.dataset import DataInput, ModelDataset
 from rslearn.train.model_context import SampleMetadata
 from rslearn.utils.geometry import PixelBounds, STGeometry
 
@@ -34,22 +34,28 @@ def get_window_patch_options(
             bottommost patches may extend beyond the provided bounds.
     """
     # We stride the patches by patch_size - overlap_size until the last patch.
+    # We handle the first patch with a special case to ensure it is always used.
     # 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(
+    cols = [bounds[0]] + list(
         range(
-            bounds[0],
+            bounds[0] + patch_size[0],
             bounds[2] - patch_size[0],
             patch_size[0] - overlap_size[0],
         )
-    ) + [bounds[2] - patch_size[0]]
-    rows = list(
+    )
+    rows = [bounds[1]] + list(
         range(
-            bounds[1],
+            bounds[1] + patch_size[1],
             bounds[3] - patch_size[1],
             patch_size[1] - overlap_size[1],
         )
-    ) + [bounds[3] - patch_size[1]]
+    )
+    # Add last patches only if the input is larger than one patch.
+    if bounds[2] - patch_size[0] > bounds[0]:
+        cols.append(bounds[2] - patch_size[0])
+    if bounds[3] - patch_size[1] > bounds[1]:
+        rows.append(bounds[3] - patch_size[1])
 
     patch_bounds: list[PixelBounds] = []
     for col in cols:
@@ -62,13 +68,17 @@ def pad_slice_protect(
     raw_inputs: dict[str, Any],
     passthrough_inputs: dict[str, Any],
     patch_size: tuple[int, int],
+    inputs: dict[str, DataInput],
 ) -> tuple[dict[str, Any], dict[str, Any]]:
     """Pad tensors in-place by patch size to protect slicing near right/bottom edges.
 
+    The padding is scaled based on each input's resolution_factor.
+
     Args:
         raw_inputs: the raw inputs to pad.
         passthrough_inputs: the passthrough inputs to pad.
-        patch_size: the size of the patches to extract.
+        patch_size: the size of the patches to extract (at window resolution).
+        inputs: the DataInput definitions, used to get resolution_factor per input.
 
     Returns:
         a tuple of (raw_inputs, passthrough_inputs).
@@ -77,8 +87,14 @@ def pad_slice_protect(
         for input_name, value in list(d.items()):
             if not isinstance(value, torch.Tensor):
                 continue
+            # Get resolution scale for this input
+            rf = inputs[input_name].resolution_factor
+            scale = rf.numerator / rf.denominator
+            # Scale the padding amount
+            scaled_pad_x = int(patch_size[0] * scale)
+            scaled_pad_y = int(patch_size[1] * scale)
             d[input_name] = torch.nn.functional.pad(
-                value, pad=(0, patch_size[0], 0, patch_size[1])
+                value, pad=(0, scaled_pad_x, 0, scaled_pad_y)
             )
     return raw_inputs, passthrough_inputs
 
@@ -123,6 +139,7 @@ class IterableAllPatchesDataset(torch.utils.data.IterableDataset):
         self.rank = rank
         self.world_size = world_size
         self.windows = self.dataset.get_dataset_examples()
+        self.inputs = dataset.inputs
 
     def set_name(self, name: str) -> None:
         """Sets dataset name.
@@ -235,8 +252,10 @@ class IterableAllPatchesDataset(torch.utils.data.IterableDataset):
 
                 # 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)
+                # the tensors later. Padding is scaled per-input based on resolution_factor.
+                pad_slice_protect(
+                    raw_inputs, passthrough_inputs, self.patch_size, self.inputs
+                )
 
                 # Now iterate over the patches and extract/yield the crops.
                 # Note that, in case user is leveraging RslearnWriter, it is important that
@@ -258,15 +277,28 @@ class IterableAllPatchesDataset(torch.utils.data.IterableDataset):
                     )
 
                     # Define a helper function to handle each input dict.
+                    # Crop coordinates are scaled based on each input's resolution_factor.
                     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.
+                                # Get resolution scale for this input
+                                rf = self.inputs[input_name].resolution_factor
+                                scale = rf.numerator / rf.denominator
+                                # Scale the crop coordinates
+                                scaled_start = (
+                                    int(start_offset[0] * scale),
+                                    int(start_offset[1] * scale),
+                                )
+                                scaled_end = (
+                                    int(end_offset[0] * scale),
+                                    int(end_offset[1] * scale),
+                                )
+                                # Crop the CHW tensor with scaled coordinates.
                                 cropped[input_name] = value[
                                     :,
-                                    start_offset[1] : end_offset[1],
-                                    start_offset[0] : end_offset[0],
+                                    scaled_start[1] : scaled_end[1],
+                                    scaled_start[0] : scaled_end[0],
                                 ].clone()
                             elif isinstance(value, list):
                                 cropped[input_name] = [
@@ -348,6 +380,7 @@ class InMemoryAllPatchesDataset(torch.utils.data.Dataset):
             round(self.patch_size[1] * overlap_ratio),
         )
         self.windows = self.dataset.get_dataset_examples()
+        self.inputs = dataset.inputs
         self.window_cache: dict[
             int, tuple[dict[str, Any], dict[str, Any], SampleMetadata]
         ] = {}
@@ -378,26 +411,41 @@ class InMemoryAllPatchesDataset(torch.utils.data.Dataset):
             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)
+        pad_slice_protect(raw_inputs, passthrough_inputs, self.patch_size, self.inputs)
 
         self.window_cache[index] = (raw_inputs, passthrough_inputs, metadata)
         return self.window_cache[index]
 
-    @staticmethod
     def _crop_input_dict(
+        self,
         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."""
+        """Crop a dictionary of inputs to the given bounds.
+
+        Crop coordinates are scaled based on each input's resolution_factor.
+        """
         cropped = {}
         for input_name, value in d.items():
             if isinstance(value, torch.Tensor):
+                # Get resolution scale for this input
+                rf = self.inputs[input_name].resolution_factor
+                scale = rf.numerator / rf.denominator
+                # Scale the crop coordinates
+                scaled_start = (
+                    int(start_offset[0] * scale),
+                    int(start_offset[1] * scale),
+                )
+                scaled_end = (
+                    int(end_offset[0] * scale),
+                    int(end_offset[1] * scale),
+                )
                 cropped[input_name] = value[
                     :,
-                    start_offset[1] : end_offset[1],
-                    start_offset[0] : end_offset[0],
+                    scaled_start[1] : scaled_end[1],
+                    scaled_start[0] : scaled_end[0],
                 ].clone()
             elif isinstance(value, list):
                 cropped[input_name] = [

rslearn/train/dataset.py

@@ -24,7 +24,7 @@ from rslearn.dataset.storage.file import FileWindowStorage
 from rslearn.dataset.window import Window, get_layer_and_group_from_dir_name
 from rslearn.log_utils import get_logger
 from rslearn.utils.feature import Feature
-from rslearn.utils.geometry import PixelBounds
+from rslearn.utils.geometry import PixelBounds, ResolutionFactor
 from rslearn.utils.mp import star_imap_unordered
 
 from .model_context import SampleMetadata
@@ -130,6 +130,10 @@ class DataInput:
     """Specification of a piece of data from a window that is needed for training.
 
     The DataInput includes which layer(s) the data can be obtained from for each window.
+
+    Note that this class is not a dataclass because jsonargparse does not play well
+    with dataclasses without enabling specialized options which we have not validated
+    will work with the rest of our code.
     """
 
     def __init__(
@@ -143,7 +147,9 @@ class DataInput:
         dtype: DType = DType.FLOAT32,
         load_all_layers: bool = False,
         load_all_item_groups: bool = False,
-    ) -> None:
+        resolution_factor: ResolutionFactor = ResolutionFactor(),
+        resampling: Resampling = Resampling.nearest,
+    ):
         """Initialize a new DataInput.
 
         Args:
@@ -166,6 +172,11 @@ class DataInput:
                 are reading from. By default, we assume the specified layer name is of
                 the form "{layer_name}.{group_idx}" and read that item group only. With
                 this option enabled, we ignore the group_idx and read all item groups.
+            resolution_factor: controls the resolution at which raster data is loaded for training.
+                By default (factor=1), data is loaded at the window resolution.
+                E.g. for a 64x64 window at 10 m/pixel with resolution_factor=1/2,
+                the resulting tensor is 32x32 (covering the same geographic area at 20 m/pixel).
+            resampling: resampling method (default nearest neighbor).
         """
         self.data_type = data_type
         self.layers = layers
@@ -176,6 +187,8 @@ class DataInput:
         self.dtype = dtype
         self.load_all_layers = load_all_layers
         self.load_all_item_groups = load_all_item_groups
+        self.resolution_factor = resolution_factor
+        self.resampling = resampling
 
 
 def read_raster_layer_for_data_input(
@@ -233,15 +246,23 @@ def read_raster_layer_for_data_input(
             + f"window {window.name} layer {layer_name} group {group_idx}"
         )
 
+    # Get the projection and bounds to read under (multiply window resolution # by
+    # the specified resolution factor).
+    final_projection = data_input.resolution_factor.multiply_projection(
+        window.projection
+    )
+    final_bounds = data_input.resolution_factor.multiply_bounds(bounds)
+
     image = torch.zeros(
-        (len(needed_bands), bounds[3] - bounds[1], bounds[2] - bounds[0]),
+        (
+            len(needed_bands),
+            final_bounds[3] - final_bounds[1],
+            final_bounds[2] - final_bounds[0],
+        ),
         dtype=get_torch_dtype(data_input.dtype),
     )
 
     for band_set, src_indexes, dst_indexes in needed_sets_and_indexes:
-        final_projection, final_bounds = band_set.get_final_projection_and_bounds(
-            window.projection, bounds
-        )
         if band_set.format is None:
             raise ValueError(f"No format specified for {layer_name}")
         raster_format = band_set.instantiate_raster_format()
@@ -249,44 +270,16 @@ def read_raster_layer_for_data_input(
             layer_name, band_set.bands, group_idx=group_idx
         )
 
-        # Previously we always read in the native projection of the data, and then
-        # zoom in or out (the resolution must be a power of two off) to match the
-        # window's resolution.
-        # However, this fails if the bounds are not multiples of the resolution factor.
-        # So we fallback to reading directly in the window projection if that is the
-        # case (which may be a bit slower).
-        is_bounds_zoomable = True
-        if band_set.zoom_offset < 0:
-            zoom_factor = 2 ** (-band_set.zoom_offset)
-            is_bounds_zoomable = (final_bounds[2] - final_bounds[0]) * zoom_factor == (
-                bounds[2] - bounds[0]
-            ) and (final_bounds[3] - final_bounds[1]) * zoom_factor == (
-                bounds[3] - bounds[1]
-            )
-
-        if is_bounds_zoomable:
-            src = raster_format.decode_raster(
-                raster_dir, final_projection, final_bounds
-            )
-
-            # Resize to patch size if needed.
-            # This is for band sets that are stored at a lower resolution.
-            # Here we assume that it is a multiple.
-            if src.shape[1:3] != image.shape[1:3]:
-                if src.shape[1] < image.shape[1]:
-                    factor = image.shape[1] // src.shape[1]
-                    src = src.repeat(repeats=factor, axis=1).repeat(
-                        repeats=factor, axis=2
-                    )
-                else:
-                    factor = src.shape[1] // image.shape[1]
-                    src = src[:, ::factor, ::factor]
-
-        else:
-            src = raster_format.decode_raster(
-                raster_dir, window.projection, bounds, resampling=Resampling.nearest
-            )
+        # TODO: previously we try to read based on band_set.zoom_offset when possible,
+        # and handle zooming in with torch.repeat (if resampling method is nearest
+        # neighbor). However, we have not benchmarked whether this actually improves
+        # data loading speed, so for simplicity, for now we let rasterio handle the
+        # resampling. If it really is much faster to handle it via torch, then it may
+        # make sense to bring back that functionality.
 
+        src = raster_format.decode_raster(
+            raster_dir, final_projection, final_bounds, resampling=Resampling.nearest
+        )
         image[dst_indexes, :, :] = torch.as_tensor(
             src[src_indexes, :, :].astype(data_input.dtype.get_numpy_dtype())
         )

rslearn/train/scheduler.py

@@ -8,6 +8,7 @@ from torch.optim.lr_scheduler import (
     CosineAnnealingLR,
     CosineAnnealingWarmRestarts,
     LRScheduler,
+    MultiStepLR,
     ReduceLROnPlateau,
 )
 
@@ -50,6 +51,20 @@ class PlateauScheduler(SchedulerFactory):
         return ReduceLROnPlateau(optimizer, **self.get_kwargs())
 
 
+@dataclass
+class MultiStepScheduler(SchedulerFactory):
+    """Step learning rate scheduler."""
+
+    milestones: list[int]
+    gamma: float | None = None
+    last_epoch: int | None = None
+
+    def build(self, optimizer: Optimizer) -> LRScheduler:
+        """Build the ReduceLROnPlateau scheduler."""
+        super().build(optimizer)
+        return MultiStepLR(optimizer, **self.get_kwargs())
+
+
 @dataclass
 class CosineAnnealingScheduler(SchedulerFactory):
     """Cosine annealing learning rate scheduler."""

rslearn/train/transforms/resize.py

@@ -0,0 +1,74 @@
+"""Resize transform."""
+
+from typing import Any
+
+import torch
+import torchvision
+from torchvision.transforms import InterpolationMode
+
+from .transform import Transform
+
+INTERPOLATION_MODES = {
+    "nearest": InterpolationMode.NEAREST,
+    "nearest_exact": InterpolationMode.NEAREST_EXACT,
+    "bilinear": InterpolationMode.BILINEAR,
+    "bicubic": InterpolationMode.BICUBIC,
+}
+
+
+class Resize(Transform):
+    """Resizes inputs to a target size."""
+
+    def __init__(
+        self,
+        target_size: tuple[int, int],
+        selectors: list[str] = [],
+        interpolation: str = "nearest",
+    ):
+        """Initialize a resize transform.
+
+        Args:
+            target_size: the (height, width) to resize to.
+            selectors: items to transform.
+            interpolation: the interpolation mode to use for resizing.
+                Must be one of "nearest", "nearest_exact", "bilinear", or "bicubic".
+        """
+        super().__init__()
+        self.target_size = target_size
+        self.selectors = selectors
+        self.interpolation = INTERPOLATION_MODES[interpolation]
+
+    def apply_resize(self, image: torch.Tensor) -> torch.Tensor:
+        """Apply resizing on the specified image.
+
+        If the image is 2D, it is unsqueezed to 3D and then squeezed
+        back after resizing.
+
+        Args:
+            image: the image to transform.
+        """
+        if image.dim() == 2:
+            image = image.unsqueeze(0)  # (H, W) -> (1, H, W)
+            result = torchvision.transforms.functional.resize(
+                image, self.target_size, self.interpolation
+            )
+            return result.squeeze(0)  # (1, H, W) -> (H, W)
+
+        return torchvision.transforms.functional.resize(
+            image, self.target_size, self.interpolation
+        )
+
+    def forward(
+        self, input_dict: dict[str, Any], target_dict: dict[str, Any]
+    ) -> tuple[dict[str, Any], dict[str, Any]]:
+        """Apply transform over the inputs and targets.
+
+        Args:
+            input_dict: the input
+            target_dict: the target
+
+        Returns:
+            transformed (input_dicts, target_dicts) tuple
+        """
+        self.apply_fn(self.apply_resize, input_dict, target_dict, self.selectors)
+        return input_dict, target_dict

rslearn/utils/geometry.py

@@ -116,6 +116,79 @@ class Projection:
 WGS84_PROJECTION = Projection(CRS.from_epsg(WGS84_EPSG), 1, 1)
 
 
+class ResolutionFactor:
+    """Multiplier for the resolution in a Projection.
+
+    The multiplier is either an integer x, or the inverse of an integer (1/x).
+
+    Factors greater than 1 increase the projection_units/pixel resolution, increasing
+    the resolution (more pixels per projection unit). Factors less than 1 make it coarser
+    (less pixels).
+    """
+
+    def __init__(self, numerator: int = 1, denominator: int = 1):
+        """Create a new ResolutionFactor.
+
+        Args:
+            numerator: the numerator of the fraction.
+            denominator: the denominator of the fraction. If set, numerator must be 1.
+        """
+        if numerator != 1 and denominator != 1:
+            raise ValueError("one of numerator or denominator must be 1")
+        if not isinstance(numerator, int) or not isinstance(denominator, int):
+            raise ValueError("numerator and denominator must be integers")
+        if numerator < 1 or denominator < 1:
+            raise ValueError("numerator and denominator must be >= 1")
+        self.numerator = numerator
+        self.denominator = denominator
+
+    def multiply_projection(self, projection: Projection) -> Projection:
+        """Multiply the projection by this factor."""
+        if self.denominator > 1:
+            return Projection(
+                projection.crs,
+                projection.x_resolution * self.denominator,
+                projection.y_resolution * self.denominator,
+            )
+        else:
+            return Projection(
+                projection.crs,
+                projection.x_resolution // self.numerator,
+                projection.y_resolution // self.numerator,
+            )
+
+    def multiply_bounds(self, bounds: PixelBounds) -> PixelBounds:
+        """Multiply the bounds by this factor.
+
+        When coarsening, the width and height of the given bounds must be a multiple of
+        the denominator.
+        """
+        if self.denominator > 1:
+            # Verify the width and height are multiples of the denominator.
+            # Otherwise the new width and height is not an integer.
+            width = bounds[2] - bounds[0]
+            height = bounds[3] - bounds[1]
+            if width % self.denominator != 0 or height % self.denominator != 0:
+                raise ValueError(
+                    f"width {width} or height {height} is not a multiple of the resolution factor {self.denominator}"
+                )
+            # TODO: an offset could be introduced by bounds not being a multiple
+            # of the denominator -> will need to decide how to handle that.
+            return (
+                bounds[0] // self.denominator,
+                bounds[1] // self.denominator,
+                bounds[2] // self.denominator,
+                bounds[3] // self.denominator,
+            )
+        else:
+            return (
+                bounds[0] * self.numerator,
+                bounds[1] * self.numerator,
+                bounds[2] * self.numerator,
+                bounds[3] * self.numerator,
+            )
+
+
 class STGeometry:
     """A spatiotemporal geometry.
 

rslearn/utils/jsonargparse.py

@@ -8,6 +8,7 @@ from rasterio.crs import CRS
 from upath import UPath
 
 from rslearn.config.dataset import LayerConfig
+from rslearn.utils.geometry import ResolutionFactor
 
 if TYPE_CHECKING:
     from rslearn.data_sources.data_source import DataSourceContext
@@ -91,6 +92,68 @@ def data_source_context_deserializer(v: dict[str, Any]) -> "DataSourceContext":
     )
 
 
+def resolution_factor_serializer(v: ResolutionFactor) -> str:
+    """Serialize ResolutionFactor for jsonargparse.
+
+    Args:
+        v: the ResolutionFactor object.
+
+    Returns:
+        the ResolutionFactor encoded to string
+    """
+    if hasattr(v, "init_args"):
+        init_args = v.init_args
+        return f"{init_args.numerator}/{init_args.denominator}"
+
+    return f"{v.numerator}/{v.denominator}"
+
+
+def resolution_factor_deserializer(v: int | str | dict) -> ResolutionFactor:
+    """Deserialize ResolutionFactor for jsonargparse.
+
+    Args:
+        v: the encoded ResolutionFactor.
+
+    Returns:
+        the decoded ResolutionFactor object
+    """
+    # Handle already-instantiated ResolutionFactor
+    if isinstance(v, ResolutionFactor):
+        return v
+
+    # Handle Namespace from class_path syntax (used during config save/validation)
+    if hasattr(v, "init_args"):
+        init_args = v.init_args
+        return ResolutionFactor(
+            numerator=init_args.numerator,
+            denominator=init_args.denominator,
+        )
+
+    # Handle dict from class_path syntax in YAML config
+    if isinstance(v, dict) and "init_args" in v:
+        init_args = v["init_args"]
+        return ResolutionFactor(
+            numerator=init_args.get("numerator", 1),
+            denominator=init_args.get("denominator", 1),
+        )
+
+    if isinstance(v, int):
+        return ResolutionFactor(numerator=v)
+    elif isinstance(v, str):
+        parts = v.split("/")
+        if len(parts) == 1:
+            return ResolutionFactor(numerator=int(parts[0]))
+        elif len(parts) == 2:
+            return ResolutionFactor(
+                numerator=int(parts[0]),
+                denominator=int(parts[1]),
+            )
+        else:
+            raise ValueError("expected resolution factor to be of the form x or 1/x")
+    else:
+        raise ValueError("expected resolution factor to be str or int")
+
+
 def init_jsonargparse() -> None:
     """Initialize custom jsonargparse serializers."""
     global INITIALIZED
@@ -100,6 +163,9 @@ def init_jsonargparse() -> None:
     jsonargparse.typing.register_type(
         datetime, datetime_serializer, datetime_deserializer
     )
+    jsonargparse.typing.register_type(
+        ResolutionFactor, resolution_factor_serializer, resolution_factor_deserializer
+    )
 
     from rslearn.data_sources.data_source import DataSourceContext
 

{rslearn-0.0.18.dist-info → rslearn-0.0.19.dist-info}/METADATA

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