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

rslearn/models/terramind.py

@@ -0,0 +1,256 @@
+"""Terramind models."""
+
+from enum import Enum
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from terratorch.registry import BACKBONE_REGISTRY
+
+from rslearn.train.model_context import ModelContext
+from rslearn.train.transforms.transform import Transform
+
+from .component import FeatureExtractor, FeatureMaps
+
+
+# TerraMind v1 provides two sizes: base and large
+class TerramindSize(str, Enum):
+    """Size of the Terramind model."""
+
+    BASE = "base"
+    LARGE = "large"
+
+
+# Pretraining image size for Terramind
+IMAGE_SIZE = 224
+# Default patch size for Terramind
+PATCH_SIZE = 16
+
+# Modalities supported by Terramind
+# S2L1C: Sentinel-2 Level 1C (Top-of-atmosphere reflectance), range: 1000 – 11000 DN
+# S2L2A: Sentinel-2 Level 2A (Bottom-of-atmosphere reflectance), range: 1000 – 11000 DN
+# S1GRD: Sentinel-1 GRD (Calibrated SAR backscatter), range: -50 – +10 dB
+# S1RTC: Sentinel-1 RTC (Radiometrically terrain corrected), range: -50 – +10 dB
+# RGB: Processed RGB images based on S2L2A, range: 0-255
+# DEM: Digital Elevation Model (Copernicus DEM, 30m), range: -400 – 8800 meters
+
+# More details in the TerraMesh paper: https://arxiv.org/pdf/2504.11172v1
+TERRAMIND_MODALITIES = ["S2L1C", "S2L2A", "S1GRD", "S1RTC", "RGB", "DEM"]
+
+# TerraMind band orders and standardization values
+PRETRAINED_BANDS = {
+    "S2L2A": {
+        "B01": [1390.458, 2106.761],
+        "B02": [1503.317, 2141.107],
+        "B03": [1718.197, 2038.973],
+        "B04": [1853.910, 2134.138],
+        "B05": [2199.100, 2085.321],
+        "B06": [2779.975, 1889.926],
+        "B07": [2987.011, 1820.257],
+        "B08": [3083.234, 1871.918],
+        "B8A": [3132.220, 1753.829],
+        "B09": [3162.988, 1797.379],
+        "B11": [2424.884, 1434.261],
+        "B12": [1857.648, 1334.311],
+    },
+    "S2L1C": {
+        "B01": [2357.089, 1624.683],
+        "B02": [2137.385, 1675.806],
+        "B03": [2018.788, 1557.708],
+        "B04": [2082.986, 1833.702],
+        "B05": [2295.651, 1823.738],
+        "B06": [2854.537, 1733.977],
+        "B07": [3122.849, 1732.131],
+        "B08": [3040.560, 1679.732],
+        "B8A": [3306.481, 1727.26],
+        "B09": [1473.847, 1024.687],
+        "B10": [506.070, 442.165],
+        "B11": [2472.825, 1331.411],
+        "B12": [1838.929, 1160.419],
+    },
+    "RGB": {
+        "Red": [87.271, 58.767],
+        "Green": [80.931, 47.663],
+        "Blue": [66.667, 42.631],
+    },
+    "S1GRD": {
+        "vv": [-12.599, 5.195],
+        "vh": [-20.293, 5.890],
+    },
+    "S1RTC": {
+        "vv": [-10.93, 4.391],
+        "vh": [-17.329, 4.459],
+    },
+    "DEM": {
+        "DEM": [670.665, 951.272],
+    },
+}
+
+
+class Terramind(FeatureExtractor):
+    """Terramind backbones."""
+
+    def __init__(
+        self,
+        model_size: TerramindSize,
+        modalities: list[str] = ["S2L2A"],
+        do_resizing: bool = False,
+    ) -> None:
+        """Initialize the Terramind model.
+
+        Args:
+            model_size: The size of the Terramind model.
+            modalities: The modalities to use.
+            do_resizing: Whether to resize the input images to the pretraining resolution.
+        """
+        super().__init__()
+
+        # Check if all modalities are valid
+        for modality in modalities:
+            if modality not in TERRAMIND_MODALITIES:
+                raise ValueError(f"Invalid modality: {modality}")
+
+        if model_size == TerramindSize.BASE:
+            self.model = BACKBONE_REGISTRY.build(
+                "terramind_v1_base", modalities=modalities, pretrained=True
+            )
+        elif model_size == TerramindSize.LARGE:
+            self.model = BACKBONE_REGISTRY.build(
+                "terramind_v1_large", modalities=modalities, pretrained=True
+            )
+        else:
+            raise ValueError(f"Invalid model size: {model_size}")
+
+        self.model_size = model_size
+        self.modalities = modalities
+        self.do_resizing = do_resizing
+
+    def forward(self, context: ModelContext) -> FeatureMaps:
+        """Forward pass for the Terramind model.
+
+        Args:
+            context: the model context. Input dicts must include modalities as keys
+                which are defined in the self.modalities list.
+
+        Returns:
+            a FeatureMaps with one feature map from the encoder, at 1/16 of the input
+                resolution.
+        """
+        model_inputs = {}
+        for modality in self.modalities:
+            # We assume the all the inputs include the same modalities
+            if modality not in context.inputs[0]:
+                continue
+            cur = torch.stack(
+                [inp[modality].single_ts_to_chw_tensor() for inp in context.inputs],
+                dim=0,
+            )  # (B, C, H, W)
+            if self.do_resizing and (
+                cur.shape[2] != IMAGE_SIZE or cur.shape[3] != IMAGE_SIZE
+            ):
+                if cur.shape[2] == 1 and cur.shape[3] == 1:
+                    new_height, new_width = PATCH_SIZE, PATCH_SIZE
+                else:
+                    new_height, new_width = IMAGE_SIZE, IMAGE_SIZE
+                cur = F.interpolate(
+                    cur,
+                    size=(new_height, new_width),
+                    mode="bilinear",
+                    align_corners=False,
+                )
+            model_inputs[modality] = cur
+
+        # By default, the patch embeddings are averaged over all modalities to reduce output tokens
+        # The output is a list of tensors (B, N, D) from each layer of the transformer
+        # We only get the last layer's output
+        image_features = self.model(model_inputs)[-1]
+        batch_size, num_patches, _ = image_features.shape
+        height, width = int(num_patches**0.5), int(num_patches**0.5)
+        return FeatureMaps(
+            [
+                rearrange(
+                    image_features,
+                    "b (h w) d -> b d h w",
+                    b=batch_size,
+                    h=height,
+                    w=width,
+                )
+            ]
+        )
+
+    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 (patch_size, depth) that corresponds
+        to the feature maps that the backbone returns.
+
+        Returns:
+            the output channels of the backbone as a list of (patch_size, depth) tuples.
+        """
+        if self.model_size == TerramindSize.BASE:
+            depth = 768
+        elif self.model_size == TerramindSize.LARGE:
+            depth = 1024
+        else:
+            raise ValueError(f"Invalid model size: {self.model_size}")
+        return [(PATCH_SIZE, depth)]
+
+
+class TerramindNormalize(Transform):
+    """Normalize inputs using Terramind normalization.
+
+    It will apply normalization to the modalities that are specified in the model configuration.
+    """
+
+    def __init__(self) -> None:
+        """Initialize a new TerramindNormalize."""
+        super().__init__()
+
+    def apply_image(
+        self, image: torch.Tensor, means: list[float], stds: list[float]
+    ) -> torch.Tensor:
+        """Normalize the specified image with Terramind normalization.
+
+        Args:
+            image: the image to normalize.
+            means: the means to use for the normalization.
+            stds: the standard deviations to use for the normalization.
+
+        Returns:
+            The normalized image.
+        """
+        images = image.float()  # (C, 1, H, W)
+        if images.shape[0] % len(means) != 0:
+            raise ValueError(
+                f"the number of image channels {images.shape[0]} is not multiple of expected number of bands {len(means)}"
+            )
+        for i in range(images.shape[0]):
+            band_idx = i % len(means)
+            images[i] = (images[i] - means[band_idx]) / stds[band_idx]
+        return images
+
+    def forward(
+        self, input_dict: dict[str, Any], target_dict: dict[str, Any]
+    ) -> tuple[dict[str, Any], dict[str, Any]]:
+        """Normalize the specified image with Terramind normalization.
+
+        Args:
+            input_dict: the input dictionary.
+            target_dict: the target dictionary.
+
+        Returns:
+            normalized (input_dicts, target_dicts) tuple
+        """
+        for modality in TERRAMIND_MODALITIES:
+            if modality not in input_dict:
+                continue
+            band_info = PRETRAINED_BANDS[modality]
+            means = [band_info[band][0] for band in band_info]
+            stds = [band_info[band][1] for band in band_info]
+            input_dict[modality].image = self.apply_image(
+                input_dict[modality].image,
+                means,
+                stds,
+            )
+        return input_dict, target_dict

rslearn/models/trunk.py

@@ -0,0 +1,139 @@
+"""Trunk module for decoder."""
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+import torch
+
+from rslearn.log_utils import get_logger
+from rslearn.models.task_embedding import BaseTaskEmbedding
+from rslearn.train.model_context import ModelOutput
+
+logger = get_logger(__name__)
+
+
+class DecoderTrunkLayer(torch.nn.Module, ABC):
+    """Trunk layer for decoder."""
+
+    def __init__(self) -> None:
+        """Initialize the DecoderTrunkLayer module."""
+        super().__init__()
+
+    @abstractmethod
+    def forward(
+        self, x: torch.Tensor, task_embedding: torch.Tensor | None = None
+    ) -> dict[str, torch.Tensor]:
+        """Forward pass.
+
+        Args:
+            x: input tensor of shape (batch_size, seq_len, dim)
+            task_embedding: task embedding tensor of shape (batch_size, dim), or None
+
+        Returns:
+            dict with key "outputs" (output tensor of shape (batch_size, seq_len, dim))
+            and optionally other keys.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def apply_auxiliary_losses(
+        self, trunk_out: dict[str, Any], outs: ModelOutput
+    ) -> None:
+        """Apply auxiliary losses in-place.
+
+        Args:
+            trunk_out: The output of the trunk.
+            outs: The output of the decoders, with key "loss_dict" containing the losses.
+        """
+        raise NotImplementedError
+
+
+class DecoderTrunk(torch.nn.Module):
+    """Trunk module for decoder, including arbitrary layers plus an optional task embedding."""
+
+    def __init__(
+        self,
+        task_embedding: BaseTaskEmbedding | None = None,
+        layers: list[DecoderTrunkLayer] | None = None,
+    ) -> None:
+        """Initialize the DecoderTrunk module.
+
+        Args:
+            task_embedding: Task-specific embedding module, or None if not using task embedding.
+            layers: List of other shared layers. The first one should expect a
+                B x T x C tensor, and the last should output a B x T x C tensor.
+                All layers must output a dict with key "outputs" (output tensor of shape
+                (B, T, C)) and optionally other keys.
+        """
+        super().__init__()
+        self.layers = torch.nn.ModuleList(layers or [])
+        self.task_embedding = task_embedding
+
+        # If we have multiple instances of the same layer class, output keys will get overwritten
+        if layers is not None:
+            types = [type(layer) for layer in layers]
+            if len(set(types)) != len(types):
+                logger.warning(
+                    "Multiple instances of the same layer class found in trunk. "
+                    "Only the keys from the last instance will be used"
+                )
+
+    def register_tasks(self, task_names: list[str]) -> None:
+        """Register tasks.
+
+        Args:
+            task_names: list of task names
+        """
+        if self.task_embedding is not None:
+            self.task_embedding.register_tasks(task_names)
+
+    def forward(
+        self,
+        features: list[torch.tensor],
+        inputs: list[dict[str, Any]],
+    ) -> dict[str, Any]:
+        """Forward pass.
+
+        Args:
+            features: The encoder features, a 1-list of B x C x H x W features.
+            inputs: The original inputs to the encoder.
+
+        Returns:
+            dict with key "outputs" (output tensor of shape (batch_size, seq_len, dim))
+            and optionally other keys from the other layers.
+        """
+        embeds = None
+        if self.task_embedding is not None:
+            embeds = self.task_embedding.compute_embeds(features, inputs)
+            features = self.task_embedding(features, inputs, embeds=embeds)
+
+        if not self.layers:
+            return {"outputs": features}
+
+        assert len(features) == 1, "DecoderTrunk only supports one feature map"
+        x = torch.einsum("bchw->bhwc", features[0])
+        x = torch.flatten(x, start_dim=1, end_dim=2)  # B x T x C, T = HW
+        out = {}
+        for layer in self.layers:
+            layer_out = layer(x, task_embedding=embeds)
+            x = layer_out.pop("outputs")  # unspecified shape
+            out.update(layer_out)
+        x = torch.einsum("btc->bct", x)  # B x C x T
+        x = x.view(*features[0].shape)  # B x C x H x W
+
+        out["outputs"] = [x]
+        return out
+
+    def apply_auxiliary_losses(
+        self, trunk_out: dict[str, Any], outs: ModelOutput
+    ) -> None:
+        """Apply auxiliary losses in-place.
+
+        Each layer handles its own auxiliary losses, assuming the loss key is `loss_dict`.
+
+        Args:
+            trunk_out: The output of the trunk.
+            outs: The output of the decoders.
+        """
+        for layer in self.layers:
+            layer.apply_auxiliary_losses(trunk_out, outs)

