rslearn 0.0.14__py3-none-any.whl → 0.0.16__py3-none-any.whl

rslearn/models/satlaspretrain.py

@@ -4,15 +4,14 @@ from typing import Any
 
 import satlaspretrain_models
 import torch
+import torch.nn.functional as F
 
 
 class SatlasPretrain(torch.nn.Module):
     """SatlasPretrain backbones."""
 
     def __init__(
-        self,
-        model_identifier: str,
-        fpn: bool = False,
+        self, model_identifier: str, fpn: bool = False, resize_to_pretrain: bool = False
     ) -> None:
         """Instantiate a new SatlasPretrain instance.
 
@@ -21,6 +20,8 @@ class SatlasPretrain(torch.nn.Module):
                 https://github.com/allenai/satlaspretrain_models
             fpn: whether to include the feature pyramid network, otherwise only the
                 Swin-v2-Transformer is used.
+            resize_to_pretrain: whether to resize inputs to the pretraining input
+                size (512 x 512)
         """
         super().__init__()
         weights_manager = satlaspretrain_models.Weights()
@@ -49,6 +50,19 @@ class SatlasPretrain(torch.nn.Module):
                 [16, 1024],
                 [32, 2048],
             ]
+        self.resize_to_pretrain = resize_to_pretrain
+
+    def maybe_resize(self, data: torch.Tensor) -> list[torch.Tensor]:
+        """Resize to pretraining sizes if resize_to_pretrain == True."""
+        if self.resize_to_pretrain:
+            return F.interpolate(
+                data,
+                size=(512, 512),
+                mode="bilinear",
+                align_corners=False,
+            )
+        else:
+            return data
 
     def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
         """Compute feature maps from the SatlasPretrain backbone.
@@ -58,7 +72,7 @@ class SatlasPretrain(torch.nn.Module):
                 process.
         """
         images = torch.stack([inp["image"] for inp in inputs], dim=0)
-        return self.model(images)
+        return self.model(self.maybe_resize(images))
 
     def get_backbone_channels(self) -> list:
         """Returns the output channels of this model when used as a backbone.

rslearn/models/terramind.py

@@ -4,6 +4,7 @@ from enum import Enum
 from typing import Any
 
 import torch
+import torch.nn.functional as F
 from einops import rearrange
 from terratorch.registry import BACKBONE_REGISTRY
 
@@ -18,6 +19,8 @@ class TerramindSize(str, Enum):
     LARGE = "large"
 
 
+# Pretraining image size for Terramind
+IMAGE_SIZE = 224
 # Default patch size for Terramind
 PATCH_SIZE = 16
 
@@ -89,12 +92,14 @@ class Terramind(torch.nn.Module):
         self,
         model_size: TerramindSize,
         modalities: list[str] = ["S2L2A"],
+        do_resizing: bool = False,
     ) -> None:
         """Initialize the Terramind model.
 
         Args:
             model_size: The size of the Terramind model.
             modalities: The modalities to use.
+            do_resizing: Whether to resize the input images to the pretraining resolution.
         """
         super().__init__()
 
@@ -116,6 +121,7 @@ class Terramind(torch.nn.Module):
 
         self.model_size = model_size
         self.modalities = modalities
+        self.do_resizing = do_resizing
 
     def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
         """Forward pass for the Terramind model.
@@ -132,6 +138,19 @@ class Terramind(torch.nn.Module):
             if modality not in inputs[0]:
                 continue
             cur = torch.stack([inp[modality] for inp in inputs], dim=0)  # (B, C, H, W)
+            if self.do_resizing and (
+                cur.shape[2] != IMAGE_SIZE or cur.shape[3] != IMAGE_SIZE
+            ):
+                if cur.shape[2] == 1 and cur.shape[3] == 1:
+                    new_height, new_width = PATCH_SIZE, PATCH_SIZE
+                else:
+                    new_height, new_width = IMAGE_SIZE, IMAGE_SIZE
+                cur = F.interpolate(
+                    cur,
+                    size=(new_height, new_width),
+                    mode="bilinear",
+                    align_corners=False,
+                )
             model_inputs[modality] = cur
 
         # By default, the patch embeddings are averaged over all modalities to reduce output tokens

rslearn/tile_stores/__init__.py

