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

rslearn/models/copernicusfm.py

@@ -0,0 +1,216 @@
+"""Copernicus FM model."""
+
+import logging
+import math
+from enum import Enum
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from upath import UPath
+
+from .copernicusfm_src.model_vit import vit_base_patch16
+
+logger = logging.getLogger(__name__)
+
+
+class CopernicusFMModality(Enum):
+    """Modality for Copernicus FM."""
+
+    SENTINEL2_L2A = "sentinel2_l2a"
+    SENTINEL1 = "sentinel1"
+
+
+MODALITY_TO_WAVELENGTH_BANDWIDTHS: dict[str, dict[str, list]] = {
+    # https://github.com/zhu-xlab/Copernicus-FM/blob/main/Copernicus-Bench/src/configs/dataset/cobench_eurosat_s2.yaml
+    CopernicusFMModality.SENTINEL2_L2A.value: {
+        "band_names": [
+            "B01",
+            "B02",
+            "B03",
+            "B04",
+            "B05",
+            "B06",
+            "B07",
+            "B08",
+            "B8A",
+            "B09",
+            "B10",
+            "B11",
+            "B12",
+        ],
+        "band_wavelengths": [
+            440,
+            490,
+            560,
+            665,
+            705,
+            740,
+            783,
+            842,
+            860,
+            940,
+            1370,
+            1610,
+            2190,
+        ],
+        "band_bandwidths": [20, 65, 35, 30, 15, 15, 20, 115, 20, 20, 30, 90, 180],
+    },
+    # https://github.com/zhu-xlab/Copernicus-FM/blob/main/Copernicus-Bench/src/configs/dataset/cobench_eurosat_s1.yaml
+    CopernicusFMModality.SENTINEL1.value: {
+        "band_names": ["vv", "vh"],
+        "band_wavelengths": [50000000, 50000000],
+        "band_bandwidths": [1e9, 1e9],
+    },
+}
+
+
+class CopernicusFM(torch.nn.Module):
+    """Wrapper for Copernicus FM to ingest Masked Helios Sample."""
+
+    image_resolution = 224
+    patch_size = 16
+    input_mode = "spectral"
+    # Don't need this as band order is provided
+    supported_modalities = [
+        CopernicusFMModality.SENTINEL2_L2A.value,
+        CopernicusFMModality.SENTINEL1.value,
+    ]
+
+    def __init__(
+        self,
+        band_order: dict[str, list[str]],
+        load_directory: str | None,
+    ) -> None:
+        """Initialize the Copernicus FM wrapper.
+
+        Args:
+            band_order: The band order for each modality
+            load_directory: The directory to load from, if None no weights are loaded
+        """
+        super().__init__()
+
+        # global_pool=True so that we initialize the fc_norm layer
+        self.band_order = band_order
+        self.model = vit_base_patch16(num_classes=10, global_pool=True)
+        if load_directory is not None:
+            check_point = torch.load(
+                UPath(load_directory) / "CopernicusFM_ViT_base_varlang_e100.pth",
+                weights_only=True,
+            )
+            if "model" in check_point:
+                state_dict = check_point["model"]
+            else:
+                state_dict = check_point
+            self.model.load_state_dict(state_dict, strict=False)
+
+        # take MODALITY_TO_WAVELENGTH_BANDWIDTHS and rearrage it so that it has the same
+        # ordering as the Helios band orders, defined by Modality.band_order
+        self.modality_to_wavelength_bandwidths = {}
+        for modality in self.supported_modalities:
+            wavelength_bandwidths = MODALITY_TO_WAVELENGTH_BANDWIDTHS[modality]
+            wavelengths = []
+            bandwidths = []
+            modality_band_order = self.band_order.get(modality, None)
+            if modality_band_order is None:
+                logger.warning(
+                    f"Band order for modality {modality} not found in band_order dictionary, unable to use this modality unless specified"
+                )
+                continue
+            for b in modality_band_order:
+                cfm_idx = wavelength_bandwidths["band_names"].index(b)
+                wavelengths.append(wavelength_bandwidths["band_wavelengths"][cfm_idx])
+                bandwidths.append(wavelength_bandwidths["band_bandwidths"][cfm_idx])
+            self.modality_to_wavelength_bandwidths[modality] = {
+                "band_bandwidths": bandwidths,
+                "band_wavelengths": wavelengths,
+            }
+
+    def _resize_data(self, data: torch.Tensor) -> torch.Tensor:
+        """Process individual modality data.
+
+        Args:
+            data: Input tensor of shape [B, C, H, W]
+
+        Returns:
+            list of tensors of shape [B, C, H, W]
+        """
+        # Get original dimensions
+        original_height = data.shape[2]
+        new_height = self.patch_size if original_height == 1 else self.image_resolution
+        data = F.interpolate(
+            data,
+            size=(new_height, new_height),
+            mode="bilinear",
+            align_corners=False,
+        )
+        return data
+
+    def prepare_input(
+        self,
+        inputs: dict[str, torch.Tensor],
+    ) -> tuple[torch.Tensor, list[int], list[int]]:
+        """Prepare input for the CopernicusFM model from MaskedHeliosSample."""
+        wavelengths: list[int] = []
+        bandwidths: list[int] = []
+        all_processed_data: list[list[torch.Tensor]] = []
+        for modality in inputs.keys():
+            if modality not in self.supported_modalities:
+                logger.debug(
+                    f"Skipping modality {modality} as it is not in the supported "
+                    f"modalities list {self.supported_modalities}"
+                )
+                continue
+
+            data = inputs[modality]
+
+            if data is None:
+                continue
+
+            all_processed_data.append(self._resize_data(data))
+            wavelengths.extend(
+                self.modality_to_wavelength_bandwidths[modality]["band_wavelengths"]
+            )
+            bandwidths.extend(
+                self.modality_to_wavelength_bandwidths[modality]["band_bandwidths"]
+            )
+
+        concatenated_processed_data = torch.cat(all_processed_data, dim=1)
+        return concatenated_processed_data, wavelengths, bandwidths
+
+    def forward(
+        self,
+        inputs: list[dict[str, torch.Tensor]],
+    ) -> torch.Tensor:
+        """Forward pass through CopernicusFM model."""
+        batch_inputs = {
+            key: torch.stack([inp[key] for inp in inputs], dim=0)
+            for key in inputs[0].keys()
+        }
+        # Prepare input
+        data, wavelengths, bandwidths = self.prepare_input(batch_inputs)
+        meta = torch.full(
+            (1, 4), float("nan"), device=data.device
+        )  # [lon, lat, delta_time, patch_token_area], assume unknown
+        # "The embed tensor contains the encoded image features, which can be used for downstream tasks."
+        _, timestep_output = self.model(
+            data,
+            meta,
+            wavelengths,
+            bandwidths,
+            None,
+            self.input_mode,
+            self.patch_size,
+        )
+        # no norm, following
+        # https://github.com/zhu-xlab/Copernicus-FM/blob/main/Copernicus-Bench/src/foundation_models/CopernicusFM/models_dwv_seg.py
+        side = math.isqrt(timestep_output.shape[1])
+        output_features = rearrange(
+            timestep_output, "b (h w) c -> b c h w ", h=side, w=side
+        )
+        return [output_features]
+
+    def get_backbone_channels(self) -> list[tuple[int, int]]:
+        """Returns the output channels of this model when used as a backbone."""
+        # TODO: load this from a constant depending on the model size
+        return [(self.patch_size, 768)]

