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

rslearn/models/clay/configs/metadata.yaml

@@ -0,0 +1,295 @@
+sentinel-2-l2a:
+  band_order:
+    - blue
+    - green
+    - red
+    - rededge1
+    - rededge2
+    - rededge3
+    - nir
+    - nir08
+    - swir16
+    - swir22
+  rgb_indices:
+    - 2
+    - 1
+    - 0
+  gsd: 10
+  bands:
+    mean:
+      blue: 1105.
+      green: 1355.
+      red: 1552.
+      rededge1: 1887.
+      rededge2: 2422.
+      rededge3: 2630.
+      nir: 2743.
+      nir08: 2785.
+      swir16: 2388.
+      swir22: 1835.
+    std:
+      blue: 1809.
+      green: 1757.
+      red: 1888.
+      rededge1: 1870.
+      rededge2: 1732.
+      rededge3: 1697.
+      nir: 1742.
+      nir08: 1648.
+      swir16: 1470.
+      swir22: 1379.
+    wavelength:
+      blue: 0.493
+      green: 0.56
+      red: 0.665
+      rededge1: 0.704
+      rededge2: 0.74
+      rededge3: 0.783
+      nir: 0.842
+      nir08: 0.865
+      swir16: 1.61
+      swir22: 2.19
+planetscope-sr:
+  band_order:
+    - coastal_blue
+    - blue
+    - green_i
+    - green
+    - yellow
+    - red
+    - rededge
+    - nir
+  rgb_indices:
+    - 5
+    - 3
+    - 1
+  gsd: 5
+  bands:
+    mean:
+      coastal_blue: 1720.
+      blue: 1715.
+      green_i: 1913.
+      green: 2088.
+      yellow: 2274.
+      red: 2290.
+      rededge: 2613.
+      nir: 3970.
+    std:
+      coastal_blue: 747.
+      blue: 698.
+      green_i: 739.
+      green: 768.
+      yellow: 849.
+      red: 868.
+      rededge: 849.
+      nir: 914.
+    wavelength:
+      coastal_blue: 0.443
+      blue: 0.490
+      green_i: 0.531
+      green: 0.565
+      yellow: 0.610
+      red: 0.665
+      rededge: 0.705
+      nir: 0.865
+landsat-c2l1:
+  band_order:
+    - red
+    - green
+    - blue
+    - nir08
+    - swir16
+    - swir22
+  rgb_indices:
+    - 0
+    - 1
+    - 2
+  gsd: 30
+  bands:
+    mean:
+      red: 10678.
+      green: 10563.
+      blue: 11083.
+      nir08: 14792.
+      swir16: 12276.
+      swir22: 10114.
+    std:
+      red: 6025.
+      green: 5411.
+      blue: 5468.
+      nir08: 6746.
+      swir16: 5897.
+      swir22: 4850.
+    wavelength:
+      red: 0.65
+      green: 0.56
+      blue: 0.48
+      nir08: 0.86
+      swir16: 1.6
+      swir22: 2.2
+landsat-c2l2-sr:
+  band_order:
+    - red
+    - green
+    - blue
+    - nir08
+    - swir16
+    - swir22
+  rgb_indices:
+    - 0
+    - 1
+    - 2
+  gsd: 30
+  bands:
+    mean:
+      red: 13705.
+      green: 13310.
+      blue: 12474.
+      nir08: 17801.
+      swir16: 14615.
+      swir22: 12701.
+    std:
+      red: 9578.
+      green: 9408.
+      blue: 10144.
+      nir08: 8277.
+      swir16: 5300.
+      swir22: 4522.
+    wavelength:
+      red: 0.65
+      green: 0.56
+      blue: 0.48
+      nir08: 0.86
+      swir16: 1.6
+      swir22: 2.2
+naip:
+  band_order:
+    - red
+    - green
+    - blue
+    - nir
+  rgb_indices:
+    - 0
+    - 1
+    - 2
+  gsd: 1.0
+  bands:
+    mean:
+      red: 110.16
+      green: 115.41
+      blue: 98.15
+      nir: 139.04
+    std:
+      red: 47.23
+      green: 39.82
+      blue: 35.43
+      nir: 49.86
+    wavelength:
+      red: 0.65
+      green: 0.56
+      blue: 0.48
+      nir: 0.842
+linz:
+  band_order:
+    - red
+    - green
+    - blue
+  rgb_indices:
+    - 0
+    - 1
+    - 2
+  gsd: 0.5
+  bands:
+    mean:
+      red: 89.96
+      green: 99.46
+      blue: 89.51
+    std:
+      red: 41.83
+      green: 36.96
+      blue: 31.45
+    wavelength:
+      red: 0.635
+      green: 0.555
+      blue: 0.465
+sentinel-1-rtc:
+  band_order:
+    - vv
+    - vh
+  gsd: 10
+  bands:
+    mean:
+      vv: -12.113
+      vh: -18.673
+    std:
+      vv: 8.314
+      vh: 8.017
+    wavelength:
+      vv: 3.5
+      vh: 4.0
+modis:
+  band_order:
+    - sur_refl_b01
+    - sur_refl_b02
+    - sur_refl_b03
+    - sur_refl_b04
+    - sur_refl_b05
+    - sur_refl_b06
+    - sur_refl_b07
+  rgb_indices:
+    - 0
+    - 3
+    - 2
+  gsd: 500
+  bands:
+    mean:
+      sur_refl_b01: 1072.
+      sur_refl_b02: 1624.
+      sur_refl_b03: 931.
+      sur_refl_b04: 1023.
+      sur_refl_b05: 1599.
+      sur_refl_b06: 1404.
+      sur_refl_b07: 1051.
+    std:
+      sur_refl_b01: 1643.
+      sur_refl_b02: 1878.
+      sur_refl_b03: 1449.
+      sur_refl_b04: 1538.
+      sur_refl_b05: 1763.
+      sur_refl_b06: 1618.
+      sur_refl_b07: 1396.
+    wavelength:
+      sur_refl_b01: .645
+      sur_refl_b02: .858
+      sur_refl_b03: .469
+      sur_refl_b04: .555
+      sur_refl_b05: 1.240
+      sur_refl_b06: 1.640
+      sur_refl_b07: 2.130
+satellogic-MSI-L1D:
+  band_order:
+    - red
+    - green
+    - blue
+    - nir
+  rgb_indices:
+    - 0
+    - 1
+    - 2
+  gsd: 1.0
+  bands:
+    mean:
+      red: 1451.54
+      green: 1456.54
+      blue: 1543.22
+      nir: 2132.68
+    std:
+      red: 995.48
+      green: 771.29
+      blue: 708.86
+      nir: 1236.71
+    wavelength:
+      red: 0.640
+      green: 0.545
+      blue: 0.480
+      nir: 0.825

