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

rslearn/data_sources/local_files.py

@@ -2,12 +2,12 @@
 
 import functools
 import json
+from collections.abc import Callable
 from typing import Any, Generic, TypeVar
 
 import fiona
 import shapely
 import shapely.geometry
-from class_registry import ClassRegistry
 from rasterio.crs import CRS
 from upath import UPath
 
@@ -23,7 +23,24 @@ from rslearn.utils.geometry import Projection, STGeometry, get_global_geometry
 from .data_source import DataSource, Item, QueryConfig
 
 logger = get_logger("__name__")
-Importers = ClassRegistry()
+_ImporterT = TypeVar("_ImporterT", bound="Importer")
+
+
+class _ImporterRegistry(dict[str, type["Importer"]]):
+    """Registry for Importer classes."""
+
+    def register(self, name: str) -> Callable[[type[_ImporterT]], type[_ImporterT]]:
+        """Decorator to register an importer class."""
+
+        def decorator(cls: type[_ImporterT]) -> type[_ImporterT]:
+            self[name] = cls
+            return cls
+
+        return decorator
+
+
+Importers = _ImporterRegistry()
+
 
 ItemType = TypeVar("ItemType", bound=Item)
 LayerConfigType = TypeVar("LayerConfigType", bound=LayerConfig)
@@ -425,7 +442,7 @@ class LocalFiles(DataSource):
         """
         self.config = config
 
-        self.importer = Importers[config.layer_type.value]
+        self.importer = Importers[config.layer_type.value]()
         self.src_dir = src_dir
 
     @staticmethod

rslearn/data_sources/planetary_computer.py

@@ -83,6 +83,10 @@ class PlanetaryComputer(DataSource, TileStore):
 
     STAC_ENDPOINT = "https://planetarycomputer.microsoft.com/api/stac/v1"
 
+    # Default threshold for recreating the STAC client to prevent memory leaks
+    # from the pystac Catalog's resolved objects cache growing unbounded
+    DEFAULT_MAX_ITEMS_PER_CLIENT = 1000
+
     def __init__(
         self,
         collection_name: str,
@@ -93,6 +97,7 @@ class PlanetaryComputer(DataSource, TileStore):
         timeout: timedelta = timedelta(seconds=10),
         skip_items_missing_assets: bool = False,
         cache_dir: UPath | None = None,
+        max_items_per_client: int | None = None,
     ):
         """Initialize a new PlanetaryComputer instance.
 
@@ -109,6 +114,9 @@ class PlanetaryComputer(DataSource, TileStore):
             cache_dir: optional directory to cache items by name, including asset URLs.
                 If not set, there will be no cache and instead STAC requests will be
                 needed each time.
+            max_items_per_client: number of STAC items to process before recreating
+                the client to prevent memory leaks from the resolved objects cache.
+                Defaults to DEFAULT_MAX_ITEMS_PER_CLIENT.
         """
         self.collection_name = collection_name
         self.asset_bands = asset_bands
@@ -118,12 +126,15 @@ class PlanetaryComputer(DataSource, TileStore):
         self.timeout = timeout
         self.skip_items_missing_assets = skip_items_missing_assets
         self.cache_dir = cache_dir
+        self.max_items_per_client = (
+            max_items_per_client or self.DEFAULT_MAX_ITEMS_PER_CLIENT
+        )
 
         if self.cache_dir is not None:
             self.cache_dir.mkdir(parents=True, exist_ok=True)
 
         self.client: pystac_client.Client | None = None
-        self.collection: pystac_client.CollectionClient | None = None
+        self._client_item_count = 0
 
     @staticmethod
     def from_config(config: RasterLayerConfig, ds_path: UPath) -> "PlanetaryComputer":
@@ -142,7 +153,12 @@ class PlanetaryComputer(DataSource, TileStore):
         if "cache_dir" in d:
             kwargs["cache_dir"] = join_upath(ds_path, d["cache_dir"])
 
-        simple_optionals = ["query", "sort_by", "sort_ascending"]
+        simple_optionals = [
+            "query",
+            "sort_by",
+            "sort_ascending",
+            "max_items_per_client",
+        ]
         for k in simple_optionals:
             if k in d:
                 kwargs[k] = d[k]
@@ -151,20 +167,40 @@ class PlanetaryComputer(DataSource, TileStore):
 
     def _load_client(
         self,
-    ) -> tuple[pystac_client.Client, pystac_client.CollectionClient]:
+    ) -> pystac_client.Client:
         """Lazily load pystac client.
 
         We don't load it when creating the data source because it takes time and caller
         may not be calling get_items. Additionally, loading it during the get_items
         call enables leveraging the retry loop functionality in
         prepare_dataset_windows.
-        """
-        if self.client is not None:
-            return self.client, self.collection
 
+        Note: We periodically recreate the client to prevent memory leaks from the
+        pystac Catalog's resolved objects cache, which grows unbounded as STAC items
+        are deserialized and cached. The cache cannot be cleared or disabled.
+        """
+        if self.client is None:
+            logger.info("Creating initial STAC client")
+            self.client = pystac_client.Client.open(self.STAC_ENDPOINT)
+            return self.client
+
+        if self._client_item_count < self.max_items_per_client:
+            return self.client
+
+        # Recreate client to clear the resolved objects cache
+        current_client = self.client
+        logger.debug(
+            "Recreating STAC client after processing %d items (threshold: %d)",
+            self._client_item_count,
+            self.max_items_per_client,
+        )
+        client_root = current_client.get_root()
+        client_root.clear_links()
+        client_root.clear_items()
+        client_root.clear_children()
+        self._client_item_count = 0
         self.client = pystac_client.Client.open(self.STAC_ENDPOINT)
-        self.collection = self.client.get_collection(self.collection_name)
-        return self.client, self.collection
+        return self.client
 
     def _stac_item_to_item(self, stac_item: pystac.Item) -> PlanetaryComputerItem:
         shp = shapely.geometry.shape(stac_item.geometry)
@@ -210,10 +246,26 @@ class PlanetaryComputer(DataSource, TileStore):
 
         # No cache or not in cache, so we need to make the STAC request.
         logger.debug("Getting STAC item {name}")
-        _, collection = self._load_client()
-        stac_item = collection.get_item(name)
+        client = self._load_client()
+
+        search_result = client.search(ids=[name], collections=[self.collection_name])
+        stac_items = list(search_result.items())
+
+        if not stac_items:
+            raise ValueError(
+                f"Item {name} not found in collection {self.collection_name}"
+            )
+        if len(stac_items) > 1:
+            raise ValueError(
+                f"Multiple items found for ID {name} in collection {self.collection_name}"
+            )
+
+        stac_item = stac_items[0]
         item = self._stac_item_to_item(stac_item)
 
