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

rslearn/data_sources/xyz_tiles.py

@@ -186,6 +186,8 @@ class XyzTiles(DataSource):
     @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"]:
@@ -207,7 +209,7 @@ class XyzTiles(DataSource):
 
     def get_items(
         self, geometries: list[STGeometry], query_config: QueryConfig
-    ) -> list[list[list[Item]]]:
+    ) -> list[list[list[XyzItem]]]:
         """Get a list of items in the data source intersecting the given geometries.
 
         In XyzTiles we treat the data source as containing a single item, i.e., the
@@ -278,7 +280,7 @@ class XyzTiles(DataSource):
     def materialize(
         self,
         window: Window,
-        item_groups: list[list[Item]],
+        item_groups: list[list[XyzItem]],
         layer_name: str,
         layer_cfg: LayerConfig,
     ) -> None:
@@ -305,13 +307,10 @@ class XyzTiles(DataSource):
             window_projection, shapely.box(*window_bounds), None
         )
         projected_geometry = window_geometry.to_projection(self.projection)
-        projected_bounds = [
-            math.floor(projected_geometry.shp.bounds[0]),
-            math.floor(projected_geometry.shp.bounds[1]),
-            math.ceil(projected_geometry.shp.bounds[2]),
-            math.ceil(projected_geometry.shp.bounds[3]),
-        ]
-        projected_raster = self.read_bounds(item.url_template, projected_bounds)
+        projected_bounds = tuple(
+            math.floor(projected_geometry.shp.bounds[i]) for i in range(4)
+        )
+        projected_raster = self.read_bounds(item.url_template, projected_bounds)  # type: ignore
 
         # Attach the transform to the raster.
         src_transform = rasterio.transform.Affine(

rslearn/dataset/add_windows.py

@@ -25,7 +25,7 @@ def add_windows_from_geometries(
     window_size: int | None = None,
     time_range: tuple[datetime, datetime] | None = None,
     use_utm: bool = False,
-):
+) -> list[Window]:
     """Create windows based on a list of STGeometry.
 
     Args:

rslearn/dataset/dataset.py

@@ -7,10 +7,13 @@ import tqdm
 from upath import UPath
 
 from rslearn.config import TileStoreConfig, load_layer_config
+from rslearn.log_utils import get_logger
 from rslearn.tile_stores import TileStore, load_tile_store
 
 from .window import Window
 
+logger = get_logger(__name__)
+
 
 class Dataset:
     """A rslearn dataset.
@@ -37,21 +40,29 @@ class Dataset:
     materialize.
     """
 
-    def __init__(self, path: UPath) -> None:
+    def __init__(self, path: UPath, disabled_layers: list[str] = []) -> None:
         """Initializes a new Dataset.
 
         Args:
             path: the root directory of the dataset
+            disabled_layers: list of layers to disable
         """
         self.path = path
 
         # Load dataset configuration.
+
         with (self.path / "config.json").open("r") as f:
             config = json.load(f)
-            self.layers = {
-                layer_name: load_layer_config(d)
-                for layer_name, d in config["layers"].items()
-            }
+            self.layers = {}
+            for layer_name, d 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.tile_store_config = TileStoreConfig.from_config(config["tile_store"])
             self.materializer_name = config.get("materialize")
 

rslearn/dataset/manage.py

@@ -1,15 +1,17 @@
 """Functions to manage datasets."""
 
 import rslearn.data_sources
-from rslearn.config import LayerConfig, LayerType
+from rslearn.config import LayerConfig, LayerType, RasterLayerConfig
 from rslearn.data_sources import DataSource, Item
+from rslearn.log_utils import get_logger
 from rslearn.tile_stores import TileStore, get_tile_store_for_layer
-from rslearn.utils import logger
 
 from .dataset import Dataset
 from .materialize import Materializers
 from .window import Window, WindowLayerData
 
