rslearn 0.0.7__py3-none-any.whl → 0.0.9__py3-none-any.whl

rslearn/dataset/materialize.py

@@ -1,10 +1,10 @@
 """Classes to implement dataset materialization."""
 
+from collections.abc import Callable
 from typing import Any, Generic, TypeVar
 
 import numpy as np
 import numpy.typing as npt
-from class_registry import ClassRegistry
 from rasterio.enums import Resampling
 
 from rslearn.config import (
@@ -25,7 +25,26 @@ from rslearn.utils.vector_format import load_vector_format
 from .remap import Remapper, load_remapper
 from .window import Window
 
-Materializers = ClassRegistry()
+_MaterializerT = TypeVar("_MaterializerT", bound="Materializer")
+
+
+class _MaterializerRegistry(dict[str, type["Materializer"]]):
+    """Registry for Materializer classes."""
+
+    def register(
+        self, name: str
+    ) -> Callable[[type[_MaterializerT]], type[_MaterializerT]]:
+        """Decorator to register a materializer class."""
+
+        def decorator(cls: type[_MaterializerT]) -> type[_MaterializerT]:
+            self[name] = cls
+            return cls
+
+        return decorator
+
+
+Materializers = _MaterializerRegistry()
+
 
 LayerConfigType = TypeVar("LayerConfigType", bound=LayerConfig)
 

rslearn/dataset/remap.py

@@ -1,18 +1,42 @@
 """Classes to remap raster values."""
 
-from typing import Any
+from collections.abc import Callable
+from typing import Any, TypeVar
 
 import numpy as np
 import numpy.typing as npt
-from class_registry import ClassRegistry
 
-Remappers = ClassRegistry()
+_RemapperT = TypeVar("_RemapperT", bound="Remapper")
+
+
+class _RemapperRegistry(dict[str, type["Remapper"]]):
+    """Registry for Remapper classes."""
+
+    def register(self, name: str) -> Callable[[type[_RemapperT]], type[_RemapperT]]:
+        """Decorator to register a remapper class."""
+
+        def decorator(cls: type[_RemapperT]) -> type[_RemapperT]:
+            self[name] = cls
+            return cls
+
+        return decorator
+
+
+Remappers = _RemapperRegistry()
 """Registry of Remapper implementations."""
 
 
 class Remapper:
     """An abstract class that remaps pixel values based on layer configuration."""
 
+    def __init__(self, config: dict[str, Any]) -> None:
+        """Initialize a Remapper.
+
+        Args:
+            config: the config dict for this remapper.
+        """
+        pass
+
     def __call__(
         self, array: npt.NDArray[Any], dtype: npt.DTypeLike
     ) -> npt.NDArray[Any]:
@@ -67,4 +91,5 @@ class LinearRemapper(Remapper):
 
 def load_remapper(config: dict[str, Any]) -> Remapper:
     """Load a remapper from a configuration dictionary."""
-    return Remappers.get(config["name"], config=config)
+    cls = Remappers[config["name"]]
+    return cls(config)

rslearn/main.py

@@ -4,6 +4,7 @@ import argparse
 import multiprocessing
 import random
 import sys
+import time
 from collections.abc import Callable
 from datetime import UTC, datetime, timedelta
 from typing import Any, TypeVar
@@ -19,8 +20,18 @@ from rslearn.const import WGS84_EPSG
 from rslearn.data_sources import Item, data_source_from_config
 from rslearn.dataset import Dataset, Window, WindowLayerData
 from rslearn.dataset.add_windows import add_windows_from_box, add_windows_from_file
+from rslearn.dataset.handler_summaries import (
+    ErrorOutcome,
+    IngestCounts,
+    IngestDatasetJobsSummary,
+    LayerIngestSummary,
+    MaterializeDatasetWindowsSummary,
+    PrepareDatasetWindowsSummary,
+    UnknownIngestCounts,
+)
 from rslearn.dataset.index import DatasetIndex
 from rslearn.dataset.manage import (
+    AttemptsCounter,
     materialize_dataset_windows,
     prepare_dataset_windows,
     retry,
@@ -287,7 +298,7 @@ def add_apply_on_windows_args(parser: argparse.ArgumentParser) -> None:
 
 
 def apply_on_windows(
-    f: Callable[[list[Window]], None],
+    f: Callable[[list[Window]], Any],
     dataset: Dataset,
     group: str | list[str] | None = None,
     names: list[str] | None = None,
@@ -367,7 +378,7 @@ def apply_on_windows(
         p.close()
 
 
-def apply_on_windows_args(f: Callable[..., None], args: argparse.Namespace) -> None:
+def apply_on_windows_args(f: Callable[..., Any], args: argparse.Namespace) -> None:
     """Call apply_on_windows with arguments passed via command-line interface."""
     dataset = Dataset(UPath(args.root), args.disabled_layers)
     apply_on_windows(
@@ -413,12 +424,12 @@ class PrepareHandler:
         """
         self.dataset = dataset
 
-    def __call__(self, windows: list[Window]) -> None:
+    def __call__(self, windows: list[Window]) -> PrepareDatasetWindowsSummary:
         """Prepares the windows from apply_on_windows."""
         logger.info(f"Running prepare on {len(windows)} windows")
         if self.dataset is None:
             raise ValueError("dataset not set")
-        prepare_dataset_windows(
+        return prepare_dataset_windows(
             self.dataset,
             windows,
             self.force,
@@ -502,14 +513,20 @@ class IngestHandler:
 
     def __call__(
         self, jobs: list[tuple[str, LayerConfig, Item, list[STGeometry]]]
-    ) -> None:
+    ) -> IngestDatasetJobsSummary:
         """Ingest the specified items.
 
         The items are computed from list of windows via IngestHandler.get_jobs.
 
         Args:
-            jobs: list of (layer_name, item, geometries) tuples to ingest.
+            jobs: list of (layer_name, layer_cfg, item, geometries) tuples to ingest.
+
+        Returns:
+            summary of the ingest jobs operation fit for telemetry purposes.
         """
+        start_time = time.monotonic()
+        layer_summaries: list[LayerIngestSummary] = []
+
         logger.info(f"Running ingest for {len(jobs)} jobs")
         import gc
 
@@ -533,6 +550,8 @@ class IngestHandler:
             layer_cfg = self.dataset.layers[layer_name]
             data_source = data_source_from_config(layer_cfg, self.dataset.path)
 
+            attempts_counter = AttemptsCounter()
+            ingest_counts: IngestCounts | UnknownIngestCounts
             try:
                 retry(
                     lambda: data_source.ingest(
@@ -544,18 +563,47 @@ class IngestHandler:
                     ),
                     retry_max_attempts=self.retry_max_attempts,
                     retry_backoff=self.retry_backoff,
+                    attempts_counter=attempts_counter,
+                )
+                ingest_counts = IngestCounts(
+                    items_ingested=len(items_and_geometries),
+                    geometries_ingested=sum(
+                        len(geometries) for _, geometries in items_and_geometries
+                    ),
                 )
             except Exception as e:
                 if not self.ignore_errors:
                     raise
 
+                ingest_counts = UnknownIngestCounts(
+                    items_attempted=len(items_and_geometries),
+                    geometries_attempted=sum(
+                        len(geometries) for _, geometries in items_and_geometries
+                    ),
+                )
                 logger.error(
                     "warning: got error while ingesting "
                     + f"{len(items_and_geometries)} items: {e}"
                 )
 
+            layer_summaries.append(
+                LayerIngestSummary(
+                    layer_name=layer_name,
+                    data_source_name=getattr(layer_cfg.data_source, "name", "N/A"),
+                    duration_seconds=time.monotonic() - start_time,
+                    ingest_counts=ingest_counts,
+                    ingest_attempts=attempts_counter.value,
+                )
+            )
+
         gc.collect()
 
+        return IngestDatasetJobsSummary(
+            duration_seconds=time.monotonic() - start_time,
+            num_jobs=len(jobs),
+            layer_summaries=layer_summaries,
+        )
+
     def _load_layer_data_for_windows(
         self, windows: list[Window], workers: int
     ) -> list[tuple[Window, dict[str, WindowLayerData]]]:
@@ -686,13 +734,16 @@ class MaterializeHandler:
         """
         self.dataset = dataset
 
-    def __call__(self, windows: list[Window]) -> None:
+    def __call__(
+        self, windows: list[Window]
+    ) -> MaterializeDatasetWindowsSummary | ErrorOutcome:
         """Materializes the windows from apply_on_windows."""
         logger.info(f"Running Materialize with {len(windows)} windows")
+        start_time = time.monotonic()
         if self.dataset is None:
             raise ValueError("dataset not set")
         try:
-            materialize_dataset_windows(
+            return materialize_dataset_windows(
                 self.dataset,
                 windows,
                 retry_max_attempts=self.retry_max_attempts,
@@ -703,6 +754,7 @@ class MaterializeHandler:
                 logger.error(f"Error materializing windows: {e}")
                 raise
             logger.warning(f"Ignoring error while materializing windows: {e}")
+            return ErrorOutcome(duration_seconds=time.monotonic() - start_time)
 
 
 @register_handler("dataset", "materialize")

rslearn/models/clay/clay.py

@@ -15,6 +15,7 @@ from huggingface_hub import hf_hub_download
 # from claymodel.module import ClayMAEModule
 from terratorch.models.backbones.clay_v15.module import ClayMAEModule
 
+from rslearn.train.transforms.normalize import Normalize
 from rslearn.train.transforms.transform import Transform
 
 
@@ -163,13 +164,36 @@ class Clay(torch.nn.Module):
 
 
 class ClayNormalize(Transform):
-    """Normalize inputs using Clay metadata."""
+    """Normalize inputs using Clay metadata.
+
+    For Sentinel-1, the intensities should be converted to decibels.
+    """
 
     def __init__(self, metadata_path: str = CLAY_METADATA_PATH) -> None:
         """Initialize ClayNormalize."""
         super().__init__()
         with open(metadata_path) as f:
-            self.metadata = yaml.safe_load(f)
+            metadata = yaml.safe_load(f)
+        normalizers = {}
+        for modality in CLAY_MODALITIES:
+            if modality not in metadata:
+                continue
+            modality_metadata = metadata[modality]
+            means = [
+                modality_metadata["bands"]["mean"][b]
+                for b in modality_metadata["band_order"]
+            ]
+            stds = [
+                modality_metadata["bands"]["std"][b]
+                for b in modality_metadata["band_order"]
+            ]
+            normalizers[modality] = Normalize(
+                mean=means,
+                std=stds,
+                selectors=[modality],
+                num_bands=len(means),
+            )
+        self.normalizers = torch.nn.ModuleDict(normalizers)
 
     def apply_image(
         self, image: torch.Tensor, means: list[float], stds: list[float]
@@ -188,17 +212,8 @@ class ClayNormalize(Transform):
         self, input_dict: dict[str, Any], target_dict: dict[str, Any]
     ) -> tuple[dict[str, Any], dict[str, Any]]:
         """Normalize the specified image with Clay normalization."""
-        for modality in CLAY_MODALITIES:
-            if modality not in input_dict or modality not in self.metadata:
+        for modality, normalizer in self.normalizers.items():
+            if modality not in input_dict:
                 continue
-            modality_metadata = self.metadata[modality]
-            means = [
-                modality_metadata["bands"]["mean"][b]
-                for b in modality_metadata["band_order"]
-            ]
-            stds = [
-                modality_metadata["bands"]["std"][b]
-                for b in modality_metadata["band_order"]
-            ]
-            input_dict[modality] = self.apply_image(input_dict[modality], means, stds)
+            input_dict, target_dict = normalizer(input_dict, target_dict)
         return input_dict, target_dict

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,166 @@
+"""DinoV3 model."""
+
+from enum import StrEnum
+from pathlib import Path
+from typing import Any
+
+import torch
+import torchvision
+from einops import rearrange
+
+from rslearn.train.transforms.normalize import Normalize
+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:
+            mean = [0.430, 0.411, 0.296]
+            std = [0.213, 0.156, 0.143]
+        else:
+            mean = [0.485, 0.456, 0.406]
+            std = [0.229, 0.224, 0.225]
+
+        self.normalize = Normalize(
+            [value * 255 for value in mean],
+            [value * 255 for value in std],
+            num_bands=3,
+        )
+
+    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
+        """
+        return self.normalize(input_dict, target_dict)