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

rslearn/dataset/window.py

@@ -1,20 +1,16 @@
 """rslearn windows."""
 
-import json
 from datetime import datetime
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
 import shapely
 from upath import UPath
 
+from rslearn.dataset.storage.storage import WindowStorage
 from rslearn.log_utils import get_logger
 from rslearn.utils import Projection, STGeometry
-from rslearn.utils.fsspec import open_atomic
 from rslearn.utils.raster_format import get_bandset_dirname
 
-if TYPE_CHECKING:
-    from .index import DatasetIndex
-
 logger = get_logger(__name__)
 
 LAYERS_DIRECTORY_NAME = "layers"
@@ -138,14 +134,13 @@ class Window:
 
     def __init__(
         self,
-        path: UPath,
+        storage: WindowStorage,
         group: str,
         name: str,
         projection: Projection,
         bounds: tuple[int, int, int, int],
         time_range: tuple[datetime, datetime] | None,
         options: dict[str, Any] = {},
-        index: "DatasetIndex | None" = None,
     ) -> None:
         """Creates a new Window instance.
 
@@ -153,23 +148,21 @@ class Window:
         stored in metadata.json.
 
         Args:
-            path: the directory of this window
+            storage: the dataset storage for the underlying rslearn dataset.
             group: the group the window belongs to
             name: the unique name for this window
             projection: the projection of the window
             bounds: the bounds of the window in pixel coordinates
             time_range: optional time range of the window
             options: additional options (?)
-            index: DatasetIndex if it is available
         """
-        self.path = path
+        self.storage = storage
         self.group = group
         self.name = name
         self.projection = projection
         self.bounds = bounds
         self.time_range = time_range
         self.options = options
-        self.index = index
 
     def get_geometry(self) -> STGeometry:
         """Computes the STGeometry corresponding to this window."""
@@ -181,29 +174,11 @@ class Window:
 
     def load_layer_datas(self) -> dict[str, WindowLayerData]:
         """Load layer datas describing items in retrieved layers from items.json."""
-        # Load from index if it is available.
-        if self.index is not None:
-            layer_datas = self.index.layer_datas.get(self.name, [])
-
-        else:
-            items_fname = self.path / "items.json"
-            if not items_fname.exists():
-                return {}
-            with items_fname.open("r") as f:
-                layer_datas = [
-                    WindowLayerData.deserialize(layer_data)
-                    for layer_data in json.load(f)
-                ]
-
-        return {layer_data.layer_name: layer_data for layer_data in layer_datas}
+        return self.storage.get_layer_datas(self.group, self.name)
 
     def save_layer_datas(self, layer_datas: dict[str, WindowLayerData]) -> None:
         """Save layer datas to items.json."""
-        json_data = [layer_data.serialize() for layer_data in layer_datas.values()]
-        items_fname = self.path / "items.json"
-        logger.info(f"Saving window items to {items_fname}")
-        with open_atomic(items_fname, "w") as f:
-            json.dump(json_data, f)
+        self.storage.save_layer_datas(self.group, self.name, layer_datas)
 
     def list_completed_layers(self) -> list[tuple[str, int]]:
         """List the layers available for this window that are completed.
@@ -211,18 +186,7 @@ class Window:
         Returns:
             a list of (layer_name, group_idx) completed layers.
         """
-        layers_directory = self.path / LAYERS_DIRECTORY_NAME
-        if not layers_directory.exists():
-            return []
-
-        completed_layers = []
-        for layer_dir in layers_directory.iterdir():
-            layer_name, group_idx = get_layer_and_group_from_dir_name(layer_dir.name)
-            if not self.is_layer_completed(layer_name, group_idx):
-                continue
-            completed_layers.append((layer_name, group_idx))
-
-        return completed_layers
+        return self.storage.list_completed_layers(self.group, self.name)
 
     def get_layer_dir(self, layer_name: str, group_idx: int = 0) -> UPath:
         """Get the directory containing materialized data for the specified layer.
@@ -235,7 +199,9 @@ class Window:
         Returns:
             the path where data is or should be materialized.
         """
