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

rslearn/train/tasks/multi_task.py

@@ -6,6 +6,7 @@ import numpy.typing as npt
 import torch
 from torchmetrics import Metric, MetricCollection
 
+from rslearn.train.model_context import SampleMetadata
 from rslearn.utils import Feature
 
 from .task import Task
@@ -30,7 +31,7 @@ class MultiTask(Task):
     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.
@@ -46,14 +47,12 @@ class MultiTask(Task):
         """
         input_dict = {}
         target_dict = {}
-        if metadata["dataset_source"] is None:
+        if metadata.dataset_source is None:
             # No multi-dataset, so always compute across all tasks
             task_iter = list(self.tasks.items())
         else:
             # Multi-dataset, so only compute for the task in this dataset
-            task_iter = [
-                (metadata["dataset_source"], self.tasks[metadata["dataset_source"]])
-            ]
+            task_iter = [(metadata.dataset_source, self.tasks[metadata.dataset_source])]
 
         for task_name, task in task_iter:
             cur_raw_inputs = {}
@@ -71,12 +70,13 @@ class MultiTask(Task):
         return input_dict, target_dict
 
     def process_output(
-        self, raw_output: Any, metadata: dict[str, Any]
+        self, raw_output: Any, metadata: SampleMetadata
     ) -> dict[str, Any]:
         """Processes an output into raster or vector data.
 
         Args:
-            raw_output: the output from prediction head.
+            raw_output: the output from prediction head. It must be a dict mapping from
+                task name to per-task output for this sample.
             metadata: metadata about the patch being read
 
         Returns:

rslearn/train/tasks/per_pixel_regression.py

@@ -8,6 +8,8 @@ import torch
 import torchmetrics
 from torchmetrics import Metric, MetricCollection
 
+from rslearn.models.component import FeatureMaps, Predictor
+from rslearn.train.model_context import ModelContext, ModelOutput, SampleMetadata
 from rslearn.utils.feature import Feature
 
 from .task import BasicTask
@@ -41,7 +43,7 @@ class PerPixelRegressionTask(BasicTask):
     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.
@@ -72,20 +74,23 @@ class PerPixelRegressionTask(BasicTask):
         }
 
     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 an HW tensor.
             metadata: metadata about the patch being read
 
         Returns:
             either raster or vector data.
         """
-        # Input could be CHW (with single channel) or just HW.
-        if len(raw_output.shape) == 2:
-            raw_output = raw_output[None, :, :]
+        if not isinstance(raw_output, torch.Tensor):
+            raise ValueError("output for PerPixelRegressionTask must be a tensor")
+        if len(raw_output.shape) != 2:
+            raise ValueError(
+                f"PerPixelRegressionTask output must be an HW tensor, but got shape {raw_output.shape}"
+            )
         return (raw_output / self.scale_factor).cpu().numpy()
 
     def visualize(
@@ -133,7 +138,7 @@ class PerPixelRegressionTask(BasicTask):
         return MetricCollection(metric_dict)
 
 
-class PerPixelRegressionHead(torch.nn.Module):
+class PerPixelRegressionHead(Predictor):
     """Head for per-pixel regression task."""
 
     def __init__(
@@ -156,24 +161,38 @@ class PerPixelRegressionHead(torch.nn.Module):
 
     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, Any]]:
+    ) -> ModelOutput:
         """Compute the regression outputs and loss from logits and targets.
 
         Args:
-            logits: BxHxW or BxCxHxW tensor.
-            inputs: original inputs (ignored).
-            targets: should contain target key that stores the regression labels.
+            intermediates: output from previous component, which must be a FeatureMaps
+                with one feature map corresponding to the logits. The channel dimension
+                size must be 1.
+            context: the model context.
+            targets: must contain values key that stores the regression labels, and
+                valid key containing mask image indicating where the labels are valid.
 
         Returns:
-            tuple of outputs and loss dict
+            tuple of outputs and loss dict. The output is a BHW tensor so that the
+                per-sample output is an HW tensor.
         """
-        assert len(logits.shape) in [3, 4]
-        if len(logits.shape) == 4:
-            assert logits.shape[1] == 1
-            logits = logits[:, 0, :, :]
+        if not isinstance(intermediates, FeatureMaps):
+            raise ValueError(
+                "the input to PerPixelRegressionHead must be a FeatureMaps"
+            )
+        if len(intermediates.feature_maps) != 1:
+            raise ValueError(
+                "the input to PerPixelRegressionHead must have one feature map"
+            )
+        if intermediates.feature_maps[0].shape[1] != 1:
+            raise ValueError(
+                f"the input to PerPixelRegressionHead must have channel dimension size 1, but got {intermediates.feature_maps[0].shape}"
+            )
+
+        logits = intermediates.feature_maps[0][:, 0, :, :]
 
         if self.use_sigmoid:
             outputs = torch.nn.functional.sigmoid(logits)
@@ -200,7 +219,10 @@ class PerPixelRegressionHead(torch.nn.Module):
             else:
                 losses["regress"] = (scores * mask).sum() / mask_total
 
-        return outputs, losses
+        return ModelOutput(
+            outputs=outputs,
+            loss_dict=losses,
+        )
 
 
 class PerPixelRegressionMetricWrapper(Metric):

rslearn/train/tasks/regression.py

@@ -10,6 +10,8 @@ import torchmetrics
 from PIL import Image, ImageDraw
 from torchmetrics import Metric, MetricCollection
 
+from rslearn.models.component import FeatureVector, Predictor
+from rslearn.train.model_context import ModelContext, ModelOutput, SampleMetadata
 from rslearn.utils.feature import Feature
 from rslearn.utils.geometry import STGeometry
 
@@ -62,7 +64,7 @@ class RegressionTask(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.
@@ -103,22 +105,26 @@ class RegressionTask(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 scalar tensor.
             metadata: metadata about the patch being read
 
         Returns:
-            either raster or vector data.
+            a list with a single Feature corresponding to the patch extent and with a
+                property containing the predicted value.
         """
+        if not isinstance(raw_output, torch.Tensor) or len(raw_output.shape) != 0:
+            raise ValueError("output for RegressionTask must be a scalar Tensor")
+
         output = raw_output.item() / self.scale_factor
         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,
             ),
             {
@@ -180,7 +186,7 @@ class RegressionTask(BasicTask):
         return MetricCollection(metric_dict)
 
 
-class RegressionHead(torch.nn.Module):
+class RegressionHead(Predictor):
     """Head for regression task."""
 
     def __init__(
@@ -199,24 +205,32 @@ class RegressionHead(torch.nn.Module):
 
     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, Any]]:
+    ) -> ModelOutput:
         """Compute the regression outputs and loss from logits and targets.
 
         Args:
-            logits: tensor that is (BatchSize, 1) or (BatchSize) in shape.
-            inputs: original inputs (ignored).
-            targets: should contain target key that stores the regression label.
+            intermediates: output from previous model component, which must be a
+                FeatureVector with channel dimension size 1 (Bx1).
+            context: the model context.
+            targets: target dicts, which each must contain a "value" key containing the
+                regression label, along with a "valid" key containing a flag indicating
+                whether each example is valid for this task.
 
         Returns:
-            tuple of outputs and loss dict
+            the model outputs. The output is a B tensor so that it is split up into a
+                scalar for each example.
         """
-        assert len(logits.shape) in [1, 2]
-        if len(logits.shape) == 2:
-            assert logits.shape[1] == 1
-            logits = logits[:, 0]
+        if not isinstance(intermediates, FeatureVector):
+            raise ValueError("the input to RegressionHead must be a FeatureVector")
+        if intermediates.feature_vector.shape[1] != 1:
+            raise ValueError(
+                f"the input to RegressionHead must have channel dimension size 1, but got shape {intermediates.feature_vector.shape}"
+            )
+
+        logits = intermediates.feature_vector[:, 0]
 
         if self.use_sigmoid:
             outputs = torch.nn.functional.sigmoid(logits)
@@ -232,9 +246,12 @@ class RegressionHead(torch.nn.Module):
             elif self.loss_mode == "l1":
                 losses["regress"] = torch.mean(torch.abs(outputs - labels) * mask)
             else:
-                assert False
+                raise ValueError(f"unknown loss mode {self.loss_mode}")
 
-        return outputs, losses
+        return ModelOutput(
+            outputs=outputs,
+            loss_dict=losses,
+        )
 
 
 class RegressionMetricWrapper(Metric):