rslearn/models/unet.py

@@ -3,9 +3,17 @@
 from typing import Any
 
 import torch
+import torch.nn.functional as F
 
+from rslearn.train.model_context import ModelContext
 
-class UNetDecoder(torch.nn.Module):
+from .component import (
+    FeatureMaps,
+    IntermediateComponent,
+)
+
+
+class UNetDecoder(IntermediateComponent):
     """UNet-style decoder.
 
     It inputs multi-scale features. Starting from last (lowest resolution) feature map,
@@ -16,20 +24,30 @@ class UNetDecoder(torch.nn.Module):
     def __init__(
         self,
         in_channels: list[tuple[int, int]],
-        out_channels: int,
+        out_channels: int | None,
         conv_layers_per_resolution: int = 1,
         kernel_size: int = 3,
-    ):
+        num_channels: dict[int, int] = {},
+        target_resolution_factor: int = 1,
+        original_size_to_interpolate: tuple[int, int] | None = None,
+    ) -> None:
         """Initialize a UNetDecoder.
 
         Args:
             in_channels: list of (downsample factor, num channels) indicating the
                 resolution (1/downsample_factor of input resolution) and number of
                 channels in each feature map of the multi-scale features.
-            out_channels: channels to output at each pixel.
+            out_channels: channels to output at each pixel, or None to skip the output
+                layer.
             conv_layers_per_resolution: number of convolutional layers to apply after
                 each up-sampling operation
             kernel_size: kernel size to use in convolutional layers