+logger = get_logger(__name__)
+
 
 def prepare_dataset_windows(
     dataset: Dataset, windows: list[Window], force: bool = False
@@ -37,7 +39,7 @@ def prepare_dataset_windows(
             if layer_name in layer_datas and not force:
                 continue
             needed_windows.append(window)
-        print(f"Preparing {len(needed_windows)} windows for layer {layer_name}")
+        logger.info(f"Preparing {len(needed_windows)} windows for layer {layer_name}")
         if len(needed_windows) == 0:
             continue
 
@@ -101,7 +103,7 @@ def ingest_dataset_windows(dataset: Dataset, windows: list[Window]) -> None:
             layer_cfg, dataset.path
         )
 
-        geometries_by_item = {}
+        geometries_by_item: dict = {}
         for window in windows:
             layer_datas = window.load_layer_datas()
             if layer_name not in layer_datas:
@@ -151,6 +153,7 @@ def is_window_ingested(
                 item = Item.deserialize(serialized_item)
 
                 if layer_cfg.layer_type == LayerType.RASTER:
+                    assert isinstance(layer_cfg, RasterLayerConfig)
                     for band_set in layer_cfg.band_sets:
                         projection, _ = band_set.get_final_projection_and_bounds(
                             window.projection, window.bounds
@@ -229,6 +232,8 @@ def materialize_window(
             item_group.append(item)
         item_groups.append(item_group)
 
+    if layer_cfg.data_source is None:
+        raise ValueError("data_source is required")
     if layer_cfg.data_source.ingest:
         if not is_window_ingested(dataset, window, check_layer_name=layer_name):
             logger.info(

rslearn/dataset/materialize.py

@@ -1,6 +1,6 @@
 """Classes to implement dataset materialization."""
 
-from typing import Any
+from typing import Any, Generic, TypeVar
 
 import numpy as np
 import numpy.typing as npt
@@ -24,8 +24,10 @@ from .window import Window
 
 Materializers = ClassRegistry()
 
+LayerConfigType = TypeVar("LayerConfigType", bound=LayerConfig)
 
-class Materializer:
+
+class Materializer(Generic[LayerConfigType]):
     """An abstract class that materializes data from a tile store."""
 
     def materialize(
@@ -33,7 +35,7 @@ class Materializer:
         tile_store: TileStore,
         window: Window,
         layer_name: str,
-        layer_cfg: LayerConfig,
+        layer_cfg: LayerConfigType,
         item_groups: list[list[Item]],
     ) -> None:
         """Materialize portions of items corresponding to this window into the dataset.
@@ -82,6 +84,8 @@ def read_raster_window_from_tiles(
     dst_row_offset = intersection[1] - bounds[1]
 
     src = ts_layer.read_raster(intersection)
+    if src is None:
+        raise ValueError(f"No raster data found for bounds {intersection}")
     src = src[src_indexes, :, :]
     if remapper:
         src = remapper(src, dst.dtype)
@@ -97,7 +101,7 @@ def read_raster_window_from_tiles(
 
 
 @Materializers.register("raster")
-class RasterMaterializer(Materializer):
+class RasterMaterializer(Materializer[RasterLayerConfig]):
     """A Materializer for raster data."""
 
     def materialize(
@@ -105,7 +109,7 @@ class RasterMaterializer(Materializer):
         tile_store: TileStore,
         window: Window,
         layer_name: str,
-        layer_cfg: LayerConfig,
+        layer_cfg: RasterLayerConfig,
         item_groups: list[list[Item]],
     ) -> None:
         """Materialize portions of items corresponding to this window into the dataset.
@@ -142,6 +146,12 @@ class RasterMaterializer(Materializer):
             if band_cfg.remap:
                 remapper = load_remapper(band_cfg.remap)
 
+            if band_cfg.format is None or band_cfg.bands is None or bounds is None:
+                raise ValueError(
+                    f"No raster format or bands specified for {layer_name} \
+                          with  {band_cfg}"
+                )
+
             raster_format = load_raster_format(
                 RasterFormatConfig(band_cfg.format["name"], band_cfg.format)
             )
@@ -182,6 +192,11 @@ class RasterMaterializer(Materializer):
                         ts_layer = layer_tile_store.get_layer(
                             (item.name, suffix, str(projection))
                         )
+                        if ts_layer is None:
+                            raise ValueError(
+                                f"No tile store layer found for {item.name} {suffix} \
+                                      {projection}"
+                            )
                         read_raster_window_from_tiles(
                             dst, ts_layer, bounds, src_indexes, dst_indexes, remapper
                         )
@@ -223,6 +238,8 @@ class VectorMaterializer(Materializer):
         projection, bounds = layer_cfg.get_final_projection_and_bounds(
             window.projection, window.bounds
         )
+        if bounds is None:
+            raise ValueError(f"No bounds specified for {layer_name}")
         vector_format = load_vector_format(layer_cfg.format)
 
         out_layer_dirs: list[UPath] = []
@@ -241,6 +258,10 @@ class VectorMaterializer(Materializer):
                 ts_layer = get_tile_store_for_layer(
                     tile_store, layer_name, layer_cfg
                 ).get_layer((item.name, str(projection)))
+                if ts_layer is None:
+                    raise ValueError(
+                        f"No tile store layer found for {item.name} {projection}"
+                    )
                 cur_features = ts_layer.read_vector(bounds)
                 features.extend(cur_features)
 

rslearn/dataset/window.py

@@ -7,9 +7,12 @@ from typing import Any
 import shapely
 from upath import UPath
 
+from rslearn.log_utils import get_logger
 from rslearn.utils import Projection, STGeometry
 from rslearn.utils.fsspec import open_atomic
 
+logger = get_logger(__name__)
+
 
 class WindowLayerData:
     """Layer data for retrieved layers specifying relevant items in the data source.
@@ -115,6 +118,7 @@ class Window:
             "options": self.options,
         }
         metadata_path = self.path / "metadata.json"
+        logger.info(f"Saving window metadata to {metadata_path}")
         with open_atomic(metadata_path, "w") as f:
             json.dump(metadata, f)
 
@@ -141,6 +145,7 @@ class Window:
         """Save layer datas to items.json."""
         json_data = [layer_data.serialize() for layer_data in layer_datas.values()]
         items_fname = self.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)
 

rslearn/log_utils.py

@@ -0,0 +1,24 @@
+"""Logging utilities."""
+
+import logging
+import os
+import sys
+
+LOG_FORMAT = "format=%(asctime)s loglevel=%(levelname)-6s logger=%(name)s %(funcName)s() L%(lineno)-4d %(message)s"
+# DETAILED_LOG_FORMAT = "format=%(asctime)s loglevel=%(levelname)-6s logger=%(name)s %(funcName)s() L%(lineno)-4d %(message)s   call_trace=%(pathname)s L%(lineno)-4d"  # noqa
+
+
+def get_logger(name: str) -> logging.Logger:
+    """Get a logger with a console handler."""
+    this_logger = logging.getLogger(name)
+    log_level = os.environ.get("RSLEARN_LOGLEVEL", "INFO")
+    if not this_logger.handlers:
+        console_handler = logging.StreamHandler(sys.stdout)
+        console_handler.setLevel(log_level)
+        console_formatter = logging.Formatter(LOG_FORMAT)
+        console_handler.setFormatter(console_formatter)
+        this_logger.addHandler(console_handler)
+
+    this_logger.setLevel(log_level)
+    this_logger.propagate = True
+    return this_logger