rslearn 0.0.11__py3-none-any.whl → 0.0.12__py3-none-any.whl

rslearn/models/copernicusfm.py

@@ -1,228 +0,0 @@
-"""Copernicus FM model."""
-
-import logging
-import math
-from enum import Enum
-from pathlib import Path
-
-import torch
-import torch.nn.functional as F
-from einops import rearrange
-from huggingface_hub import hf_hub_download
-
-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],
-    },
-}
-
-HF_REPO_ID = "wangyi111/Copernicus-FM"
-HF_REPO_REVISION = "e1db406d517a122c8373802e1c130c5fc4789f84"
-HF_FILENAME = "CopernicusFM_ViT_base_varlang_e100.pth"
-
-
-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]],
-        cache_dir: str | Path | None = None,
-    ) -> None:
-        """Initialize the Copernicus FM wrapper.
-
-        Args:
-            band_order: The band order for each modality that will be used. The bands
-                can be provided in any order, and any subset can be used.
-            cache_dir: The directory to cache the weights. If None, a default directory
-                managed by huggingface_hub is used. The weights are downloaded from
-                Hugging Face (https://huggingface.co/wangyi111/Copernicus-FM).
-        """
-        super().__init__()
-
-        # Make sure all keys in band_order are in supported_modalities.
-        for modality_name in band_order.keys():
-            if modality_name in self.supported_modalities:
-                continue
-            raise ValueError(
-                f"band_order contains unsupported modality {modality_name}"
-            )
-
-        # global_pool=True so that we initialize the fc_norm layer
-        self.model = vit_base_patch16(num_classes=10, global_pool=True)
-
-        # Load weights, downloading if needed.
-        local_fname = hf_hub_download(
-            repo_id=HF_REPO_ID,
-            revision=HF_REPO_REVISION,
-            filename=HF_FILENAME,
-            local_dir=cache_dir,
-        )  # nosec
-        state_dict = torch.load(local_fname, weights_only=True)
-        self.model.load_state_dict(state_dict, strict=False)
-
-        # take MODALITY_TO_WAVELENGTH_BANDWIDTHS and rearrange it so that it has the same
-        # ordering as the user-provided band order.
-        self.modality_to_wavelength_bandwidths = {}
-        for modality in self.supported_modalities:
-            if modality not in band_order:
-                continue
-
-            wavelength_bandwidths = MODALITY_TO_WAVELENGTH_BANDWIDTHS[modality]
-            wavelengths = []
-            bandwidths = []
-            for b in band_order[modality]:
-                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

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

rslearn/models/copernicusfm_src/aurora/area.py

@@ -1,50 +0,0 @@
-"""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

@@ -1,134 +0,0 @@
-# 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)