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

rslearn/models/copernicusfm_src/util/pos_embed.py

@@ -0,0 +1,216 @@
+# type: ignore
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# Position embedding utils
+# --------------------------------------------------------
+
+import numpy as np
+import torch
+
+
+# --------------------------------------------------------
+# 2D sine-cosine position embedding
+# References:
+# Transformer: https://github.com/tensorflow/models/blob/master/official/nlp/transformer/model_utils.py
+# MoCo v3: https://github.com/facebookresearch/moco-v3
+# --------------------------------------------------------
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
+    """grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token:
+        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    # omega = np.arange(embed_dim // 2, dtype=np.float) # numpy deprecated in 1.20
+    omega = np.arange(embed_dim // 2, dtype=float)
+
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out)  # (M, D/2)
+    emb_cos = np.cos(out)  # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+# --------------------------------------------------------
+# Interpolate position embeddings for high-resolution
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# --------------------------------------------------------
+def interpolate_pos_embed(model, checkpoint_model):
+    if "pos_embed" in checkpoint_model:
+        pos_embed_checkpoint = checkpoint_model["pos_embed"]
+        embedding_size = pos_embed_checkpoint.shape[-1]
+        num_patches = model.patch_embed.num_patches
+        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
+        # height (== width) for the checkpoint position embedding
+        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+        # height (== width) for the new position embedding
+        new_size = int(num_patches**0.5)
+        # class_token and dist_token are kept unchanged
+        if orig_size != new_size:
+            print(
+                "Position interpolate from %dx%d to %dx%d"
+                % (orig_size, orig_size, new_size, new_size)
+            )
+            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+            # only the position tokens are interpolated
+            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+            pos_tokens = pos_tokens.reshape(
+                -1, orig_size, orig_size, embedding_size
+            ).permute(0, 3, 1, 2)
+            pos_tokens = torch.nn.functional.interpolate(
+                pos_tokens,
+                size=(new_size, new_size),
+                mode="bicubic",
+                align_corners=False,
+            )
+            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+            checkpoint_model["pos_embed"] = new_pos_embed
+
+
+def interpolate_pos_embed_ofa(model, checkpoint_model):
+    if "pos_embed" in checkpoint_model:
+        pos_embed_dict = checkpoint_model["pos_embed"]
+
+        for key, pos_embed in pos_embed_dict.items():
+            pos_embed_checkpoint = pos_embed
+            embedding_size = pos_embed_checkpoint.shape[-1]
+            num_patches = model.patch_embed[key].num_patches
+            num_extra_tokens = model.pos_embed[key].shape[-2] - num_patches
+            # height (== width) for the checkpoint position embedding
+            orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+            # height (== width) for the new position embedding
+            new_size = int(num_patches**0.5)
+            # class_token and dist_token are kept unchanged
+            if orig_size != new_size:
+                print(
+                    "Position interpolate from %dx%d to %dx%d"
+                    % (orig_size, orig_size, new_size, new_size)
+                )
+                extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+                # only the position tokens are interpolated
+                pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+                pos_tokens = pos_tokens.reshape(
+                    -1, orig_size, orig_size, embedding_size
+                ).permute(0, 3, 1, 2)
+                pos_tokens = torch.nn.functional.interpolate(
+                    pos_tokens,
+                    size=(new_size, new_size),
+                    mode="bicubic",
+                    align_corners=False,
+                )
+                pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+                new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+                checkpoint_model["pos_embed"][key] = new_pos_embed
+
+
+def get_2d_sincos_pos_embed_with_resolution(
+    embed_dim, grid_size, res, cls_token=False, device="cpu"
+):
+    """grid_size: int of the grid height and width
+    res: array of size n, representing the resolution of a pixel (say, in meters),
+
+    Return:
+    pos_embed: [n,grid_size*grid_size, embed_dim] or [n,1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    # res = torch.FloatTensor(res).to(device)
+    res = res.to(device)
+    grid_h = torch.arange(grid_size, dtype=torch.float32, device=device)
+    grid_w = torch.arange(grid_size, dtype=torch.float32, device=device)
+    grid = torch.meshgrid(
+        grid_w, grid_h, indexing="xy"
+    )  # here h goes first,direction reversed for numpy
+    grid = torch.stack(grid, dim=0)  # 2 x h x w
+
+    # grid = grid.reshape([2, 1, grid_size, grid_size])
+    grid = torch.einsum("chw,n->cnhw", grid, res)  # 2 x n x h x w
+    _, n, h, w = grid.shape
+    pos_embed = get_2d_sincos_pos_embed_from_grid_torch(
+        embed_dim, grid
+    )  #  # (nxH*W, D/2)
+    pos_embed = pos_embed.reshape(n, h * w, embed_dim)
+    if cls_token:
+        pos_embed = torch.cat(
+            [
+                torch.zeros(
+                    [n, 1, embed_dim], dtype=torch.float32, device=pos_embed.device
+                ),
+                pos_embed,
+            ],
+            dim=1,
+        )
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid_torch(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid_torch(
+        embed_dim // 2, grid[0]
+    )  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid_torch(
+        embed_dim // 2, grid[1]
+    )  # (H*W, D/2)
+
+    emb = torch.cat([emb_h, emb_w], dim=1)  # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid_torch(embed_dim, pos):
+    """embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    old_shape = pos
+    omega = torch.arange(embed_dim // 2, dtype=torch.float32, device=pos.device)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = torch.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = torch.sin(out)  # (M, D/2)
+    emb_cos = torch.cos(out)  # (M, D/2)
+
+    emb = torch.cat([emb_sin, emb_cos], dim=1)  # (M, D)
+    return emb

rslearn/models/panopticon.py

@@ -0,0 +1,167 @@
+"""Wrapper for the Panopticon model."""
+
+import math
+from enum import StrEnum
+from importlib import resources
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+import yaml
+from einops import rearrange, repeat
+from torch import nn
+
+from rslearn.log_utils import get_logger
+
+logger = get_logger(__name__)
+
+
+class PanopticonModalities(StrEnum):
+    """Modalities supported by Panopticon.
+
+    These are the keys needed to load the yaml file from panopticon_data/sensors
+    """
+
+    SENTINEL2 = "sentinel2"
+    LANDSAT8 = "landsat8"
+    SENTINEL1 = "sentinel1"
+    # Add more modalities as needed
+
+
+class Panopticon(nn.Module):
+    """Class containing the Panopticon model that can ingest MaskedHeliosSample objects."""
+
+    patch_size: int = 14
+    base_image_size: int = 224
+
+    def __init__(
+        self,
+        band_order: dict[str, list[str]],
+        torchhub_id: str = "panopticon_vitb14",
+    ):
+        """Initialize the Panopticon wrapper.
+
+        Args:
+            band_order: The band order for the panopticon model, must match the specified order in the data config
+            torchhub_id: The torch hub model ID for panopticon
+        """
+        super().__init__()
+        # Load the panopticon model
+        self._load_model(torchhub_id)
+        self.output_dim = self.model.embed_dim
+        self.band_order = band_order
+        self.supported_modalities = list(band_order.keys())
+
+    def _load_model(self, torchhub_id: str) -> None:
+        """Load the panopticon model from torch hub."""
+        import time
+
+        # Hack to get around https://discuss.pytorch.org/t/torch-hub-load-gives-httperror-rate-limit-exceeded/124769
+        torch.hub._validate_not_a_forked_repo = lambda a, b, c: True
+        for attempt in range(2):
+            try:
+                self.model = torch.hub.load(  # nosec B614
+                    "panopticon-FM/panopticon",
+                    torchhub_id,
+                )
+                break
+            except Exception as e:
+                logger.warning(
+                    f"Error loading panopticon model: {e}. Retrying in 5 seconds..."
+                )
+                time.sleep(5)
+        else:
+            raise RuntimeError(
+                f"Failed to load panopticon model {torchhub_id} after retrying."
+            )
+
+    def _process_modality_data(self, data: torch.Tensor) -> torch.Tensor:
+        """Process individual modality data.
+
+        Args:
+            data: Input tensor of shape [B, C, H, W]
+
+        Returns:
+            Processed tensor of shape [B, C, H, W]
+        """
+        original_height = data.shape[2]
+        new_height = self.patch_size if original_height == 1 else self.base_image_size
+
+        data = F.interpolate(
+            data,
+            size=(new_height, new_height),
+            mode="bilinear",
+            align_corners=False,
+        )
+        return data
+
+    def _create_channel_ids(
+        self, modality: str, batch_size: int, device: torch.device
+    ) -> torch.Tensor:
+        """Create channel IDs for the panopticon model."""
+        with resources.open_text(
+            "rslearn.models.panopticon_data.sensors", f"{modality}.yaml"
+        ) as f:
+            sensor_config = yaml.safe_load(f)
+
+        band_order = self.band_order[modality]
+        chn_ids = [
+            sensor_config["bands"][band.upper()]["gaussian"]["mu"]
+            for band in band_order
+        ]
+        chn_ids = torch.tensor(chn_ids, dtype=torch.float32, device=device)
+        chn_ids = repeat(chn_ids, "c -> b c", b=batch_size)
+        return chn_ids
+
+    def prepare_input(
+        self, input_data: dict[str, torch.Tensor]
+    ) -> dict[str, torch.Tensor]:
+        """Prepare input for the panopticon model from MaskedHeliosSample."""
+        channel_ids_list: list[torch.Tensor] = []
+        processed_data_list: list[torch.Tensor] = []
+        for modality in self.supported_modalities:
+            if modality not in input_data.keys():
+                logger.debug(f"Modality {modality} not found in input data")
+                continue
+            data = input_data[modality]
+            device = data.device
+            processed_data = self._process_modality_data(data)
+            processed_data_list.append(processed_data)
+            batch_size = processed_data.shape[0]
+            chn_ids = self._create_channel_ids(modality, batch_size, device)
+            channel_ids_list.append(chn_ids)
+
+        processed_data = torch.cat(processed_data_list, dim=1)
+        chn_ids = torch.cat(channel_ids_list, dim=1)
+        return {
+            "imgs": processed_data,
+            "chn_ids": chn_ids,
+        }
+
+    def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
+        """Forward pass through the panopticon model."""
+        batch_inputs = {
+            key: torch.stack([inp[key] for inp in inputs], dim=0)
+            for key in inputs[0].keys()
+        }
+        panopticon_inputs = self.prepare_input(batch_inputs)
+        output_features = self.model.forward_features(panopticon_inputs)[
+            "x_norm_patchtokens"
+        ]
+
+        num_tokens = output_features.shape[1]
+        height = int(math.sqrt(num_tokens))
+        output_features = rearrange(
+            output_features, "b (h w) d -> b d h w", h=height, w=height
+        )
+        return [output_features]
+
+    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 (downsample_factor, depth) that corresponds
+        to the feature maps that the backbone returns. For example, an element [2, 32]
+        indicates that the corresponding feature map is 1/2 the input resolution and
+        has 32 channels.
+        """
+        return [(self.patch_size, self.output_dim)]

rslearn/models/presto/__init__.py

@@ -0,0 +1,5 @@
+"""Presto."""
+
+from .presto import Presto
+
+__all__ = ["Presto"]

rslearn/models/presto/presto.py

@@ -0,0 +1,247 @@
+"""Presto wrapper to ingest Masked Helios Samples."""
+
+import logging
+import tempfile
+from typing import Any
+
+import torch
+from einops import rearrange, repeat
+from huggingface_hub import hf_hub_download
+from torch import nn
+from upath import UPath
+
+from rslearn.models.presto.single_file_presto import (
+    ERA5_BANDS,
+    NUM_DYNAMIC_WORLD_CLASSES,
+    PRESTO_ADD_BY,
+    PRESTO_BANDS,
+    PRESTO_DIV_BY,
+    PRESTO_S1_BANDS,
+    PRESTO_S2_BANDS,
+    SRTM_BANDS,
+)
+from rslearn.models.presto.single_file_presto import Presto as SFPresto
+
+logger = logging.getLogger(__name__)
+
+INPUT_PRESTO_BANDS = [b for b in PRESTO_BANDS if b != "B09"]
+INPUT_PRESTO_S2_BANDS = [b for b in PRESTO_S2_BANDS if b != "B09"]
+
+PRESTO_S1_SUBTRACT_VALUE = -25.0
+PRESTO_S1_DIV_VALUE = 25.0
+PRESTO_S2_SUBTRACT_VALUE = 0.0
+PRESTO_S2_DIV_VALUE = 1e4
+
+HF_HUB_ID = "nasaharvest/presto"
+MODEL_FILENAME = "default_model.pt"
+
+
+class Presto(nn.Module):
+    """Presto."""
+
+    input_keys = [
+        "s1",
+        "s2",
+        "era5",
+        "srtm",
+        "dynamic_world",
+        "latlon",
+    ]
+
+    def __init__(
+        self,
+        pretrained_path: str | UPath | None = None,
+        pixel_batch_size: int = 128,
+    ):
+        """Initialize the Presto wrapper.
+
+        Args:
+            pretrained_path: The directory to load from
+            pixel_batch_size: If the input has a h,w dimension >1, this is
+                flattened into a batch dimension (b h w) before being passed
+                to the model (since Presto is designed for pixel timeseries).
+        """
+        super().__init__()
+
+        if pretrained_path is None:
+            pretrained_path = UPath(tempfile.gettempdir(), "rslearn_cache", "presto")
+        if not (UPath(pretrained_path) / MODEL_FILENAME).exists():
+            _ = hf_hub_download(
+                local_dir=UPath(pretrained_path),
+                repo_id=HF_HUB_ID,
+                filename=MODEL_FILENAME,
+                # pin the model to a specific hugging face commit
+                revision="1b97f885969da4e2d5834ca8c92707c737911464",
+            )
+
+        model = SFPresto.construct()
+        model.load_state_dict(
+            torch.load(
+                UPath(pretrained_path) / MODEL_FILENAME,
+                map_location="cpu",
+                weights_only=True,
+            )
+        )
+        self.pixel_batch_size = pixel_batch_size
+        self.model = model.encoder
+        self.month = 6  # default month
+
+    def construct_presto_input(
+        self,
+        s1: torch.Tensor | None = None,
+        s1_bands: torch.Tensor | None = None,
+        s2: torch.Tensor | None = None,
+        s2_bands: torch.Tensor | None = None,
+        era5: torch.Tensor | None = None,
+        era5_bands: torch.Tensor | None = None,
+        srtm: torch.Tensor | None = None,
+        srtm_bands: torch.Tensor | None = None,
+        dynamic_world: torch.Tensor | None = None,
+        months: torch.Tensor | None = None,
+        normalize: bool = True,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Inputs are paired into a tensor input <X> and a list <X>_bands, which describes <X>.
+
+        <X> should have shape (b, num_timesteps, h, w len(<X>_bands)), with the following bands for
+        each input:
+
+        s1: ["VV", "VH"]
+        s2: ["B2", "B3", "B4", "B5", "B6", "B7", "B8", "B8A", "B9", "B11", "B12"]
+        era5: ["temperature_2m", "total_precipitation"]
+            "temperature_2m": Temperature of air at 2m above the surface of land,
+                sea or in-land waters in Kelvin (K)
+            "total_precipitation": Accumulated liquid and frozen water, including rain and snow,
+                that falls to the Earth's surface. Measured in metres (m)
+        srtm: ["elevation", "slope"]
+
+        dynamic_world is a 1d input of shape (num_timesteps,) representing the dynamic world classes
+            of each timestep for that pixel
+        """
+        bs = [x.shape[0] for x in [s1, s2, era5, srtm] if x is not None]
+        hs = [x.shape[2] for x in [s1, s2, era5, srtm] if x is not None]
+        ws = [x.shape[3] for x in [s1, s2, era5, srtm] if x is not None]
+
+        assert len(bs) > 0
+        assert len(set(bs)) == 1
+        assert len(set(hs)) == 1
+        assert len(set(ws)) == 1
+        b, h, w = bs[0], hs[0], ws[0]
+
+        # these values will be initialized as
+        # we iterate through the data
+        x: torch.Tensor | None = None
+        mask: torch.Tensor | None = None
+        t: int | None = None
+
+        for band_group in [
+            (s1, s1_bands),
+            (s2, s2_bands),
+            (era5, era5_bands),
+            (srtm, srtm_bands),
+        ]:
+            data, input_bands = band_group
+            if data is not None:
+                assert input_bands is not None
+            else:
+                continue
+
+            m_t = data.shape[1] // len(input_bands)
+            if t is None:
+                t = m_t
+            else:
+                if t != m_t:
+                    raise ValueError("inconsistent values for t")
+
+            data = rearrange(data, "b (t c) h w -> b t h w c", t=m_t)
+            if x is None:
+                x = torch.zeros(b, t, h, w, len(INPUT_PRESTO_BANDS))
+            if mask is None:
+                mask = torch.ones(b, t, h, w, len(INPUT_PRESTO_BANDS))
+
+            # construct a mapping from the input bands to the presto input bands
+            input_to_output_mapping = [
+                INPUT_PRESTO_BANDS.index(val) for val in input_bands
+            ]
+            x[:, :, :, :, input_to_output_mapping] = data
+            mask[:, :, :, :, input_to_output_mapping] = 0
+
+        assert x is not None
+        assert mask is not None
+        assert t is not None
+
+        if dynamic_world is None:
+            dynamic_world = torch.ones(b, t, h, w) * NUM_DYNAMIC_WORLD_CLASSES
+
+        if months is None:
+            months = torch.ones((b, t), device=x.device) * self.month
+        else:
+            assert months.shape[-1] == t
+
+        if normalize:
+            x = (x + PRESTO_ADD_BY) / PRESTO_DIV_BY
+        return x, mask, dynamic_world.long(), months.long()
+
+    def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
+        """Compute feature maps from the Presto backbone.
+
+        Inputs:
+            inputs
+        """
+        stacked_inputs = {}
+        latlons: torch.Tensor | None = None
+        for key in inputs[0].keys():
+            # assume all the keys in an input are consistent
+            if key in self.input_keys:
+                if key == "latlon":
+                    latlons = torch.stack([inp[key] for inp in inputs], dim=0)
+                else:
+                    stacked_inputs[key] = torch.stack(
+                        [inp[key] for inp in inputs], dim=0
+                    )
+
+        (
+            x,
+            mask,
+            dynamic_world,
+            months,
+        ) = self.construct_presto_input(
+            **stacked_inputs,
+            s1_bands=PRESTO_S1_BANDS,
+            s2_bands=INPUT_PRESTO_S2_BANDS,
+            era5_bands=ERA5_BANDS,
+            srtm_bands=SRTM_BANDS,
+            normalize=True,
+        )
+        b, _, h, w, _ = x.shape
+
+        output_features = torch.zeros(
+            b * h * w, self.model.embedding_size, device=x.device
+        )
+
+        x = rearrange(x, "b t h w d -> (b h w) t d")
+        mask = rearrange(mask, "b t h w d -> (b h w) t d")
+        dynamic_world = rearrange(dynamic_world, "b t h w -> (b h w) t")
+        months = repeat(months, "b t -> (b h w) t", h=h, w=w)
+        if latlons is not None:
+            latlons = rearrange(latlons, "b c h w -> (b h w) c")
+
+        for batch_idx in range(0, b * h * w, self.pixel_batch_size):
+            x_b = x[batch_idx : batch_idx + self.pixel_batch_size]
+            mask_b = mask[batch_idx : batch_idx + self.pixel_batch_size]
+            dw = dynamic_world[batch_idx : batch_idx + self.pixel_batch_size]
+            months_b = months[batch_idx : batch_idx + self.pixel_batch_size]
+            if latlons is not None:
+                l_b = latlons[batch_idx : batch_idx + self.pixel_batch_size]
+            else:
+                l_b = None
+            output_b = self.model(
+                x=x_b,
+                dynamic_world=dw,
+                mask=mask_b,
+                month=months_b,
+                latlons=l_b,
+                eval_task=True,
+            )
+            output_features[batch_idx : batch_idx + self.pixel_batch_size] = output_b
+
+        return [rearrange(output_features, "(b h w) d -> b d h w", h=h, w=w, b=b)]