+        # Track items processed for client recreation threshold (after deserialization)
+        self._client_item_count += 1
+
         # Finally we cache it if cache_dir is set.
         if cache_fname is not None:
             with cache_fname.open("w") as f:
@@ -233,7 +285,7 @@ class PlanetaryComputer(DataSource, TileStore):
         Returns:
             List of groups of items that should be retrieved for each geometry.
         """
-        client, _ = self._load_client()
+        client = self._load_client()
 
         groups = []
         for geometry in geometries:
@@ -247,7 +299,9 @@ class PlanetaryComputer(DataSource, TileStore):
                 datetime=wgs84_geometry.time_range,
                 query=self.query,
             )
-            stac_items = [item for item in result.item_collection()]
+            stac_items = [item for item in result.items()]
+            # Track items processed for client recreation threshold (after deserialization)
+            self._client_item_count += len(stac_items)
             logger.debug("STAC search yielded %d items", len(stac_items))
 
             if self.skip_items_missing_assets:
@@ -580,7 +634,13 @@ class Sentinel2(PlanetaryComputer):
         if "cache_dir" in d:
             kwargs["cache_dir"] = join_upath(ds_path, d["cache_dir"])
 
-        simple_optionals = ["harmonize", "query", "sort_by", "sort_ascending"]
+        simple_optionals = [
+            "harmonize",
+            "query",
+            "sort_by",
+            "sort_ascending",
+            "max_items_per_client",
+        ]
         for k in simple_optionals:
             if k in d:
                 kwargs[k] = d[k]
@@ -756,7 +816,12 @@ class Sentinel1(PlanetaryComputer):
         if "cache_dir" in d:
             kwargs["cache_dir"] = join_upath(ds_path, d["cache_dir"])
 
-        simple_optionals = ["query", "sort_by", "sort_ascending"]
+        simple_optionals = [
+            "query",
+            "sort_by",
+            "sort_ascending",
+            "max_items_per_client",
+        ]
         for k in simple_optionals:
             if k in d:
                 kwargs[k] = d[k]

rslearn/dataset/manage.py

@@ -396,9 +396,9 @@ def materialize_window(
         )
 
         if dataset.materializer_name:
-            materializer = Materializers[dataset.materializer_name]
+            materializer = Materializers[dataset.materializer_name]()
         else:
-            materializer = Materializers[layer_cfg.layer_type.value]
+            materializer = Materializers[layer_cfg.layer_type.value]()
 
         retry(
             fn=lambda: materializer.materialize(

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/models/dinov3.py

@@ -7,8 +7,8 @@ from typing import Any
 import torch
 import torchvision
 from einops import rearrange
-from torchvision.transforms import v2
 
+from rslearn.train.transforms.normalize import Normalize
 from rslearn.train.transforms.transform import Transform
 
 
@@ -139,15 +139,17 @@ class DinoV3Normalize(Transform):
         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),
-            )
+            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),
-            )
+            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]
@@ -161,5 +163,4 @@ class DinoV3Normalize(Transform):
         Returns:
             normalized (input_dicts, target_dicts) tuple
         """
-        input_dict["image"] = self.normalize(input_dict["image"] / 255.0)
-        return input_dict, target_dict
+        return self.normalize(input_dict, target_dict)

rslearn/models/galileo/galileo.py

@@ -2,6 +2,7 @@
 
 import math
 import tempfile
+from contextlib import nullcontext
 from enum import StrEnum
 from typing import Any, cast
 
@@ -63,6 +64,11 @@ pretrained_weights: dict[GalileoSize, str] = {
 
 DEFAULT_NORMALIZER = Normalizer()
 
+AUTOCAST_DTYPE_MAP = {
+    "bfloat16": torch.bfloat16,
+    "float32": torch.float32,
+}
+
 
 class GalileoModel(nn.Module):
     """Galileo backbones."""
@@ -85,6 +91,7 @@ class GalileoModel(nn.Module):
         size: GalileoSize,
         patch_size: int = 4,
         pretrained_path: str | UPath | None = None,
+        autocast_dtype: str | None = "bfloat16",
     ) -> None:
         """Initialize the Galileo model.
 
@@ -93,6 +100,7 @@ class GalileoModel(nn.Module):
             patch_size: The patch size to use.
             pretrained_path: the local path to the pretrained weights. Otherwise it is
                 downloaded and cached in temp directory.
+            autocast_dtype: which dtype to use for autocasting, or set None to disable.
         """
         super().__init__()
         if pretrained_path is None:
@@ -128,8 +136,14 @@ class GalileoModel(nn.Module):
             idx for idx, key in enumerate(SPACE_TIME_BANDS_GROUPS_IDX) if "S1" in key
         ]
 
+        self.size = size
         self.patch_size = patch_size
 
+        if autocast_dtype is not None:
+            self.autocast_dtype = AUTOCAST_DTYPE_MAP[autocast_dtype]
+        else:
+            self.autocast_dtype = None
+
     @staticmethod
     def to_cartesian(
         lat: float | np.ndarray | torch.Tensor, lon: float | np.ndarray | torch.Tensor
@@ -484,18 +498,31 @@ class GalileoModel(nn.Module):
             patch_size = h
         else:
             patch_size = self.patch_size
-        outputs = self.model(
-            s_t_x=galileo_input.s_t_x,
-            s_t_m=galileo_input.s_t_m,
-            sp_x=galileo_input.sp_x,
-            sp_m=galileo_input.sp_m,
-            t_x=galileo_input.t_x,
-            t_m=galileo_input.t_m,
-            st_x=galileo_input.st_x,
-            st_m=galileo_input.st_m,
-            months=galileo_input.months,
-            patch_size=patch_size,
-        )
+
+        # Decide context based on self.autocast_dtype.
+        device = galileo_input.s_t_x.device
+        if self.autocast_dtype is None:
+            context = nullcontext()
+        else:
+            assert device is not None
+            context = torch.amp.autocast(
+                device_type=device.type, dtype=self.autocast_dtype
+            )
+
+        with context:
+            outputs = self.model(
+                s_t_x=galileo_input.s_t_x,
+                s_t_m=galileo_input.s_t_m,
+                sp_x=galileo_input.sp_x,
+                sp_m=galileo_input.sp_m,
+                t_x=galileo_input.t_x,
+                t_m=galileo_input.t_m,
+                st_x=galileo_input.st_x,
+                st_m=galileo_input.st_m,
+                months=galileo_input.months,
+                patch_size=patch_size,
+            )
+
         if h == patch_size:
             # only one spatial patch, so we can just take an average
             # of all the tokens to output b c_g 1 1
@@ -515,3 +542,22 @@ class GalileoModel(nn.Module):
                     "b h w c_g d -> b c_g d h w",
                 ).mean(dim=1)
             ]