+            num_channels: override number of output channels to use at different
+                downsample factors.
+            target_resolution_factor: output features at 1/target_resolution_factor
+                relative to the input resolution. The default is 1 which outputs pixel
+                level features.
+            original_size_to_interpolate: the original size to interpolate the output to.
         """
         super().__init__()
 
@@ -52,7 +70,7 @@ class UNetDecoder(torch.nn.Module):
             ]
         )
         channels_by_factor = {factor: channels for factor, channels in in_channels}
-        while cur_factor > 1:
+        while cur_factor > target_resolution_factor:
             # Add upsampling layer.
             cur_layers.append(torch.nn.Upsample(scale_factor=2))
             cur_factor //= 2
@@ -62,28 +80,39 @@ class UNetDecoder(torch.nn.Module):
             # concatenating with.
             if cur_factor in channels_by_factor:
                 layers.append(torch.nn.Sequential(*cur_layers))
+                # Number of output channels for this layer can be configured
+                # per-resolution by the user, otherwise we default to the feature map
+                # channels at the corresponding downsample factor.
+                cur_out_channels = num_channels.get(
+                    cur_factor, channels_by_factor[cur_factor]
+                )
                 cur_layers = [
                     torch.nn.Conv2d(
                         in_channels=cur_channels + channels_by_factor[cur_factor],
-                        out_channels=channels_by_factor[cur_factor],
+                        out_channels=cur_out_channels,
                         kernel_size=kernel_size,
                         padding="same",
                     ),
                     torch.nn.ReLU(inplace=True),
                 ]
-                cur_channels = channels_by_factor[cur_factor]
+                cur_channels = cur_out_channels
             else:
+                # Since there is no feature map at the next downsample factor, the
+                # default is to keep the same number of channels (but the user can
+                # still override it with num_channels).
+                cur_out_channels = num_channels.get(cur_factor, cur_channels)
                 cur_layers.extend(
                     [
                         torch.nn.Conv2d(
                             in_channels=cur_channels,
-                            out_channels=cur_channels,
+                            out_channels=cur_out_channels,
                             kernel_size=kernel_size,
                             padding="same",
                         ),
                         torch.nn.ReLU(inplace=True),
                     ]
                 )
+                cur_channels = cur_out_channels
 
             # Add remaining conv layers.
             for _ in range(conv_layers_per_resolution - 1):
@@ -99,30 +128,47 @@ class UNetDecoder(torch.nn.Module):
                     ]
                 )
 
-        cur_layers.append(
-            torch.nn.Conv2d(
-                in_channels=cur_channels,
-                out_channels=out_channels,
-                kernel_size=kernel_size,
-                padding="same",
-            ),
-        )
+        if out_channels is not None:
+            cur_layers.append(
+                torch.nn.Conv2d(
+                    in_channels=cur_channels,
+                    out_channels=out_channels,
+                    kernel_size=kernel_size,
+                    padding="same",
+                ),
+            )
         layers.append(torch.nn.Sequential(*cur_layers))
         self.layers = torch.nn.ModuleList(layers)
+        self.original_size_to_interpolate = original_size_to_interpolate
 
-    def forward(self, in_features: list[torch.Tensor], inputs: list[dict[str, Any]]):
+    def _resize(self, features: torch.Tensor) -> torch.Tensor:
+        """Interpolate the features to the original size."""
+        return F.interpolate(
+            features,
+            size=self.original_size_to_interpolate,
+            mode="bilinear",
+            align_corners=False,
+        )
+
+    def forward(self, intermediates: Any, context: ModelContext) -> FeatureMaps:
         """Compute output from multi-scale feature map.
 
         Args:
