Skip to content

rslearn.models.terramind

terramind

Terramind models.

TerramindSize

Bases: StrEnum

Size of the Terramind model.

Source code in rslearn/models/terramind.py 
18
19
20
21
22
class TerramindSize(StrEnum):
    """Size of the Terramind model."""

    BASE = "base"
    LARGE = "large"

Terramind

Bases: FeatureExtractor

Terramind backbones.

Source code in rslearn/models/terramind.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
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)]

forward 

forward(context: ModelContext) -> FeatureMaps

Forward pass for the Terramind model.

Parameters:

Name Type Description Default
context ModelContext

the model context. Input dicts must include modalities as keys which are defined in the self.modalities list.

 required

Returns:

Type Description
FeatureMaps

a FeatureMaps with one feature map from the encoder, at 1/16 of the input resolution.
Source code in rslearn/models/terramind.py 
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
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,
            )
        ]
    )

get_backbone_channels 

get_backbone_channels() -> 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:

Type Description
list

the output channels of the backbone as a list of (patch_size, depth) tuples.
Source code in rslearn/models/terramind.py 
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
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)]

TerramindNormalize

Bases: Transform

Normalize inputs using Terramind normalization.

It will apply normalization to the modalities that are specified in the model configuration.

Source code in rslearn/models/terramind.py 
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
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

apply_image 

apply_image(image: Tensor, means: list[float], stds: list[float]) -> Tensor

Normalize the specified image with Terramind normalization.

Parameters:

Name Type Description Default
image Tensor

the image to normalize.

 required
means list[float]

the means to use for the normalization.

 required
stds list[float]

the standard deviations to use for the normalization.

 required

Returns:

Type Description
Tensor

The normalized image.
Source code in rslearn/models/terramind.py 
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
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

forward 

forward(input_dict: dict[str, Any], target_dict: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]

Normalize the specified image with Terramind normalization.

Parameters:

Name Type Description Default
input_dict dict[str, Any]

the input dictionary.

 required
target_dict dict[str, Any]

the target dictionary.

 required

Returns:

Type Description
tuple[dict[str, Any], dict[str, Any]]

normalized (input_dicts, target_dicts) tuple
Source code in rslearn/models/terramind.py 
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
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