rslearn 0.0.1__py3-none-any.whl → 0.0.21__py3-none-any.whl

rslearn/models/resize_features.py

@@ -0,0 +1,59 @@
+"""The ResizeFeatures module."""
+
+from typing import Any
+
+import torch
+
+from rslearn.train.model_context import ModelContext
+
+from .component import (
+    FeatureMaps,
+    IntermediateComponent,
+)
+
+
+class ResizeFeatures(IntermediateComponent):
+    """Resize input features to new sizes."""
+
+    def __init__(
+        self,
+        out_sizes: list[tuple[int, int]],
+        mode: str = "bilinear",
+    ):
+        """Initialize a ResizeFeatures.
+
+        Args:
+            out_sizes: the output sizes of the feature maps. There must be one entry
+                for each input feature map.
+            mode: mode to pass to torch.nn.Upsample, e.g. "bilinear" (default) or
+                "nearest".
+        """
+        super().__init__()
+        layers = []
+        for size in out_sizes:
+            layers.append(
+                torch.nn.Upsample(
+                    size=size,
+                    mode=mode,
+                )
+            )
+        self.layers = torch.nn.ModuleList(layers)
+
+    def forward(self, intermediates: Any, context: ModelContext) -> FeatureMaps:
+        """Resize the input feature maps to new sizes.
+
+        Args:
+            intermediates: the outputs from the previous component, which must be a FeatureMaps.
+            context: the model context.
+
+        Returns:
+            resized feature maps
+        """
+        if not isinstance(intermediates, FeatureMaps):
+            raise ValueError("input to ResizeFeatures must be a FeatureMaps")
+
+        feat_maps = intermediates.feature_maps
+        resized_feat_maps = [
+            self.layers[idx](feat_map) for idx, feat_map in enumerate(feat_maps)
+        ]
+        return FeatureMaps(resized_feat_maps)

rslearn/models/sam2_enc.py

@@ -1,14 +1,15 @@
 """SegmentAnything2 encoders."""
 
-from typing import Any
-
 import torch
-import torch.nn as nn
 from sam2.build_sam import build_sam2
 from upath import UPath
 
+from rslearn.train.model_context import ModelContext
+
+from .component import FeatureExtractor, FeatureMaps
+
 
-class SAM2Encoder(nn.Module):
+class SAM2Encoder(FeatureExtractor):
     """SAM2's image encoder."""
 
     def __init__(self, model_identifier: str) -> None:
@@ -84,18 +85,21 @@ class SAM2Encoder(nn.Module):
         del self.model.obj_ptr_proj
         del self.model.image_encoder.neck
 
-    def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
+    def forward(self, context: ModelContext) -> FeatureMaps:
         """Extract multi-scale features from a batch of images.
 
         Args:
-            inputs: List of dictionaries, each containing the input image under the key 'image'.
+            context: the model context. Input dicts must have a key 'image' containing
+                the input for the SAM2 image encoder.
 
         Returns:
-            List[torch.Tensor]: Multi-scale feature tensors from the encoder.
+            feature maps from the encoder.
         """
-        images = torch.stack([inp["image"] for inp in inputs], dim=0)
+        images = torch.stack(
+            [inp["image"].single_ts_to_chw_tensor() for inp in context.inputs], dim=0
+        )
         features = self.encoder(images)
-        return features
+        return FeatureMaps(features)
 
     def get_backbone_channels(self) -> list[list[int]]:
         """Returns the output channels of the encoder at different scales.

rslearn/models/satlaspretrain.py

@@ -1,19 +1,20 @@
 """SatlasPretrain models."""
 
-from typing import Any
-
 import satlaspretrain_models
 import torch
+import torch.nn.functional as F
+
+from rslearn.train.model_context import ModelContext
 
+from .component import FeatureExtractor, FeatureMaps
 
-class SatlasPretrain(torch.nn.Module):
+
+class SatlasPretrain(FeatureExtractor):
     """SatlasPretrain backbones."""
 
     def __init__(
-        self,
-        model_identifier: str,
-        fpn: bool = False,
-    ):
+        self, model_identifier: str, fpn: bool = False, resize_to_pretrain: bool = False
+    ) -> None:
         """Instantiate a new SatlasPretrain instance.
 
         Args:
@@ -21,11 +22,13 @@ class SatlasPretrain(torch.nn.Module):
                 https://github.com/allenai/satlaspretrain_models
             fpn: whether to include the feature pyramid network, otherwise only the
                 Swin-v2-Transformer is used.
