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

rslearn/models/anysat.py

@@ -0,0 +1,215 @@
+"""AnySat model.
+
+This code loads the AnySat model from torch hub. See
+https://github.com/gastruc/AnySat for applicable license and copyright information.
+"""
+
+from datetime import datetime
+
+import torch
+from einops import rearrange
+
+from rslearn.train.model_context import ModelContext
+
+from .component import FeatureExtractor, FeatureMaps
+
+# AnySat github: https://github.com/gastruc/AnySat
+# Modalities and expected resolutions (meters)
+MODALITY_RESOLUTIONS: dict[str, float] = {
+    "aerial": 0.2,
+    "aerial-flair": 0.2,
+    "spot": 1,
+    "naip": 1.25,
+    "s2": 10,
+    "s1-asc": 10,
+    "s1": 10,
+    "alos": 30,
+    "l7": 30,
+    "l8": 10,  # L8 must be upsampled to 10 m in AnySat
+    "modis": 250,
+}
+
+# Modalities and expected band names
+MODALITY_BANDS: dict[str, list[str]] = {
+    "aerial": ["R", "G", "B", "NiR"],
+    "aerial-flair": ["R", "G", "B", "NiR", "Elevation"],
+    "spot": ["R", "G", "B"],
+    "naip": ["R", "G", "B", "NiR"],
+    "s2": ["B2", "B3", "B4", "B5", "B6", "B7", "B8", "B8a", "B11", "B12"],
+    "s1-asc": ["VV", "VH"],
+    "s1": ["VV", "VH", "Ratio"],
+    "alos": ["HH", "HV", "Ratio"],
+    "l7": ["B1", "B2", "B3", "B4", "B5", "B7"],
+    "l8": ["B8", "B1", "B2", "B3", "B4", "B5", "B6", "B7", "B9", "B10", "B11"],
+    "modis": ["B1", "B2", "B3", "B4", "B5", "B6", "B7"],
+}
+
+# Modalities that require *_dates* input
+TIME_SERIES_MODALITIES = {"s2", "s1-asc", "s1", "alos", "l7", "l8", "modis"}
+
+
+class AnySat(FeatureExtractor):
+    """AnySat backbone (outputs one feature map)."""
+
+    def __init__(
+        self,
+        modalities: list[str],
+        patch_size_meters: int,
+        output: str = "patch",
+        output_modality: str | None = None,
+        hub_repo: str = "gastruc/anysat",
+        pretrained: bool = True,
+        force_reload: bool = False,
+        flash_attn: bool = False,
+    ) -> None:
+        """Initialize an AnySat model.
+
+        Args:
+            modalities: list of modalities to use as input (1 or more).
+            patch_size_meters: patch size in meters (must be multiple of 10). Avoid having more than 1024 patches per tile
+                ie, the height/width in meters should be <= 32 * patch_size_meters.
+            dates: dict mapping time-series modalities to list of dates (day number in a year, 0-255).
+            output: 'patch' (default) or 'dense'. Use 'patch' for classification tasks,
+                'dense' for segmentation tasks.
+            output_modality: required if output='dense', specifies which modality to use
+                for the dense output (one of the input modalities).
+            hub_repo: torch.hub repository to load AnySat from.
+            pretrained: whether to load pretrained weights.
+            force_reload: whether to force re-download of the model.
+            flash_attn: whether to use flash attention (if available).
+        """
+        super().__init__()
+
+        if not modalities:
+            raise ValueError("At least one modality must be specified.")
+        for m in modalities:
+            if m not in MODALITY_RESOLUTIONS:
+                raise ValueError(f"Invalid modality: {m}")
+
+        if patch_size_meters % 10 != 0:
+            raise ValueError(
+                "In AnySat, `patch_size` is in meters and must be a multiple of 10."
+            )
+
+        output = output.lower()
+        if output not in {"patch", "dense"}:
+            raise ValueError("`output` must be 'patch' or 'dense'.")
+        if output == "dense" and output_modality is None:
+            raise ValueError("`output_modality` is required when output='dense'.")
+
+        self.modalities = modalities
+        self.patch_size_meters = int(patch_size_meters)
+        self.output = output
+        self.output_modality = output_modality
+
+        self.model = torch.hub.load(  # nosec B614
+            hub_repo,
+            "anysat",
+            pretrained=pretrained,
+            force_reload=force_reload,
+            flash_attn=flash_attn,
+        )
+        self._embed_dim = 768  # base width, 'dense' returns 2x
+
+    @staticmethod
+    def time_ranges_to_doy(
+        time_ranges: list[tuple[datetime, datetime]],
+        device: torch.device,
+    ) -> torch.Tensor:
+        """Turn the time ranges stored in a RasterImage to timestamps accepted by AnySat.
+
+        AnySat uses the doy with each timestamp, so we take the midpoint
+        the time range. For some inputs (e.g. Sentinel 2) we take an image from a specific
+        time so that start_time == end_time == mid_time.
+        """
+        doys = [(t[0] + ((t[1] - t[0]) / 2)).timetuple().tm_yday for t in time_ranges]
+        return torch.tensor(doys, dtype=torch.int32, device=device)
+
+    def forward(self, context: ModelContext) -> FeatureMaps:
+        """Forward pass for the AnySat 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 at the configured patch size.
+        """
+        inputs = context.inputs
+
+        batch: dict[str, torch.Tensor] = {}
+        spatial_extent: tuple[float, float] | None = None
+
+        for modality in self.modalities:
+            if modality not in inputs[0]:
+                raise ValueError(f"Modality '{modality}' not present in inputs.")
+
+            cur = torch.stack(
+                [inp[modality].image for inp in inputs], dim=0
+            )  # (B, C, T, H, W)
+
+            if modality in TIME_SERIES_MODALITIES:
+                num_bands = cur.shape[1]
+                cur = rearrange(cur, "b c t h w -> b t c h w")
+                H, W = cur.shape[-2], cur.shape[-1]
+
+                if inputs[0][modality].timestamps is None:
+                    raise ValueError(
+                        f"Require timestamps for time series modality {modality}"
+                    )
+                timestamps = torch.stack(
+                    [
+                        self.time_ranges_to_doy(inp[modality].timestamps, cur.device)  # type: ignore
+                        for inp in inputs
+                    ],
+                    dim=0,
+                )
+                batch[f"{modality}_dates"] = timestamps
+            else:
+                # take the first (assumed only) timestep
+                cur = cur[:, :, 0]
+                num_bands = cur.shape[1]
+                H, W = cur.shape[-2], cur.shape[-1]
+
+            if num_bands != len(MODALITY_BANDS[modality]):
+                raise ValueError(
+                    f"Modality '{modality}' expected {len(MODALITY_BANDS[modality])} bands, "
+                    f"got {num_bands} (shape {tuple(cur.shape)})"
+                )
+
+            batch[modality] = cur
+
+            # Ensure same spatial extent across all modalities (H*res, W*res)
+            extent = (
+                H * MODALITY_RESOLUTIONS[modality],
+                W * MODALITY_RESOLUTIONS[modality],
+            )
+            if spatial_extent is None:
+                spatial_extent = extent
+            elif spatial_extent != extent:
+                raise ValueError(
+                    "All modalities must share the same spatial extent (H*res, W*res)."
+                )
+
+        kwargs = {"patch_size": self.patch_size_meters, "output": self.output}
+        if self.output == "dense":
+            kwargs["output_modality"] = self.output_modality
+
+        features = self.model(batch, **kwargs)
+        return FeatureMaps([rearrange(features, "b h w d -> b d h w")])
+
+    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.output == "patch":
+            return [(self.patch_size_meters // 10, 768)]
+        elif self.output == "dense":
+            return [(1, 1536)]
+        else:
+            raise ValueError(f"invalid output type: {self.output}")

rslearn/models/attention_pooling.py

@@ -0,0 +1,177 @@
+"""An attention pooling layer."""
+
+import math
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from torch import nn
+
+from rslearn.models.component import (
+    FeatureMaps,
+    IntermediateComponent,
+    TokenFeatureMaps,
+)
+from rslearn.train.model_context import ModelContext
+
+
+class SimpleAttentionPool(IntermediateComponent):
+    """Simple Attention Pooling.
+
+    Given a token feature map of shape BCHWN,
+    learn an attention layer which aggregates over
+    the N dimension.
+
+    This is done simply by learning a mapping D->1 which is the weight
+    which should be assigned to each token during averaging:
+
+    output = sum [feat_token * W(feat_token) for feat_token in feat_tokens]
+    """
+
+    def __init__(self, in_dim: int, hidden_linear: bool = False) -> None:
+        """Initialize the simple attention pooling layer.
+
+        Args:
+            in_dim: the encoding dimension D
+            hidden_linear: whether to apply an additional linear transformation D -> D
+                to the feat tokens. If this is True, a ReLU activation is applied
+                after the first linear transformation.
+        """
+        super().__init__()
+        if hidden_linear:
+            self.hidden_linear = nn.Linear(in_features=in_dim, out_features=in_dim)
+        else:
+            self.hidden_linear = None
+        self.linear = nn.Linear(in_features=in_dim, out_features=1)
+
+    def forward_for_map(self, feat_tokens: torch.Tensor) -> torch.Tensor:
+        """Attention pooling for a single feature map (BCHWN tensor)."""
+        B, D, H, W, N = feat_tokens.shape
+        feat_tokens = rearrange(feat_tokens, "b d h w n -> (b h w) n d")
+        if self.hidden_linear is not None:
+            feat_tokens = torch.nn.functional.relu(self.hidden_linear(feat_tokens))
+        attention_scores = torch.nn.functional.softmax(self.linear(feat_tokens), dim=1)
+        feat_tokens = (attention_scores * feat_tokens).sum(dim=1)
+        return rearrange(feat_tokens, "(b h w) d -> b d h w", b=B, h=H, w=W)
+
+    def forward(self, intermediates: Any, context: ModelContext) -> FeatureMaps:
+        """Forward pass for attention pooling linear probe.
+
+        Args:
+            intermediates: the output from the previous component, which must be a TokenFeatureMaps.
+                We pool over the final dimension in the TokenFeatureMaps. If multiple maps
+                are passed, we apply the same linear layers to all of them.
+            context: the model context.
+            feat_tokens (torch.Tensor): Input feature tokens of shape (B, C, H, W, N).
+
+        Returns:
+            torch.Tensor:
+                - output, attentioned pool over the last dimension (B, C, H, W)
+        """
+        if not isinstance(intermediates, TokenFeatureMaps):
+            raise ValueError("input to Attention Pool must be a TokenFeatureMaps")
+
+        features = []
+        for feat in intermediates.feature_maps:
+            features.append(self.forward_for_map(feat))
+        return FeatureMaps(features)
+
+
+class AttentionPool(IntermediateComponent):
+    """Attention Pooling.
+
+    Given a feature map of shape BCHWN,
+    learn an attention layer which aggregates over
+    the N dimension.
+
+    We do this by learning a query token, and applying a standard
+    attention mechanism against this learned query token.
+    """
+
+    def __init__(self, in_dim: int, num_heads: int, linear_on_kv: bool = True) -> None:
+        """Initialize the attention pooling layer.
+
+        Args:
+            in_dim: the encoding dimension D
+            num_heads: the number of heads to use
+            linear_on_kv: Whether to apply a linear layer on the input tokens
+            to create the key and value tokens.
+        """
+        super().__init__()
+        self.query_token: nn.Parameter = nn.Parameter(torch.empty(in_dim))
+        if linear_on_kv:
+            self.k_linear = nn.Linear(in_dim, in_dim)
+            self.v_linear = nn.Linear(in_dim, in_dim)
+        else:
+            self.k_linear = None
+            self.v_linear = None
+        if in_dim % num_heads != 0:
+            raise ValueError(
+                f"in_dim must be divisible by num_heads. Got {in_dim} and {num_heads}."
+            )
+        self.num_heads = num_heads
+        self.init_weights()
+
+    def init_weights(self) -> None:
+        """Initialize weights for the probe."""
+        nn.init.trunc_normal_(self.query_token, std=0.02)
+
+    def forward_for_map(self, feat_tokens: torch.Tensor) -> torch.Tensor:
+        """Attention pooling for a single feature map (BCHWN tensor)."""
+        B, D, H, W, N = feat_tokens.shape
+        feat_tokens = rearrange(feat_tokens, "b d h w n -> (b h w) n d")
+        collapsed_dim = B * H * W
+        q = self.query_token.expand(collapsed_dim, 1, -1)
+        q = q.reshape(
+            collapsed_dim, 1, self.num_heads, D // self.num_heads
+        )  # [B, 1, head, D_head]
+        q = rearrange(q, "b h n d -> b n h d")
+        if self.k_linear is not None:
+            assert self.v_linear is not None
+            k = self.k_linear(feat_tokens).reshape(
+                collapsed_dim, N, self.num_heads, D // self.num_heads
+            )
+            v = self.v_linear(feat_tokens).reshape(
+                collapsed_dim, N, self.num_heads, D // self.num_heads
+            )
+        else:
+            k = feat_tokens.reshape(
+                collapsed_dim, N, self.num_heads, D // self.num_heads
+            )
+            v = feat_tokens.reshape(
+                collapsed_dim, N, self.num_heads, D // self.num_heads
+            )
+        k = rearrange(k, "b n h d -> b h n d")
+        v = rearrange(v, "b n h d -> b h n d")
+
+        # Compute attention scores
+        attn_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(
+            D // self.num_heads
+        )
+        attn_weights = F.softmax(attn_scores, dim=-1)
+        x = torch.matmul(attn_weights, v)  # [B, head, 1, D_head]
+        return x.reshape(B, D, H, W)
+
+    def forward(self, intermediates: Any, context: ModelContext) -> FeatureMaps:
+        """Forward pass for attention pooling linear probe.
+
+        Args:
+            intermediates: the output from the previous component, which must be a TokenFeatureMaps.
+                We pool over the final dimension in the TokenFeatureMaps. If multiple feature
+                maps are passed, we apply the same attention weights (query token and linear k, v layers)
+                to all the maps.
+            context: the model context.
+            feat_tokens (torch.Tensor): Input feature tokens of shape (B, C, H, W, N).
+
+        Returns:
+            torch.Tensor:
+                - output, attentioned pool over the last dimension (B, C, H, W)
+        """
+        if not isinstance(intermediates, TokenFeatureMaps):
+            raise ValueError("input to Attention Pool must be a TokenFeatureMaps")
+
+        features = []
+        for feat in intermediates.feature_maps:
+            features.append(self.forward_for_map(feat))
+        return FeatureMaps(features)

rslearn/models/clay/clay.py

@@ -0,0 +1,231 @@
+"""Clay models."""
+
+from __future__ import annotations
+
+import math
+from enum import Enum
+from importlib.resources import files
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+import yaml
+from einops import rearrange
+from huggingface_hub import hf_hub_download
+
+# from claymodel.module import ClayMAEModule
+from terratorch.models.backbones.clay_v15.module import ClayMAEModule
+
+from rslearn.models.component import FeatureExtractor, FeatureMaps
+from rslearn.train.model_context import ModelContext
+from rslearn.train.transforms.normalize import Normalize
+from rslearn.train.transforms.transform import Transform
+
+
+class ClaySize(str, Enum):
+    """Size of the Clay model."""
+
+    BASE = "base"
+    LARGE = "large"
+
+
+PATCH_SIZE = 8
+CLAY_MODALITIES = ["sentinel-2-l2a", "sentinel-1-rtc", "landsat-c2l1", "naip"]
+CONFIG_DIR = files("rslearn.models.clay.configs")
+CLAY_METADATA_PATH = str(CONFIG_DIR / "metadata.yaml")
+DEFAULT_IMAGE_RESOLUTION = 128  # image resolution during pretraining
+
+
+def get_clay_checkpoint_path(
+    filename: str = "v1.5/clay-v1.5.ckpt",
+    repo_id: str = "made-with-clay/Clay",
+) -> str:
+    """Return a cached local path to the Clay ckpt from the Hugging Face Hub."""
+    return hf_hub_download(repo_id=repo_id, filename=filename)  # nosec B615
+
+
+class Clay(FeatureExtractor):
+    """Clay backbones."""
+
+    def __init__(
+        self,
+        model_size: ClaySize,
+        modality: str = "sentinel-2-l2a",
+        checkpoint_path: str | None = None,
+        metadata_path: str = CLAY_METADATA_PATH,
+        do_resizing: bool = False,
+    ) -> None:
+        """Initialize the Clay model.
+
+        Args:
+            model_size: The size of the Clay model.
+            modality: The modality to use (subset of CLAY_MODALITIES).
+            checkpoint_path: Path to clay-v1.5.ckpt, if None, fetch from HF Hub.
+            metadata_path: Path to metadata.yaml.
+            do_resizing: Whether to resize the image to the input resolution.
+        """
+        super().__init__()
+
+        # Clay only supports single modality input
+        if modality not in CLAY_MODALITIES:
+            raise ValueError(f"Invalid modality: {modality}")
+
+        ckpt = checkpoint_path or get_clay_checkpoint_path()
+        if model_size == ClaySize.LARGE:
+            self.model = ClayMAEModule.load_from_checkpoint(
+                checkpoint_path=ckpt,
+                model_size="large",
+                metadata_path=metadata_path,
+                dolls=[16, 32, 64, 128, 256, 768, 1024],
+                doll_weights=[1, 1, 1, 1, 1, 1, 1],
+                mask_ratio=0.0,
+                shuffle=False,
+            )
+        elif model_size == ClaySize.BASE:
+            # Failed to load Base model in Clay v1.5
+            raise ValueError("Clay BASE model currently not supported in v1.5.")
+            self.model = ClayMAEModule.load_from_checkpoint(
+                checkpoint_path=ckpt,
+                model_size="base",
+                metadata_path=metadata_path,
+                dolls=[16, 32, 64, 128, 256, 768],
+                doll_weights=[1, 1, 1, 1, 1, 1],
+                mask_ratio=0.0,
+                shuffle=False,
+            )
+        else:
+            raise ValueError(f"Invalid model size: {model_size}")
+
+        with open(metadata_path) as f:
+            self.metadata = yaml.safe_load(f)
+
+        self.model_size = model_size
+        self.modality = modality
+        self.do_resizing = do_resizing
+
+    def _resize_image(self, image: torch.Tensor, original_hw: int) -> torch.Tensor:
+        """Resize the image to the input resolution."""
+        new_hw = self.patch_size if original_hw == 1 else DEFAULT_IMAGE_RESOLUTION
+        return F.interpolate(
+            image, size=(new_hw, new_hw), mode="bilinear", align_corners=False
+        )
+
+    def forward(self, context: ModelContext) -> FeatureMaps:
+        """Forward pass for the Clay model.
+
+        Args:
+            context: the model context. Input dicts must include `self.modality` as a key
+
+        Returns:
+            a FeatureMaps consisting of one feature map, computed by Clay.
+        """
+        param = next(self.model.parameters())
+        device = param.device
+
+        chips = torch.stack(
+            [inp[self.modality] for inp in context.inputs], dim=0
+        )  # (B, C, H, W)
+        if self.do_resizing:
+            chips = self._resize_image(chips, chips.shape[2])
+        order = self.metadata[self.modality]["band_order"]
+        wavelengths = []
+        for band in self.metadata[self.modality]["band_order"]:
+            wavelengths.append(
+                self.metadata[self.modality]["bands"]["wavelength"][band] * 1000
+            )  # Convert to nm
+        # Check channel count matches Clay expectation
+        if chips.shape[1] != len(order):
+            raise ValueError(
+                f"Channel count {chips.shape[1]} does not match expected {len(order)} for {self.modality}"
+            )
+
+        # Time & latlon zeros are valid per Clay doc
+        # https://clay-foundation.github.io/model/getting-started/basic_use.html
+        datacube = {
+            "platform": self.modality,
+            "time": torch.zeros(chips.shape[0], 4).to(device),
+            "latlon": torch.zeros(chips.shape[0], 4).to(device),
+            "pixels": chips.to(device),
+            "gsd": torch.tensor(self.metadata[self.modality]["gsd"]).to(device),
+            "waves": torch.tensor(wavelengths).to(device),
+        }
+
+        tokens, *_ = self.model.model.encoder(datacube)  # (B, 1 + N, D)
+
+        # Remove CLS token
+        spatial = tokens[:, 1:, :]  # (B, N, D)
+        n_tokens = spatial.shape[1]
+        side = int(math.isqrt(n_tokens))
+        if chips.shape[2] != side * PATCH_SIZE or chips.shape[3] != side * PATCH_SIZE:
+            raise ValueError(
+                f"Input spatial size {(chips.shape[2], chips.shape[3])} is not compatible with patch size {PATCH_SIZE}"
+            )
+
+        features = rearrange(spatial, "b (h w) d -> b d h w", h=side, w=side)
+        return FeatureMaps([features])
+
+    def get_backbone_channels(self) -> list:
+        """Return output channels of this model when used as a backbone."""
+        if self.model_size == ClaySize.LARGE:
+            depth = 1024
+        elif self.model_size == ClaySize.BASE:
+            depth = 768
+        else:
+            raise ValueError(f"Invalid model size: {self.model_size}")
+        return [(PATCH_SIZE, depth)]
+
+
+class ClayNormalize(Transform):
+    """Normalize inputs using Clay metadata.
+
+    For Sentinel-1, the intensities should be converted to decibels.
+    """
+
+    def __init__(self, metadata_path: str = CLAY_METADATA_PATH) -> None:
+        """Initialize ClayNormalize."""
+        super().__init__()
+        with open(metadata_path) as f:
+            metadata = yaml.safe_load(f)
+        normalizers = {}
+        for modality in CLAY_MODALITIES:
+            if modality not in metadata:
+                continue
+            modality_metadata = metadata[modality]
+            means = [
+                modality_metadata["bands"]["mean"][b]
+                for b in modality_metadata["band_order"]
+            ]
+            stds = [
+                modality_metadata["bands"]["std"][b]
+                for b in modality_metadata["band_order"]
+            ]
+            normalizers[modality] = Normalize(
+                mean=means,
+                std=stds,
+                selectors=[modality],
+                num_bands=len(means),
+            )
+        self.normalizers = torch.nn.ModuleDict(normalizers)
+
+    def apply_image(
+        self, image: torch.Tensor, means: list[float], stds: list[float]
+    ) -> torch.Tensor:
+        """Normalize the specified image with Clay normalization."""
+        x = image.float()
+        if x.shape[0] != len(means):
+            raise ValueError(
+                f"channel count {x.shape[0]} does not match provided band stats {len(means)}"
+            )
+        for c in range(x.shape[0]):
+            x[c] = (x[c] - means[c]) / stds[c]
+        return x
+
+    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 Clay normalization."""
+        for modality, normalizer in self.normalizers.items():
+            if modality not in input_dict:
+                continue
+            input_dict, target_dict = normalizer(input_dict, target_dict)
+        return input_dict, target_dict