+
+    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 (patch_size, depth) that corresponds
+        to the feature maps that the backbone returns.
+
+        Returns:
+            the output channels of the backbone as a list of (patch_size, depth) tuples.
+        """
+        if self.size == GalileoSize.BASE:
+            depth = 768
+        elif self.model_size == GalileoSize.TINY:
+            depth = 192
+        elif self.model_size == GalileoSize.NANO:
+            depth = 128
+        else:
+            raise ValueError(f"Invalid model size: {self.size}")
+        return [(self.patch_size, depth)]

rslearn/models/galileo/single_file_galileo.py

@@ -1469,7 +1469,13 @@ class Encoder(GalileoBase):
             # we take the inverse of the mask because a value
             # of True indicates the value *should* take part in
             # attention
-            x = blk(x=x, y=None, attn_mask=~new_m.bool())
+            temp_mask = ~new_m.bool()
+            if temp_mask.all():
+                # if all the tokens are used in attention we can pass a None mask
+                # to the attention block
+                temp_mask = None
+
+            x = blk(x=x, y=None, attn_mask=temp_mask)
 
         if exit_ids_seq is not None:
             assert exited_tokens is not None

rslearn/models/olmoearth_pretrain/__init__.py

@@ -0,0 +1 @@
+"""OlmoEarth model architecture."""

rslearn/models/olmoearth_pretrain/model.py

@@ -0,0 +1,203 @@
+"""OlmoEarth model wrapper for fine-tuning in rslearn."""
+
+import json
+from contextlib import nullcontext
+from typing import Any
+
+import torch
+from einops import rearrange
+from olmo_core.config import Config
+from olmo_core.distributed.checkpoint import load_model_and_optim_state
+from olmoearth_pretrain.data.constants import Modality
+from olmoearth_pretrain.nn.flexihelios import Encoder, TokensAndMasks
+from olmoearth_pretrain.train.masking import MaskedOlmoEarthSample, MaskValue
+from upath import UPath
+
+from rslearn.log_utils import get_logger
+
+logger = get_logger(__name__)
+
+MODALITY_NAMES = [
+    "sentinel2_l2a",
+    "sentinel1",
+    "worldcover",
+    "openstreetmap_raster",
+    "landsat",
+]
+
+AUTOCAST_DTYPE_MAP = {
+    "bfloat16": torch.bfloat16,
+    "float16": torch.float16,
+    "float32": torch.float32,
+}
+
+
+class OlmoEarth(torch.nn.Module):
+    """A wrapper to support the OlmoEarth model."""
+
+    def __init__(
+        self,
+        # TODO: we should accept model ID instead of checkpoint_path once we are closer
+        # to being ready for release.
+        checkpoint_path: str,
+        selector: list[str | int] = [],
+        forward_kwargs: dict[str, Any] = {},
+        random_initialization: bool = False,
+        embedding_size: int | None = None,
+        patch_size: int | None = None,
+        autocast_dtype: str | None = "bfloat16",
+    ):
+        """Create a new OlmoEarth model.
+
+        Args:
+            checkpoint_path: the checkpoint directory to load. It should contain
+                config.json file as well as model_and_optim folder.
+            selector: an optional sequence of attribute names or list indices to select
+                the sub-module that should be applied on the input images.
+            forward_kwargs: additional arguments to pass to forward pass besides the
+                 MaskedOlmoEarthSample.
+            random_initialization: whether to skip loading the checkpoint so the
+                weights are randomly initialized. In this case, the checkpoint is only
+                used to define the model architecture.
+            embedding_size: optional embedding size to report via
+                get_backbone_channels.
+            patch_size: optional patch size to report via get_backbone_channels.
+            autocast_dtype: which dtype to use for autocasting, or set None to disable.
+        """
+        super().__init__()
+        _checkpoint_path = UPath(checkpoint_path)
+        self.forward_kwargs = forward_kwargs
+        self.embedding_size = embedding_size
+        self.patch_size = patch_size
+
+        if autocast_dtype is not None:
+            self.autocast_dtype = AUTOCAST_DTYPE_MAP[autocast_dtype]
+        else:
+            self.autocast_dtype = None
+
+        # Load the model config and initialize it.
+        # We avoid loading the train module here because it depends on running within
+        # olmo_core.
+        with (_checkpoint_path / "config.json").open() as f:
+            config_dict = json.load(f)
+            model_config = Config.from_dict(config_dict["model"])
+
+        model = model_config.build()
+
+        # Load the checkpoint.
+        if not random_initialization:
+            train_module_dir = _checkpoint_path / "model_and_optim"
+            if train_module_dir.exists():
+                load_model_and_optim_state(str(train_module_dir), model)
+                logger.info(f"loaded OlmoEarth encoder from {train_module_dir}")
+            else:
+                logger.info(f"could not find OlmoEarth encoder at {train_module_dir}")
+        else:
+            logger.info("skipping loading OlmoEarth encoder")
+
+        # Select just the portion of the model that we actually want to use.
+        for part in selector:
+            if isinstance(part, str):
+                model = getattr(model, part)
+            else:
+                model = model[part]
+        self.model = model
+
+    def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
+        """Compute feature maps from the OlmoEarth backbone.
+
+        Inputs:
+            inputs: input dicts. It should include keys corresponding to the modalities
+                that should be passed to the OlmoEarth model.
+        """
+        kwargs = {}
+        present_modalities = []
+        device = None
+        # Handle the case where some modalities are multitemporal and some are not.
+        # We assume all multitemporal modalities have the same number of timesteps.
+        max_timesteps = 1
+        for modality in MODALITY_NAMES:
+            if modality not in inputs[0]:
+                continue
+            present_modalities.append(modality)
+            cur = torch.stack([inp[modality] for inp in inputs], dim=0)
+            device = cur.device
+            # Check if it's single or multitemporal, and reshape accordingly
+            num_bands = Modality.get(modality).num_bands
+            num_timesteps = cur.shape[1] // num_bands
+            max_timesteps = max(max_timesteps, num_timesteps)
+            cur = rearrange(cur, "b (t c) h w -> b h w t c", t=num_timesteps)
+            kwargs[modality] = cur
+            # Create mask array which is BHWTS (without channels but with band sets).
+            num_band_sets = len(Modality.get(modality).band_sets)
+            mask_shape = cur.shape[0:4] + (num_band_sets,)
+            mask = (
+                torch.ones(mask_shape, dtype=torch.int32, device=device)
+                * MaskValue.ONLINE_ENCODER.value
+            )
+            kwargs[f"{modality}_mask"] = mask
+
+        # Timestamps is required.
+        # Note that only months (0 to 11) are used in OlmoEarth position encoding.
+        # For now, we assign same timestamps to all inputs, but later we should handle varying timestamps per input.
+        timestamps = torch.zeros(
+            (len(inputs), max_timesteps, 3), dtype=torch.int32, device=device
+        )
+        timestamps[:, :, 0] = 1  # day
+        timestamps[:, :, 1] = torch.arange(max_timesteps, device=device)[
+            None, :
+        ]  # month
+        timestamps[:, :, 2] = 2024  # year
+        kwargs["timestamps"] = timestamps
+
+        sample = MaskedOlmoEarthSample(**kwargs)
+
+        # Decide context based on self.autocast_dtype.
+        if self.autocast_dtype is None:
+            context = nullcontext()
+        else:
+            assert device is not None
+            context = torch.amp.autocast(
+                device_type=device.type, dtype=self.autocast_dtype
+            )
+
+        with context:
+            # Currently we assume the provided model always returns a TokensAndMasks object.
+            tokens_and_masks: TokensAndMasks
+            if isinstance(self.model, Encoder):
+                # Encoder has a fast_pass argument to indicate mask is not needed.
+                tokens_and_masks = self.model(
+                    sample, fast_pass=True, **self.forward_kwargs
+                )["tokens_and_masks"]
+            else:
+                # Other models like STEncoder do not have this option supported.
+                tokens_and_masks = self.model(sample, **self.forward_kwargs)[
+                    "tokens_and_masks"
+                ]
+
+        # Apply temporal/modality pooling so we just have one feature per patch.
+        features = []
+        for modality in present_modalities:
+            modality_features = getattr(tokens_and_masks, modality)
+            # Pool over band sets and timesteps (BHWTSC -> BHWC).
+            pooled = modality_features.mean(dim=[3, 4])
+            # We want BHWC -> BCHW.
+            pooled = rearrange(pooled, "b h w c -> b c h w")
+            features.append(pooled)
+        # Pool over the modalities, so we get one BCHW feature map.
+        pooled = torch.stack(features, dim=0).mean(dim=0)
+        return [pooled]
+
+    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.
+
+        Returns:
+            the output channels of the backbone as a list of (downsample_factor, depth)
+            tuples.
+        """
+        return [(self.patch_size, self.embedding_size)]