rslearn/models/copernicusfm_src/__init__.py

@@ -0,0 +1 @@
+# mypy: ignore-errors

rslearn/models/copernicusfm_src/aurora/area.py

@@ -0,0 +1,50 @@
+"""Copyright (c) Microsoft Corporation. Licensed under the MIT license."""
+
+import torch
+
+__all__ = ["area", "radius_earth"]
+
+
+# float: Radius of the earth in kilometers.
+radius_earth = 6378137 / 1000
+
+
+def area(polygon: torch.Tensor) -> torch.Tensor:
+    """Compute the area of a polygon specified by latitudes and longitudes in degrees.
+
+    This function is a PyTorch port of the PyPI package `area`. In particular, it is heavily
+    inspired by the following file:
+
+        https://github.com/scisco/area/blob/9d9549d6ebffcbe4bffe11b71efa2d406d1c9fe9/area/__init__.py
+
+    Args:
+        polygon (:class:`torch.Tensor`): Polygon of the shape `(*b, n, 2)` where `b` is an optional
+            multidimensional batch size, `n` is the number of points of the polygon, and 2
+            concatenates first latitudes and then longitudes. The polygon does not have be closed.
+
+    Returns:
+        :class:`torch.Tensor`: Area in square kilometers.
+    """
+    # Be sure to close the loop.
+    polygon = torch.cat((polygon, polygon[..., -1:, :]), axis=-2)
+
+    area = torch.zeros(polygon.shape[:-2], dtype=polygon.dtype, device=polygon.device)
+    n = polygon.shape[-2]  # Number of points of the polygon
+
+    rad = torch.deg2rad  # Convert degrees to radians.
+
+    if n > 2:
+        for i in range(n):
+            i_lower = i
+            i_middle = (i + 1) % n
+            i_upper = (i + 2) % n
+
+            lon_lower = polygon[..., i_lower, 1]
+            lat_middle = polygon[..., i_middle, 0]
+            lon_upper = polygon[..., i_upper, 1]
+
+            area = area + (rad(lon_upper) - rad(lon_lower)) * torch.sin(rad(lat_middle))
+
+    area = area * radius_earth * radius_earth / 2
+
+    return torch.abs(area)