-        return get_window_layer_dir(self.path, layer_name, group_idx)
+        return get_window_layer_dir(
+            self.storage.get_window_root(self.group, self.name), layer_name, group_idx
+        )
 
     def is_layer_completed(self, layer_name: str, group_idx: int = 0) -> bool:
         """Check whether the specified layer is completed.
@@ -250,14 +216,9 @@ class Window:
         Returns:
             whether the layer is completed
         """
-        # Use the index to speed up the completed check if it is available.
-        if self.index is not None:
-            return (layer_name, group_idx) in self.index.completed_layers.get(
-                self.name, []
-            )
-
-        layer_dir = self.get_layer_dir(layer_name, group_idx)
-        return (layer_dir / "completed").exists()
+        return self.storage.is_layer_completed(
+            self.group, self.name, layer_name, group_idx
+        )
 
     def mark_layer_completed(self, layer_name: str, group_idx: int = 0) -> None:
         """Mark the specified layer completed.
@@ -272,8 +233,7 @@ class Window:
             layer_name: the layer name.
             group_idx: the index of the group within the layer.
         """
-        layer_dir = self.get_layer_dir(layer_name, group_idx)
-        (layer_dir / "completed").touch()
+        self.storage.mark_layer_completed(self.group, self.name, layer_name, group_idx)
 
     def get_raster_dir(
         self, layer_name: str, bands: list[str], group_idx: int = 0
@@ -289,7 +249,12 @@ class Window:
         Returns:
             the directory containing the raster.
         """
-        return get_window_raster_dir(self.path, layer_name, bands, group_idx)
+        return get_window_raster_dir(
+            self.storage.get_window_root(self.group, self.name),
+            layer_name,
+            bands,
+            group_idx,
+        )
 
     def get_metadata(self) -> dict[str, Any]:
         """Returns the window metadata dictionary."""
@@ -308,18 +273,14 @@ class Window:
 
     def save(self) -> None:
         """Save the window metadata to its root directory."""
-        self.path.mkdir(parents=True, exist_ok=True)
-        metadata_path = self.path / "metadata.json"
-        logger.debug(f"Saving window metadata to {metadata_path}")
-        with open_atomic(metadata_path, "w") as f:
-            json.dump(self.get_metadata(), f)
+        self.storage.create_or_update_window(self)
 
     @staticmethod
-    def from_metadata(path: UPath, metadata: dict[str, Any]) -> "Window":
-        """Create a Window from its path and metadata dictionary.
+    def from_metadata(storage: WindowStorage, metadata: dict[str, Any]) -> "Window":
+        """Create a Window from the WindowStorage and the window's metadata dictionary.
 
         Args:
-            path: the root directory of the window.
+            storage: the WindowStorage for the underlying dataset.
             metadata: the window metadata.
 
         Returns:
@@ -334,7 +295,7 @@ class Window:
         )
 
         return Window(
-            path=path,
+            storage=storage,
             group=metadata["group"],
             name=metadata["name"],
             projection=Projection.deserialize(metadata["projection"]),
@@ -350,21 +311,6 @@ class Window:
             options=metadata["options"],
         )
 
-    @staticmethod
-    def load(path: UPath) -> "Window":
-        """Load a Window from a UPath.
-
-        Args:
-            path: the root directory of the window
-
-        Returns:
-            the Window
-        """
-        metadata_fname = path / "metadata.json"
-        with metadata_fname.open("r") as f:
-            metadata = json.load(f)
-        return Window.from_metadata(path, metadata)
-
     @staticmethod
     def get_window_root(ds_path: UPath, group: str, name: str) -> UPath:
         """Gets the root directory of a window.

rslearn/main.py

@@ -27,13 +27,13 @@ from rslearn.dataset.handler_summaries import (
     PrepareDatasetWindowsSummary,
     UnknownIngestCounts,
 )
-from rslearn.dataset.index import DatasetIndex
 from rslearn.dataset.manage import (
     AttemptsCounter,
     materialize_dataset_windows,
     prepare_dataset_windows,
     retry,
 )
+from rslearn.dataset.storage.file import FileWindowStorage
 from rslearn.log_utils import get_logger
 from rslearn.tile_stores import get_tile_store_with_layer
 from rslearn.utils import Projection, STGeometry
@@ -315,7 +315,8 @@ def apply_on_windows(
         load_workers: optional different number of workers to use for loading the
             windows. If set, workers controls the number of workers to process the
             jobs, while load_workers controls the number of workers to use for reading
-            windows from the rslearn dataset.
+            windows from the rslearn dataset. Workers is only passed if the window
+            storage is FileWindowStorage.
         batch_size: if workers > 0, the maximum number of windows to pass to the
             function.
         jobs_per_process: optional, terminate processes after they have handled this
@@ -336,11 +337,14 @@ def apply_on_windows(
     else:
         groups = group
 
-    if load_workers is None:
-        load_workers = workers
-    windows = dataset.load_windows(
-        groups=groups, names=names, workers=load_workers, show_progress=True
-    )
+    # Load the windows. We pass workers and show_progress if it is FileWindowStorage.
+    kwargs: dict[str, Any] = {}
+    if isinstance(dataset.storage, FileWindowStorage):
+        if load_workers is None:
+            load_workers = workers
+        kwargs["workers"] = load_workers
+        kwargs["show_progress"] = True
+    windows = dataset.load_windows(groups=groups, names=names, **kwargs)
     logger.info(f"found {len(windows)} windows")
 
     if hasattr(f, "get_jobs"):
@@ -798,35 +802,6 @@ def dataset_materialize() -> None:
     apply_on_windows_args(fn, args)
 
 
-@register_handler("dataset", "build_index")
-def dataset_build_index() -> None:
-    """Handler for the rslearn dataset build_index command."""
-    parser = argparse.ArgumentParser(
-        prog="rslearn dataset build_index",
-        description=("rslearn dataset build_index: " + "create a dataset index file"),
-    )
-    parser.add_argument(
-        "--root",
-        type=str,
-        required=True,
-        help="Dataset path",
-    )
-    parser.add_argument(
-        "--workers",
-        type=int,
-        default=16,
-        help="Number of workers",
-    )
-    args = parser.parse_args(args=sys.argv[3:])
-    ds_path = UPath(args.root)
-    dataset = Dataset(ds_path)
-    index = DatasetIndex.build_index(
-        dataset=dataset,
-        workers=args.workers,
-    )
-    index.save_index(ds_path)
-
-
 @register_handler("model", "fit")
 def model_fit() -> None:
     """Handler for rslearn model fit."""

rslearn/models/anysat.py

@@ -4,11 +4,13 @@ This code loads the AnySat model from torch hub. See
 https://github.com/gastruc/AnySat for applicable license and copyright information.
 """
 
-from typing import Any
-
 import torch
 from einops import rearrange
 
+from rslearn.train.model_context import ModelContext
+
+from .component import FeatureExtractor, FeatureMaps
+
 # AnySat github: https://github.com/gastruc/AnySat
 # Modalities and expected resolutions (meters)
 MODALITY_RESOLUTIONS: dict[str, float] = {
@@ -44,7 +46,7 @@ MODALITY_BANDS: dict[str, list[str]] = {
 TIME_SERIES_MODALITIES = {"s2", "s1-asc", "s1", "alos", "l7", "l8", "modis"}
 
 
-class AnySat(torch.nn.Module):
+class AnySat(FeatureExtractor):
     """AnySat backbone (outputs one feature map)."""
 
     def __init__(
@@ -117,17 +119,17 @@ class AnySat(torch.nn.Module):
         )
         self._embed_dim = 768  # base width, 'dense' returns 2x
 
-    def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
+    def forward(self, context: ModelContext) -> FeatureMaps:
         """Forward pass for the AnySat model.
 
         Args:
-            inputs: input dicts that must include modalities as keys which are defined in the self.modalities list
+            context: the model context. Input dicts must include modalities as keys
+                which are defined in the self.modalities list
 
         Returns:
-            List[torch.Tensor]: Single-scale feature tensors from the encoder.
+            a FeatureMaps with one feature map at the configured patch size.
         """
-        if not inputs:
-            raise ValueError("empty inputs")
+        inputs = context.inputs
 
         batch: dict[str, torch.Tensor] = {}
         spatial_extent: tuple[float, float] | None = None
@@ -192,7 +194,7 @@ class AnySat(torch.nn.Module):
             kwargs["output_modality"] = self.output_modality
 
         features = self.model(batch, **kwargs)
-        return [rearrange(features, "b h w d -> b d h w")]
+        return FeatureMaps([rearrange(features, "b h w d -> b d h w")])
 
     def get_backbone_channels(self) -> list:
         """Returns the output channels of this model when used as a backbone.

rslearn/models/clay/clay.py

@@ -16,6 +16,8 @@ from huggingface_hub import hf_hub_download
 # from claymodel.module import ClayMAEModule
 from terratorch.models.backbones.clay_v15.module import ClayMAEModule
 
+from rslearn.models.component import FeatureExtractor, FeatureMaps
+from rslearn.train.model_context import ModelContext
 from rslearn.train.transforms.normalize import Normalize
 from rslearn.train.transforms.transform import Transform
 
@@ -42,7 +44,7 @@ def get_clay_checkpoint_path(
     return hf_hub_download(repo_id=repo_id, filename=filename)  # nosec B615
 
 
-class Clay(torch.nn.Module):
+class Clay(FeatureExtractor):
     """Clay backbones."""
 
     def __init__(
@@ -108,23 +110,20 @@ class Clay(torch.nn.Module):
             image, size=(new_hw, new_hw), mode="bilinear", align_corners=False
         )
 
-    def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
+    def forward(self, context: ModelContext) -> FeatureMaps:
         """Forward pass for the Clay model.
 
         Args:
-            inputs: input dicts that must include `self.modality` as a key
+            context: the model context. Input dicts must include `self.modality` as a key
 
         Returns:
-            List[torch.Tensor]: Single-scale feature tensors from the encoder.
+            a FeatureMaps consisting of one feature map, computed by Clay.
         """
-        if self.modality not in inputs[0]:
-            raise ValueError(f"Missing modality {self.modality} in inputs.")
-
         param = next(self.model.parameters())
         device = param.device
 
         chips = torch.stack(
-            [inp[self.modality] for inp in inputs], dim=0
+            [inp[self.modality] for inp in context.inputs], dim=0
         )  # (B, C, H, W)
         if self.do_resizing:
             chips = self._resize_image(chips, chips.shape[2])
@@ -163,7 +162,7 @@ class Clay(torch.nn.Module):
             )
 
         features = rearrange(spatial, "b (h w) d -> b d h w", h=side, w=side)
-        return [features]
+        return FeatureMaps([features])
 
     def get_backbone_channels(self) -> list:
         """Return output channels of this model when used as a backbone."""

rslearn/models/clip.py

@@ -1,12 +1,13 @@
 """OpenAI CLIP models."""
 
-from typing import Any
-
-import torch
 from transformers import AutoModelForZeroShotImageClassification, AutoProcessor
 
+from rslearn.train.model_context import ModelContext
+
+from .component import FeatureExtractor, FeatureMaps
+
 
-class CLIP(torch.nn.Module):
+class CLIP(FeatureExtractor):
     """CLIP image encoder."""
 
     def __init__(
@@ -31,17 +32,17 @@ class CLIP(torch.nn.Module):
         self.height = crop_size["height"] // stride[0]
         self.width = crop_size["width"] // stride[1]
 
-    def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
+    def forward(self, context: ModelContext) -> FeatureMaps:
         """Compute outputs from the backbone.
 
-        Inputs:
-            inputs: input dicts that must include "image" key containing the image to
-                process. The images should have values 0-255.
+        Args:
+            context: the model context. Input dicts must include "image" key containing
+                the image to process. The images should have values 0-255.
 
         Returns:
-            list of feature maps. The ViT produces features at one scale, so the list
-                contains a single Bx24x24x1024 feature map.
+            a FeatureMaps with one feature map from the ViT, which is always Bx24x24x1024.
         """
+        inputs = context.inputs
         device = inputs[0]["image"].device
         clip_inputs = self.processor(
             images=[inp["image"].cpu().numpy().transpose(1, 2, 0) for inp in inputs],
@@ -55,8 +56,10 @@ class CLIP(torch.nn.Module):
         batch_size = image_features.shape[0]
 
         # 576x1024 -> HxWxC
-        return [
-            image_features.reshape(
-                batch_size, self.height, self.width, self.num_features
-            ).permute(0, 3, 1, 2)
-        ]
+        return FeatureMaps(
+            [
+                image_features.reshape(
+                    batch_size, self.height, self.width, self.num_features
+                ).permute(0, 3, 1, 2)
+            ]
+        )

rslearn/models/component.py

@@ -0,0 +1,99 @@
+"""Model component API."""
+
+import abc
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+from rslearn.train.model_context import ModelContext, ModelOutput
+
+
+class FeatureExtractor(torch.nn.Module, abc.ABC):
+    """A feature extractor that performs initial processing of the inputs.
+
+    The FeatureExtractor is the first component in the encoders list for
+    SingleTaskModel and MultiTaskModel.
+    """
+
+    @abc.abstractmethod
+    def forward(self, context: ModelContext) -> Any:
+        """Extract an initial intermediate from the model context.
+
+        Args:
+            context: the model context.
+
+        Returns:
+            any intermediate to pass to downstream components. Oftentimes this is a
+                FeatureMaps.
+        """
+        raise NotImplementedError
+
+
+class IntermediateComponent(torch.nn.Module, abc.ABC):
+    """An intermediate component in the model.
+
+    In SingleTaskModel and MultiTaskModel, modules after the first module
+    in the encoders list are IntermediateComponents, as are modules before the last
+    module in the decoders list(s).
+    """
+
+    @abc.abstractmethod
+    def forward(self, intermediates: Any, context: ModelContext) -> Any:
+        """Process the given intermediate into another intermediate.
+
+        Args:
+            intermediates: the output from the previous component (either a
+                FeatureExtractor or another IntermediateComponent).
+            context: the model context.
+
+        Returns:
+            any intermediate to pass to downstream components.
+        """
+        raise NotImplementedError
+
+
+class Predictor(torch.nn.Module, abc.ABC):
+    """A predictor that computes task-specific outputs and a loss dict.
+
+    In SingleTaskModel and MultiTaskModel, the last module(s) in the decoders list(s)
+    are Predictors.
+    """
+
+    @abc.abstractmethod
+    def forward(
+        self,
+        intermediates: Any,
+        context: ModelContext,
+        targets: list[dict[str, torch.Tensor]] | None = None,
+    ) -> ModelOutput:
+        """Compute task-specific outputs and loss dict.
+
+        Args:
+            intermediates: the output from the previous component.
+            context: the model context.
+            targets: the training targets, or None during prediction.
+
+        Returns:
+            a tuple of the task-specific outputs (which should be compatible with the
+                configured Task) and loss dict. The loss dict maps from a name for each
+                loss to a scalar tensor.
+        """
+        raise NotImplementedError
+
+
+@dataclass
+class FeatureMaps:
+    """An intermediate output type for multi-resolution feature maps."""
+
+    # List of BxCxHxW feature maps at different scales, ordered from highest resolution
+    # (most fine-grained) to lowest resolution (coarsest).
+    feature_maps: list[torch.Tensor]
+
+
+@dataclass
+class FeatureVector:
+    """An intermediate output type for a flat feature vector."""
+
+    # Flat BxC feature vector.
+    feature_vector: torch.Tensor

rslearn/models/concatenate_features.py

@@ -4,8 +4,12 @@ from typing import Any
 
 import torch
 
+from rslearn.train.model_context import ModelContext
 
-class ConcatenateFeatures(torch.nn.Module):
+from .component import FeatureMaps, IntermediateComponent
+
+
+class ConcatenateFeatures(IntermediateComponent):
     """Concatenate feature map with additional raw data inputs."""
 
     def __init__(
@@ -55,26 +59,32 @@ class ConcatenateFeatures(torch.nn.Module):
 
         self.conv_layers = torch.nn.Sequential(*conv_layers)
 
-    def forward(
-        self, features: list[torch.Tensor], inputs: list[dict[str, Any]]
-    ) -> list[torch.Tensor]:
+    def forward(self, intermediates: Any, context: ModelContext) -> FeatureMaps:
         """Concatenate the feature map with the raw data inputs.
 
         Args:
-            features: list of feature maps at different resolutions.
-            inputs: original inputs.
+            intermediates: the previous output, which must be a FeatureMaps.
+            context: the model context. The input dicts must have a key matching the
+                configured key.
 
         Returns:
             concatenated feature maps.
         """
-        if not features:
-            raise ValueError("Expected at least one feature map, got none.")
+        if (
+            not isinstance(intermediates, FeatureMaps)
+            or len(intermediates.feature_maps) == 0
+        ):
+            raise ValueError(
+                "Expected input to be FeatureMaps with at least one feature map"
+            )
 
-        add_data = torch.stack([input_data[self.key] for input_data in inputs], dim=0)
+        add_data = torch.stack(
+            [input_data[self.key] for input_data in context.inputs], dim=0
+        )
         add_features = self.conv_layers(add_data)
 
         new_features: list[torch.Tensor] = []
-        for feature_map in features:
+        for feature_map in intermediates.feature_maps:
             # Shape of feature map: BCHW
             feat_h, feat_w = feature_map.shape[2], feature_map.shape[3]
 
@@ -90,4 +100,4 @@ class ConcatenateFeatures(torch.nn.Module):
 
             new_features.append(torch.cat([feature_map, resized_add_features], dim=1))
 
-        return new_features
+        return FeatureMaps(new_features)

rslearn/models/conv.py

@@ -4,8 +4,12 @@ from typing import Any
 
 import torch
 
+from rslearn.train.model_context import ModelContext
 
-class Conv(torch.nn.Module):
+from .component import FeatureMaps, IntermediateComponent
+
+
+class Conv(IntermediateComponent):
     """A single convolutional layer.
 
     It inputs a set of feature maps; the conv layer is applied to each feature map
@@ -38,19 +42,22 @@ class Conv(torch.nn.Module):
         )
         self.activation = activation
 
-    def forward(self, features: list[torch.Tensor], inputs: Any) -> list[torch.Tensor]:
-        """Compute flat output vector from multi-scale feature map.
+    def forward(self, intermediates: Any, context: ModelContext) -> FeatureMaps:
+        """Apply conv layer on each feature map.
 
         Args:
-            features: list of feature maps at different resolutions.
-            inputs: original inputs (ignored).
+            intermediates: the previous output, which must be a FeatureMaps.
+            context: the model context.
 
         Returns:
-            flat feature vector
+            the resulting feature maps after applying the same Conv2d on each one.
         """
+        if not isinstance(intermediates, FeatureMaps):
+            raise ValueError("input to Conv must be FeatureMaps")
+
         new_features = []
-        for feat_map in features:
+        for feat_map in intermediates.feature_maps:
             feat_map = self.layer(feat_map)
             feat_map = self.activation(feat_map)
             new_features.append(feat_map)
-        return new_features
+        return FeatureMaps(new_features)