rslearn/models/olmoearth_pretrain/norm.py

@@ -0,0 +1,84 @@
+"""Normalization transforms."""
+
+import json
+from typing import Any
+
+from olmoearth_pretrain.data.normalize import load_computed_config
+
+from rslearn.log_utils import get_logger
+from rslearn.train.transforms.transform import Transform
+
+logger = get_logger(__file__)
+
+
+class OlmoEarthNormalize(Transform):
+    """Normalize using OlmoEarth JSON config.
+
+    For Sentinel-1 data, the values should be converted to decibels before being passed
+    to this transform.
+    """
+
+    def __init__(
+        self,
+        band_names: dict[str, list[str]],
+        std_multiplier: float | None = 2,
+        config_fname: str | None = None,
+    ) -> None:
+        """Initialize a new OlmoEarthNormalize.
+
+        Args:
+            band_names: map from modality name to the list of bands in that modality in
+                the order they are being loaded. Note that this order must match the
+                expected order for the OlmoEarth model.
+            std_multiplier: the std multiplier matching the one used for the model
+                training in OlmoEarth.
+            config_fname: load the normalization configuration from this file, instead
+                of getting it from OlmoEarth.
+        """
+        super().__init__()
+        self.band_names = band_names
+        self.std_multiplier = std_multiplier
+
+        if config_fname is None:
+            self.norm_config = load_computed_config()
+        else:
+            logger.warning(
+                f"Loading normalization config from {config_fname}. This argument is deprecated and will be removed in a future version."
+            )
+            with open(config_fname) as f:
+                self.norm_config = json.load(f)
+
+    def forward(
+        self, input_dict: dict[str, Any], target_dict: dict[str, Any]
+    ) -> tuple[dict[str, Any], dict[str, Any]]:
+        """Apply normalization over the inputs and targets.
+
+        Args:
+            input_dict: the input
+            target_dict: the target
+
+        Returns:
+            normalized (input_dicts, target_dicts) tuple
+        """
+        for modality_name, cur_band_names in self.band_names.items():
+            band_norms = self.norm_config[modality_name]
+            image = input_dict[modality_name]
+            # Keep a set of indices to make sure that we normalize all of them.
+            needed_band_indices = set(range(image.shape[0]))
+            num_timesteps = image.shape[0] // len(cur_band_names)
+
+            for band, norm_dict in band_norms.items():
+                # If multitemporal, normalize each timestep separately.
+                for t in range(num_timesteps):
+                    band_idx = cur_band_names.index(band) + t * len(cur_band_names)
+                    min_val = norm_dict["mean"] - self.std_multiplier * norm_dict["std"]
+                    max_val = norm_dict["mean"] + self.std_multiplier * norm_dict["std"]
+                    image[band_idx] = (image[band_idx] - min_val) / (max_val - min_val)
+                    needed_band_indices.remove(band_idx)
+
+            if len(needed_band_indices) > 0:
+                raise ValueError(
+                    f"for modality {modality_name}, bands {needed_band_indices} were unexpectedly not normalized"
+                )
+
+        return input_dict, target_dict

rslearn/models/pooling_decoder.py

@@ -76,3 +76,46 @@ class PoolingDecoder(torch.nn.Module):
         features = torch.amax(features, dim=(2, 3))
         features = self.fc_layers(features)
         return self.output_layer(features)
+
+
+class SegmentationPoolingDecoder(PoolingDecoder):
+    """Like PoolingDecoder, but copy output to all pixels.
+
+    This allows for the model to produce a global output while still being compatible
+    with SegmentationTask. This only makes sense for very small windows, since the
+    output probabilities will be the same at all pixels. The main use case is to train
+    for a classification-like task on small windows, but still produce a raster during
+    inference on large windows.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        image_key: str = "image",
+        **kwargs: Any,
+    ):
+        """Create a new SegmentationPoolingDecoder.
+
+        Args:
+            in_channels: input channels (channels in the last feature map passed to
+                this module)
+            out_channels: channels for the output flat feature vector
+            image_key: the key in inputs for the image from which the expected width
+                and height is derived.
+            kwargs: other arguments to pass to PoolingDecoder.
+        """
+        super().__init__(in_channels=in_channels, out_channels=out_channels, **kwargs)
+        self.image_key = image_key
+
+    def forward(
+        self, features: list[torch.Tensor], inputs: list[dict[str, Any]]
+    ) -> torch.Tensor:
+        """Extend PoolingDecoder forward to upsample the output to a segmentation mask.
+
+        This only works when all of the pixels have the same segmentation target.
+        """
+        output_probs = super().forward(features, inputs)
+        # BC -> BCHW
+        h, w = inputs[0][self.image_key].shape[1:3]
+        return output_probs[:, :, None, None].repeat([1, 1, h, w])

rslearn/models/presto/presto.py

@@ -248,3 +248,14 @@ class Presto(nn.Module):
             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)]
+
+    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 (patch_size, depth) that corresponds
+        to the feature maps that the backbone returns.
+
+        Returns:
+            the output channels of the backbone as a list of (patch_size, depth) tuples.
+        """
+        return [(1, 128)]