rslearn/models/clip.py

@@ -0,0 +1,68 @@
+"""OpenAI CLIP models."""
+
+from transformers import AutoModelForZeroShotImageClassification, AutoProcessor
+
+from rslearn.train.model_context import ModelContext
+
+from .component import FeatureExtractor, FeatureMaps
+
+
+class CLIP(FeatureExtractor):
+    """CLIP image encoder."""
+
+    def __init__(
+        self,
+        model_name: str,
+    ):
+        """Instantiate a new CLIP instance.
+
+        Args:
+            model_name: the model name like "openai/clip-vit-large-patch14-336".
+        """
+        super().__init__()
+
+        self.processor = AutoProcessor.from_pretrained(model_name)  # nosec
+        model = AutoModelForZeroShotImageClassification.from_pretrained(model_name)  # nosec
+        self.encoder = model.vision_model
+
+        # Get number of features and token map size from encoder attributes.
+        self.num_features = self.encoder.post_layernorm.normalized_shape[0]
+        crop_size = self.processor.image_processor.crop_size
+        stride = self.encoder.embeddings.patch_embedding.stride
+        self.height = crop_size["height"] // stride[0]
+        self.width = crop_size["width"] // stride[1]
+
+    def forward(self, context: ModelContext) -> FeatureMaps:
+        """Compute outputs from the backbone.
+
+        Args:
+            context: the model context. Input dicts must include "image" key containing
+                the image to process. The images should have values 0-255.
+
+        Returns:
+            a FeatureMaps with one feature map from the ViT, which is always Bx24x24x1024.
+        """
+        inputs = context.inputs
+        device = inputs[0]["image"].image.device
+        clip_inputs = self.processor(
+            images=[
+                inp["image"].single_ts_to_chw_tensor().cpu().numpy().transpose(1, 2, 0)
+                for inp in inputs
+            ],
+            return_tensors="pt",
+            padding=True,
+        )
+        pixel_values = clip_inputs["pixel_values"].to(device)
+        output = self.encoder(pixel_values=pixel_values)
+        # Ignore class token output which is before the patch tokens.
+        image_features = output.last_hidden_state[:, 1:, :]
+        batch_size = image_features.shape[0]
+
+        # 576x1024 -> HxWxC
+        return FeatureMaps(
+            [
+                image_features.reshape(
+                    batch_size, self.height, self.width, self.num_features
+                ).permute(0, 3, 1, 2)
+            ]
+        )

