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

rslearn/models/use_croma.py

@@ -0,0 +1,508 @@
+"""CROMA code, copied from https://github.com/antofuller/CROMA.
+
+The code is released under:
+
+MIT License
+Copyright (c) 2023 Anthony Fuller
+"""
+
+import itertools
+import math
+import warnings
+
+import torch
+from einops import rearrange
+from torch import einsum, nn
+
+
+class PretrainedCROMA(nn.Module):
+    """Pre-trained CROMA model."""
+
+    def __init__(
+        self,
+        pretrained_path: str = "CROMA_base.pt",
+        size: str = "base",
+        modality: str = "both",
+        image_resolution: int = 120,
+    ):
+        """Create a new PretrainedCROMA.
+
+        NOTE: image_resolution is not the spatial, spectral, or temporal resolution. It is the height and width of the image, in pixels.
+        E.g., CROMA was pretrained on 120x120px images, hence image_resolution is 120 by default
+        """
+        super().__init__()
+        # check types
+        assert isinstance(pretrained_path, str), (
+            f"pretrained_path must be a string, not {type(pretrained_path)}"
+        )
+        assert isinstance(size, str), f"size must be a string, not {type(size)}"
+        assert isinstance(modality, str), (
+            f"modality must be a string, not {type(modality)}"
+        )
+        assert isinstance(image_resolution, int), (
+            f"image_resolution must be an int, not {type(image_resolution)}"
+        )
+
+        # check values
+        assert size in [
+            "base",
+            "large",
+        ], f"size must be either base or large, not {size}"
+        assert image_resolution % 8 == 0, (
+            f"image_resolution must be a multiple of 8, not {image_resolution}"
+        )
+        assert modality in [
+            "both",
+            "SAR",
+            "optical",
+        ], f"modality must be either both, SAR, or optical, not {modality}"
+
+        # warn the user if the path contains a different size than the size parameter
+        if size == "base" and "large" in pretrained_path:
+            warnings.warn(
+                "The size is set to base, but the word large appears in the pretrained path!"
+            )
+        elif size == "large" and "base" in pretrained_path:
+            warnings.warn(
+                "The size is set to large, but the word base appears in the pretrained path!"
+            )
+
+        if size == "base":
+            self.encoder_dim = 768
+            self.encoder_depth = 12
+            self.num_heads = 16
+            self.patch_size = 8
+        else:
+            # large by default
+            self.encoder_dim = 1024
+            self.encoder_depth = 24
+            self.num_heads = 16
+            self.patch_size = 8
+
+        self.modality = modality
+        self.num_patches = int((image_resolution / 8) ** 2)
+        self.s1_channels = 2  # fixed at 2 SAR backscatter channels
+        self.s2_channels = 12  # fixed at 12 multispectral optical channels
+        self.attn_bias = get_2dalibi(
+            num_heads=self.num_heads, num_patches=self.num_patches
+        )
+
+        if modality in ["SAR", "both"]:
+            print("Initializing SAR encoder")
+            self.s1_encoder = ViT(
+                dim=self.encoder_dim,
+                depth=int(self.encoder_depth / 2),
+                in_channels=self.s1_channels,
+            )
+            self.GAP_FFN_s1 = nn.Sequential(
+                nn.LayerNorm(self.encoder_dim),
+                nn.Linear(
+                    self.encoder_dim, int(4 * self.encoder_dim)
+                ),  # (BSZ, num_patches, inner_dim)
+                nn.GELU(),  # (BSZ, num_patches, inner_dim)
+                nn.Linear(
+                    int(4 * self.encoder_dim), self.encoder_dim
+                ),  # (BSZ, num_patches, dim)
+            )
+
+            # load weights
+            self.s1_encoder.load_state_dict(
+                torch.load(pretrained_path, weights_only=True)["s1_encoder"]
+            )
+            self.GAP_FFN_s1.load_state_dict(
+                torch.load(pretrained_path, weights_only=True)["s1_GAP_FFN"]
+            )
+
+        if modality in ["optical", "both"]:
+            print("Initializing optical encoder")
+            self.s2_encoder = ViT(
+                dim=self.encoder_dim,
+                depth=self.encoder_depth,
+                in_channels=self.s2_channels,
+            )
+            self.GAP_FFN_s2 = nn.Sequential(
+                nn.LayerNorm(self.encoder_dim),
+                nn.Linear(
+                    self.encoder_dim, int(4 * self.encoder_dim)
+                ),  # (BSZ, num_patches, inner_dim)
+                nn.GELU(),  # (BSZ, num_patches, inner_dim)
+                nn.Linear(
+                    int(4 * self.encoder_dim), self.encoder_dim
+                ),  # (BSZ, num_patches, dim)
+            )
+
+            # load weights
+            self.s2_encoder.load_state_dict(
+                torch.load(pretrained_path, weights_only=True)["s2_encoder"]
+            )
+            self.GAP_FFN_s2.load_state_dict(
+                torch.load(pretrained_path, weights_only=True)["s2_GAP_FFN"]
+            )
+
+        if modality == "both":
+            print("Initializing joint SAR-optical encoder")
+            self.cross_encoder = BaseTransformerCrossAttn(
+                dim=self.encoder_dim,
+                depth=int(self.encoder_depth / 2),
+                num_heads=self.num_heads,
+            )
+
+            # load weights
+            self.cross_encoder.load_state_dict(
+                torch.load(pretrained_path, weights_only=True)["joint_encoder"]
+            )
+
+    def forward(
+        self,
+        SAR_images: torch.Tensor | None = None,
+        optical_images: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
+        """Forward pass through PretrainedCROMA."""
+        return_dict = {}
+        if self.modality in ["SAR", "both"]:
+            assert SAR_images is not None, (
+                f"Modality is set to {self.modality}, but SAR_images are None"
+            )
+            SAR_encodings = self.s1_encoder(
+                imgs=SAR_images, attn_bias=self.attn_bias.to(SAR_images.device)
+            )  # (bsz, num_patches, encoder_dim)
+            SAR_GAP = self.GAP_FFN_s1(SAR_encodings.mean(dim=1))  # (bsz, encoder_dim)
+            return_dict["SAR_encodings"] = SAR_encodings
+            return_dict["SAR_GAP"] = SAR_GAP
+
+        if self.modality in ["optical", "both"]:
+            assert optical_images is not None, (
+                f"Modality is set to {self.modality}, but optical_images are None"
+            )
+            optical_encodings = self.s2_encoder(
+                imgs=optical_images, attn_bias=self.attn_bias.to(optical_images.device)
+            )  # (bsz, num_patches, encoder_dim)
+            optical_GAP = self.GAP_FFN_s2(
+                optical_encodings.mean(dim=1)
+            )  # (bsz, encoder_dim)
+            return_dict["optical_encodings"] = optical_encodings
+            return_dict["optical_GAP"] = optical_GAP
+
+        if self.modality == "both":
+            assert SAR_images is not None
+            assert optical_images is not None
+            joint_encodings = self.cross_encoder(
+                x=SAR_encodings,
+                context=optical_encodings,
+                relative_position_bias=self.attn_bias.to(optical_images.device),
+            )  # (bsz, num_patches, encoder_dim)
+            joint_GAP = joint_encodings.mean(dim=1)  # (bsz, encoder_dim)
+            return_dict["joint_encodings"] = joint_encodings
+            return_dict["joint_GAP"] = joint_GAP
+
+        return return_dict
+
+
+def get_2dalibi(num_heads: int, num_patches: int) -> torch.Tensor:
+    """Get 2D bias initialization for attention layer.
+
+    Args:
+        num_heads: the number of heads in the attention layer.
+        num_patches: the total number of patches, which should be a square.
+    """
+    # inspired by: https://github.com/ofirpress/attention_with_linear_biases
+    points = list(
+        itertools.product(
+            range(int(math.sqrt(num_patches))), range(int(math.sqrt(num_patches)))
+        )
+    )
+
+    def get_slopes(n: int) -> list[float]:
+        def get_slopes_power_of_2(n: int) -> list[float]:
+            start = 2 ** (-(2 ** -(math.log2(n) - 3)))
+            ratio = start
+            return [start * ratio**i for i in range(n)]
+
+        if math.log2(n).is_integer():
+            return get_slopes_power_of_2(n)
+        else:
+            closest_power_of_2 = 2 ** math.floor(math.log2(n))
+            return (
+                get_slopes_power_of_2(closest_power_of_2)
+                + get_slopes(2 * closest_power_of_2)[0::2][: n - closest_power_of_2]
+            )
+
+    slopes = torch.Tensor(get_slopes(num_heads)).unsqueeze(1)
+    idxs = []
+    for p1 in points:
+        for p2 in points:
+            dist = math.sqrt((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2)
+            idxs.append(dist * slopes * -1)
+    all_bias = torch.cat(idxs, dim=1)
+    return all_bias.view(1, num_heads, num_patches, num_patches)
+
+
+class FFN(nn.Module):
+    """Feed-forward network block."""
+
+    def __init__(
+        self,
+        dim: int,
+        mult: int = 4,
+        dropout: float = 0.0,
+    ):
+        """Create a new FFN.
+
+        Args:
+            dim: the input dimension.
+            mult: the MLP factor (how much larger the hidden dimension should be).
+            dropout: the dropout rate.
+        """
+        super().__init__()
+        inner_dim = int(dim * mult)
+
+        self.net = nn.Sequential(
+            nn.Linear(dim, inner_dim),  # (BSZ, num_patches, inner_dim)
+            nn.GELU(),  # (BSZ, num_patches, inner_dim)
+            nn.Dropout(dropout),  # (BSZ, num_patches, inner_dim)
+            nn.Linear(inner_dim, dim),  # (BSZ, num_patches, dim)
+        )
+        self.input_norm = nn.LayerNorm(dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass through the FFN."""
+        x = self.input_norm(x)  # (BSZ, num_patches, dim)
+        return self.net(x)  # (BSZ, num_patches, dim)
+
+
+class Attention(nn.Module):
+    """Attention block."""
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        dropout: float = 0.0,
+    ):
+        """Create a new Attention."""
+        super().__init__()
+        self.num_heads = num_heads
+        assert dim % num_heads == 0, "dim must be evenly divisible by num_heads"
+        dim_head = int(dim / num_heads)
+        self.scale = dim_head**-0.5
+
+        self.to_qkv = nn.Linear(dim, dim * 3, bias=False)
+        self.to_out = nn.Linear(dim, dim)
+        self.input_norm = nn.LayerNorm(dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(
+        self, x: torch.Tensor, relative_position_bias: torch.Tensor
+    ) -> torch.Tensor:
+        """Forward pass through the Attention."""
+        x = self.input_norm(x)  # (BSZ, num_patches, dim)
+        q, k, v = self.to_qkv(x).chunk(3, dim=-1)  # (BSZ, num_patches, dim)
+        q, k, v = map(
+            lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.num_heads), (q, k, v)
+        )  # (BSZ, num_heads, num_patches, dim_head)
+
+        attention_scores = (
+            einsum("b h i d, b h j d -> b h i j", q, k) * self.scale
+        )  # (BSZ, num_heads, num_patches, num_patches)
+        attention_scores = (
+            attention_scores + relative_position_bias
+        )  # (BSZ, num_heads, num_patches, num_patches)
+
+        attn = attention_scores.softmax(
+            dim=-1
+        )  # (BSZ, num_heads, num_patches, num_patches)
+        attn = self.dropout(attn)  # (BSZ, num_heads, num_patches, num_patches)
+
+        out = einsum(
+            "b h i j, b h j d -> b h i d", attn, v
+        )  # (BSZ, num_heads, num_patches, dim_head)
+        out = rearrange(out, "b h n d -> b n (h d)")  # (BSZ, num_patches, dim)
+        return self.to_out(out)  # (BSZ, num_patches, dim)
+
+
+class CrossAttention(nn.Module):
+    """Cross-attention block."""
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        dropout: float = 0.0,
+    ):
+        """Create a new CrossAttention."""
+        super().__init__()
+        self.num_heads = num_heads
+        assert dim % num_heads == 0, "dim must be evenly divisible by num_heads"
+        dim_head = int(dim / num_heads)
+        self.scale = dim_head**-0.5
+
+        self.to_q = nn.Linear(dim, dim, bias=False)
+        self.to_k = nn.Linear(dim, dim, bias=False)
+        self.to_v = nn.Linear(dim, dim, bias=False)
+
+        self.to_out = nn.Linear(dim, dim)
+        self.input_norm = nn.LayerNorm(dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        context: torch.Tensor,
+        relative_position_bias: torch.Tensor,
+    ) -> torch.Tensor:
+        """Forward pass through the CrossAttention."""
+        x = self.input_norm(x)  # (BSZ, num_patches, dim)
+        context = self.input_norm(context)  # (BSZ, num_patches, dim)
+
+        q = self.to_q(x)  # (BSZ, num_patches, dim)
+        k = self.to_k(context)  # (BSZ, num_patches, dim)
+        v = self.to_v(context)  # (BSZ, num_patches, dim)
+
+        q, k, v = map(
+            lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.num_heads), (q, k, v)
+        )  # (BSZ, num_heads, num_patches, dim_head)
+
+        attention_scores = (
+            einsum("b h i d, b h j d -> b h i j", q, k) * self.scale
+        )  # (BSZ, num_heads, num_patches, num_patches)
+        attention_scores = (
+            attention_scores + relative_position_bias
+        )  # (BSZ, num_heads, num_patches, num_patches)
+
+        attn = attention_scores.softmax(
+            dim=-1
+        )  # (BSZ, num_heads, num_patches, num_patches)
+        attn = self.dropout(attn)  # (BSZ, num_heads, num_patches, num_patches)
+
+        out = einsum(
+            "b h i j, b h j d -> b h i d", attn, v
+        )  # (BSZ, num_heads, num_patches, dim_head)
+        out = rearrange(out, "b h n d -> b n (h d)")  # (BSZ, num_patches, dim)
+        return self.to_out(out)  # (BSZ, num_patches, dim)
+
+
+class BaseTransformer(nn.Module):
+    """Base transformer class."""
+
+    def __init__(
+        self,
+        dim: int,
+        depth: int,
+        num_heads: int = 8,
+        attn_dropout: float = 0.0,
+        ff_dropout: float = 0.0,
+        ff_mult: int = 4,
+        final_norm: bool = True,
+    ):
+        """Create a new BaseTransformer."""
+        super().__init__()
+        self.final_norm = final_norm
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        Attention(dim=dim, num_heads=num_heads, dropout=attn_dropout),
+                        FFN(dim=dim, mult=ff_mult, dropout=ff_dropout),
+                    ]
+                )
+            )
+
+        if self.final_norm:
+            self.norm_out = nn.LayerNorm(dim)
+
+    def forward(
+        self, x: torch.Tensor, relative_position_bias: torch.Tensor
+    ) -> torch.Tensor:
+        """Forward pass through the BaseTransformer."""
+        for self_attn, ffn in self.layers:
+            x = self_attn(x, relative_position_bias) + x  # (BSZ, num_patches, dim)
+            x = ffn(x) + x  # (BSZ, num_patches, dim)
+
+        if self.final_norm:
+            return self.norm_out(x)
+        else:
+            return x
+
+
+class BaseTransformerCrossAttn(nn.Module):
+    """Base transformer class for cross attention."""
+
+    def __init__(
+        self,
+        dim: int,
+        depth: int,
+        num_heads: int = 8,
+        attn_dropout: float = 0.0,
+        ff_dropout: float = 0.0,
+        ff_mult: int = 4,
+    ):
+        """Create a new BaseTransformerCrossAttn."""
+        super().__init__()
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        Attention(dim=dim, num_heads=num_heads, dropout=attn_dropout),
+                        CrossAttention(
+                            dim=dim, num_heads=num_heads, dropout=attn_dropout
+                        ),
+                        FFN(dim=dim, mult=ff_mult, dropout=ff_dropout),
+                    ]
+                )
+            )
+
+        self.norm_out = nn.LayerNorm(dim)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        context: torch.Tensor,
+        relative_position_bias: torch.Tensor,
+    ) -> torch.Tensor:
+        """Forward pass through the BaseTransformerCrossAttn."""
+        for self_attn, cross_attn, ffn in self.layers:
+            x = self_attn(x, relative_position_bias) + x  # (BSZ, num_patches, dim)
+            x = (
+                cross_attn(x, context, relative_position_bias) + x
+            )  # (BSZ, num_patches, dim)
+            x = ffn(x) + x  # (BSZ, num_patches, dim)
+
+        x = self.norm_out(x)
+        return x  # (BSZ, num_patches, dim)
+
+
+class ViT(nn.Module):
+    """ViT model."""
+
+    def __init__(self, dim: int, depth: int, in_channels: int):
+        """Create a new ViT."""
+        super().__init__()
+        self.depth = depth
+        self.in_channels = in_channels
+        self.dim = dim
+        self.num_heads = 16  # always 16, for base and large models
+        self.patch_size = 8  # always 8, for base and large models
+
+        pixels_per_patch = int(self.patch_size * self.patch_size * in_channels)
+        self.linear_input = nn.Linear(pixels_per_patch, self.dim)
+        self.transformer = BaseTransformer(
+            dim=self.dim,
+            depth=self.depth,
+            num_heads=self.num_heads,
+        )
+
+    def forward(self, imgs: torch.Tensor, attn_bias: torch.Tensor) -> torch.Tensor:
+        """Forward pass through the ViT."""
+        x = rearrange(
+            imgs,
+            "b c (h i) (w j) -> b (h w) (c i j)",
+            i=self.patch_size,
+            j=self.patch_size,
+        )
+        # x is shape -> (bsz, num_patches, self.channels*self.patch_size*self.patch_size)
+
+        x = self.linear_input(x)  # (bsz, num_patches, dim)
+        x = self.transformer(x, relative_position_bias=attn_bias)
+        return x

rslearn/template_params.py

@@ -0,0 +1,26 @@
+"""Template parameter substitution utilities for rslearn configuration files."""
+
+import os
+import re
+
+
+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/JSON 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)

rslearn/tile_stores/__init__.py

@@ -1,37 +1,60 @@
 """Tile stores that store ingested raster and vector data before materialization."""
 
-from upath import UPath
+from typing import Any
 
-from rslearn.config import TileStoreConfig
+import jsonargparse
+from upath import UPath
 
-from .file import FileTileStore
-from .tile_store import (
-    LayerMetadata,
-    PrefixedTileStore,
-    TileStore,
-    TileStoreLayer,
-    get_tile_store_for_layer,
-)
+from rslearn.config import LayerConfig
+from rslearn.utils.jsonargparse import init_jsonargparse
 
-registry = {"file": FileTileStore}
+from .default import DefaultTileStore
+from .tile_store import TileStore, TileStoreWithLayer
 
 
-def load_tile_store(config: TileStoreConfig, ds_path: UPath) -> TileStore:
+def load_tile_store(config: dict[str, Any], ds_path: UPath) -> TileStore:
     """Load a tile store from a configuration.
 
     Args:
         config: the tile store configuration.
         ds_path: the dataset root path.
+
+    Returns:
+        the TileStore
+    """
+    init_jsonargparse()
+    parser = jsonargparse.ArgumentParser()
+    parser.add_argument("--tile_store", type=TileStore)
+    cfg = parser.parse_object({"tile_store": config})
+    tile_store = parser.instantiate_classes(cfg).tile_store
+    tile_store.set_dataset_path(ds_path)
+    return tile_store
+
+
+def get_tile_store_with_layer(
+    tile_store: TileStore, layer_name: str, layer_cfg: LayerConfig
+) -> TileStoreWithLayer:
+    """Get the TileStoreWithLayer for the specified layer.
+
+    Uses alias of the layer if it is set, otherwise just the layer name.
+
+    Args:
+        tile_store: the tile store.
+        layer_name: the layer name.
+        layer_cfg: the layer configuration which can specify an alias.
+
+    Returns:
+        corresponding TileStoreWithLayer
     """
-    return registry[config.name].from_config(config, ds_path)
+    if layer_cfg.alias is not None:
+        return TileStoreWithLayer(tile_store, layer_cfg.alias)
+    return TileStoreWithLayer(tile_store, layer_name)
 
 
 __all__ = (
-    "FileTileStore",
-    "LayerMetadata",
-    "PrefixedTileStore",
+    "DefaultTileStore",
     "TileStore",
-    "TileStoreLayer",
+    "TileStoreWithLayer",
     "load_tile_store",
-    "get_tile_store_for_layer",
+    "get_tile_store_with_layer",
 )