+            resize_to_pretrain: whether to resize inputs to the pretraining input
+                size (512 x 512)
         """
         super().__init__()
         weights_manager = satlaspretrain_models.Weights()
         self.model = weights_manager.get_pretrained_model(
-            model_identifier=model_identifier, fpn=fpn
+            model_identifier=model_identifier, fpn=fpn, device="cpu"
         )
 
         if "SwinB" in model_identifier:
@@ -49,21 +52,38 @@ class SatlasPretrain(torch.nn.Module):
                 [16, 1024],
                 [32, 2048],
             ]
+        self.resize_to_pretrain = resize_to_pretrain
+
+    def maybe_resize(self, data: torch.Tensor) -> list[torch.Tensor]:
+        """Resize to pretraining sizes if resize_to_pretrain == True."""
+        if self.resize_to_pretrain:
+            return F.interpolate(
+                data,
+                size=(512, 512),
+                mode="bilinear",
+                align_corners=False,
+            )
+        else:
+            return data
 
-    def forward(
-        self, inputs: list[dict[str, Any]], targets: list[dict[str, Any]] = None
-    ):
+    def forward(self, context: ModelContext) -> FeatureMaps:
         """Compute feature maps from the SatlasPretrain backbone.
 
-        Inputs:
-            inputs: input dicts that must include "image" key containing the image to
-                process.
-            targets: target dicts that are ignored
+        Args:
+            context: the model context. Input dicts must contain an "image" key
+                containing the image input to the model.
+
+        Returns:
+            multi-resolution feature maps computed by the model.
         """
-        images = torch.stack([inp["image"] for inp in inputs], dim=0)
-        return self.model(images)
+        # take the first (assumed to be only) timestep
+        images = torch.stack(
+            [inp["image"].single_ts_to_chw_tensor() for inp in context.inputs], dim=0
+        )
+        feature_maps = self.model(self.maybe_resize(images))
+        return FeatureMaps(feature_maps)
 
-    def get_backbone_channels(self):
+    def get_backbone_channels(self) -> list:
         """Returns the output channels of this model when used as a backbone.
 
         The output channels is a list of (downsample_factor, depth) that corresponds

rslearn/models/simple_time_series.py

@@ -3,12 +3,17 @@
 from typing import Any
 
 import torch
+from einops import rearrange
 
+from rslearn.train.model_context import ModelContext, RasterImage
 
-class SimpleTimeSeries(torch.nn.Module):
-    """SimpleTimeSeries wraps another encoder and applies it on an image time series.
+from .component import FeatureExtractor, FeatureMaps
 
-    It independently applies the other encoder on each image in the time series to
+
+class SimpleTimeSeries(FeatureExtractor):
+    """SimpleTimeSeries wraps another FeatureExtractor and applies it on an image time series.
+
+    It independently applies the other FeatureExtractor on each image in the time series to
     extract feature maps. It then provides a few ways to combine the features into one
     final feature map:
     - Temporal max pooling.
@@ -19,17 +24,21 @@ class SimpleTimeSeries(torch.nn.Module):
 
     def __init__(
         self,
-        encoder: torch.nn.Module,
-        image_channels: int,
+        encoder: FeatureExtractor,
+        image_channels: int | None = None,
         op: str = "max",
         groups: list[list[int]] | None = None,
         num_layers: int | None = None,
-    ):
+        image_key: str = "image",
+        backbone_channels: list[tuple[int, int]] | None = None,
+        image_keys: dict[str, int] | None = None,
+    ) -> None:
         """Create a new SimpleTimeSeries.
 
         Args:
-            encoder: the underlying encoder. It must provide get_backbone_channels
-                function that returns the output channels.
+            encoder: the underlying FeatureExtractor. It must provide get_backbone_channels
+                function that returns the output channels, or backbone_channels must be set.
+                It must output a FeatureMaps.
             image_channels: the number of channels per image of the time series. The
                 input should have multiple images concatenated on the channel axis, so
                 this parameter is used to distinguish the different images.
@@ -42,76 +51,101 @@ class SimpleTimeSeries(torch.nn.Module):
                 combined before features and the combined after features. groups is a
                 list of sets, and each set is a list of image indices.
             num_layers: the number of layers for convrnn, conv3d, and conv1d ops.
+            image_key: the key to access the images.
+            backbone_channels: manually specify the backbone channels. Can be set if
+                the encoder does not provide get_backbone_channels function.
+            image_keys: as an alternative to setting image_channels, map from the key
+                in input dict to the number of channels per timestep for that modality.
+                This way SimpleTimeSeries can be used with multimodal inputs. One of
+                image_channels or image_keys must be specified.
         """