rslearn/models/copernicusfm_src/aurora/fourier.py

@@ -0,0 +1,134 @@
+# type: ignore
+"""Copyright (c) Microsoft Corporation. Licensed under the MIT license."""
+
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from .area import area, radius_earth
+
+__all__ = [
+    "FourierExpansion",
+    "pos_expansion",
+    "scale_expansion",
+    "lead_time_expansion",
+    "levels_expansion",
+    "absolute_time_expansion",
+]
+
+
+class FourierExpansion(nn.Module):
+    """A Fourier series-style expansion into a high-dimensional space.
+
+    Attributes:
+        lower (float): Lower wavelength.
+        upper (float): Upper wavelength.
+        assert_range (bool): Assert that the encoded tensor is within the specified wavelength
+            range.
+    """
+
+    def __init__(self, lower: float, upper: float, assert_range: bool = True) -> None:
+        """Initialise.
+
+        Args:
+            lower (float): Lower wavelength.
+            upper (float): Upper wavelength.
+            assert_range (bool, optional): Assert that the encoded tensor is within the specified
+                wavelength range. Defaults to `True`.
+        """
+        super().__init__()
+        self.lower = lower
+        self.upper = upper
+        self.assert_range = assert_range
+
+    def forward(self, x: torch.Tensor, d: int) -> torch.Tensor:
+        """Perform the expansion.
+
+        Adds a dimension of length `d` to the end of the shape of `x`.
+
+        Args:
+            x (:class:`torch.Tensor`): Input to expand of shape `(..., n)`. All elements of `x` must
+                lie within `[self.lower, self.upper]` if `self.assert_range` is `True`.
+            d (int): Dimensionality. Must be a multiple of two.
+
+        Raises:
+            AssertionError: If `self.assert_range` is `True` and not all elements of `x` are not
+                within `[self.lower, self.upper]`.
+            ValueError: If `d` is not a multiple of two.
+
+        Returns:
+            torch.Tensor: Fourier series-style expansion of `x` of shape `(..., n, d)`.
+        """
+        # If the input is not within the configured range, the embedding might be ambiguous!
+        in_range = torch.logical_and(
+            self.lower <= x.abs(), torch.all(x.abs() <= self.upper)
+        )
+        in_range_or_zero = torch.all(
+            torch.logical_or(in_range, x == 0)
+        )  # Allow zeros to pass through.
+        if self.assert_range and not in_range_or_zero:
+            raise AssertionError(
+                f"The input tensor is not within the configured range"
+                f" `[{self.lower}, {self.upper}]`."
+            )
+
+        # We will use half of the dimensionality for `sin` and the other half for `cos`.
+        if not (d % 2 == 0):
+            raise ValueError("The dimensionality must be a multiple of two.")
+
+        # Always perform the expansion with `float64`s to avoid numerical accuracy shenanigans.
+        x = x.double()
+
+        wavelengths = torch.logspace(
+            math.log10(self.lower),
+            math.log10(self.upper),
+            d // 2,
+            base=10,
+            device=x.device,
+            dtype=x.dtype,
+        )
+        prod = torch.einsum("...i,j->...ij", x, 2 * np.pi / wavelengths)
+        encoding = torch.cat((torch.sin(prod), torch.cos(prod)), dim=-1)
+
+        return encoding.float()  # Cast to `float32` to avoid incompatibilities.
+
+
+# Determine a reasonable smallest value for the scale embedding by assuming a smallest delta in
+# latitudes and longitudes.
+_delta = 0.01  # Reasonable smallest delta in latitude and longitude
+_min_patch_area: float = area(
+    torch.tensor(
+        [
+            # The smallest patches will be at the poles. Just use the north pole.
+            [90, 0],
+            [90, _delta],
+            [90 - _delta, _delta],
+            [90 - _delta, 0],
+        ],
+        dtype=torch.float64,
+    )
+).item()
+_area_earth = 4 * np.pi * radius_earth * radius_earth
+
+pos_expansion = FourierExpansion(_delta, 720)
+
+
+scale_expansion = FourierExpansion(_min_patch_area, _area_earth)
+
+
+lead_time_expansion = FourierExpansion(1 / 60, 24 * 7 * 3)
+
+levels_expansion = FourierExpansion(0.01, 1e5)
+
+absolute_time_expansion = FourierExpansion(1, 24 * 365.25, assert_range=False)
+
+### new for SSL4EO-S ###
+# min wavelength: ultraviolet light (100 nm)
+# max wavelength: radio waves (1 m)
+spectrum_central_expansion = FourierExpansion(1e-7, 1)
+
+# min bandwidth: 10nm
+# max bandwidth: 1m
+spectrum_width_expansion = FourierExpansion(1e-7, 1)