rslearn/models/prithvi.py

@@ -173,6 +173,17 @@ class PrithviV2(nn.Module):
             features, num_timesteps
         )
 
+    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 (patch_size, depth) that corresponds
+        to the feature maps that the backbone returns.
+
+        Returns:
+            the output channels of the backbone as a list of (patch_size, depth) tuples.
+        """
+        return [(1, 1024)]
+
 
 class PrithviNormalize(Transform):
     """Normalize inputs using Prithvi normalization.

rslearn/models/registry.py

@@ -1,5 +1,22 @@
 """Model registry."""
 
-from class_registry import ClassRegistry
+from collections.abc import Callable
+from typing import Any, TypeVar
 
-Models = ClassRegistry()
+_ModelT = TypeVar("_ModelT")
+
+
+class _ModelRegistry(dict[str, type[Any]]):
+    """Registry for Model classes."""
+
+    def register(self, name: str) -> Callable[[type[_ModelT]], type[_ModelT]]:
+        """Decorator to register a model class."""
+
+        def decorator(cls: type[_ModelT]) -> type[_ModelT]:
+            self[name] = cls
+            return cls
+
+        return decorator
+
+
+Models = _ModelRegistry()

rslearn/tile_stores/default.py

@@ -130,10 +130,12 @@ class DefaultTileStore(TileStore):
         """
         raster_dir = self._get_raster_dir(layer_name, item_name, bands)
         for fname in raster_dir.iterdir():
-            # Ignore completed sentinel files as well as temporary files created by
+            # Ignore completed sentinel files, bands files, as well as temporary files created by
             # open_atomic (in case this tile store is on local filesystem).
             if fname.name == COMPLETED_FNAME:
                 continue
+            if fname.name == BANDS_FNAME:
+                continue
             if ".tmp." in fname.name:
                 continue
             return fname

rslearn/train/transforms/transform.py

@@ -54,7 +54,7 @@ def read_selector(
         the item specified by the selector
     """
     d, selector = get_dict_and_subselector(input_dict, target_dict, selector)
-    parts = selector.split("/")
+    parts = selector.split("/") if selector else []
     cur = d
     for part in parts:
         cur = cur[part]
