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

rslearn/models/copernicusfm.py

@@ -3,11 +3,12 @@
 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 upath import UPath
+from huggingface_hub import hf_hub_download
 
 from .copernicusfm_src.model_vit import vit_base_patch16
 
@@ -64,6 +65,10 @@ MODALITY_TO_WAVELENGTH_BANDWIDTHS: dict[str, dict[str, list]] = {
     },
 }
 
+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."""
@@ -80,44 +85,51 @@ class CopernicusFM(torch.nn.Module):
     def __init__(
         self,
         band_order: dict[str, list[str]],
-        load_directory: str | None,
+        cache_dir: str | Path | None = 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
+            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.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
+
+        # 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 = []
-            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:
+            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])

rslearn/models/dinov3.py

@@ -0,0 +1,165 @@
+"""DinoV3 model."""
+
+from enum import StrEnum
+from pathlib import Path
+from typing import Any
+
+import torch
+import torchvision
+from einops import rearrange
+from torchvision.transforms import v2
+
+from rslearn.train.transforms.transform import Transform
+
+
+class DinoV3Models(StrEnum):
+    """Names for different DinoV3 images on torch hub."""
+
+    SMALL_WEB = "dinov3_vits16"
+    SMALL_PLUS_WEB = "dinov3_vits16plus"
+    BASE_WEB = "dinov3_vitb16"
+    LARGE_WEB = "dinov3_vitl16"
+    HUGE_PLUS_WEB = "dinov3_vith16plus"
+    FULL_7B_WEB = "dinov3_vit7b16"
+    LARGE_SATELLITE = "dinov3_vitl16_sat"
+    FULL_7B_SATELLITE = "dinov3_vit7b16_sat"
+
+
+DINOV3_PTHS: dict[str, str] = {
+    DinoV3Models.LARGE_SATELLITE: "dinov3_vitl16_pretrain_sat493m-eadcf0ff.pth",
+    DinoV3Models.FULL_7B_SATELLITE: "dinov3_vit7b16_pretrain_sat493m-a6675841.pth",
+    DinoV3Models.BASE_WEB: "dinov3_vitb16_pretrain_lvd1689m-73cec8be.pth",
+    DinoV3Models.LARGE_WEB: "dinov3_vitl16_pretrain_lvd1689m-8aa4cbdd.pth",
+    DinoV3Models.HUGE_PLUS_WEB: "dinov3_vith16plus_pretrain_lvd1689m-7c1da9a5.pth",
+    DinoV3Models.FULL_7B_WEB: "dinov3_vit7b16_pretrain_lvd1689m-a955f4.pth",
+}
+
+
+class DinoV3(torch.nn.Module):
+    """DinoV3 Backbones.
+
+    Must have the pretrained weights downloaded in checkpoint_dir for them to be loaded.
+    See https://github.com/facebookresearch/dinov3?tab=readme-ov-file#pretrained-models
+
+    Only takes RGB as input. Expects normalized data (use the below normalizer).
+
+    Uses patch size 16. The input is resized to 256x256; when applying DinoV3 on
+    segmentation or detection tasks with inputs larger than 256x256, it may be best to
+    train and predict on 256x256 crops (using SplitConfig.patch_size argument).
+    """
+
+    image_size: int = 256
+    patch_size: int = 16
+    output_dim: int = 1024
+
+    def _load_model(self, size: str, checkpoint_dir: str | None) -> torch.nn.Module:
+        model_name = size.replace("_sat", "")
+        if checkpoint_dir is not None:
+            weights = str(Path(checkpoint_dir) / DINOV3_PTHS[size])
+            return torch.hub.load(
+                "facebookresearch/dinov3",
+                model_name,
+                weights=weights,
+            )  # nosec
+        return torch.hub.load("facebookresearch/dinov3", model_name, pretrained=False)  # nosec
+
+    def __init__(
+        self,
+        checkpoint_dir: str | None,
+        size: str = DinoV3Models.LARGE_SATELLITE,
+        use_cls_token: bool = False,
+        do_resizing: bool = True,
+    ) -> None:
+        """Instantiate a new DinoV3 instance.
+
+        Args:
+            checkpoint_dir: the local path to the pretrained weight dir. If None, we load the architecture
+                only (randomly initialized).
+            size: the model size, see class for various models.
+            use_cls_token: use pooled class token (for classification), otherwise returns spatial feature map.
+            do_resizing: whether to resize inputs to 256x256. Default true.
+        """
+        super().__init__()
+        self.size = size
+        self.checkpoint_dir = checkpoint_dir
+        self.use_cls_token = use_cls_token
+        self.do_resizing = do_resizing
+        self.model = self._load_model(size, checkpoint_dir)
+
+    def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
+        """Forward pass for the dinov3 model.
+
+        Args:
+            inputs: input dicts that must include "image" key.
+
+        Returns:
+            List[torch.Tensor]: Single-scale feature tensors from the encoder.
+        """
+        cur = torch.stack([inp["image"] for inp in inputs], dim=0)  # (B, C, H, W)
+
+        if self.do_resizing and (
+            cur.shape[2] != self.image_size or cur.shape[3] != self.image_size
+        ):
+            cur = torchvision.transforms.functional.resize(
+                cur,
+                [self.image_size, self.image_size],
+            )
+
+        if self.use_cls_token:
+            features = self.model(cur)
+        else:
+            features = self.model.forward_features(cur)["x_norm_patchtokens"]
+            batch_size, num_patches, _ = features.shape
+            height, width = int(num_patches**0.5), int(num_patches**0.5)
+            features = rearrange(features, "b (h w) d -> b d h w", h=height, w=width)
+
+        return [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)]
+
+
+class DinoV3Normalize(Transform):
+    """Normalize inputs using DinoV3 normalization.
+
+    Normalize "image" key in input according to Dino statistics from pretraining. Satellite pretraining has slightly different normalizing than the base image model so set 'satellite' depending on what pretrained model you are using.
+
+    Input "image" should be RGB-like image between 0-255.
+    """
+
+    def __init__(self, satellite: bool = True):
+        """Initialize a new DinoV3Normalize."""
+        super().__init__()
+        self.satellite = satellite
+        if satellite:
+            self.normalize = v2.Normalize(
+                mean=(0.430, 0.411, 0.296),
+                std=(0.213, 0.156, 0.143),
+            )
+        else:
+            self.normalize = v2.Normalize(
+                mean=(0.485, 0.456, 0.406),
+                std=(0.229, 0.224, 0.225),
+            )
+
+    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 DinoV3 normalization.
+
+        Args:
+            input_dict: the input dictionary.
+            target_dict: the target dictionary.
+
+        Returns:
+            normalized (input_dicts, target_dicts) tuple
+        """
+        input_dict["image"] = self.normalize(input_dict["image"] / 255.0)
+        return input_dict, target_dict

rslearn/models/galileo/__init__.py

@@ -0,0 +1,5 @@
+"""Galileo model."""
+
+from .galileo import GalileoModel, GalileoSize
+
+__all__ = ["GalileoModel", "GalileoSize"]