rslearn 0.0.17__py3-none-any.whl → 0.0.19__py3-none-any.whl

rslearn/arg_parser.py

@@ -1,6 +1,5 @@
 """Custom Lightning ArgumentParser with environment variable substitution support."""
 
-import os
 from typing import Any
 
 from jsonargparse import Namespace
@@ -21,11 +20,7 @@ class RslearnArgumentParser(LightningArgumentParser):
     def parse_string(
         self,
         cfg_str: str,
-        cfg_path: str | os.PathLike = "",
-        ext_vars: dict | None = None,
-        env: bool | None = None,
-        defaults: bool = True,
-        with_meta: bool | None = None,
+        *args: Any,
         **kwargs: Any,
     ) -> Namespace:
         """Pre-processes string for environment variable substitution before parsing."""
@@ -33,6 +28,4 @@ class RslearnArgumentParser(LightningArgumentParser):
         substituted_cfg_str = substitute_env_vars_in_string(cfg_str)
 
         # Call the parent method with the substituted config
-        return super().parse_string(
-            substituted_cfg_str, cfg_path, ext_vars, env, defaults, with_meta, **kwargs
-        )
+        return super().parse_string(substituted_cfg_str, *args, **kwargs)

rslearn/config/__init__.py

@@ -10,6 +10,7 @@ from .dataset import (
     LayerType,
     QueryConfig,
     SpaceMode,
+    StorageConfig,
     TimeMode,
 )
 
@@ -23,5 +24,6 @@ __all__ = [
     "LayerType",
     "QueryConfig",
     "SpaceMode",
+    "StorageConfig",
     "TimeMode",
 ]

rslearn/config/dataset.py

@@ -25,12 +25,13 @@ from rasterio.enums import Resampling
 from upath import UPath
 
 from rslearn.log_utils import get_logger
-from rslearn.utils import PixelBounds, Projection
+from rslearn.utils.geometry import PixelBounds, Projection, ResolutionFactor
 from rslearn.utils.raster_format import RasterFormat
 from rslearn.utils.vector_format import VectorFormat
 
 if TYPE_CHECKING:
     from rslearn.data_sources.data_source import DataSource
+    from rslearn.dataset.storage.storage import WindowStorageFactory
 
 logger = get_logger("__name__")
 
@@ -132,7 +133,11 @@ class BandSetConfig(BaseModel):
     bands.
     """
 
-    dtype: DType = Field(description="Pixel value type to store the data under")
+    model_config = ConfigDict(extra="forbid")
+
+    dtype: DType = Field(
+        description="Pixel value type to store the data under. This is used during dataset materialize and model predict."
+    )
     bands: list[str] = Field(
         default_factory=lambda: [],
         description="List of band names in this BandSetConfig. One of bands or num_bands must be set.",
@@ -210,22 +215,12 @@ class BandSetConfig(BaseModel):
         Returns:
             tuple of updated projection and bounds with zoom offset applied
         """
-        if self.zoom_offset == 0:
-            return projection, bounds
-        projection = Projection(
-            projection.crs,
-            projection.x_resolution / (2**self.zoom_offset),
-            projection.y_resolution / (2**self.zoom_offset),
-        )
-        if self.zoom_offset > 0:
-            zoom_factor = 2**self.zoom_offset
-            bounds = tuple(x * zoom_factor for x in bounds)  # type: ignore
+        if self.zoom_offset >= 0:
+            factor = ResolutionFactor(numerator=2**self.zoom_offset)
         else:
-            bounds = tuple(
-                x // (2 ** (-self.zoom_offset))
-                for x in bounds  # type: ignore
-            )
-        return projection, bounds
+            factor = ResolutionFactor(denominator=2 ** (-self.zoom_offset))
+
+        return (factor.multiply_projection(projection), factor.multiply_bounds(bounds))
 
     @field_validator("format", mode="before")
     @classmethod
@@ -329,7 +324,7 @@ class TimeMode(StrEnum):
 class QueryConfig(BaseModel):
     """A configuration for querying items in a data source."""
 