@@ -22,17 +22,6 @@ def load_tile_store(config: dict[str, Any], ds_path: UPath) -> TileStore:
     Returns:
         the TileStore
     """
-    if config is None:
-        tile_store = DefaultTileStore()
-        tile_store.set_dataset_path(ds_path)
-        return tile_store
-
-    # Backwards compatability.
-    if "name" in config and "root_dir" in config and config["name"] == "file":
-        tile_store = DefaultTileStore(config["root_dir"])
-        tile_store.set_dataset_path(ds_path)
-        return tile_store
-
     init_jsonargparse()
     parser = jsonargparse.ArgumentParser()
     parser.add_argument("--tile_store", type=TileStore)

rslearn/train/dataset.py

@@ -17,9 +17,7 @@ from rasterio.warp import Resampling
 import rslearn.train.transforms.transform
 from rslearn.config import (
     DType,
-    RasterFormatConfig,
-    RasterLayerConfig,
-    VectorLayerConfig,
+    LayerConfig,
 )
 from rslearn.dataset.dataset import Dataset
 from rslearn.dataset.window import Window, get_layer_and_group_from_dir_name
@@ -28,8 +26,6 @@ from rslearn.train.tasks import Task
 from rslearn.utils.feature import Feature
 from rslearn.utils.geometry import PixelBounds
 from rslearn.utils.mp import star_imap_unordered
-from rslearn.utils.raster_format import load_raster_format
-from rslearn.utils.vector_format import load_vector_format
 
 from .transforms import Sequential
 
@@ -185,7 +181,7 @@ def read_raster_layer_for_data_input(
     bounds: PixelBounds,
     layer_name: str,
     group_idx: int,
-    layer_config: RasterLayerConfig,
+    layer_config: LayerConfig,
     data_input: DataInput,
 ) -> torch.Tensor:
     """Read a raster layer for a DataInput.
@@ -246,9 +242,7 @@ def read_raster_layer_for_data_input(
         )
         if band_set.format is None:
             raise ValueError(f"No format specified for {layer_name}")
-        raster_format = load_raster_format(
-            RasterFormatConfig(band_set.format["name"], band_set.format)
-        )
+        raster_format = band_set.instantiate_raster_format()
         raster_dir = window.get_raster_dir(
             layer_name, band_set.bands, group_idx=group_idx
         )
