Allow dataset config to be passed for dataset init (#400)

* used with precedence over `path / config.json`
* bit of validation logic moved out into pydantic model

Author: cmwilhelm

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rslearn"
-version = "0.0.18"
+version = "0.0.19"
 description = "A library for developing remote sensing datasets and models"
 authors = [
     { name = "OlmoEarth Team" },

rslearn/config/dataset.py

@@ -635,3 +635,12 @@ class DatasetConfig(BaseModel):
         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/dataset.py

@@ -23,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/
@@ -40,37 +40,43 @@ 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
-            self.storage = (
-                config.storage.instantiate_window_storage_factory().get_storage(
-                    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,

rslearn/main.py

@@ -380,7 +380,7 @@ def apply_on_windows(
 
 def apply_on_windows_args(f: Callable[..., Any], args: argparse.Namespace) -> None:
     """Call apply_on_windows with arguments passed via command-line interface."""
-    dataset = Dataset(UPath(args.root), args.disabled_layers)
+    dataset = Dataset(UPath(args.root), disabled_layers=args.disabled_layers)
     apply_on_windows(
         f=f,
         dataset=dataset,

tests/unit/dataset/test_dataset.py

@@ -7,6 +7,7 @@
 import pytest
 from upath import UPath
 
+from rslearn.config import BandSetConfig, DatasetConfig, DType, LayerConfig, LayerType
 from rslearn.dataset import Dataset
 
 
@@ -110,3 +111,25 @@ def test_no_template_variables(self) -> None:
                 dataset.tile_store_config["init_args"]["path_suffix"]
                 == "/static/path/to/tiles"
             )
+
+    def test_load_from_config_object(self) -> None:
+        """Test that Dataset can be initialized with a pre-built DatasetConfig object."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            dataset_path = Path(temp_dir)
+
+            # Create a DatasetConfig object directly (no config.json file needed)
+            config = DatasetConfig(
+                layers={
+                    "images": LayerConfig(
+                        type=LayerType.RASTER,
+                        band_sets=[
+                            BandSetConfig(dtype=DType.UINT8, bands=["R", "G", "B"])
+                        ],
+                    ),
+                    "labels": LayerConfig(type=LayerType.VECTOR),
+                },
+            )
+
+            dataset = Dataset(UPath(dataset_path), dataset_config=config)
+
+            assert set(dataset.layers.keys()) == {"images", "labels"}