+        if (image_channels is None and image_keys is None) or (
+            image_channels is not None and image_keys is not None
+        ):
+            raise ValueError(
+                "exactly one of image_channels and image_keys must be specified"
+            )
+
         super().__init__()
         self.encoder = encoder
         self.image_channels = image_channels
         self.op = op
         self.groups = groups
+        self.image_key = image_key
+        self.image_keys = image_keys
 
-        out_channels = self.encoder.get_backbone_channels()
+        if backbone_channels is not None:
+            out_channels = backbone_channels
+        else:
+            out_channels = self.encoder.get_backbone_channels()
         if self.groups:
             self.num_groups = len(self.groups)
         else:
             self.num_groups = 1
 
-        if self.op == "convrnn":
-            rnn_kernel_size = 3
-            self.rnn_layers = []
-            for _, count in out_channels:
-                cur_layer = [
-                    torch.nn.Sequential(
-                        torch.nn.Conv2d(
-                            2 * count, count, rnn_kernel_size, padding="same"
-                        ),
-                        torch.nn.ReLU(inplace=True),
-                    )
-                ]
-                for _ in range(num_layers - 1):
-                    cur_layer.append(
+        if self.op in ["convrnn", "conv3d", "conv1d"]:
+            if num_layers is None:
+                raise ValueError(f"num_layers must be specified for {self.op} op")
+
+            if self.op == "convrnn":
+                rnn_kernel_size = 3
+                rnn_layers = []
+                for _, count in out_channels:
+                    cur_layer = [
                         torch.nn.Sequential(
                             torch.nn.Conv2d(
-                                count, count, rnn_kernel_size, padding="same"
+                                2 * count, count, rnn_kernel_size, padding="same"
                             ),
                             torch.nn.ReLU(inplace=True),
                         )
-                    )
-                cur_layer = torch.nn.Sequential(*cur_layer)
-                self.rnn_layers.append(cur_layer)
-            self.rnn_layers = torch.nn.ModuleList(self.rnn_layers)
-
-        elif self.op == "conv3d":
-            self.conv3d_layers = []
-            for _, count in out_channels:
-                cur_layer = [
-                    torch.nn.Sequential(
-                        torch.nn.Conv3d(count, count, 3, padding=1, stride=(2, 1, 1)),
-                        torch.nn.ReLU(inplace=True),
-                    )
-                    for _ in range(num_layers)
-                ]
-                cur_layer = torch.nn.Sequential(*cur_layer)
-                self.conv3d_layers.append(cur_layer)
-            self.conv3d_layers = torch.nn.ModuleList(self.conv3d_layers)
-
-        elif self.op == "conv1d":
-            self.conv1d_layers = []
-            for _, count in out_channels:
-                cur_layer = [
-                    torch.nn.Sequential(
-                        torch.nn.Conv1d(count, count, 3, padding=1, stride=2),
-                        torch.nn.ReLU(inplace=True),
-                    )
-                    for _ in range(num_layers)
-                ]
-                cur_layer = torch.nn.Sequential(*cur_layer)
-                self.conv1d_layers.append(cur_layer)
-            self.conv1d_layers = torch.nn.ModuleList(self.conv1d_layers)
+                    ]
+                    for _ in range(num_layers - 1):
+                        cur_layer.append(
+                            torch.nn.Sequential(
+                                torch.nn.Conv2d(
+                                    count, count, rnn_kernel_size, padding="same"
+                                ),
+                                torch.nn.ReLU(inplace=True),
+                            )
+                        )
+                    cur_layer = torch.nn.Sequential(*cur_layer)
+                    rnn_layers.append(cur_layer)
+                self.rnn_layers = torch.nn.ModuleList(rnn_layers)
+
+            elif self.op == "conv3d":
+                conv3d_layers = []
+                for _, count in out_channels:
+                    cur_layer = [
+                        torch.nn.Sequential(
+                            torch.nn.Conv3d(
+                                count, count, 3, padding=1, stride=(2, 1, 1)
+                            ),
+                            torch.nn.ReLU(inplace=True),
+                        )
+                        for _ in range(num_layers)
+                    ]
+                    cur_layer = torch.nn.Sequential(*cur_layer)
+                    conv3d_layers.append(cur_layer)
+                self.conv3d_layers = torch.nn.ModuleList(conv3d_layers)
+
+            elif self.op == "conv1d":
+                conv1d_layers = []
+                for _, count in out_channels:
+                    cur_layer = [
+                        torch.nn.Sequential(
+                            torch.nn.Conv1d(count, count, 3, padding=1, stride=2),
+                            torch.nn.ReLU(inplace=True),
+                        )
+                        for _ in range(num_layers)
+                    ]
+                    cur_layer = torch.nn.Sequential(*cur_layer)
+                    conv1d_layers.append(cur_layer)
+                self.conv1d_layers = torch.nn.ModuleList(conv1d_layers)
 
         else:
             assert self.op in ["max", "mean"]
 