rslearn/models/component.py

@@ -0,0 +1,111 @@
+"""Model component API."""
+
+import abc
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+from rslearn.train.model_context import ModelContext, ModelOutput
+
+
+class FeatureExtractor(torch.nn.Module, abc.ABC):
+    """A feature extractor that performs initial processing of the inputs.
+
+    The FeatureExtractor is the first component in the encoders list for
+    SingleTaskModel and MultiTaskModel.
+    """
+
+    @abc.abstractmethod
+    def forward(self, context: ModelContext) -> Any:
+        """Extract an initial intermediate from the model context.
+
+        Args:
+            context: the model context.
+
+        Returns:
+            any intermediate to pass to downstream components. Oftentimes this is a
+                FeatureMaps.
+        """
+        raise NotImplementedError
+
+
+class IntermediateComponent(torch.nn.Module, abc.ABC):
+    """An intermediate component in the model.
+
+    In SingleTaskModel and MultiTaskModel, modules after the first module
+    in the encoders list are IntermediateComponents, as are modules before the last
+    module in the decoders list(s).
+    """
+
+    @abc.abstractmethod
+    def forward(self, intermediates: Any, context: ModelContext) -> Any:
+        """Process the given intermediate into another intermediate.
+
+        Args:
+            intermediates: the output from the previous component (either a
+                FeatureExtractor or another IntermediateComponent).
+            context: the model context.
+
+        Returns:
+            any intermediate to pass to downstream components.
+        """
+        raise NotImplementedError
+
+
+class Predictor(torch.nn.Module, abc.ABC):
+    """A predictor that computes task-specific outputs and a loss dict.
+
+    In SingleTaskModel and MultiTaskModel, the last module(s) in the decoders list(s)
+    are Predictors.
+    """
+
+    @abc.abstractmethod
+    def forward(
+        self,
+        intermediates: Any,
+        context: ModelContext,
+        targets: list[dict[str, torch.Tensor]] | None = None,
+    ) -> ModelOutput:
+        """Compute task-specific outputs and loss dict.
+
+        Args:
+            intermediates: the output from the previous component.
+            context: the model context.
+            targets: the training targets, or None during prediction.
+
+        Returns:
+            a tuple of the task-specific outputs (which should be compatible with the
+                configured Task) and loss dict. The loss dict maps from a name for each
+                loss to a scalar tensor.
+        """
+        raise NotImplementedError
+
+
+@dataclass
+class FeatureMaps:
+    """An intermediate output type for multi-resolution feature maps."""
+
+    # List of BxCxHxW feature maps at different scales, ordered from highest resolution
+    # (most fine-grained) to lowest resolution (coarsest).
+    feature_maps: list[torch.Tensor]
+
+
+@dataclass
+class TokenFeatureMaps:
+    """An intermediate output type for multi-resolution BCHWN feature maps with a token dimension.
+
+    Unlike `FeatureMaps`, these include an additional dimension for unpooled tokens.
+    """
+
+    # List of BxCxHxWxN feature maps at different scales, ordered from highest resolution
+    # (most fine-grained) to lowest resolution (coarsest).
+    feature_maps: list[torch.Tensor]
+
+
+@dataclass
+class FeatureVector:
+    """An intermediate output type for a flat feature vector."""
+
+    # Flat BxC feature vector.
+    feature_vector: torch.Tensor

rslearn/models/concatenate_features.py

