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

rslearn/train/prediction_writer.py

@@ -1,6 +1,7 @@
 """rslearn PredictionWriter implementation."""
 
-from collections.abc import Sequence
+import json
+from collections.abc import Iterable, Sequence
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
@@ -12,11 +13,15 @@ from lightning.pytorch.callbacks import BasePredictionWriter
 from upath import UPath
 
 from rslearn.config import (
+    DatasetConfig,
     LayerConfig,
     LayerType,
+    StorageConfig,
 )
-from rslearn.dataset import Dataset, Window
+from rslearn.dataset import Window
+from rslearn.dataset.storage.storage import WindowStorage
 from rslearn.log_utils import get_logger
+from rslearn.train.model_context import SampleMetadata
 from rslearn.utils.array import copy_spatial_array
 from rslearn.utils.feature import Feature
 from rslearn.utils.geometry import PixelBounds
@@ -27,6 +32,7 @@ from rslearn.utils.raster_format import (
 from rslearn.utils.vector_format import VectorFormat
 
 from .lightning_module import RslearnLightningModule
+from .model_context import ModelOutput
 from .tasks.task import Task
 
 logger = get_logger(__name__)
@@ -43,12 +49,18 @@ class PendingPatchOutput:
 class PatchPredictionMerger:
     """Base class for merging predictions from multiple patches."""
 
-    def merge(self, window: Window, outputs: Sequence[PendingPatchOutput]) -> Any:
+    def merge(
+        self,
+        window: Window,
+        outputs: Sequence[PendingPatchOutput],
+        layer_config: LayerConfig,
+    ) -> Any:
         """Merge the outputs.
 
         Args:
             window: the window we are merging the outputs for.
             outputs: the outputs to process.
+            layer_config: the output layer configuration.
 
         Returns:
             the merged outputs.
@@ -60,7 +72,10 @@ class VectorMerger(PatchPredictionMerger):
     """Merger for vector data that simply concatenates the features."""
 
     def merge(
-        self, window: Window, outputs: Sequence[PendingPatchOutput]
+        self,
+        window: Window,
+        outputs: Sequence[PendingPatchOutput],
+        layer_config: LayerConfig,
     ) -> list[Feature]:
         """Concatenate the vector features."""
         return [feat for output in outputs for feat in output.output]
@@ -83,18 +98,20 @@ class RasterMerger(PatchPredictionMerger):
         self.downsample_factor = downsample_factor
 
     def merge(
-        self, window: Window, outputs: Sequence[PendingPatchOutput]
+        self,
+        window: Window,
+        outputs: Sequence[PendingPatchOutput],
+        layer_config: LayerConfig,
     ) -> npt.NDArray:
         """Merge the raster outputs."""
         num_channels = outputs[0].output.shape[0]
-        dtype = outputs[0].output.dtype
         merged_image = np.zeros(
             (
                 num_channels,
                 (window.bounds[3] - window.bounds[1]) // self.downsample_factor,
                 (window.bounds[2] - window.bounds[0]) // self.downsample_factor,
             ),
-            dtype=dtype,
+            dtype=layer_config.band_sets[0].dtype.get_numpy_dtype(),
         )
 
         # Ensure the outputs are sorted by height then width.
@@ -148,6 +165,7 @@ class RslearnWriter(BasePredictionWriter):
         merger: PatchPredictionMerger | None = None,
         output_path: str | Path | None = None,
         layer_config: LayerConfig | None = None,
+        storage_config: StorageConfig | None = None,
     ):
         """Create a new RslearnWriter.
 
@@ -163,28 +181,24 @@ class RslearnWriter(BasePredictionWriter):
             layer_config: optional layer configuration. If provided, this config will be
                 used instead of reading from the dataset config, allowing usage without
                 requiring dataset config at the output path.
+            storage_config: optional storage configuration, needed similar to layer_config
+                if there is no dataset config.
         """
         super().__init__(write_interval="batch")
         self.output_layer = output_layer
         self.selector = selector or []
-        self.path = UPath(path, **path_options or {})
-        self.output_path = (
+        ds_upath = UPath(path, **path_options or {})
+        output_upath = (
             UPath(output_path, **path_options or {})
             if output_path is not None
-            else None
+            else ds_upath
         )
 
-        # Handle dataset and layer config
-        self.layer_config: LayerConfig
-        if layer_config:
-            self.layer_config = layer_config
-        else:
-            dataset = Dataset(self.path)
-            if self.output_layer not in dataset.layers:
-                raise KeyError(
-                    f"Output layer '{self.output_layer}' not found in dataset layers."
-                )
-            self.layer_config = dataset.layers[self.output_layer]
+        self.layer_config, self.dataset_storage = (
+            self._get_layer_config_and_dataset_storage(
+                ds_upath, output_upath, layer_config, storage_config
+            )
+        )
 
         self.format: RasterFormat | VectorFormat
         if self.layer_config.type == LayerType.RASTER:
@@ -207,11 +221,73 @@ class RslearnWriter(BasePredictionWriter):
         # patches of each window need to be reconstituted.
         self.pending_outputs: dict[str, list[PendingPatchOutput]] = {}
 
+    def _get_layer_config_and_dataset_storage(
+        self,
+        ds_upath: UPath,
+        output_upath: UPath,
+        layer_config: LayerConfig | None,
+        storage_config: StorageConfig | None,
+    ) -> tuple[LayerConfig, WindowStorage]:
+        """Get the layer config and dataset storage to use.
+
+        This is a helper function for the init method.
+
+        If layer_config is set, we use that. If storage_config is set, we use it to
+        instantiate a WindowStorage using the output_upath.
+
+        If one of them is not set, we load the config from the ds_upath. Otherwise, we
+        avoid reading the dataset config; this way, RslearnWriter can be used with
+        output directories that do not contain the dataset config, as long as
+        layer_config and storage_config are both provided.
+
+        Args:
+            ds_upath: the dataset path, where a dataset config can be loaded from if
+                layer_config or storage_config is not provided.
+            output_upath: the output directory, which could be different from the
+                dataset path.
+            layer_config: optional LayerConfig to provide.
+            storage_config: optional StorageConfig to provide.
+
+        Returns:
+            a tuple (layer_config, dataset_storage)
+        """
+        dataset_storage: WindowStorage | None = None
+
+        # Instantiate the WindowStorage from the storage_config if provided.
+        if storage_config:
+            dataset_storage = (
+                storage_config.instantiate_window_storage_factory().get_storage(
+                    output_upath
+                )
+            )
+
+        if not layer_config or not dataset_storage:
+            # Need to load dataset config since one of LayerConfig/StorageConfig is missing.
+            # We use DatasetConfig.model_validate instead of initializing the Dataset
+            # because we want to get a WindowStorage that has the dataset path set to
+            # output_upath instead of ds_upath.
+            with (ds_upath / "config.json").open() as f:
+                dataset_config = DatasetConfig.model_validate(json.load(f))
+
+            if not layer_config:
+                if self.output_layer not in dataset_config.layers:
+                    raise KeyError(
+                        f"Output layer '{self.output_layer}' not found in dataset layers."
+                    )
+                layer_config = dataset_config.layers[self.output_layer]
+
+            if not dataset_storage:
+                dataset_storage = dataset_config.storage.instantiate_window_storage_factory().get_storage(
+                    output_upath
+                )
+
+        return (layer_config, dataset_storage)
+
     def write_on_batch_end(
         self,
         trainer: Trainer,
         pl_module: LightningModule,
-        prediction: dict[str, Sequence],
+        prediction: ModelOutput,
         batch_indices: Sequence[int] | None,
         batch: tuple[list, list, list],
         batch_idx: int,
@@ -232,13 +308,13 @@ class RslearnWriter(BasePredictionWriter):
         assert isinstance(pl_module, RslearnLightningModule)
         task = pl_module.task
         _, _, metadatas = batch
-        self.process_output_batch(task, prediction["outputs"], metadatas)
+        self.process_output_batch(task, prediction.outputs, metadatas)
 
     def process_output_batch(
         self,
         task: Task,
-        prediction: Sequence,
-        metadatas: Sequence,
+        prediction: Iterable[Any],
+        metadatas: Iterable[SampleMetadata],
     ) -> None:
         """Write a prediction batch with simplified API.
 
@@ -263,25 +339,19 @@ class RslearnWriter(BasePredictionWriter):
             for k in self.selector:
                 output = output[k]
 
-            # Use custom output_path if provided, otherwise use dataset path
-            window_base_path = (
-                self.output_path if self.output_path is not None else self.path
-            )
             window = Window(
-                path=Window.get_window_root(
-                    window_base_path, metadata["group"], metadata["window_name"]
-                ),
-                group=metadata["group"],
-                name=metadata["window_name"],
-                projection=metadata["projection"],
-                bounds=metadata["window_bounds"],
-                time_range=metadata["time_range"],
+                storage=self.dataset_storage,
+                group=metadata.window_group,
+                name=metadata.window_name,
+                projection=metadata.projection,
+                bounds=metadata.window_bounds,
+                time_range=metadata.time_range,
             )
             self.process_output(
                 window,
-                metadata["patch_idx"],
-                metadata["num_patches"],
-                metadata["bounds"],
+                metadata.patch_idx,
+                metadata.num_patches_in_window,
+                metadata.patch_bounds,
                 output,
             )
 
@@ -320,7 +390,7 @@ class RslearnWriter(BasePredictionWriter):
 
         # Merge outputs from overlapped patches if merger is set.
         logger.debug(f"Merging and writing for window {window.name}")
-        merged_output = self.merger.merge(window, pending_output)
+        merged_output = self.merger.merge(window, pending_output, self.layer_config)
 
         if self.layer_config.type == LayerType.RASTER:
             raster_dir = window.get_raster_dir(

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/tasks/classification.py

@@ -15,6 +15,8 @@ from torchmetrics.classification import (
     MulticlassRecall,
 )
 
+from rslearn.models.component import FeatureVector, Predictor
+from rslearn.train.model_context import ModelContext, ModelOutput, SampleMetadata
 from rslearn.utils import Feature, STGeometry
 
 from .task import BasicTask
@@ -98,7 +100,7 @@ class ClassificationTask(BasicTask):
     def process_inputs(
         self,
         raw_inputs: dict[str, torch.Tensor | list[Feature]],
-        metadata: dict[str, Any],
+        metadata: SampleMetadata,
         load_targets: bool = True,
     ) -> tuple[dict[str, Any], dict[str, Any]]:
         """Processes the data into targets.
@@ -154,17 +156,25 @@ class ClassificationTask(BasicTask):
         }
 
     def process_output(
-        self, raw_output: Any, metadata: dict[str, Any]
-    ) -> npt.NDArray[Any] | list[Feature]:
+        self, raw_output: Any, metadata: SampleMetadata
+    ) -> list[Feature]:
         """Processes an output into raster or vector data.
 
         Args:
-            raw_output: the output from prediction head.
+            raw_output: the output from prediction head, which must be a tensor
+                containing output probabilities (one dimension).
             metadata: metadata about the patch being read
 
         Returns:
-            either raster or vector data.
+            a list with one Feature corresponding to the input patch extent with a
+                property name containing the predicted class. It will have another
+                property containing the probabilities if prob_property was set.
         """
+        if not isinstance(raw_output, torch.Tensor) or len(raw_output.shape) != 1:
+            raise ValueError(
+                "expected output for ClassificationTask to be a Tensor with one dimension"
+            )
+
         probs = raw_output.cpu().numpy()
         if len(self.classes) == 2 and self.positive_class_threshold != 0.5:
             positive_class_prob = probs[self.positive_class_id]
@@ -184,8 +194,8 @@ class ClassificationTask(BasicTask):
 
         feature = Feature(
             STGeometry(
-                metadata["projection"],
-                shapely.Point(metadata["bounds"][0], metadata["bounds"][1]),
+                metadata.projection,
+                shapely.Point(metadata.patch_bounds[0], metadata.patch_bounds[1]),
                 None,
             ),
             {
@@ -265,25 +275,31 @@ class ClassificationTask(BasicTask):
         return MetricCollection(metrics)
 
 
-class ClassificationHead(torch.nn.Module):
+class ClassificationHead(Predictor):
     """Head for classification task."""
 
     def forward(
         self,
-        logits: torch.Tensor,
-        inputs: list[dict[str, Any]],
+        intermediates: Any,
+        context: ModelContext,
         targets: list[dict[str, Any]] | None = None,
-    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+    ) -> ModelOutput:
         """Compute the classification outputs and loss from logits and targets.
 
         Args:
-            logits: tensor that is (BatchSize, NumClasses) in shape.
-            inputs: original inputs (ignored).
-            targets: should contain class key that stores the class label.
+            intermediates: output from the previous model component, it should be a
+                FeatureVector with a tensor that is (BatchSize, NumClasses) in shape.
+            context: the model context.
+            targets: must contain "class" key that stores the class label, along with
+                "valid" key indicating whether the label is valid for each example.
 
         Returns:
             tuple of outputs and loss dict
         """
+        if not isinstance(intermediates, FeatureVector):
+            raise ValueError("the input to ClassificationHead must be a FeatureVector")
+
+        logits = intermediates.feature_vector
         outputs = torch.nn.functional.softmax(logits, dim=1)
 
         losses = {}
@@ -298,7 +314,10 @@ class ClassificationHead(torch.nn.Module):
             )
             losses["cls"] = torch.mean(loss)
 
-        return outputs, losses
+        return ModelOutput(
+            outputs=outputs,
+            loss_dict=losses,
+        )
 
 
 class ClassificationMetric(Metric):

rslearn/train/tasks/detection.py

@@ -12,6 +12,7 @@ import torchmetrics.classification
 import torchvision
 from torchmetrics import Metric, MetricCollection
 
+from rslearn.train.model_context import SampleMetadata
 from rslearn.utils import Feature, STGeometry
 
 from .task import BasicTask
@@ -127,7 +128,7 @@ class DetectionTask(BasicTask):
     def process_inputs(
         self,
         raw_inputs: dict[str, torch.Tensor | list[Feature]],
-        metadata: dict[str, Any],
+        metadata: SampleMetadata,
         load_targets: bool = True,
     ) -> tuple[dict[str, Any], dict[str, Any]]:
         """Processes the data into targets.
@@ -144,6 +145,8 @@ class DetectionTask(BasicTask):
         if not load_targets:
             return {}, {}
 
+        bounds = metadata.patch_bounds
+
         boxes = []
         class_labels = []
         valid = 1
@@ -186,39 +189,33 @@ class DetectionTask(BasicTask):
             else:
                 box = [int(val) for val in shp.bounds]
 
-            if box[0] >= metadata["bounds"][2] or box[2] <= metadata["bounds"][0]:
+            if box[0] >= bounds[2] or box[2] <= bounds[0]:
                 continue
-            if box[1] >= metadata["bounds"][3] or box[3] <= metadata["bounds"][1]:
+            if box[1] >= bounds[3] or box[3] <= bounds[1]:
                 continue
 
             if self.exclude_by_center:
                 center_col = (box[0] + box[2]) // 2
                 center_row = (box[1] + box[3]) // 2
-                if (
-                    center_col <= metadata["bounds"][0]
-                    or center_col >= metadata["bounds"][2]
-                ):
+                if center_col <= bounds[0] or center_col >= bounds[2]:
                     continue
-                if (
-                    center_row <= metadata["bounds"][1]
-                    or center_row >= metadata["bounds"][3]
-                ):
+                if center_row <= bounds[1] or center_row >= bounds[3]:
                     continue
 
             if self.clip_boxes:
                 box = [
-                    np.clip(box[0], metadata["bounds"][0], metadata["bounds"][2]),
-                    np.clip(box[1], metadata["bounds"][1], metadata["bounds"][3]),
-                    np.clip(box[2], metadata["bounds"][0], metadata["bounds"][2]),
-                    np.clip(box[3], metadata["bounds"][1], metadata["bounds"][3]),
+                    np.clip(box[0], bounds[0], bounds[2]),
+                    np.clip(box[1], bounds[1], bounds[3]),
+                    np.clip(box[2], bounds[0], bounds[2]),
+                    np.clip(box[3], bounds[1], bounds[3]),
                 ]
 
             # Convert to relative coordinates.
             box = [
-                box[0] - metadata["bounds"][0],
-                box[1] - metadata["bounds"][1],
-                box[2] - metadata["bounds"][0],
-                box[3] - metadata["bounds"][1],
+                box[0] - bounds[0],
+                box[1] - bounds[1],
+                box[2] - bounds[0],
+                box[3] - bounds[1],
             ]
 
             boxes.append(box)
@@ -238,16 +235,12 @@ class DetectionTask(BasicTask):
             "valid": torch.tensor(valid, dtype=torch.int32),
             "boxes": boxes,
             "labels": class_labels,
-            "width": torch.tensor(
-                metadata["bounds"][2] - metadata["bounds"][0], dtype=torch.float32
-            ),
-            "height": torch.tensor(
-                metadata["bounds"][3] - metadata["bounds"][1], dtype=torch.float32
-            ),
+            "width": torch.tensor(bounds[2] - bounds[0], dtype=torch.float32),
+            "height": torch.tensor(bounds[3] - bounds[1], dtype=torch.float32),
         }
 
     def process_output(
-        self, raw_output: Any, metadata: dict[str, Any]
+        self, raw_output: Any, metadata: SampleMetadata
     ) -> npt.NDArray[Any] | list[Feature]:
         """Processes an output into raster or vector data.
 
@@ -267,12 +260,12 @@ class DetectionTask(BasicTask):
         features = []
         for box, class_id, score in zip(boxes, class_ids, scores):
             shp = shapely.box(
-                metadata["bounds"][0] + float(box[0]),
-                metadata["bounds"][1] + float(box[1]),
-                metadata["bounds"][0] + float(box[2]),
-                metadata["bounds"][1] + float(box[3]),
+                metadata.patch_bounds[0] + float(box[0]),
+                metadata.patch_bounds[1] + float(box[1]),
+                metadata.patch_bounds[0] + float(box[2]),
+                metadata.patch_bounds[1] + float(box[3]),
             )
-            geom = STGeometry(metadata["projection"], shp, None)
+            geom = STGeometry(metadata.projection, shp, None)
             properties: dict[str, Any] = {
                 "score": float(score),
             }

rslearn/train/tasks/embedding.py

@@ -6,6 +6,8 @@ import numpy.typing as npt
 import torch
 from torchmetrics import MetricCollection
 
+from rslearn.models.component import FeatureMaps
+from rslearn.train.model_context import ModelContext, ModelOutput, SampleMetadata
 from rslearn.utils import Feature
 
 from .task import Task
@@ -21,7 +23,7 @@ class EmbeddingTask(Task):
     def process_inputs(
         self,
         raw_inputs: dict[str, torch.Tensor],
-        metadata: dict[str, Any],
+        metadata: SampleMetadata,
         load_targets: bool = True,
     ) -> tuple[dict[str, Any], dict[str, Any]]:
         """Processes the data into targets.
@@ -38,17 +40,22 @@ class EmbeddingTask(Task):
         return {}, {}
 
     def process_output(
-        self, raw_output: Any, metadata: dict[str, Any]
+        self, raw_output: Any, metadata: SampleMetadata
     ) -> npt.NDArray[Any] | list[Feature]:
         """Processes an output into raster or vector data.
 
         Args:
-            raw_output: the output from prediction head.
+            raw_output: the output from prediction head, which must be a CxHxW tensor.
             metadata: metadata about the patch being read
 
         Returns:
             either raster or vector data.
         """
+        if not isinstance(raw_output, torch.Tensor) or len(raw_output.shape) != 3:
+            raise ValueError(
+                "output for EmbeddingTask must be a tensor with three dimensions"
+            )
+
         # Just convert the raw output to numpy array that can be saved to GeoTIFF.
         return raw_output.cpu().numpy()
 
@@ -76,41 +83,38 @@ class EmbeddingTask(Task):
         return MetricCollection({})
 
 
-class EmbeddingHead(torch.nn.Module):
+class EmbeddingHead:
     """Head for embedding task.
 
-    This picks one feature map from the input list of feature maps to output. It also
-    returns a dummy loss.
+    It just adds a dummy loss to act as a Predictor.
     """
 
-    def __init__(self, feature_map_index: int | None = 0):
-        """Create a new EmbeddingHead.
-
-        Args:
-            feature_map_index: the index of the feature map to choose from the input
-                list of multi-scale feature maps (default 0). If the input is already
-                a single feature map, then set to None.
-        """
-        super().__init__()
-        self.feature_map_index = feature_map_index
-
     def forward(
         self,
-        features: torch.Tensor,
-        inputs: list[dict[str, Any]],
+        intermediates: Any,
+        context: ModelContext,
         targets: list[dict[str, Any]] | None = None,
-    ) -> tuple[torch.Tensor, dict[str, Any]]:
-        """Select the desired feature map and return it along with a dummy loss.
+    ) -> ModelOutput:
+        """Return the feature map along with a dummy loss.
 
         Args:
-            features: list of BCHW feature maps (or one feature map, if feature_map_index is None).
-            inputs: original inputs (ignored).
-            targets: should contain classes key that stores the per-pixel class labels.
+            intermediates: output from the previous model component, which must be a
+                FeatureMaps consisting of a single feature map.
+            context: the model context.
+            targets: the targets (ignored).
 
         Returns:
-            tuple of outputs and loss dict
+            model output with the feature map that was input to this component along
+                with a dummy loss.
         """
-        if self.feature_map_index is not None:
-            features = features[self.feature_map_index]
-
-        return features, {"loss": 0}
+        if not isinstance(intermediates, FeatureMaps):
+            raise ValueError("input to EmbeddingHead must be a FeatureMaps")
+        if len(intermediates.feature_maps) != 1:
+            raise ValueError(
+                f"input to EmbeddingHead must have one feature map, but got {len(intermediates.feature_maps)}"
+            )
+
+        return ModelOutput(
+            outputs=intermediates.feature_maps[0],
+            loss_dict={"loss": 0},
+        )