-    def get_backbone_channels(self):
+    def get_backbone_channels(self) -> list:
         """Returns the output channels of this model when used as a backbone.
 
         The output channels is a list of (downsample_factor, depth) that corresponds
@@ -128,27 +162,105 @@ class SimpleTimeSeries(torch.nn.Module):
             out_channels.append((downsample_factor, depth * self.num_groups))
         return out_channels
 
+    def _get_batched_images(
+        self, input_dicts: list[dict[str, Any]], image_key: str, image_channels: int
+    ) -> list[RasterImage]:
+        """Collect and reshape images across input dicts.
+
+        The BTCHW image time series are reshaped to (B*T)CHW so they can be passed to
+        the forward pass of a per-image (unitemporal) model.
+        """
+        images = torch.stack(
+            [input_dict[image_key].image for input_dict in input_dicts], dim=0
+        )  # B, C, T, H, W
+        timestamps = [input_dict[image_key].timestamps for input_dict in input_dicts]
+        # if image channels is not equal to the actual number of channels, then
+        # then every N images should be batched together. For example, if the
+        # number of input channels c == 2, and image_channels == 4, then we
+        # want to pass 2 timesteps to the model.
+        # TODO is probably to make this behaviour clearer but lets leave it like
+        # this for now to not break things.
+        num_timesteps = images.shape[1] // image_channels
+        batched_timesteps = images.shape[2] // num_timesteps
+        images = rearrange(
+            images,
+            "b c (b_t k_t) h w -> (b b_t) c k_t h w",
+            b_t=batched_timesteps,
+            k_t=num_timesteps,
+        )
+        if timestamps[0] is None:
+            new_timestamps = [None] * images.shape[0]
+        else:
+            # we also need to split the timestamps
+            new_timestamps = []
+            for t in timestamps:
+                for i in range(batched_timesteps):
+                    new_timestamps.append(
+                        t[i * num_timesteps : (i + 1) * num_timesteps]
+                    )
+        return [
+            RasterImage(image=image, timestamps=timestamps)
+            for image, timestamps in zip(images, new_timestamps)
+        ]  # C, T, H, W
+
     def forward(
-        self, inputs: list[dict[str, Any]], targets: list[dict[str, Any]] = None
-    ):
+        self,
+        context: ModelContext,
+    ) -> FeatureMaps:
         """Compute outputs from the backbone.
 
-        Inputs:
-            inputs: input dicts that must include "image" key containing the image time
+        Args:
+            context: the model context. Input dicts must include "image" key containing the image time
                 series to process (with images concatenated on the channel dimension).
