rslearn 0.0.1__py3-none-any.whl → 0.0.21__py3-none-any.whl

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/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)

rslearn/dataset/storage/storage.py

@@ -0,0 +1,140 @@
+"""Abstract classes for window metadata storage."""
+
+import abc
+from typing import TYPE_CHECKING
+
+from upath import UPath
+
+if TYPE_CHECKING:
+    from rslearn.dataset.window import Window, WindowLayerData
+
+
+class WindowStorage(abc.ABC):
+    """An abstract class for the storage backend for window metadata.
+
+    This is instantiated by a WindowStorageFactory for a specific rslearn dataset.
+
+    Window metadata includes the location and time range of windows (metadata.json),
+    the window layer datas (items.json), and the completed (materialized) layers. It
+    excludes the actual materialized data. All operations involving window metadata go
+    through the WindowStorage, including enumerating windows, creating new windows, and
+    updating window layer datas during `rslearn dataset prepare` or the completed
+    layers during `rslearn dataset materialize`.
+    """
+
+    @abc.abstractmethod
+    def get_window_root(self, group: str, name: str) -> UPath:
+        """Get the path where the window should be stored."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def get_windows(
+        self,
+        groups: list[str] | None = None,
+        names: list[str] | None = None,
+    ) -> 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
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def create_or_update_window(self, window: "Window") -> None:
+        """Create or update the window.
+
+        An existing window is only updated if there is one with the same name and group.
+
+        If there is a window with the same name but a different group, the behavior is
+        undefined.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def get_layer_datas(self, group: str, name: str) -> dict[str, "WindowLayerData"]:
+        """Get the window layer datas for the specified window.
+
+        Args:
+            group: the window group.
+            name: the window name.
+
+        Returns:
+            a dict mapping from the layer name to the layer data for that layer, if one
+                was previously saved.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def save_layer_datas(
+        self, group: str, name: str, layer_datas: dict[str, "WindowLayerData"]
+    ) -> None:
+        """Set the window layer datas for the specified window."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def list_completed_layers(self, group: str, name: str) -> list[tuple[str, int]]:
+        """List the layers available for this window that are completed.
+
+        Args:
+            group: the window group.
+            name: the window name.
+
+        Returns:
+            a list of (layer_name, group_idx) completed layers.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def is_layer_completed(
+        self, group: str, name: str, layer_name: str, group_idx: int = 0
+    ) -> bool:
+        """Check whether the specified layer is completed in the given window.
+
+        Completed means there is data in the layer and the data has been written
+        (materialized).
+
+        Args:
+            group: the window group.
+            name: the window name.
+            layer_name: the layer name.
+            group_idx: the index of the group within the layer.
+
+        Returns:
+            whether the layer is completed.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def mark_layer_completed(
+        self, group: str, name: str, layer_name: str, group_idx: int = 0
+    ) -> None:
+        """Mark the specified layer completed for the given window.
+
+        This must be done after the contents of the layer have been written. If a layer
+        has multiple groups, the caller should wait until the contents of all groups
+        have been written before marking them completed; this is because, when
+        materializing a window, we skip materialization if the first group
+        (group_idx=0) is marked completed.
+
+        Args:
+            group: the window group.
+            name: the window name.
+            layer_name: the layer name.
+            group_idx: the index of the group within the layer.
+        """
+        raise NotImplementedError
+
+
+class WindowStorageFactory(abc.ABC):
+    """An abstract class for a configurable storage backend for window metadata.
+
+    The dataset config includes a StorageConfig that configures a WindowStorageFactory,
+    which in turn creates a WindowStorage given a dataset path.
+    """
+
+    @abc.abstractmethod
+    def get_storage(self, ds_path: UPath) -> WindowStorage:
+        """Get a WindowStorage for the given dataset path."""
+        raise NotImplementedError

rslearn/dataset/window.py

@@ -1,14 +1,79 @@
 """rslearn windows."""
 
-import json
 from datetime import datetime
 from typing import Any
 
 import shapely
 from upath import UPath
 
+from rslearn.dataset.storage.storage import WindowStorage
+from rslearn.log_utils import get_logger
 from rslearn.utils import Projection, STGeometry
