rslearn 0.0.6__py3-none-any.whl → 0.0.8__py3-none-any.whl

rslearn/models/anysat.py

@@ -0,0 +1,207 @@
+"""AnySat model."""
+
+from typing import Any
+
+import torch
+from einops import rearrange
+
+# 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(torch.nn.Module):
+    """AnySat backbone (outputs one feature map)."""
+
+    def __init__(
+        self,
+        modalities: list[str],
+        patch_size_meters: int,
+        dates: dict[str, list[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 not all(m in TIME_SERIES_MODALITIES for m in dates.keys()):
+            raise ValueError("`dates` keys must be time-series modalities only.")
+        for m in modalities:
+            if m in TIME_SERIES_MODALITIES and m not in dates:
+                raise ValueError(
+                    f"Missing required dates for time-series 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.dates = dates
+        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
+
+    def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
+        """Forward pass for the AnySat model.
+
+        Args:
+            inputs: input dicts that must include modalities as keys which are defined in the self.modalities list
+
+        Returns:
+            List[torch.Tensor]: Single-scale feature tensors from the encoder.
+        """
+        if not inputs:
+            raise ValueError("empty 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] for inp in inputs], dim=0
+            )  # (B, C, H, W) or (B, T*C, H, W)
+
+            if modality in TIME_SERIES_MODALITIES:
+                num_dates = len(self.dates[modality])
+                num_bands = cur.shape[1] // num_dates
+                cur = rearrange(
+                    cur, "b (t c) h w -> b t c h w", t=num_dates, c=num_bands
+                )
+                H, W = cur.shape[-2], cur.shape[-1]
+            else:
+                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)."
+                )
+
+        # Add *_dates
+        to_add = {}
+        for modality, x in list(batch.items()):
+            if modality in TIME_SERIES_MODALITIES:
+                B, T = x.shape[0], x.shape[1]
+                d = torch.as_tensor(
+                    self.dates[modality], dtype=torch.long, device=x.device
+                )
+                if d.ndim != 1 or d.numel() != T:
+                    raise ValueError(
+                        f"dates for '{modality}' must be 1D length {T}, got {tuple(d.shape)}"
+                    )
+                to_add[f"{modality}_dates"] = d.unsqueeze(0).repeat(B, 1)
+
+        batch.update(to_add)
+
+        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 [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/clay/clay.py

@@ -0,0 +1,219 @@
+"""Clay models."""
+
+from __future__ import annotations
+
+import math
+from enum import Enum
+from importlib.resources import files
+from typing import Any
+
+import torch
+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.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")
+
+
+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(torch.nn.Module):
+    """Clay backbones."""
+
+    def __init__(
+        self,
+        model_size: ClaySize,
+        modality: str = "sentinel-2-l2a",
+        checkpoint_path: str | None = None,
+        metadata_path: str = CLAY_METADATA_PATH,
+    ) -> 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.
+        """
+        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
+
+    def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
+        """Forward pass for the Clay model.
+
+        Args:
+            inputs: input dicts that must include `self.modality` as a key
+
+        Returns:
+            List[torch.Tensor]: Single-scale feature tensors from the encoder.
+        """
+        if self.modality not in inputs[0]:
+            raise ValueError(f"Missing modality {self.modality} in inputs.")
+
+        param = next(self.model.parameters())
+        device = param.device
+
+        chips = torch.stack(
+            [inp[self.modality] for inp in inputs], dim=0
+        )  # (B, C, H, W)
+
+        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 [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

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