-            targets: target dicts that are ignored unless
+
+        Returns:
+            the FeatureMaps aggregated temporally.
         """
         # First get features of each image.
         # To do so, we need to split up each grouped image into its component images (which have had their channels stacked).
-        images = torch.stack([inp["image"] for inp in inputs], dim=0)
-        n_batch = images.shape[0]
-        n_images = images.shape[1] // self.image_channels
-        n_height = images.shape[2]
-        n_width = images.shape[3]
-        batched_images = images.reshape(
-            n_batch * n_images, self.image_channels, n_height, n_width
+        batched_inputs: list[dict[str, Any]] | None = None
+        n_batch = len(context.inputs)
+        n_images: int | None = None
+
+        if self.image_keys is not None:
+            for image_key, image_channels in self.image_keys.items():
+                batched_images = self._get_batched_images(
+                    context.inputs, image_key, image_channels
+                )
+
+                if batched_inputs is None:
+                    batched_inputs = [{} for _ in batched_images]
+                    n_images = len(batched_images) // n_batch
+                elif n_images != len(batched_images) // n_batch:
+                    raise ValueError(
+                        "expected all modalities to have the same number of timesteps"
+                    )
+
+                for i, image in enumerate(batched_images):
+                    batched_inputs[i][image_key] = image
+
+        else:
+            assert self.image_channels is not None
+            batched_images = self._get_batched_images(
+                context.inputs, self.image_key, self.image_channels
+            )
+            batched_inputs = [{self.image_key: image} for image in batched_images]
+            n_images = len(batched_images) // n_batch
+
+        assert n_images is not None
+        # Now we can apply the underlying FeatureExtractor.
+        # Its output must be a FeatureMaps.
+        assert batched_inputs is not None
+        encoder_output = self.encoder(
+            ModelContext(
+                inputs=batched_inputs,
+                metadatas=context.metadatas,
+            )
         )
-        batched_inputs = [{"image": image} for image in batched_images]
+        if not isinstance(encoder_output, FeatureMaps):
+            raise ValueError(
+                "output of underlying FeatureExtractor in SimpleTimeSeries must be a FeatureMaps"
+            )
         all_features = [
             feat_map.reshape(
                 n_batch,
@@ -157,9 +269,8 @@ class SimpleTimeSeries(torch.nn.Module):
                 feat_map.shape[2],
                 feat_map.shape[3],
             )
-            for feat_map in self.encoder(batched_inputs)
+            for feat_map in encoder_output.feature_maps
         ]
-
         # Groups defaults to flattening all the feature maps.
         groups = self.groups
         if not groups:
@@ -171,13 +282,13 @@ class SimpleTimeSeries(torch.nn.Module):
         for feature_idx in range(len(all_features)):
             aggregated_features = []
             for group in groups:
-                group_features = []
+                group_features_list = []
                 for image_idx in group:
-                    group_features.append(
+                    group_features_list.append(
                         all_features[feature_idx][:, image_idx, :, :, :]
                     )
                 # Resulting group features are (depth, batch, C, height, width).
-                group_features = torch.stack(group_features, dim=0)
+                group_features = torch.stack(group_features_list, dim=0)
 
                 if self.op == "max":
                     group_features = torch.amax(group_features, dim=0)
@@ -213,7 +324,7 @@ class SimpleTimeSeries(torch.nn.Module):
                         .permute(0, 3, 1, 2)
                     )
                 else:
-                    raise Exception(f"unknown aggregation op {self.op}")
+                    raise ValueError(f"unknown aggregation op {self.op}")
 
                 aggregated_features.append(group_features)
 
@@ -222,4 +333,4 @@ class SimpleTimeSeries(torch.nn.Module):
 
             output_features.append(aggregated_features)
 
-        return output_features
+        return FeatureMaps(output_features)

rslearn/models/singletask.py

@@ -4,6 +4,10 @@ from typing import Any
 
 import torch
 
+from rslearn.train.model_context import ModelContext, ModelOutput
+
+from .component import FeatureExtractor, IntermediateComponent, Predictor
+
 
 class SingleTaskModel(torch.nn.Module):
     """Standard model wrapper.
@@ -14,34 +18,41 @@ class SingleTaskModel(torch.nn.Module):
     outputs and targets from the last module (which also receives the targets).
     """
 
-    def __init__(self, encoder: list[torch.nn.Module], decoder: list[torch.nn.Module]):
+    def __init__(
+        self,
+        encoder: list[FeatureExtractor | IntermediateComponent],
+        decoder: list[IntermediateComponent | Predictor],
+    ):
         """Initialize a new SingleTaskModel.
 
         Args:
-            encoder: modules to compute intermediate feature representations.
-            decoder: modules to compute outputs and loss.
+            encoder: modules to compute intermediate feature representations. The first
+                module must be a FeatureExtractor, and following modules must be
+                IntermediateComponents.
+            decoder: modules to compute outputs and loss. The last module must be a
+                Predictor, while the previous modules must be IntermediateComponents.
         """
         super().__init__()
-        self.encoder = torch.nn.Sequential(*encoder)
+        self.encoder = torch.nn.ModuleList(encoder)
         self.decoder = torch.nn.ModuleList(decoder)
 
     def forward(
         self,
-        inputs: list[dict[str, Any]],
+        context: ModelContext,
         targets: list[dict[str, Any]] | None = None,
-    ) -> tuple[list[Any], dict[str, torch.Tensor]]:
+    ) -> ModelOutput:
         """Apply the sequence of modules on the inputs.
 
         Args:
-            inputs: list of input dicts
+            context: the model context.
             targets: optional list of target dicts
 
         Returns:
-            tuple (outputs, loss_dict) from the last module.
+            the model output.
         """
-        features = self.encoder(inputs)
-        cur = features
+        cur = self.encoder[0](context)
+        for module in self.encoder[1:]:
+            cur = module(cur, context)
         for module in self.decoder[:-1]:
-            cur = module(cur, inputs)
-
-        return self.decoder[-1](cur, inputs, targets)
+            cur = module(cur, context)
+        return self.decoder[-1](cur, context, targets)