-            in_features: list of feature maps at different resolutions.
-            inputs: original inputs (ignored).
+            intermediates: the output from the previous model component, which must be a FeatureMaps.
+            context: the model context.
 
         Returns:
-            output image
+            output FeatureMaps consisting of one map. The embedding size is equal to the
+                configured out_channels.
         """
+        if not isinstance(intermediates, FeatureMaps):
+            raise ValueError("input to UNetDecoder must be a FeatureMaps")
+
         # Reverse the features since we will pass them in from lowest resolution to highest.
-        in_features = list(reversed(in_features))
+        in_features = list(reversed(intermediates.feature_maps))
         cur_features = self.layers[0](in_features[0])
         for in_feat, layer in zip(in_features[1:], self.layers[1:]):
             cur_features = layer(torch.cat([cur_features, in_feat], dim=1))
-        return cur_features
+        if self.original_size_to_interpolate is not None:
+            cur_features = self._resize(cur_features)
+        return FeatureMaps([cur_features])

rslearn/models/upsample.py

@@ -0,0 +1,48 @@
+"""An upsampling layer."""
+
+from typing import Any
+
+import torch
+
+from rslearn.train.model_context import ModelContext
+
+from .component import (
+    FeatureMaps,
+    IntermediateComponent,
+)
+
+
+class Upsample(IntermediateComponent):
+    """Upsamples each input feature map by the same factor."""
+
+    def __init__(
+        self,
+        scale_factor: int,
+        mode: str = "bilinear",
+    ):
+        """Initialize an Upsample.
+
+        Args:
+            scale_factor: the upsampling factor, e.g. 2 to double the size.
+            mode: "nearest" or "bilinear".
+        """
+        super().__init__()
+        self.layer = torch.nn.Upsample(scale_factor=scale_factor, mode=mode)
+
+    def forward(self, intermediates: Any, context: ModelContext) -> FeatureMaps:
+        """Upsample each feature map by scale_factor.
+
+        Args:
+            intermediates: the output from the previous component, which must be a FeatureMaps.
+            context: the model context.
+
+        Returns:
+            upsampled feature maps.
+        """
+        if not isinstance(intermediates, FeatureMaps):
+            raise ValueError("input to Upsample must be a FeatureMaps")
+
+        upsampled_feat_maps = [
+            self.layer(feat_map) for feat_map in intermediates.feature_maps
+        ]
+        return FeatureMaps(upsampled_feat_maps)