rslearn 0.0.5__py3-none-any.whl → 0.0.7__py3-none-any.whl

rslearn/arg_parser.py

@@ -1,33 +1,12 @@
 """Custom Lightning ArgumentParser with environment variable substitution support."""
 
 import os
-import re
 from typing import Any
 
 from jsonargparse import Namespace
 from lightning.pytorch.cli import LightningArgumentParser
 
-
-def substitute_env_vars_in_string(content: str) -> str:
-    """Substitute environment variables in a string.
-
-    Replaces ${VAR_NAME} patterns with os.getenv(VAR_NAME, "") values.
-    This works on raw string content before YAML parsing.
-
-    Args:
-        content: The string content containing template variables
-
-    Returns:
-        The string with environment variables substituted
-    """
-    pattern = r"\$\{([^}]+)\}"
-
-    def replace_variable(match_obj: re.Match[str]) -> str:
-        var_name = match_obj.group(1)
-        env_value = os.getenv(var_name, "")
-        return env_value if env_value is not None else ""
-
-    return re.sub(pattern, replace_variable, content)
+from rslearn.template_params import substitute_env_vars_in_string
 
 
 class RslearnArgumentParser(LightningArgumentParser):

rslearn/dataset/dataset.py

@@ -8,6 +8,7 @@ from upath import UPath
 
 from rslearn.config import load_layer_config
 from rslearn.log_utils import get_logger
+from rslearn.template_params import substitute_env_vars_in_string
 from rslearn.tile_stores import TileStore, load_tile_store
 
 from .index import DatasetIndex
@@ -52,7 +53,9 @@ class Dataset:
 
         # Load dataset configuration.
         with (self.path / "config.json").open("r") as f:
-            config = json.load(f)
+            config_content = f.read()
+            config_content = substitute_env_vars_in_string(config_content)
+            config = json.loads(config_content)
             self.layers = {}
             for layer_name, d in config["layers"].items():
                 # Layer names must not contain period, since we use period to

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,204 @@
+"""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.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."""
+
+    def __init__(self, metadata_path: str = CLAY_METADATA_PATH) -> None:
+        """Initialize ClayNormalize."""
+        super().__init__()
+        with open(metadata_path) as f:
+            self.metadata = yaml.safe_load(f)
+
+    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 in CLAY_MODALITIES:
+            if modality not in input_dict or modality not in self.metadata:
+                continue
+            modality_metadata = self.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"]
+            ]
+            input_dict[modality] = self.apply_image(input_dict[modality], means, stds)
+        return input_dict, target_dict