-from rslearn.utils.fsspec import open_atomic
+from rslearn.utils.raster_format import get_bandset_dirname
+
+logger = get_logger(__name__)
+
+LAYERS_DIRECTORY_NAME = "layers"
+
+
+def get_window_layer_dir(
+    window_path: UPath, layer_name: str, group_idx: int = 0
+) -> UPath:
+    """Get the directory containing materialized data for the specified layer.
+
+    Args:
+        window_path: the window directory.
+        layer_name: the layer name.
+        group_idx: the index of the group within the layer to get the directory
+            for (default 0).
+
+    Returns:
+        the path where data is or should be materialized.
+    """
+    if group_idx == 0:
+        folder_name = layer_name
+    else:
+        folder_name = f"{layer_name}.{group_idx}"
+    return window_path / LAYERS_DIRECTORY_NAME / folder_name
+
+
+def get_layer_and_group_from_dir_name(layer_dir_name: str) -> tuple[str, int]:
+    """Get the layer name and group index from the layer directory name.
+
+    Args:
+        layer_dir_name: the name of the layer folder.
+
+    Returns:
+        a tuple (layer_name, group_idx)
+    """
+    if "." in layer_dir_name:
+        parts = layer_dir_name.split(".")
+        if len(parts) != 2:
+            raise ValueError(
+                f"expected layer directory name {layer_dir_name} to only contain one '.'"
+            )
+        return (parts[0], int(parts[1]))
+    else:
+        return (layer_dir_name, 0)
+
+
+def get_window_raster_dir(
+    window_path: UPath, layer_name: str, bands: list[str], group_idx: int = 0
+) -> UPath:
+    """Get the directory where the raster is materialized.
+
+    Args:
+        window_path: the window directory.
+        layer_name: the layer name
+        bands: the bands in the raster. It should match a band set defined for this
+            layer.
+        group_idx: the index of the group within the layer.
+
+    Returns:
+        the directory containing the raster.
+    """
+    dirname = get_bandset_dirname(bands)
+    return get_window_layer_dir(window_path, layer_name, group_idx) / dirname
 
 
 class WindowLayerData:
@@ -69,7 +134,7 @@ class Window:
 
     def __init__(
         self,
-        path: UPath,
+        storage: WindowStorage,
         group: str,
         name: str,
         projection: Projection,
@@ -83,7 +148,7 @@ class Window:
         stored in metadata.json.
 
         Args:
-            path: the directory of this window
+            storage: the dataset storage for the underlying rslearn dataset.
             group: the group the window belongs to
             name: the unique name for this window
             projection: the projection of the window
@@ -91,7 +156,7 @@ class Window:
             time_range: optional time range of the window
             options: additional options (?)
         """
-        self.path = path
+        self.storage = storage
         self.group = group
         self.name = name
         self.projection = projection
@@ -99,25 +164,6 @@ class Window:
         self.time_range = time_range
         self.options = options
 
-    def save(self) -> None:
-        """Save the window metadata to its root directory."""
-        self.path.mkdir(parents=True, exist_ok=True)
-        metadata = {
-            "group": self.group,
-            "name": self.name,
-            "projection": self.projection.serialize(),
-            "bounds": self.bounds,
-            "time_range": (
-                [self.time_range[0].isoformat(), self.time_range[1].isoformat()]
-                if self.time_range
-                else None
-            ),
-            "options": self.options,
-        }
-        metadata_path = self.path / "metadata.json"
-        with open_atomic(metadata_path, "w") as f:
-            json.dump(metadata, f)
-
     def get_geometry(self) -> STGeometry:
         """Computes the STGeometry corresponding to this window."""
         return STGeometry(
@@ -128,41 +174,132 @@ class Window:
 
     def load_layer_datas(self) -> dict[str, WindowLayerData]:
         """Load layer datas describing items in retrieved layers from items.json."""
-        items_fname = self.path / "items.json"
-        if not items_fname.exists():
-            return {}
-        with items_fname.open("r") 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}
+        return self.storage.get_layer_datas(self.group, self.name)
 
     def save_layer_datas(self, layer_datas: dict[str, WindowLayerData]) -> None:
         """Save layer datas to items.json."""
-        json_data = [layer_data.serialize() for layer_data in layer_datas.values()]
-        items_fname = self.path / "items.json"
-        with open_atomic(items_fname, "w") as f:
-            json.dump(json_data, f)
+        self.storage.save_layer_datas(self.group, self.name, layer_datas)
+
+    def list_completed_layers(self) -> list[tuple[str, int]]:
+        """List the layers available for this window that are completed.
+
+        Returns:
+            a list of (layer_name, group_idx) completed layers.
+        """
+        return self.storage.list_completed_layers(self.group, self.name)
+
+    def get_layer_dir(self, layer_name: str, group_idx: int = 0) -> UPath:
+        """Get the directory containing materialized data for the specified layer.
+
+        Args:
+            layer_name: the layer name.
+            group_idx: the index of the group within the layer to get the directory
+                for (default 0).
+
+        Returns:
+            the path where data is or should be materialized.
+        """
+        return get_window_layer_dir(
+            self.storage.get_window_root(self.group, self.name), layer_name, group_idx
+        )
+
+    def is_layer_completed(self, layer_name: str, group_idx: int = 0) -> bool:
+        """Check whether the specified layer is completed.
+
+        Completed means there is data in the layer and the data has been written
+        (materialized).
+
+        Args:
+            layer_name: the layer name.
+            group_idx: the index of the group within the layer.
+
+        Returns:
+            whether the layer is completed
+        """
+        return self.storage.is_layer_completed(
+            self.group, self.name, layer_name, group_idx
+        )
+
+    def mark_layer_completed(self, layer_name: str, group_idx: int = 0) -> None:
+        """Mark the specified layer completed.
+
+        This must be done after the contents of the layer have been written. If a layer
+        has multiple groups, the caller should wait until the contents of all groups
+        have been written before marking them completed; this is because, when
+        materializing a window, we skip materialization if the first group
+        (group_idx=0) is marked completed.
+
+        Args:
+            layer_name: the layer name.
+            group_idx: the index of the group within the layer.
+        """
+        self.storage.mark_layer_completed(self.group, self.name, layer_name, group_idx)
+
+    def get_raster_dir(
+        self, layer_name: str, bands: list[str], group_idx: int = 0
+    ) -> UPath:
+        """Get the directory where the raster is materialized.
+
+        Args:
+            layer_name: the layer name
+            bands: the bands in the raster. It should match a band set defined for this
+                layer.
+            group_idx: the index of the group within the layer.
+
+        Returns:
+            the directory containing the raster.
+        """
+        return get_window_raster_dir(
+            self.storage.get_window_root(self.group, self.name),
+            layer_name,
+            bands,
+            group_idx,
+        )
+
+    def get_metadata(self) -> dict[str, Any]:
+        """Returns the window metadata dictionary."""
+        return {
+            "group": self.group,
+            "name": self.name,
+            "projection": self.projection.serialize(),
+            "bounds": self.bounds,
+            "time_range": (
+                [self.time_range[0].isoformat(), self.time_range[1].isoformat()]
+                if self.time_range
+                else None
+            ),
+            "options": self.options,
+        }
+
+    def save(self) -> None:
+        """Save the window metadata to its root directory."""
+        self.storage.create_or_update_window(self)
 
     @staticmethod
-    def load(path: UPath) -> "Window":
-        """Load a Window from a UPath.
+    def from_metadata(storage: WindowStorage, metadata: dict[str, Any]) -> "Window":
+        """Create a Window from the WindowStorage and the window's metadata dictionary.
 
         Args:
-            path: the root directory of the window
+            storage: the WindowStorage for the underlying dataset.
+            metadata: the window metadata.
 
         Returns:
             the Window
         """
-        metadata_fname = path / "metadata.json"
-        with metadata_fname.open("r") as f:
-            metadata = json.load(f)
+        # Ensure bounds is converted from list to tuple.
+        bounds = (
+            metadata["bounds"][0],
+            metadata["bounds"][1],
+            metadata["bounds"][2],
+            metadata["bounds"][3],
+        )
+
         return Window(
-            path=path,
+            storage=storage,
             group=metadata["group"],
             name=metadata["name"],
             projection=Projection.deserialize(metadata["projection"]),
-            bounds=metadata["bounds"],
+            bounds=bounds,
             time_range=(
                 (
                     datetime.fromisoformat(metadata["time_range"][0]),