@@ -349,7 +343,6 @@ def read_data_input(
         images: list[torch.Tensor] = []
         for layer_name, group_idx in layers_to_read:
             layer_config = dataset.layers[layer_name]
-            assert isinstance(layer_config, RasterLayerConfig)
             images.append(
                 read_raster_layer_for_data_input(
                     window, bounds, layer_name, group_idx, layer_config, data_input
@@ -363,8 +356,7 @@ def read_data_input(
         features: list[Feature] = []
         for layer_name, group_idx in layers_to_read:
             layer_config = dataset.layers[layer_name]
-            assert isinstance(layer_config, VectorLayerConfig)
-            vector_format = load_vector_format(layer_config.format)
+            vector_format = layer_config.instantiate_vector_format()
             layer_dir = window.get_layer_dir(layer_name, group_idx=group_idx)
             cur_features = vector_format.decode_vector(
                 layer_dir, window.projection, window.bounds

rslearn/train/prediction_writer.py

@@ -12,10 +12,8 @@ from lightning.pytorch.callbacks import BasePredictionWriter
 from upath import UPath
 
 from rslearn.config import (
+    LayerConfig,
     LayerType,
-    RasterFormatConfig,
-    RasterLayerConfig,
-    VectorLayerConfig,
 )
 from rslearn.dataset import Dataset, Window
 from rslearn.log_utils import get_logger
@@ -25,9 +23,8 @@ from rslearn.utils.geometry import PixelBounds
 from rslearn.utils.raster_format import (
     RasterFormat,
     adjust_projection_and_bounds_for_array,
-    load_raster_format,
 )
-from rslearn.utils.vector_format import VectorFormat, load_vector_format
+from rslearn.utils.vector_format import VectorFormat
 
 from .lightning_module import RslearnLightningModule
 from .tasks.task import Task
@@ -150,7 +147,7 @@ class RslearnWriter(BasePredictionWriter):
         selector: list[str] | None = None,
         merger: PatchPredictionMerger | None = None,
         output_path: str | Path | None = None,
-        layer_config: RasterLayerConfig | VectorLayerConfig | None = None,
+        layer_config: LayerConfig | None = None,
     ):
         """Create a new RslearnWriter.
 
@@ -178,43 +175,31 @@ class RslearnWriter(BasePredictionWriter):
         )
 
         # Handle dataset and layer config
-        self.layer_config: RasterLayerConfig | VectorLayerConfig
+        self.layer_config: LayerConfig
         if layer_config:
             self.layer_config = layer_config
-            self.dataset = None if self.output_path else Dataset(self.path)
         else:
-            self.dataset = Dataset(self.path)
-            if self.output_layer not in self.dataset.layers:
+            dataset = Dataset(self.path)
+            if self.output_layer not in dataset.layers:
                 raise KeyError(
                     f"Output layer '{self.output_layer}' not found in dataset layers."
                 )
-            raw_layer_config = self.dataset.layers[self.output_layer]
-            # Type narrowing to ensure compatibility
-            if isinstance(raw_layer_config, (RasterLayerConfig | VectorLayerConfig)):
-                self.layer_config = raw_layer_config
-            else:
-                raise ValueError(
-                    f"Layer config must be RasterLayerConfig or VectorLayerConfig, got {type(raw_layer_config)}"
-                )
+            self.layer_config = dataset.layers[self.output_layer]
 
         self.format: RasterFormat | VectorFormat
-        if self.layer_config.layer_type == LayerType.RASTER:
-            assert isinstance(self.layer_config, RasterLayerConfig)
+        if self.layer_config.type == LayerType.RASTER:
             band_cfg = self.layer_config.band_sets[0]
-            self.format = load_raster_format(
-                RasterFormatConfig(band_cfg.format["name"], band_cfg.format)
-            )
-        elif self.layer_config.layer_type == LayerType.VECTOR:
-            assert isinstance(self.layer_config, VectorLayerConfig)
-            self.format = load_vector_format(self.layer_config.format)
+            self.format = band_cfg.instantiate_raster_format()
+        elif self.layer_config.type == LayerType.VECTOR:
+            self.format = self.layer_config.instantiate_vector_format()
         else:
-            raise ValueError(f"invalid layer type {self.layer_config.layer_type}")
+            raise ValueError(f"invalid layer type {self.layer_config.type}")
 
         if merger is not None:
             self.merger = merger
-        elif self.layer_config.layer_type == LayerType.RASTER:
+        elif self.layer_config.type == LayerType.RASTER:
             self.merger = RasterMerger()
-        elif self.layer_config.layer_type == LayerType.VECTOR:
+        elif self.layer_config.type == LayerType.VECTOR:
             self.merger = VectorMerger()
 
         # Map from window name to pending data to write.
@@ -337,8 +322,7 @@ class RslearnWriter(BasePredictionWriter):
         logger.debug(f"Merging and writing for window {window.name}")
         merged_output = self.merger.merge(window, pending_output)
 
-        if self.layer_config.layer_type == LayerType.RASTER:
-            assert isinstance(self.layer_config, RasterLayerConfig)
+        if self.layer_config.type == LayerType.RASTER:
             raster_dir = window.get_raster_dir(
                 self.output_layer, self.layer_config.band_sets[0].bands
             )
@@ -351,7 +335,7 @@ class RslearnWriter(BasePredictionWriter):
             )
             self.format.encode_raster(raster_dir, projection, bounds, merged_output)
 
-        elif self.layer_config.layer_type == LayerType.VECTOR:
+        elif self.layer_config.type == LayerType.VECTOR:
             layer_dir = window.get_layer_dir(self.output_layer)
             assert isinstance(self.format, VectorFormat)
             self.format.encode_vector(layer_dir, merged_output)

rslearn/train/tasks/classification.py

@@ -26,7 +26,7 @@ class ClassificationTask(BasicTask):
     def __init__(
         self,
         property_name: str,
-        classes: list,  # TODO: Should this be a list of str or int or can it be both?
+        classes: list[str],
         filters: list[tuple[str, str]] = [],
         read_class_id: bool = False,
         allow_invalid: bool = False,
@@ -176,6 +176,7 @@ class ClassificationTask(BasicTask):
             # For multiclass classification or when using the default threshold
             class_idx = probs.argmax().item()
 
+        value: str | int
         if not self.read_class_id:
             value = self.classes[class_idx]  # type: ignore
         else:

rslearn/utils/fsspec.py

@@ -156,3 +156,23 @@ def open_rasterio_upath_writer(
         with path.open("wb") as f:
             with rasterio.open(f, "w", **kwargs) as raster:
                 yield raster
+
+
+def get_relative_suffix(base_dir: UPath, fname: UPath) -> str:
+    """Get the suffix of fname relative to base_dir.
+
+    Args:
+        base_dir: the base directory.
+        fname: a filename within the base directory.
+
+    Returns:
+        the suffix on base_dir that would yield the given filename.
+    """
+    if not fname.path.startswith(base_dir.path):
+        raise ValueError(
+            f"filename {fname.path} must start with base directory {base_dir.path}"
+        )
+    suffix = fname.path[len(base_dir.path) :]
+    if suffix.startswith("/"):
+        suffix = suffix[1:]
+    return suffix

rslearn/utils/jsonargparse.py

@@ -1,7 +1,18 @@
 """Custom serialization for jsonargparse."""
 
+from datetime import datetime
+from typing import TYPE_CHECKING, Any
+
 import jsonargparse
 from rasterio.crs import CRS
+from upath import UPath
+
+from rslearn.config.dataset import LayerConfig
+
+if TYPE_CHECKING:
+    from rslearn.data_sources.data_source import DataSourceContext
+
+INITIALIZED = False
 
 
 def crs_serializer(v: CRS) -> str:
@@ -28,6 +39,74 @@ def crs_deserializer(v: str) -> CRS:
     return CRS.from_string(v)
 
 
+def datetime_serializer(v: datetime) -> str:
+    """Serialize datetime for jsonargparse.
+
+    Args:
+        v: the datetime object.
+
+    Returns:
+        the datetime encoded to string
+    """
+    return v.isoformat()
+
+
+def datetime_deserializer(v: str) -> datetime:
+    """Deserialize datetime for jsonargparse.
+
+    Args:
+        v: the encoded datetime.
+
+    Returns:
+        the decoded datetime object
+    """
+    return datetime.fromisoformat(v)
+
+
+def data_source_context_serializer(v: "DataSourceContext") -> dict[str, Any]:
+    """Serialize DataSourceContext for jsonargparse."""
+    x = {
+        "ds_path": (str(v.ds_path) if v.ds_path is not None else None),
+        "layer_config": (
+            v.layer_config.model_dump(mode="json")
+            if v.layer_config is not None
+            else None
+        ),
+    }
+    return x
+
+
+def data_source_context_deserializer(v: dict[str, Any]) -> "DataSourceContext":
+    """Deserialize DataSourceContext for jsonargparse."""
+    # We lazily import these to avoid cyclic dependency.
+    from rslearn.data_sources.data_source import DataSourceContext
+
+    return DataSourceContext(
+        ds_path=(UPath(v["ds_path"]) if v["ds_path"] is not None else None),
+        layer_config=(
+            LayerConfig.model_validate(v["layer_config"])
+            if v["layer_config"] is not None
+            else None
+        ),
+    )
+
+
 def init_jsonargparse() -> None:
     """Initialize custom jsonargparse serializers."""
+    global INITIALIZED
+    if INITIALIZED:
+        return
     jsonargparse.typing.register_type(CRS, crs_serializer, crs_deserializer)
+    jsonargparse.typing.register_type(
+        datetime, datetime_serializer, datetime_deserializer
+    )
+
+    from rslearn.data_sources.data_source import DataSourceContext
+
+    jsonargparse.typing.register_type(
+        DataSourceContext,
+        data_source_context_serializer,
+        data_source_context_deserializer,
+    )
+
+    INITIALIZED = True

rslearn/utils/raster_format.py

@@ -2,8 +2,7 @@
 
 import hashlib
 import json
-from collections.abc import Callable
-from typing import Any, BinaryIO, TypeVar
+from typing import Any, BinaryIO
 
 import affine
 import numpy as np
@@ -14,34 +13,12 @@ from rasterio.crs import CRS
 from rasterio.enums import Resampling
 from upath import UPath
 
-from rslearn.config import RasterFormatConfig
 from rslearn.const import TILE_SIZE
 from rslearn.log_utils import get_logger
 from rslearn.utils.fsspec import open_rasterio_upath_reader, open_rasterio_upath_writer
 
 from .geometry import PixelBounds, Projection
 
-_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__)
 
 
@@ -219,7 +196,6 @@ class RasterFormat:
         raise NotImplementedError
 
 
-@RasterFormats.register("image_tile")
 class ImageTileRasterFormat(RasterFormat):
     """A RasterFormat that stores data in image tiles corresponding to grid cells.
 
@@ -468,7 +444,6 @@ class ImageTileRasterFormat(RasterFormat):
         )
 
 
-@RasterFormats.register("geotiff")
 class GeotiffRasterFormat(RasterFormat):
     """A raster format that uses one big, tiled GeoTIFF with small block size."""
 
@@ -623,7 +598,6 @@ class GeotiffRasterFormat(RasterFormat):
         return GeotiffRasterFormat(**kwargs)
 
 
-@RasterFormats.register("single_image")
 class SingleImageRasterFormat(RasterFormat):
     """A raster format that produces a single image called image.png/jpg.
 
@@ -775,17 +749,3 @@ class SingleImageRasterFormat(RasterFormat):
         if "format" in config:
             kwargs["format"] = config["format"]
         return SingleImageRasterFormat(**kwargs)
-
-
-def load_raster_format(config: RasterFormatConfig) -> RasterFormat:
-    """Loads a RasterFormat from a RasterFormatConfig.
-
-    Args:
-        config: the RasterFormatConfig configuration object specifying the
-            RasterFormat.
-
-    Returns:
-        the loaded RasterFormat implementation
-    """
-    cls = RasterFormats[config.name]
-    return cls.from_config(config.name, config.config_dict)

rslearn/utils/vector_format.py

@@ -1,15 +1,13 @@
 """Classes for writing vector data to a UPath."""
 
 import json
-from collections.abc import Callable
 from enum import Enum
-from typing import Any, TypeVar
+from typing import Any
 
 import shapely
 from rasterio.crs import CRS
 from upath import UPath
 
-from rslearn.config import VectorFormatConfig
 from rslearn.const import WGS84_PROJECTION
 from rslearn.log_utils import get_logger
 from rslearn.utils.fsspec import open_atomic
@@ -18,25 +16,6 @@ from .feature import Feature
 from .geometry import PixelBounds, Projection, STGeometry, safely_reproject_and_clip
 
 logger = get_logger(__name__)
-_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:
@@ -85,7 +64,6 @@ class VectorFormat:
         raise NotImplementedError
 
 
-@VectorFormats.register("tile")
 class TileVectorFormat(VectorFormat):
     """TileVectorFormat stores data in GeoJSON files corresponding to grid cells.
 
@@ -275,7 +253,6 @@ class GeojsonCoordinateMode(Enum):
     WGS84 = "wgs84"
 
 
-@VectorFormats.register("geojson")
 class GeojsonVectorFormat(VectorFormat):
     """A vector format that uses one big GeoJSON."""
 
@@ -429,17 +406,3 @@ class GeojsonVectorFormat(VectorFormat):
         if "coordinate_mode" in config:
             kwargs["coordinate_mode"] = GeojsonCoordinateMode(config["coordinate_mode"])
         return GeojsonVectorFormat(**kwargs)
-
-
-def load_vector_format(config: VectorFormatConfig) -> VectorFormat:
-    """Loads a VectorFormat from a VectorFormatConfig.
-
-    Args:
-        config: the VectorFormatConfig configuration object specifying the
-            VectorFormat.
-
-    Returns:
-        the loaded VectorFormat implementation
-    """
-    cls = VectorFormats[config.name]
-    return cls.from_config(config.name, config.config_dict)

{rslearn-0.0.14.dist-info → rslearn-0.0.16.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rslearn
-Version: 0.0.14
+Version: 0.0.16
 Summary: A library for developing remote sensing datasets and models
 Author: OlmoEarth Team
 License:                                  Apache License