rslearn/train/tasks/segmentation.py

@@ -8,7 +8,8 @@ import torch
 import torchmetrics.classification
 from torchmetrics import Metric, MetricCollection
 
-from rslearn.utils import Feature
+from rslearn.models.component import FeatureMaps, Predictor
+from rslearn.train.model_context import ModelContext, ModelOutput, SampleMetadata
 
 from .task import BasicTask
 
@@ -108,7 +109,7 @@ class SegmentationTask(BasicTask):
     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.
@@ -148,17 +149,20 @@ class SegmentationTask(BasicTask):
         }
 
     def process_output(
-        self, raw_output: Any, metadata: dict[str, Any]
-    ) -> npt.NDArray[Any] | list[Feature]:
+        self, raw_output: Any, metadata: SampleMetadata
+    ) -> npt.NDArray[Any]:
         """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 CHW tensor.
             metadata: metadata about the patch being read
 
         Returns:
-            either raster or vector data.
+            CHW numpy array with one channel, containing the predicted class IDs.
         """
+        if not isinstance(raw_output, torch.Tensor) or len(raw_output.shape) != 3:
+            raise ValueError("the output for SegmentationTask must be a CHW tensor")
+
         if self.prob_scales is not None:
             raw_output = (
                 raw_output
@@ -166,7 +170,7 @@ class SegmentationTask(BasicTask):
                     self.prob_scales, device=raw_output.device, dtype=raw_output.dtype
                 )[:, None, None]
             )
-        classes = raw_output.argmax(dim=0).cpu().numpy().astype(np.uint8)
+        classes = raw_output.argmax(dim=0).cpu().numpy()
         return classes[None, :, :]
 
     def visualize(
@@ -258,25 +262,36 @@ class SegmentationTask(BasicTask):
         return MetricCollection(metrics)
 
 
-class SegmentationHead(torch.nn.Module):
+class SegmentationHead(Predictor):
     """Head for segmentation 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, Any]]:
+    ) -> ModelOutput:
         """Compute the segmentation outputs from logits and targets.
 
         Args:
-            logits: tensor that is (BatchSize, NumClasses, Height, Width) in shape.
-            inputs: original inputs (ignored).
-            targets: should contain classes key that stores the per-pixel class labels.
+            intermediates: a FeatureMaps with a single feature map containing the
+                segmentation logits.
+            context: the model context
+            targets: list of target dicts, where each target dict must contain a key
+                "classes" containing the per-pixel class labels, along with "valid"
+                containing a mask indicating where the example is valid.
 
         Returns:
             tuple of outputs and loss dict
         """
+        if not isinstance(intermediates, FeatureMaps):
+            raise ValueError("input to SegmentationHead must be a FeatureMaps")
+        if len(intermediates.feature_maps) != 1:
+            raise ValueError(
+                f"input to SegmentationHead must have one feature map, but got {len(intermediates.feature_maps)}"
+            )
+
+        logits = intermediates.feature_maps[0]
         outputs = torch.nn.functional.softmax(logits, dim=1)
 
         losses = {}
@@ -295,7 +310,10 @@ class SegmentationHead(torch.nn.Module):
                 # the summed mask loss be zero.
                 losses["cls"] = torch.sum(per_pixel_loss * mask)
 
-        return outputs, losses
+        return ModelOutput(
+            outputs=outputs,
+            loss_dict=losses,
+        )
 
 
 class SegmentationMetric(Metric):

rslearn/train/tasks/task.py

@@ -7,6 +7,7 @@ import numpy.typing as npt
 import torch
 from torchmetrics import MetricCollection
 
+from rslearn.train.model_context import SampleMetadata
 from rslearn.utils import Feature
 
 
@@ -21,7 +22,7 @@ class Task:
     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.
@@ -38,7 +39,7 @@ class Task:
         raise NotImplementedError
 
     def process_output(
-        self, raw_output: Any, metadata: dict[str, Any]
+        self, raw_output: Any, metadata: SampleMetadata
     ) -> npt.NDArray[Any] | list[Feature] | dict[str, Any]:
         """Processes an output into raster or vector data.
 

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.17.dist-info → rslearn-0.0.19.dist-info}/METADATA

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