@@ -0,0 +1,103 @@
+"""Concatenate feature map with features from input data."""
+
+from typing import Any
+
+import torch
+
+from rslearn.train.model_context import ModelContext
+
+from .component import FeatureMaps, IntermediateComponent
+
+
+class ConcatenateFeatures(IntermediateComponent):
+    """Concatenate feature map with additional raw data inputs."""
+
+    def __init__(
+        self,
+        key: str,
+        in_channels: int | None = None,
+        conv_channels: int = 64,
+        out_channels: int | None = None,
+        num_conv_layers: int = 1,
+        kernel_size: int = 3,
+        final_relu: bool = False,
+    ):
+        """Create a new ConcatenateFeatures.
+
+        Args:
+            key: the key of the input_dict to concatenate.
+            in_channels: number of input channels of the additional features.
+            conv_channels: number of channels of the convolutional layers.
+            out_channels: number of output channels of the additional features.
+            num_conv_layers: number of convolutional layers to apply to the additional features.
+            kernel_size: kernel size of the convolutional layers.
+            final_relu: whether to apply a ReLU activation to the final output, default False.
+        """
+        super().__init__()
+        self.key = key
+
+        if num_conv_layers > 0:
+            if in_channels is None or out_channels is None:
+                raise ValueError(
+                    "in_channels and out_channels must be specified if num_conv_layers > 0"
+                )
+
+        conv_layers = []
+        for i in range(num_conv_layers):
+            conv_in = in_channels if i == 0 else conv_channels
+            conv_out = out_channels if i == num_conv_layers - 1 else conv_channels
+            conv_layers.append(
+                torch.nn.Conv2d(
+                    in_channels=conv_in,
+                    out_channels=conv_out,
+                    kernel_size=kernel_size,
+                    padding="same",
+                )
+            )
+            if i < num_conv_layers - 1 or final_relu:
+                conv_layers.append(torch.nn.ReLU(inplace=True))
+
+        self.conv_layers = torch.nn.Sequential(*conv_layers)
+
+    def forward(self, intermediates: Any, context: ModelContext) -> FeatureMaps:
+        """Concatenate the feature map with the raw data inputs.
+
+        Args:
+            intermediates: the previous output, which must be a FeatureMaps.
+            context: the model context. The input dicts must have a key matching the
+                configured key.
+
+        Returns:
+            concatenated feature maps.
+        """
+        if (
+            not isinstance(intermediates, FeatureMaps)
+            or len(intermediates.feature_maps) == 0
+        ):
+            raise ValueError(
+                "Expected input to be FeatureMaps with at least one feature map"
+            )
+
+        add_data = torch.stack(
+            [input_data[self.key] for input_data in context.inputs], dim=0
+        )
+        add_features = self.conv_layers(add_data)
+
+        new_features: list[torch.Tensor] = []
+        for feature_map in intermediates.feature_maps:
+            # Shape of feature map: BCHW
+            feat_h, feat_w = feature_map.shape[2], feature_map.shape[3]
+
+            resized_add_features = add_features
+            # Resize additional features to match each feature map size if needed
+            if add_features.shape[2] != feat_h or add_features.shape[3] != feat_w:
+                resized_add_features = torch.nn.functional.interpolate(
+                    add_features,
+                    size=(feat_h, feat_w),
+                    mode="bilinear",
+                    align_corners=False,
+                )
+
+            new_features.append(torch.cat([feature_map, resized_add_features], dim=1))
+
+        return FeatureMaps(new_features)

rslearn/models/conv.py

@@ -0,0 +1,63 @@
+"""A single convolutional layer."""
+
+from typing import Any
+
+import torch
+
+from rslearn.train.model_context import ModelContext
+
+from .component import FeatureMaps, IntermediateComponent
+
+
+class Conv(IntermediateComponent):
+    """A single convolutional layer.
+
+    It inputs a set of feature maps; the conv layer is applied to each feature map
+    independently, and list of outputs is returned.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        padding: str | int = "same",
+        stride: int = 1,
+        activation: torch.nn.Module = torch.nn.ReLU(inplace=True),
+    ):
+        """Initialize a Conv.
+
+        Args:
+            in_channels: number of input channels.
+            out_channels: number of output channels.
+            kernel_size: kernel size, see torch.nn.Conv2D.
+            padding: padding to apply, see torch.nn.Conv2D.
+            stride: stride to apply, see torch.nn.Conv2D.
+            activation: activation to apply after convolution
+        """
+        super().__init__()
+
+        self.layer = torch.nn.Conv2d(
+            in_channels, out_channels, kernel_size, padding=padding, stride=stride
+        )
+        self.activation = activation
+
+    def forward(self, intermediates: Any, context: ModelContext) -> FeatureMaps:
+        """Apply conv layer on each feature map.
+
+        Args:
+            intermediates: the previous output, which must be a FeatureMaps.
+            context: the model context.
+
+        Returns:
+            the resulting feature maps after applying the same Conv2d on each one.
+        """
+        if not isinstance(intermediates, FeatureMaps):
+            raise ValueError("input to Conv must be FeatureMaps")
+
+        new_features = []
+        for feat_map in intermediates.feature_maps:
+            feat_map = self.layer(feat_map)
+            feat_map = self.activation(feat_map)
+            new_features.append(feat_map)
+        return FeatureMaps(new_features)