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

rslearn/data_sources/worldcover.py

@@ -9,7 +9,8 @@ import requests
 from fsspec.implementations.local import LocalFileSystem
 from upath import UPath
 
-from rslearn.config import LayerConfig
+from rslearn.config import LayerType
+from rslearn.data_sources import DataSourceContext
 from rslearn.data_sources.local_files import LocalFiles
 from rslearn.log_utils import get_logger
 from rslearn.utils.fsspec import get_upath_local, join_upath, open_atomic
@@ -52,8 +53,8 @@ class WorldCover(LocalFiles):
 
     def __init__(
         self,
-        config: LayerConfig,
-        worldcover_dir: UPath,
+        worldcover_dir: str,
+        context: DataSourceContext = DataSourceContext(),
     ) -> None:
         """Create a new WorldCover.
 
@@ -64,18 +65,19 @@ class WorldCover(LocalFiles):
                 high performance, this should be a local directory; if the dataset is
                 remote, prefix with a protocol ("file://") to use a local directory
                 instead of a path relative to the dataset path.
+            context: the data source context.
         """
-        tif_dir = self.download_worldcover_data(worldcover_dir)
-        super().__init__(config, tif_dir)
-
-    @staticmethod
-    def from_config(config: LayerConfig, ds_path: UPath) -> "LocalFiles":
-        """Creates a new LocalFiles instance from a configuration dictionary."""
-        if config.data_source is None:
-            raise ValueError("LocalFiles data source requires a data source config")
-        d = config.data_source.config_dict
-        return WorldCover(
-            config=config, worldcover_dir=join_upath(ds_path, d["worldcover_dir"])
+        if context.ds_path is not None:
+            worldcover_upath = join_upath(context.ds_path, worldcover_dir)
+        else:
+            worldcover_upath = UPath(worldcover_dir)
+
+        tif_dir = self.download_worldcover_data(worldcover_upath)
+
+        super().__init__(
+            src_dir=tif_dir,
+            layer_type=LayerType.RASTER,
+            context=context,
         )
 
     def download_worldcover_data(self, worldcover_dir: UPath) -> UPath:

rslearn/data_sources/worldpop.py

@@ -6,10 +6,10 @@ from html.parser import HTMLParser
 from urllib.parse import urljoin
 
 import requests
-import requests.auth
 from upath import UPath
 
-from rslearn.config import LayerConfig
+from rslearn.config import LayerType
+from rslearn.data_sources import DataSourceContext
 from rslearn.data_sources.local_files import LocalFiles
 from rslearn.log_utils import get_logger
 from rslearn.utils.fsspec import join_upath, open_atomic
@@ -59,33 +59,30 @@ class WorldPop(LocalFiles):
 
     def __init__(
         self,
-        config: LayerConfig,
-        worldpop_dir: UPath,
+        worldpop_dir: str,
         timeout: timedelta = timedelta(seconds=30),
+        context: DataSourceContext = DataSourceContext(),
     ):
         """Create a new WorldPop.
 
         Args:
-            config: configuration for this layer. It should specify a single band
-                called B1 which will contain the population counts.
             worldpop_dir: the directory to extract the WorldPop GeoTIFF files. For
                 high performance, this should be a local directory; if the dataset is
                 remote, prefix with a protocol ("file://") to use a local directory
                 instead of a path relative to the dataset path.
             timeout: timeout for HTTP requests.
+            context: the data source context.
         """
-        worldpop_dir.mkdir(parents=True, exist_ok=True)
-        self.download_worldpop_data(worldpop_dir, timeout)
-        super().__init__(config, worldpop_dir)
-
-    @staticmethod
-    def from_config(config: LayerConfig, ds_path: UPath) -> "LocalFiles":
-        """Creates a new LocalFiles instance from a configuration dictionary."""
-        if config.data_source is None:
-            raise ValueError("LocalFiles data source requires a data source config")
-        d = config.data_source.config_dict
-        return WorldPop(
-            config=config, worldpop_dir=join_upath(ds_path, d["worldpop_dir"])
+        if context.ds_path is not None:
+            worldpop_upath = join_upath(context.ds_path, worldpop_dir)
+        else:
+            worldpop_upath = UPath(worldpop_dir)
+        worldpop_upath.mkdir(parents=True, exist_ok=True)
+        self.download_worldpop_data(worldpop_upath, timeout)
+        super().__init__(
+            src_dir=worldpop_upath,
+            layer_type=LayerType.RASTER,
+            context=context,
         )
 
     def download_worldpop_data(self, worldpop_dir: UPath, timeout: timedelta) -> None:

rslearn/data_sources/xyz_tiles.py

@@ -14,9 +14,8 @@ import shapely
 from PIL import Image
 from rasterio.crs import CRS
 from rasterio.enums import Resampling
-from upath import UPath
 
-from rslearn.config import LayerConfig, QueryConfig, RasterLayerConfig
+from rslearn.config import LayerConfig, QueryConfig
 from rslearn.dataset import Window
 from rslearn.dataset.materialize import RasterMaterializer
 from rslearn.tile_stores import TileStore, TileStoreWithLayer
@@ -24,7 +23,7 @@ from rslearn.utils import PixelBounds, Projection, STGeometry
 from rslearn.utils.array import copy_spatial_array
 from rslearn.utils.raster_format import get_transform_from_projection_and_bounds
 
-from .data_source import DataSource, Item
+from .data_source import DataSource, DataSourceContext, Item
 from .utils import match_candidate_items_to_window
 
 WEB_MERCATOR_EPSG = 3857
@@ -96,11 +95,12 @@ class XyzTiles(DataSource, TileStore):
         url_templates: list[str],
         time_ranges: list[tuple[datetime, datetime]],
         zoom: int,
-        crs: CRS = CRS.from_epsg(WEB_MERCATOR_EPSG),
+        crs: str | CRS = CRS.from_epsg(WEB_MERCATOR_EPSG),
         total_units: float = WEB_MERCATOR_UNITS,
         offset: float = WEB_MERCATOR_UNITS / 2,
         tile_size: int = 256,
         band_names: list[str] = ["R", "G", "B"],
+        context: DataSourceContext = DataSourceContext(),
     ):
         """Initialize an XyzTiles instance.
 
@@ -121,16 +121,22 @@ class XyzTiles(DataSource, TileStore):
             offset: offset added to projection units when converting to tile positions.
             tile_size: size in pixels of each tile. Tiles must be square.
             band_names: what to name the bands that we read.
+            context: the data source context.
         """
         self.url_templates = url_templates
         self.time_ranges = time_ranges
         self.zoom = zoom
-        self.crs = crs
         self.total_units = total_units
         self.offset = offset
         self.tile_size = tile_size
         self.band_names = band_names
 
+        # Convert to CRS if needed.
+        if isinstance(crs, str):
+            self.crs = CRS.from_string(crs)
+        else:
+            self.crs = crs
+
         # Compute total number of pixels (a function of the zoom level and tile size).
         self.total_pixels = tile_size * (2**zoom)
         # Compute pixel size (resolution).
@@ -153,30 +159,6 @@ class XyzTiles(DataSource, TileStore):
             item = Item(url_template, geometry)
             self.items.append(item)
 
-    @staticmethod
-    def from_config(config: LayerConfig, ds_path: UPath) -> "XyzTiles":
-        """Creates a new XyzTiles instance from a configuration dictionary."""
-        if config.data_source is None:
-            raise ValueError("data_source is required")
-        d = config.data_source.config_dict
-        time_ranges = []
-        for str1, str2 in d["time_ranges"]:
-            time1 = datetime.fromisoformat(str1)
-            time2 = datetime.fromisoformat(str2)
-            time_ranges.append((time1, time2))
-        kwargs = dict(
-            url_templates=d["url_templates"], zoom=d["zoom"], time_ranges=time_ranges
-        )
-        if "crs" in d:
-            kwargs["crs"] = CRS.from_string(d["crs"])
-        if "total_units" in d:
-            kwargs["total_units"] = d["total_units"]
-        if "offset" in d:
-            kwargs["offset"] = d["offset"]
-        if "tile_size" in d:
-            kwargs["tile_size"] = d["tile_size"]
-        return XyzTiles(**kwargs)
-
     def get_items(
         self, geometries: list[STGeometry], query_config: QueryConfig
     ) -> list[list[list[Item]]]:
@@ -381,7 +363,6 @@ class XyzTiles(DataSource, TileStore):
             layer_name: the name of this layer
             layer_cfg: the config of this layer
         """
-        assert isinstance(layer_cfg, RasterLayerConfig)
         RasterMaterializer().materialize(
             TileStoreWithLayer(self, layer_name),
             window,

rslearn/dataset/dataset.py

@@ -6,7 +6,7 @@ import multiprocessing
 import tqdm
 from upath import UPath
 
-from rslearn.config import load_layer_config
+from rslearn.config import DatasetConfig
 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
@@ -55,19 +55,19 @@ class Dataset:
         with (self.path / "config.json").open("r") as f:
             config_content = f.read()
             config_content = substitute_env_vars_in_string(config_content)
-            config = json.loads(config_content)
+            config = DatasetConfig.model_validate(json.loads(config_content))
+
             self.layers = {}
-            for layer_name, d in config["layers"].items():
+            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] = load_layer_config(d)
+                self.layers[layer_name] = layer_config
 
-            self.tile_store_config = config.get("tile_store", None)
-            self.materializer_name = config.get("materialize")
+            self.tile_store_config = config.tile_store
 
     def _get_index(self) -> DatasetIndex | None:
         index_fname = self.path / DatasetIndex.FNAME

rslearn/dataset/manage.py

@@ -6,11 +6,9 @@ from collections.abc import Callable
 from datetime import timedelta
 from typing import Any
 
-import rslearn.data_sources
 from rslearn.config import (
     LayerConfig,
     LayerType,
-    RasterLayerConfig,
 )
 from rslearn.data_sources import DataSource, Item
 from rslearn.dataset.handler_summaries import (
@@ -24,7 +22,7 @@ from rslearn.log_utils import get_logger
 from rslearn.tile_stores import TileStore, get_tile_store_with_layer
 
 from .dataset import Dataset
-from .materialize import Materializers
+from .materialize import Materializer, RasterMaterializer, VectorMaterializer
 from .window import Window, WindowLayerData
 
 logger = get_logger(__name__)
@@ -124,12 +122,24 @@ def prepare_dataset_windows(
             )
             continue
         data_source_cfg = layer_cfg.data_source
+        min_matches = data_source_cfg.query_config.min_matches
 
         # Get windows that need to be prepared for this layer.
+        # Also track which windows are skipped vs previously rejected.
         needed_windows = []
+        windows_skipped = 0
+        windows_rejected = 0
         for window in windows:
             layer_datas = window.load_layer_datas()
             if layer_name in layer_datas and not force:
+                # Window already has layer data - check if it was previously rejected
+                layer_data = layer_datas[layer_name]
+                if len(layer_data.serialized_item_groups) == 0 and min_matches > 0:
+                    # Previously rejected due to min_matches
+                    windows_rejected += 1
+                else:
+                    # Successfully prepared previously
+                    windows_skipped += 1
                 continue
             needed_windows.append(window)
         logger.info(f"Preparing {len(needed_windows)} windows for layer {layer_name}")
@@ -138,11 +148,11 @@ def prepare_dataset_windows(
             layer_summaries.append(
                 LayerPrepareSummary(
                     layer_name=layer_name,
-                    data_source_name=data_source_cfg.name,
+                    data_source_name=data_source_cfg.class_path,
                     duration_seconds=time.monotonic() - layer_start_time,
                     windows_prepared=0,
-                    windows_skipped=len(windows),
-                    windows_rejected=0,
+                    windows_skipped=windows_skipped,
+                    windows_rejected=windows_rejected,
                     get_items_attempts=0,
                 )
             )
@@ -150,9 +160,7 @@ def prepare_dataset_windows(
 
         # Create data source after checking for at least one window so it can be fast
         # if there are no windows to prepare.
-        data_source = rslearn.data_sources.data_source_from_config(
-            layer_cfg, dataset.path
-        )
+        data_source = layer_cfg.instantiate_data_source(dataset.path)
 
         # Get STGeometry for each window.
         geometries = []
@@ -184,8 +192,6 @@ def prepare_dataset_windows(
         )
 
         windows_prepared = 0
-        windows_rejected = 0
-        min_matches = data_source_cfg.query_config.min_matches
         for window, result in zip(needed_windows, results):
             layer_datas = window.load_layer_datas()
             layer_datas[layer_name] = WindowLayerData(
@@ -202,12 +208,10 @@ def prepare_dataset_windows(
             else:
                 windows_prepared += 1
 
-        windows_skipped = len(windows) - len(needed_windows)
-
         layer_summaries.append(
             LayerPrepareSummary(
                 layer_name=layer_name,
-                data_source_name=data_source_cfg.name,
+                data_source_name=data_source_cfg.class_path,
                 duration_seconds=time.monotonic() - layer_start_time,
                 windows_prepared=windows_prepared,
                 windows_skipped=windows_skipped,
@@ -250,9 +254,7 @@ def ingest_dataset_windows(
         if not layer_cfg.data_source.ingest:
             continue
 
-        data_source = rslearn.data_sources.data_source_from_config(
-            layer_cfg, dataset.path
-        )
+        data_source = layer_cfg.instantiate_data_source(dataset.path)
 
         geometries_by_item: dict = {}
         for window in windows:
@@ -314,8 +316,7 @@ def is_window_ingested(
             for serialized_item in group:
                 item = Item.deserialize(serialized_item)
 
-                if layer_cfg.layer_type == LayerType.RASTER:
-                    assert isinstance(layer_cfg, RasterLayerConfig)
+                if layer_cfg.type == LayerType.RASTER:
                     for band_set in layer_cfg.band_sets:
                         # Make sure that layers exist containing each configured band.
                         # And that those layers are marked completed.
@@ -409,10 +410,13 @@ def materialize_window(
             f"Materializing {len(item_groups)} item groups in layer {layer_name} from tile store"
         )
 
-        if dataset.materializer_name:
-            materializer = Materializers[dataset.materializer_name]()
+        materializer: Materializer
+        if layer_cfg.type == LayerType.RASTER:
+            materializer = RasterMaterializer()
+        elif layer_cfg.type == LayerType.VECTOR:
+            materializer = VectorMaterializer()
         else:
-            materializer = Materializers[layer_cfg.layer_type.value]()
+            raise ValueError(f"unknown layer type {layer_cfg.type}")
 
         retry(
             fn=lambda: materializer.materialize(
@@ -483,10 +487,8 @@ def materialize_dataset_windows(
         if not layer_cfg.data_source:
             total_skipped = len(windows)
         else:
-            data_source_name = layer_cfg.data_source.name
-            data_source = rslearn.data_sources.data_source_from_config(
-                layer_cfg, dataset.path
-            )
+            data_source_name = layer_cfg.data_source.class_path
+            data_source = layer_cfg.instantiate_data_source(dataset.path)
 
             for window in windows:
                 window_summary = materialize_window(

rslearn/dataset/materialize.py

@@ -1,7 +1,6 @@
 """Classes to implement dataset materialization."""
 
-from collections.abc import Callable
-from typing import Any, Generic, TypeVar
+from typing import Any
 
 import numpy as np
 import numpy.typing as npt
@@ -11,45 +10,17 @@ from rslearn.config import (
     BandSetConfig,
     CompositingMethod,
     LayerConfig,
-    RasterFormatConfig,
-    RasterLayerConfig,
-    VectorLayerConfig,
 )
 from rslearn.data_sources.data_source import ItemType
 from rslearn.tile_stores import TileStoreWithLayer
 from rslearn.utils.feature import Feature
 from rslearn.utils.geometry import PixelBounds, Projection
-from rslearn.utils.raster_format import load_raster_format
-from rslearn.utils.vector_format import load_vector_format
 
 from .remap import Remapper, load_remapper
 from .window import Window
 
-_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)
-
-
-class Materializer(Generic[LayerConfigType]):
+class Materializer:
     """An abstract class that materializes data from a tile store."""
 
     def materialize(
@@ -57,7 +28,7 @@ class Materializer(Generic[LayerConfigType]):
         tile_store: TileStoreWithLayer,
         window: Window,
         layer_name: str,
-        layer_cfg: LayerConfigType,
+        layer_cfg: LayerConfig,
         item_groups: list[list[ItemType]],
     ) -> None:
         """Materialize portions of items corresponding to this window into the dataset.
@@ -473,7 +444,7 @@ def build_composite(
     group: list[ItemType],
     compositing_method: CompositingMethod,
     tile_store: TileStoreWithLayer,
-    layer_cfg: RasterLayerConfig,
+    layer_cfg: LayerConfig,
     band_cfg: BandSetConfig,
     projection: Projection,
     bounds: PixelBounds,
@@ -503,13 +474,12 @@ def build_composite(
         band_dtype=band_cfg.dtype.value,
         tile_store=tile_store,
         projection=projection,
-        resampling_method=layer_cfg.resampling_method,
+        resampling_method=layer_cfg.resampling_method.get_rasterio_resampling(),
         remapper=remapper,
     )
 
 
-@Materializers.register("raster")
-class RasterMaterializer(Materializer[RasterLayerConfig]):
+class RasterMaterializer(Materializer):
     """A Materializer for raster data."""
 
     def materialize(
@@ -517,7 +487,7 @@ class RasterMaterializer(Materializer[RasterLayerConfig]):
         tile_store: TileStoreWithLayer,
         window: Window,
         layer_name: str,
-        layer_cfg: RasterLayerConfig,
+        layer_cfg: LayerConfig,
         item_groups: list[list[ItemType]],
     ) -> None:
         """Materialize portions of items corresponding to this window into the dataset.
@@ -529,8 +499,6 @@ class RasterMaterializer(Materializer[RasterLayerConfig]):
             layer_cfg: the configuration of the layer to materialize
             item_groups: the items associated with this window and layer
         """
-        assert isinstance(layer_cfg, RasterLayerConfig)
-
         for band_cfg in layer_cfg.band_sets:
             # band_cfg could specify zoom_offset and maybe other parameters that affect
             # projection/bounds, so use the corrected projection/bounds.
@@ -543,9 +511,7 @@ class RasterMaterializer(Materializer[RasterLayerConfig]):
             if band_cfg.remap:
                 remapper = load_remapper(band_cfg.remap)
 
-            raster_format = load_raster_format(
-                RasterFormatConfig(band_cfg.format["name"], band_cfg.format)
-            )
+            raster_format = band_cfg.instantiate_raster_format()
 
             for group_id, group in enumerate(item_groups):
                 composite = build_composite(
@@ -569,7 +535,6 @@ class RasterMaterializer(Materializer[RasterLayerConfig]):
             window.mark_layer_completed(layer_name, group_id)
 
 
-@Materializers.register("vector")
 class VectorMaterializer(Materializer):
     """A Materializer for vector data."""
 
@@ -590,8 +555,7 @@ class VectorMaterializer(Materializer):
             layer_cfg: the configuration of the layer to materialize
             item_groups: the items associated with this window and layer
         """
-        assert isinstance(layer_cfg, VectorLayerConfig)
-        vector_format = load_vector_format(layer_cfg.format)
+        vector_format = layer_cfg.instantiate_vector_format()
 
         for group_id, group in enumerate(item_groups):
             features: list[Feature] = []