-    model_config = ConfigDict(frozen=True)
+    model_config = ConfigDict(frozen=True, extra="forbid")
 
     space_mode: SpaceMode = Field(
         default=SpaceMode.MOSAIC,
@@ -363,7 +358,7 @@ class QueryConfig(BaseModel):
 class DataSourceConfig(BaseModel):
     """Configuration for a DataSource in a dataset layer."""
 
-    model_config = ConfigDict(frozen=True)
+    model_config = ConfigDict(frozen=True, extra="forbid")
 
     class_path: str = Field(description="Class path for the data source.")
     init_args: dict[str, Any] = Field(
@@ -469,7 +464,7 @@ class CompositingMethod(StrEnum):
 class LayerConfig(BaseModel):
     """Configuration of a layer in a dataset."""
 
-    model_config = ConfigDict(frozen=True)
+    model_config = ConfigDict(frozen=True, extra="forbid")
 
     type: LayerType = Field(description="The LayerType (raster or vector).")
     data_source: DataSourceConfig | None = Field(
@@ -592,11 +587,60 @@ class LayerConfig(BaseModel):
         return vector_format
 
 
+class StorageConfig(BaseModel):
+    """Configuration for the WindowStorageFactory (window metadata storage backend)."""
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    class_path: str = Field(
+        default="rslearn.dataset.storage.file.FileWindowStorageFactory",
+        description="Class path for the WindowStorageFactory.",
+    )
+    init_args: dict[str, Any] = Field(
+        default_factory=lambda: {},
+        description="jsonargparse init args for the WindowStorageFactory.",
+    )
+
+    def instantiate_window_storage_factory(self) -> "WindowStorageFactory":
+        """Instantiate the WindowStorageFactory specified by this config."""
+        from rslearn.dataset.storage.storage import WindowStorageFactory
+        from rslearn.utils.jsonargparse import init_jsonargparse
+
+        init_jsonargparse()
+        parser = jsonargparse.ArgumentParser()
+        parser.add_argument("--wsf", type=WindowStorageFactory)
+        cfg = parser.parse_object(
+            {
+                "wsf": dict(
+                    class_path=self.class_path,
+                    init_args=self.init_args,
+                )
+            }
+        )
+        wsf = parser.instantiate_classes(cfg).wsf
+        return wsf
+
+
 class DatasetConfig(BaseModel):
     """Overall dataset configuration."""
 
+    model_config = ConfigDict(extra="forbid")
+
     layers: dict[str, LayerConfig] = Field(description="Layers in the dataset.")
     tile_store: dict[str, Any] = Field(
         default={"class_path": "rslearn.tile_stores.default.DefaultTileStore"},
         description="jsonargparse configuration for the TileStore.",
     )
+    storage: StorageConfig = Field(
+        default_factory=lambda: StorageConfig(),
+        description="jsonargparse configuration for the WindowStorageFactory.",
+    )
+
+    @field_validator("layers", mode="after")
+    @classmethod
+    def layer_names_validator(cls, v: dict[str, LayerConfig]) -> dict[str, LayerConfig]:
+        """Ensure layer names don't contain periods, since we use periods to distinguish different materialized groups within a layer."""
+        for layer_name in v.keys():
+            if "." in layer_name:
+                raise ValueError(f"layer names must not contain periods: {layer_name}")
+        return v

rslearn/dataset/add_windows.py

@@ -131,7 +131,7 @@ def add_windows_from_geometries(
                     f"_{time_range[0].isoformat()}_{time_range[1].isoformat()}"
                 )
         window = Window(
-            path=dataset.path / "windows" / group / cur_window_name,
+            storage=dataset.storage,
             group=group,
             name=cur_window_name,
             projection=cur_projection,

rslearn/dataset/dataset.py

@@ -1,9 +1,8 @@
 """rslearn dataset class."""
 
 import json
-import multiprocessing
+from typing import Any
 
-import tqdm
 from upath import UPath
 
 from rslearn.config import DatasetConfig
@@ -11,7 +10,6 @@ from rslearn.log_utils import get_logger
 from rslearn.template_params import substitute_env_vars_in_string
 from rslearn.tile_stores import TileStore, load_tile_store
 
-from .index import DatasetIndex
 from .window import Window
 
 logger = get_logger(__name__)
@@ -25,7 +23,7 @@ class Dataset:
     .. code-block:: none
 
         dataset/
-            config.json
+            config.json  # optional, if config provided as runtime object
             windows/
                 group1/
                     epsg:3857_10_623565_1528020/
@@ -42,106 +40,58 @@ class Dataset:
     materialize.
     """
 
-    def __init__(self, path: UPath, disabled_layers: list[str] = []) -> None:
+    def __init__(
+        self,
+        path: UPath,
+        disabled_layers: list[str] = [],
+        dataset_config: DatasetConfig | None = None,
+    ) -> None:
         """Initializes a new Dataset.
 
         Args:
             path: the root directory of the dataset
             disabled_layers: list of layers to disable
+            dataset_config: optional dataset configuration to use instead of loading from the dataset directory
         """
         self.path = path
 
-        # Load dataset configuration.
-        with (self.path / "config.json").open("r") as f:
-            config_content = f.read()
-            config_content = substitute_env_vars_in_string(config_content)
-            config = DatasetConfig.model_validate(json.loads(config_content))
-
-            self.layers = {}
-            for layer_name, layer_config in config.layers.items():
-                # Layer names must not contain period, since we use period to
-                # distinguish different materialized groups within a layer.
-                assert "." not in layer_name, "layer names must not contain periods"
-                if layer_name in disabled_layers:
-                    logger.warning(f"Layer {layer_name} is disabled")
-                    continue
-                self.layers[layer_name] = layer_config
-
-            self.tile_store_config = config.tile_store
-
-    def _get_index(self) -> DatasetIndex | None:
-        index_fname = self.path / DatasetIndex.FNAME
-        if not index_fname.exists():
-            return None
-        return DatasetIndex.load_index(self.path)
+        if dataset_config is None:
+            # Load dataset configuration from the dataset directory.
+            with (self.path / "config.json").open("r") as f:
+                config_content = f.read()
+                config_content = substitute_env_vars_in_string(config_content)
+                dataset_config = DatasetConfig.model_validate(
+                    json.loads(config_content)
+                )
+
+        self.layers = {}
+        for layer_name, layer_config in dataset_config.layers.items():
+            if layer_name in disabled_layers:
+                logger.warning(f"Layer {layer_name} is disabled")
+                continue
+            self.layers[layer_name] = layer_config
+
+        self.tile_store_config = dataset_config.tile_store
+        self.storage = (
+            dataset_config.storage.instantiate_window_storage_factory().get_storage(
+                self.path
+            )
+        )
 
     def load_windows(
         self,
         groups: list[str] | None = None,
         names: list[str] | None = None,
-        show_progress: bool = False,
-        workers: int = 0,
-        no_index: bool = False,
+        **kwargs: Any,
     ) -> list[Window]:
         """Load the windows in the dataset.
 
         Args:
             groups: an optional list of groups to filter loading
             names: an optional list of window names to filter loading
-            show_progress: whether to show tqdm progress bar
-            workers: number of parallel workers, default 0 (use main thread only to load windows)
-            no_index: don't use the dataset index even if it exists.
+            kwargs: optional keyword arguments to pass to WindowStorage.get_windows.
         """
-        # Load from index if it exists.
-        # We never use the index if names is set since loading the index will likely be
-        # slower than loading a few windows.
-        if not no_index and names is None:
-            dataset_index = self._get_index()
-            if dataset_index is not None:
-                return dataset_index.get_windows(groups=groups, names=names)
-
-        # Avoid directory does not exist errors later.
-        if not (self.path / "windows").exists():
-            return []
-
-        window_dirs = []
-        if not groups:
-            groups = []
-            for p in (self.path / "windows").iterdir():
-                groups.append(p.name)
-        for group in groups:
-            group_dir = self.path / "windows" / group
-            if not group_dir.exists():
-                logger.warning(
-                    f"Skipping group directory {group_dir} since it does not exist"
-                )
-                continue
-            if names:
-                cur_names = names
-            else:
-                cur_names = []
-                for p in group_dir.iterdir():
-                    cur_names.append(p.name)
-
-            for window_name in cur_names:
-                window_dir = group_dir / window_name
-                window_dirs.append(window_dir)
-
-        if workers == 0:
-            windows = [Window.load(window_dir) for window_dir in window_dirs]
-        else:
-            p = multiprocessing.Pool(workers)
-            outputs = p.imap_unordered(Window.load, window_dirs)
-            if show_progress:
-                outputs = tqdm.tqdm(
-                    outputs, total=len(window_dirs), desc="Loading windows"
-                )
-            windows = []
-            for window in outputs:
-                windows.append(window)
-            p.close()
-
-        return windows
+        return self.storage.get_windows(groups=groups, names=names, **kwargs)
 
     def get_tile_store(self) -> TileStore:
         """Get the tile store associated with this dataset.

rslearn/dataset/materialize.py

@@ -161,7 +161,7 @@ def build_first_valid_composite(
     nodata_vals: list[Any],
     bands: list[str],
     bounds: PixelBounds,
-    band_dtype: Any,
+    band_dtype: npt.DTypeLike,
     tile_store: TileStoreWithLayer,
     projection: Projection,
     remapper: Remapper | None,
@@ -233,7 +233,7 @@ def read_and_stack_raster_windows(
     projection: Projection,
     nodata_vals: list[Any],
     remapper: Remapper | None,
-    band_dtype: Any,
+    band_dtype: npt.DTypeLike,
     resampling_method: Resampling = Resampling.bilinear,
 ) -> npt.NDArray[np.generic]:
     """Create a stack of extent aligned raster windows.
@@ -326,7 +326,7 @@ def build_mean_composite(
     nodata_vals: list[Any],
     bands: list[str],
     bounds: PixelBounds,
-    band_dtype: Any,
+    band_dtype: npt.DTypeLike,
     tile_store: TileStoreWithLayer,
     projection: Projection,
     remapper: Remapper | None,
@@ -383,7 +383,7 @@ def build_median_composite(
     nodata_vals: list[Any],
     bands: list[str],
     bounds: PixelBounds,
-    band_dtype: Any,
+    band_dtype: npt.DTypeLike,
     tile_store: TileStoreWithLayer,
     projection: Projection,
     remapper: Remapper | None,
@@ -471,7 +471,7 @@ def build_composite(
         nodata_vals=nodata_vals,
         bands=band_cfg.bands,
         bounds=bounds,
-        band_dtype=band_cfg.dtype.value,
+        band_dtype=band_cfg.dtype.get_numpy_dtype(),
         tile_store=tile_store,
         projection=projection,
         resampling_method=layer_cfg.resampling_method.get_rasterio_resampling(),

rslearn/dataset/storage/__init__.py

@@ -0,0 +1 @@
+"""Storage backends for rslearn window metadata."""

rslearn/dataset/storage/file.py

@@ -0,0 +1,202 @@
+"""The default file-based window storage backend."""
+
+import json
+import multiprocessing
+
+import tqdm
+from typing_extensions import override
+from upath import UPath
+
+from rslearn.dataset.window import (
+    LAYERS_DIRECTORY_NAME,
+    Window,
+    WindowLayerData,
+    get_layer_and_group_from_dir_name,
+    get_window_layer_dir,
+)
+from rslearn.log_utils import get_logger
+from rslearn.utils.fsspec import open_atomic
+from rslearn.utils.mp import star_imap_unordered
+
+from .storage import WindowStorage, WindowStorageFactory
+
+logger = get_logger(__name__)
+
+
+def load_window(storage: "FileWindowStorage", window_dir: UPath) -> Window:
+    """Load the window from its directory by reading metadata.json.
+
+    Args:
+        storage: the underlying FileWindowStorage.
+        window_dir: the path where the window is stored.
+
+    Returns:
+        the window object.
+    """
+    metadata_fname = window_dir / "metadata.json"
+    with metadata_fname.open() as f:
+        metadata = json.load(f)
+    return Window.from_metadata(storage, metadata)
+
+
+class FileWindowStorage(WindowStorage):
+    """The default file-backed window storage."""
+
+    def __init__(self, path: UPath):
+        """Create a new FileWindowStorage.
+
+        Args:
+            path: the path to the dataset.
+        """
+        self.path = path
+
+    @override
+    def get_window_root(self, group: str, name: str) -> UPath:
+        return Window.get_window_root(self.path, group, name)
+
+    @override
+    def get_windows(
+        self,
+        groups: list[str] | None = None,
+        names: list[str] | None = None,
+        show_progress: bool = False,
+        workers: int = 0,
+    ) -> list["Window"]:
+        """Load the windows in the dataset.
+
+        Args:
+            groups: an optional list of groups to filter loading
+            names: an optional list of window names to filter loading
+            show_progress: whether to show tqdm progress bar
+            workers: number of parallel workers, default 0 (use main thread only to load windows)
+        """
+        # Avoid directory does not exist errors later.
+        if not (self.path / "windows").exists():
+            return []
+
+        window_dirs = []
+        if not groups:
+            groups = []
+            for p in (self.path / "windows").iterdir():
+                groups.append(p.name)
+        for group in groups:
+            group_dir = self.path / "windows" / group
+            if not group_dir.exists():
+                logger.warning(
+                    f"Skipping group directory {group_dir} since it does not exist"
+                )
+                continue
+            if names:
+                cur_names = names
+            else:
+                cur_names = []
+                for p in group_dir.iterdir():
+                    cur_names.append(p.name)
+
+            for window_name in cur_names:
+                window_dir = group_dir / window_name
+                window_dirs.append(window_dir)
+
+        if workers == 0:
+            windows = [load_window(self, window_dir) for window_dir in window_dirs]
+        else:
+            p = multiprocessing.Pool(workers)
+            outputs = star_imap_unordered(
+                p,
+                load_window,
+                [
+                    dict(storage=self, window_dir=window_dir)
+                    for window_dir in window_dirs
+                ],
+            )
+            if show_progress:
+                outputs = tqdm.tqdm(
+                    outputs, total=len(window_dirs), desc="Loading windows"
+                )
+            windows = []
+            for window in outputs:
+                windows.append(window)
+            p.close()
+
+        return windows
+
+    @override
+    def create_or_update_window(self, window: Window) -> None:
+        window_path = self.get_window_root(window.group, window.name)
+        window_path.mkdir(parents=True, exist_ok=True)
+        metadata_path = window_path / "metadata.json"
+        logger.debug(f"Saving window metadata to {metadata_path}")
+        with open_atomic(metadata_path, "w") as f:
+            json.dump(window.get_metadata(), f)
+
+    @override
+    def get_layer_datas(self, group: str, name: str) -> dict[str, "WindowLayerData"]:
+        window_path = self.get_window_root(group, name)
+        items_fname = window_path / "items.json"
+        if not items_fname.exists():
+            return {}
+
+        with items_fname.open() as f:
+            layer_datas = [
+                WindowLayerData.deserialize(layer_data) for layer_data in json.load(f)
+            ]
+
+        return {layer_data.layer_name: layer_data for layer_data in layer_datas}
+
+    @override
+    def save_layer_datas(
+        self, group: str, name: str, layer_datas: dict[str, "WindowLayerData"]
+    ) -> None:
+        window_path = self.get_window_root(group, name)
+        json_data = [layer_data.serialize() for layer_data in layer_datas.values()]
+        items_fname = window_path / "items.json"
+        logger.info(f"Saving window items to {items_fname}")
+        with open_atomic(items_fname, "w") as f:
+            json.dump(json_data, f)
+
+    @override
+    def list_completed_layers(self, group: str, name: str) -> list[tuple[str, int]]:
+        window_path = self.get_window_root(group, name)
+        layers_directory = window_path / LAYERS_DIRECTORY_NAME
+        if not layers_directory.exists():
+            return []
+
+        completed_layers = []
+        for layer_dir in layers_directory.iterdir():
+            layer_name, group_idx = get_layer_and_group_from_dir_name(layer_dir.name)
+            if not self.is_layer_completed(group, name, layer_name, group_idx):
+                continue
+            completed_layers.append((layer_name, group_idx))
+
+        return completed_layers
+
+    @override
+    def is_layer_completed(
+        self, group: str, name: str, layer_name: str, group_idx: int = 0
+    ) -> bool:
+        window_path = self.get_window_root(group, name)
+        layer_dir = get_window_layer_dir(
+            window_path,
+            layer_name,
+            group_idx,
+        )
+        return (layer_dir / "completed").exists()
+
+    @override
+    def mark_layer_completed(
+        self, group: str, name: str, layer_name: str, group_idx: int = 0
+    ) -> None:
+        window_path = self.get_window_root(group, name)
+        layer_dir = get_window_layer_dir(window_path, layer_name, group_idx)
+        # We assume the directory exists because the layer should be materialized before
+        # being marked completed.
+        (layer_dir / "completed").touch()
+
+
+class FileWindowStorageFactory(WindowStorageFactory):
+    """Factory class for FileWindowStorage."""
+
+    @override
+    def get_storage(self, ds_path: UPath) -> FileWindowStorage:
+        """Get a FileWindowStorage for the given dataset path."""
+        return FileWindowStorage(ds_path)