@@ -76,11 +76,28 @@ def write_selector(
         v: the value to write
     """
     d, selector = get_dict_and_subselector(input_dict, target_dict, selector)
-    parts = selector.split("/")
-    cur = d
-    for part in parts[:-1]:
-        cur = cur[part]
-    cur[parts[-1]] = v
+    if selector:
+        parts = selector.split("/")
+        cur = d
+        for part in parts[:-1]:
+            cur = cur[part]
+        cur[parts[-1]] = v
+    else:
+        # If the selector references the input or target dictionary directly, then we
+        # have a special case where instead of overwriting with v, we replace the keys
+        # with those in v. v must be a dictionary here, not a tensor, since otherwise
+        # it wouldn't match the type of the input or target dictionary.
+        if not isinstance(v, dict):
+            raise ValueError(
+                "when directly specifying the input or target dict, expected the value to be a dict"
+            )
+        if d == v:
+            # This may happen if the writer did not make a copy of the dictionary. In
+            # this case the code below would not update d correctly since it would also
+            # clear v.
+            return
+        d.clear()
+        d.update(v)
 
 
 class Transform(torch.nn.Module):

rslearn/utils/raster_format.py

@@ -2,13 +2,13 @@
 
 import hashlib
 import json
-from typing import Any, BinaryIO
+from collections.abc import Callable
+from typing import Any, BinaryIO, TypeVar
 
 import affine
 import numpy as np
 import numpy.typing as npt
 import rasterio
-from class_registry import ClassRegistry
 from PIL import Image
 from rasterio.crs import CRS
 from rasterio.enums import Resampling
@@ -21,7 +21,27 @@ from rslearn.utils.fsspec import open_rasterio_upath_reader, open_rasterio_upath
 
 from .geometry import PixelBounds, Projection
 
-RasterFormats = ClassRegistry()
+_RasterFormatT = TypeVar("_RasterFormatT", bound="RasterFormat")
+
+
+class _RasterFormatRegistry(dict[str, type["RasterFormat"]]):
+    """Registry for RasterFormat classes."""
+
+    def register(
+        self, name: str
+    ) -> Callable[[type[_RasterFormatT]], type[_RasterFormatT]]:
+        """Decorator to register a raster format class."""
+
+        def decorator(cls: type[_RasterFormatT]) -> type[_RasterFormatT]:
+            self[name] = cls
+            return cls
+
+        return decorator
+
+
+RasterFormats = _RasterFormatRegistry()
+
+
 logger = get_logger(__name__)
 
 
@@ -147,6 +167,19 @@ class RasterFormat:
         """
         raise NotImplementedError
 
+    @staticmethod
+    def from_config(name: str, config: dict[str, Any]) -> "RasterFormat":
+        """Create a RasterFormat from a config dict.
+
+        Args:
+            name: the name of this format
+            config: the config dict
+
+        Returns:
+            the RasterFormat instance
+        """
+        raise NotImplementedError
+
 
 @RasterFormats.register("image_tile")
 class ImageTileRasterFormat(RasterFormat):
@@ -716,5 +749,5 @@ def load_raster_format(config: RasterFormatConfig) -> RasterFormat:
     Returns:
         the loaded RasterFormat implementation
     """
-    cls = RasterFormats.get_class(config.name)
+    cls = RasterFormats[config.name]
     return cls.from_config(config.name, config.config_dict)

rslearn/utils/vector_format.py

@@ -1,11 +1,11 @@
 """Classes for writing vector data to a UPath."""
 
 import json
+from collections.abc import Callable
 from enum import Enum
-from typing import Any
+from typing import Any, TypeVar
 
 import shapely
-from class_registry import ClassRegistry
 from rasterio.crs import CRS
 from upath import UPath
 
@@ -18,7 +18,25 @@ from .feature import Feature
 from .geometry import PixelBounds, Projection, STGeometry, safely_reproject_and_clip
 
 logger = get_logger(__name__)
-VectorFormats = ClassRegistry()
+_VectorFormatT = TypeVar("_VectorFormatT", bound="VectorFormat")
+
+
+class _VectorFormatRegistry(dict[str, type["VectorFormat"]]):
+    """Registry for VectorFormat classes."""
+
+    def register(
+        self, name: str
+    ) -> Callable[[type[_VectorFormatT]], type[_VectorFormatT]]:
+        """Decorator to register a vector format class."""
+
+        def decorator(cls: type[_VectorFormatT]) -> type[_VectorFormatT]:
+            self[name] = cls
+            return cls
+
+        return decorator
+
+
+VectorFormats = _VectorFormatRegistry()
 
 
 class VectorFormat:
@@ -53,6 +71,19 @@ class VectorFormat:
         """
         raise NotImplementedError
 
+    @staticmethod
+    def from_config(name: str, config: dict[str, Any]) -> "VectorFormat":
+        """Create a VectorFormat from a config dict.
+
+        Args:
+            name: the name of this format
+            config: the config dict
+
+        Returns:
+            the VectorFormat instance
+        """
+        raise NotImplementedError
+
 
 @VectorFormats.register("tile")
 class TileVectorFormat(VectorFormat):
@@ -410,5 +441,5 @@ def load_vector_format(config: VectorFormatConfig) -> VectorFormat:
     Returns:
         the loaded VectorFormat implementation
     """
-    cls = VectorFormats.get_class(config.name)
+    cls = VectorFormats[config.name]
     return cls.from_config(config.name, config.config_dict)

{rslearn-0.0.8.dist-info → rslearn-0.0.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rslearn
-Version: 0.0.8
+Version: 0.0.11
 Summary: A library for developing remote sensing datasets and models
 Author: OlmoEarth Team
 License:                                  Apache License
@@ -212,7 +212,6 @@ Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: boto3>=1.39
-Requires-Dist: class_registry>=2.1
 Requires-Dist: fiona>=1.10
 Requires-Dist: fsspec>=2025.9.0
 Requires-Dist: jsonargparse>=4.35.0
@@ -244,6 +243,7 @@ Requires-Dist: planetary_computer>=1.0; extra == "extra"
 Requires-Dist: pycocotools>=2.0; extra == "extra"
 Requires-Dist: pystac_client>=0.9; extra == "extra"
 Requires-Dist: rtree>=1.4; extra == "extra"
+Requires-Dist: termcolor>=3.0; extra == "extra"
 Requires-Dist: satlaspretrain_models>=0.3; extra == "extra"
 Requires-Dist: scipy>=1.16; extra == "extra"
 Requires-Dist: terratorch>=1.0.2; extra == "extra"
@@ -284,6 +284,7 @@ Quick links:
 - [Examples](docs/Examples.md) contains more examples, including customizing different
   stages of rslearn with additional code.
 - [DatasetConfig](docs/DatasetConfig.md) documents the dataset configuration file.
+- [ModelConfig](docs/ModelConfig.md) documents the model configuration file.
 
 
 Setup

{rslearn-0.0.8.dist-info → rslearn-0.0.11.dist-info}/RECORD

@@ -20,11 +20,11 @@ rslearn/data_sources/eurocrops.py,sha256=bH_ul45RvNLLvVKU9JjZH1XMenAmgn7Lakq0pXj
 rslearn/data_sources/gcp_public_data.py,sha256=kr9stYo7ZCvz8s4E3wmoY-yAGZoLa_9RCwjS-Q5k9dM,36128
 rslearn/data_sources/geotiff.py,sha256=sFUp919chaX4j6lQytNp__xnMLlDI3Ac3rfB6F8sgZ0,45
 rslearn/data_sources/google_earth_engine.py,sha256=hpkt74ly2lEwjRrDp8FBmGvB3MEw_mQ38Av4rQOR3_w,24246
-rslearn/data_sources/local_files.py,sha256=NnmXWbOoKHxZFzwCiWsybWfPrFoPTpoZbjoK57xhoqc,19049
+rslearn/data_sources/local_files.py,sha256=d08m6IzrUN_80VfvgpHahMJrv-n6_CI6EIocp6kyDRs,19490
 rslearn/data_sources/openstreetmap.py,sha256=qUSMFiIA_laJkO3meBXf9TmSI7OBD-o3i4JxqllUv3Q,19232
 rslearn/data_sources/planet.py,sha256=F2JoLaQ5Cb3k1cTm0hwSWTL2TPfbaAUMXZ8q4Dy7UlA,10109
 rslearn/data_sources/planet_basemap.py,sha256=wuWM9dHSJMdINfyWb78Zk9i-KvJHTrf9J0Q2gyEyiiA,10450
-rslearn/data_sources/planetary_computer.py,sha256=Wchr-OmAffuVteUW6VRofIqFpE-cJqI6GMTU9v1i8pw,28089
+rslearn/data_sources/planetary_computer.py,sha256=uHNYxvnMkmo8zbqIiDRdnkz8LQ7TSs6K39Y1AXjboDI,30392
 rslearn/data_sources/raster_source.py,sha256=b8wo55GhVLxXwx1WYLzeRAlzD_ZkE_P9tnvUOdnsfQE,689
 rslearn/data_sources/usda_cdl.py,sha256=2_V11AhPRgLEGd4U5Pmx3UvE2HWBPbsFXhUIQVRVFeE,7138
 rslearn/data_sources/usgs_landsat.py,sha256=31GmOUfmxwTE6MTiVI4psb-ciVmunuA8cfvqDuvTHPE,19312
@@ -39,9 +39,9 @@ rslearn/dataset/add_windows.py,sha256=pwCEvwLE1jQCoqQxw6CJ-sP46ayWppFa2hGYIB6VVk
 rslearn/dataset/dataset.py,sha256=bjf9nI55j-MF0bIQWSNPjNbpfqnLK4jy-96TAcwO0MM,5214
 rslearn/dataset/handler_summaries.py,sha256=wGnbBpjLWTxVn3UT7j7nPoHlYsaWb9_MVJ5DhU0qWXY,2581
 rslearn/dataset/index.py,sha256=Wni5m6h4gisRB54fPLnCfUrRTEsJ5EvwS0fs9sYc2wg,6025
-rslearn/dataset/manage.py,sha256=E2PF29MWuYu8XMcX8PYG4itCdgxaUV768R5NCHmJ-e8,18058
-rslearn/dataset/materialize.py,sha256=emsxgkBKWC-ybZDtudsYOsxx7yxgmdKy5YrysgqzNac,21556
-rslearn/dataset/remap.py,sha256=HF5Kn5z6dEPbCpSxaBmYVDSzCuBjG2w8csPsk0RRlv4,2088
+rslearn/dataset/manage.py,sha256=mkdBHo1RFGxMx8f9zBT_VmRO_6y8Qb2KfWPPziKWYkg,18062
+rslearn/dataset/materialize.py,sha256=-z47svc_JqGhzkp8kq5Hd9fykWNqFEUCQezo887TWBw,22056
+rslearn/dataset/remap.py,sha256=6MaImsY02GNACpvRM81RvWmjZWRfAHxo_R3Ox6XLF6A,2723
 rslearn/dataset/window.py,sha256=I5RqZ12jlIXhohw4qews1x_I4tSDpml709DZRtLiN24,12546
 rslearn/models/__init__.py,sha256=_vWoF9d2Slah8-6XhYhdU4SRsy_CNxXjCGQTD2yvu3Q,22
 rslearn/models/anysat.py,sha256=3BnaiS1sYB4SnV6qRjHksiz_r9vUuZeGPUO2XUziFA0,7810
@@ -49,7 +49,7 @@ rslearn/models/clip.py,sha256=u5aqYnVB4Jag7o1h8EzPDAc1t2BAHeALA9FcUwP5tfo,2238
 rslearn/models/conv.py,sha256=fWyByeswIOKKzyPmP3erYUlZaKEV0huWHA4CyKTBbfY,1703
 rslearn/models/copernicusfm.py,sha256=3AiORuUre9sZYwydbrDgShwKtxeTLmExp7WQmJtBylg,7842
 rslearn/models/croma.py,sha256=cOazTp3l2PNJltKrmPqD5Gy4pi3CI03-X9G4T10cX2k,9529
-rslearn/models/dinov3.py,sha256=KgR-vbmb3-6Qn4AC7ybGL6Ri3_f4jf2XOaKxAR9uMnU,6150
+rslearn/models/dinov3.py,sha256=GKk5qXZPCEporATJdjaSWsDTfWDlAGRWBplFUJN5nRM,6146
 rslearn/models/faster_rcnn.py,sha256=uaxX6-E1f0BibaA9sorEg3be83C7kTdTc39pC5jRqwE,8286
 rslearn/models/fpn.py,sha256=s3cz29I14FaSuvBvLOcwCrqVsaRBxG5GjLlqap4WgPc,1603
 rslearn/models/module_wrapper.py,sha256=H2zb-8Au4t31kawW_4JEKHsaXFjpYDawb31ZEauKcxU,2728
@@ -57,9 +57,9 @@ rslearn/models/molmo.py,sha256=mVrARBhZciMzOgOOjGB5AHlPIf2iO9IBSJmdyKSl1L8,2061
 rslearn/models/multitask.py,sha256=j2Kiwj_dUiUp_CIUr25bS8HiyeoFlr1PGqjTfpgIGLc,14672
 rslearn/models/panopticon.py,sha256=woNEs53wVc5D-NxbSDEPRZ_mYe8vllnuldmADjvhfDQ,5806
 rslearn/models/pick_features.py,sha256=y8e4tJFhyG7ZuVSElWhQ5-Aer4ZKJCEH9wLGJU7WqGI,1551
-rslearn/models/pooling_decoder.py,sha256=jZfEQCfthfa21C9sEjgFHUcfhHMVlvG7_nDMw_1FLaE,2727
-rslearn/models/prithvi.py,sha256=bGY2Tf_V9tlJGDtDSCMkoFMJdFXUCW6cyj4r8FER8Pw,39557
-rslearn/models/registry.py,sha256=0j6jY-rmB_wKckHEUeCYWXtjviDOEr5c7WhIMxxDwBk,90
+rslearn/models/pooling_decoder.py,sha256=unr2fSE_QmJHPi3dKtopqMtb1Kn-2h94LgwwAVP9vZg,4437
+rslearn/models/prithvi.py,sha256=SVM3ypJlVTkXQ69pPhB4UeJr87VnmADTCuyV365dbkU,39961
+rslearn/models/registry.py,sha256=yCcrOvLkbn07Xtln1j7hAB_kmGw0MGsiR2TloJq9Bmk,504
 rslearn/models/resize_features.py,sha256=asKXWrLHIBrU6GaAV0Ory9YuK7IK104XjhkB4ljzI3A,1289
 rslearn/models/sam2_enc.py,sha256=gNlPokr7eNxO2KvnzDMXNxYM2WRO0YkQPjR4110n6cw,3508
 rslearn/models/satlaspretrain.py,sha256=YpjXl-uClhTZMDmyhN64Fg3AszzT-ymZgJB0fO9RyoY,2419
@@ -91,8 +91,11 @@ rslearn/models/detr/position_encoding.py,sha256=8FFoBT-Jtgqk7D4qDBTbVLOeAdmjdjtJ
 rslearn/models/detr/transformer.py,sha256=aK4HO7AkCZn7xGHP3Iq91w2iFPVshugOILYAjVjroCw,13971
 rslearn/models/detr/util.py,sha256=NMHhHbkIo7PoBUVbDqa2ZknJBTswmaxFCGHrPtFKnGg,676
 rslearn/models/galileo/__init__.py,sha256=QQa0C29nuPRva0KtGiMHQ2ZB02n9SSwj_wqTKPz18NM,112
-rslearn/models/galileo/galileo.py,sha256=b8LMyn5lN-WHmCNTJUy-ykjL0s5_-jg9RTO1CwJVra4,19533
-rslearn/models/galileo/single_file_galileo.py,sha256=hmtBGdlFfX9JceZLguntSlXtuHElOafxwu1yIrhrEJY,55920
+rslearn/models/galileo/galileo.py,sha256=jUHA64YvVC3Fz5fevc_9dFJfZaINODRDrhSGLIiOZcw,21115
+rslearn/models/galileo/single_file_galileo.py,sha256=l5tlmmdr2eieHNH-M7rVIvcptkv0Fuk3vKXFW691ezA,56143
+rslearn/models/olmoearth_pretrain/__init__.py,sha256=AjRvbjBdadCdPh-EdvySH76sVAQ8NGQaJt11Tsn1D5I,36
+rslearn/models/olmoearth_pretrain/model.py,sha256=F-B1ym9UZuTPJ0OY15Jwb1TkNtr_EtAUlqI-tr_Z2uo,8352
+rslearn/models/olmoearth_pretrain/norm.py,sha256=rHjFyWkpNLYMx9Ow7TsU-jGm9Sjx7FVf0p4R__ohx2c,3266
 rslearn/models/panopticon_data/sensors/drone.yaml,sha256=xqWS-_QMtJyRoWXJm-igoSur9hAmCFdqkPin8DT5qpw,431
 rslearn/models/panopticon_data/sensors/enmap.yaml,sha256=b2j6bSgYR2yKR9DRm3SPIzSVYlHf51ny_p-1B4B9sB4,13431
 rslearn/models/panopticon_data/sensors/goes.yaml,sha256=o00aoWCYqam0aB1rPmXq1MKe8hsKak_qyBG7BPL27Sc,152
@@ -106,10 +109,10 @@ rslearn/models/panopticon_data/sensors/sentinel2.yaml,sha256=qYJ92x-GHO0ZdCrTtCj
 rslearn/models/panopticon_data/sensors/superdove.yaml,sha256=QpIRyopdV4hAez_EIsDwhGFT4VtTk7UgzQveyc8t8fc,795
 rslearn/models/panopticon_data/sensors/wv23.yaml,sha256=SWYSlkka6UViKAz6YI8aqwQ-Ayo-S5kmNa9rO3iGW6o,1172
 rslearn/models/presto/__init__.py,sha256=eZrB-XKi_vYqZhpyAOwppJi4dRuMtYVAdbq7KRygze0,64
-rslearn/models/presto/presto.py,sha256=Vr4DBIP5Twnop5pZ21BhChbOwWP7KI-GMdg2XILkj5c,8774
+rslearn/models/presto/presto.py,sha256=8mZnc0jk_r_JikybHQNyyHg6t7JNPmoPmgoivyNf-U8,9177
 rslearn/models/presto/single_file_presto.py,sha256=Kbwp8V7pO8HHM2vlCPpjekQiFiDryW8zQkWmt1g05BY,30381
 rslearn/tile_stores/__init__.py,sha256=o_tWVKu6UwFzZbO9jn_3cmIDqc_Q3qDd6tA9If0T_Qk,2050
-rslearn/tile_stores/default.py,sha256=P78LQm9OXNjd6Igkh-YqUFDmHSBN1Uk092xwqkuOPSs,14973
+rslearn/tile_stores/default.py,sha256=PYaDNvBxhJTDKJGw0EjDTSE1OKajR7_iJpMbOjj-mE8,15054
 rslearn/tile_stores/tile_store.py,sha256=9AeYduDYPp_Ia2NMlq6osptpz_AFGIOQcLJrqZ_m-z0,10469
 rslearn/train/__init__.py,sha256=fnJyY4aHs5zQqbDKSfXsJZXY_M9fbTsf7dRYaPwZr2M,30
 rslearn/train/data_module.py,sha256=K-nQgnOZn-KGq_G2pVOQFtWRrlWih0212i_bkXZ2bEE,23515
@@ -140,7 +143,7 @@ rslearn/train/transforms/normalize.py,sha256=uyv2hE5hw5B2kCRHa4JIx0tfowm-C7bgumw
 rslearn/train/transforms/pad.py,sha256=EDswS9KYRSloM3DQlbCz6S0WYqFQJvI433qMqTtqrZw,4686
 rslearn/train/transforms/select_bands.py,sha256=uDfD9G8Z4VTt88QZsjj1FB20QEmzSefhKf7uDXYn77M,2441
 rslearn/train/transforms/sentinel1.py,sha256=FrLaYZs2AjqWQCun8DTFtgo1l0xLxqaFKtDNIehtpDg,1913
-rslearn/train/transforms/transform.py,sha256=8Q-dPrmDr0tJ9ZOwjWqWK8kbnKi4uLxEnS9Nwf6BVJk,3594
+rslearn/train/transforms/transform.py,sha256=n1Qzqix2dVvej-Q7iPzHeOQbqH79IBlvqPoymxhNVpE,4446
 rslearn/utils/__init__.py,sha256=GNvdTUmXakiEMnLdje7k1fe5aC7SFVqP757kbpN6Fzw,558
 rslearn/utils/array.py,sha256=JwZi7o0uj-dftREzJmqrRVR2joIwBikm3Er9KeHVIZU,2402
 rslearn/utils/feature.py,sha256=lsg0WThZDJzo1mrbaL04dXYI5G3x-n5FG9aEjj7uUaI,1649
@@ -150,15 +153,15 @@ rslearn/utils/get_utm_ups_crs.py,sha256=kUrcyjCK7KWvuP1XR-nURPeRqYeRO-3L8QUJ1QTF
 rslearn/utils/grid_index.py,sha256=hRmrtgpqN1pLa-djnZtgSXqKJlbgGyttGnCEmPLD0zo,2347
 rslearn/utils/jsonargparse.py,sha256=JcTKQoZ6jgwag-kSeTIEVBO9AsRj0X1oEJBsoaCazH4,658
 rslearn/utils/mp.py,sha256=XYmVckI5TOQuCKc49NJyirDJyFgvb4AI-gGypG2j680,1399
-rslearn/utils/raster_format.py,sha256=3hjBR_0Gl1M1b0Vjx4UsQBgAVziCh34VAVWeW32k3Ds,25167
+rslearn/utils/raster_format.py,sha256=dBTSa8l6Ms9Ndbx9Krgqm9z4RU7j2hwLBkw2w-KibU4,26009
 rslearn/utils/rtree_index.py,sha256=j0Zwrq3pXuAJ-hKpiRFQ7VNtvO3fZYk-Em2uBPAqfx4,6460
 rslearn/utils/spatial_index.py,sha256=eomJAUgzmjir8j9HZnSgQoJHwN9H0wGTjmJkMkLLfsU,762
 rslearn/utils/sqlite_index.py,sha256=YGOJi66544e6JNtfSft6YIlHklFdSJO2duxQ4TJ2iu4,2920
 rslearn/utils/time.py,sha256=2ilSLG94_sxLP3y5RSV5L5CG8CoND_dbdzYEHVtN-I8,387
-rslearn/utils/vector_format.py,sha256=C0-qlSXFFnhUU0S7eiOn4-4d5-pCkwU2jZ-q0Ldikk8,15299
-rslearn-0.0.8.dist-info/licenses/LICENSE,sha256=_99ZWPoLdlUbqZoSC5DF4ihiNwl5rTEmBaq2fACecdg,11352
-rslearn-0.0.8.dist-info/METADATA,sha256=hdHA3W7Hb4bQPPZG6SscWr1AOOVhALL0F08tEMt7neQ,36206
-rslearn-0.0.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-rslearn-0.0.8.dist-info/entry_points.txt,sha256=doTBQ57NT7nq-dgYGgTTw6mafcGWb_4PWYtYR4rGm50,46
-rslearn-0.0.8.dist-info/top_level.txt,sha256=XDKo90WBH8P9RQumHxo0giLJsoufT4r9odv-WE6Ahk4,8
-rslearn-0.0.8.dist-info/RECORD,,
+rslearn/utils/vector_format.py,sha256=EIChYCL6GLOILS2TO2JBkca1TuaWsSubWv6iRS3P2ds,16139
+rslearn-0.0.11.dist-info/licenses/LICENSE,sha256=_99ZWPoLdlUbqZoSC5DF4ihiNwl5rTEmBaq2fACecdg,11352
+rslearn-0.0.11.dist-info/METADATA,sha256=jwB0ZZ-oLa1Y_1iuZRKCQoB4i3kOFDJ0xSeMTJP7zww,36297
+rslearn-0.0.11.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+rslearn-0.0.11.dist-info/entry_points.txt,sha256=doTBQ57NT7nq-dgYGgTTw6mafcGWb_4PWYtYR4rGm50,46
+rslearn-0.0.11.dist-info/top_level.txt,sha256=XDKo90WBH8P9RQumHxo0giLJsoufT4r9odv-WE6Ahk4,8
+rslearn-0.0.11.dist-info/RECORD,,