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

rslearn/data_sources/local_files.py

@@ -1,52 +1,61 @@
 """Data source for raster or vector data in local files."""
 
-from typing import Any
+import functools
+import json
+from typing import Any, Generic, TypeVar
 
 import fiona
-import rasterio
 import shapely
 import shapely.geometry
-from class_registry import ClassRegistry
 from rasterio.crs import CRS
 from upath import UPath
 
 import rslearn.data_sources.utils
-from rslearn.config import LayerConfig, LayerType, VectorLayerConfig
-from rslearn.const import SHAPEFILE_AUX_EXTENSIONS, WGS84_PROJECTION
-from rslearn.tile_stores import LayerMetadata, PrefixedTileStore, TileStore
-from rslearn.utils import Feature, Projection, STGeometry
-from rslearn.utils.fsspec import get_upath_local, join_upath
+from rslearn.config import LayerType
+from rslearn.const import SHAPEFILE_AUX_EXTENSIONS
+from rslearn.log_utils import get_logger
+from rslearn.tile_stores import TileStoreWithLayer
+from rslearn.utils.feature import Feature
+from rslearn.utils.fsspec import (
+    get_relative_suffix,
+    get_upath_local,
+    join_upath,
+    open_rasterio_upath_reader,
+)
+from rslearn.utils.geometry import Projection, STGeometry, get_global_geometry
 
-from .data_source import DataSource, Item, QueryConfig
-from .raster_source import get_needed_projections, ingest_raster
+from .data_source import DataSource, DataSourceContext, Item, QueryConfig
 
-Importers = ClassRegistry()
+logger = get_logger("__name__")
 
 
-class Importer:
-    """An abstract class for importing data from local files."""
+ItemType = TypeVar("ItemType", bound=Item)
+ImporterType = TypeVar("ImporterType", bound="Importer")
 
-    def list_items(self, config: LayerConfig, src_dir: UPath) -> list[Item]:
+SOURCE_NAME = "rslearn.data_sources.local_files.LocalFiles"
+
+
+class Importer(Generic[ItemType]):
+    """An abstract base class for importing data from local files."""
+
+    def list_items(self, src_dir: UPath) -> list[ItemType]:
         """Extract a list of Items from the source directory.
 
         Args:
-            config: the configuration of the layer.
             src_dir: the source directory.
         """
         raise NotImplementedError
 
     def ingest_item(
         self,
-        config: LayerConfig,
-        tile_store: TileStore,
-        item: Item,
+        tile_store: TileStoreWithLayer,
+        item: ItemType,
         cur_geometries: list[STGeometry],
-    ):
+    ) -> None:
         """Ingest the specified local file item.
 
         Args:
-            config: the configuration of the layer.
-            tile_store: the TileStore to ingest the data into.
+            tile_store: the tile store to ingest the data into.
             item: the Item to ingest
             cur_geometries: the geometries where the item is needed.
         """
@@ -58,7 +67,7 @@ class RasterItemSpec:
 
     def __init__(
         self,
-        fnames: list[UPath],
+        fnames: list[str],
         bands: list[list[str]] | None = None,
         name: str | None = None,
     ):
@@ -73,25 +82,6 @@ class RasterItemSpec:
         self.bands = bands
         self.name = name
 
-    @staticmethod
-    def from_config(src_dir: UPath, d: dict[str, Any]) -> "RasterItemSpec":
-        """Decode a dict into a RasterItemSpec.
-
-        Args:
-            src_dir: the source directory.
-            d: the configuration dict.
-
-        Returns:
-            the RasterItemSpec.
-        """
-        kwargs = dict(
-            fnames=[join_upath(src_dir, suffix) for suffix in d["fnames"]],
-            bands=d["bands"],
-        )
-        if "name" in d:
-            kwargs["name"] = d["name"]
-        return RasterItemSpec(**kwargs)
-
     def serialize(self) -> dict[str, Any]:
         """Serializes the RasterItemSpec to a JSON-encodable dictionary."""
         return {
@@ -104,7 +94,7 @@ class RasterItemSpec:
     def deserialize(d: dict[str, Any]) -> "RasterItemSpec":
         """Deserializes a RasterItemSpec from a JSON-decoded dictionary."""
         return RasterItemSpec(
-            fnames=[UPath(s) for s in d["fnames"]],
+            fnames=[s for s in d["fnames"]],
             bands=d["bands"],
             name=d["name"],
         )
@@ -113,29 +103,37 @@ class RasterItemSpec:
 class RasterItem(Item):
     """An item corresponding to a local file."""
 
-    def __init__(self, name: str, geometry: STGeometry, spec: RasterItemSpec):
-        """Creates a new LocalFileItem.
+    def __init__(
+        self, name: str, geometry: STGeometry, src_dir: str, spec: RasterItemSpec
+    ):
+        """Creates a new RasterItem.
 
         Args:
             name: unique name of the item
             geometry: the spatial and temporal extent of the item
+            src_dir: the source directory.
             spec: the RasterItemSpec that specifies the filename(s) and bands.
         """
         super().__init__(name, geometry)
+        self.src_dir = src_dir
         self.spec = spec
 
     def serialize(self) -> dict:
         """Serializes the item to a JSON-encodable dictionary."""
         d = super().serialize()
+        d["src_dir"] = str(self.src_dir)
         d["spec"] = self.spec.serialize()
         return d
 
     @staticmethod
-    def deserialize(d: dict) -> Item:
+    def deserialize(d: dict) -> "RasterItem":
         """Deserializes an item from a JSON-decoded dictionary."""
         item = super(RasterItem, RasterItem).deserialize(d)
+        src_dir = d["src_dir"]
         spec = RasterItemSpec.deserialize(d["spec"])
-        return RasterItem(name=item.name, geometry=item.geometry, spec=spec)
+        return RasterItem(
+            name=item.name, geometry=item.geometry, src_dir=src_dir, spec=spec
+        )
 
 
 class VectorItem(Item):
@@ -159,7 +157,7 @@ class VectorItem(Item):
         return d
 
     @staticmethod
-    def deserialize(d: dict) -> Item:
+    def deserialize(d: dict) -> "VectorItem":
         """Deserializes an item from a JSON-decoded dictionary."""
         item = super(VectorItem, VectorItem).deserialize(d)
         return VectorItem(
@@ -167,119 +165,146 @@ class VectorItem(Item):
         )
 
 
-@Importers.register("raster")
 class RasterImporter(Importer):
     """An Importer for raster data."""
 
-    def list_items(self, config: LayerConfig, src_dir: UPath) -> list[Item]:
+    def __init__(self, item_specs: list[RasterItemSpec] | None = None):
+        """Create a new RasterImporter.
+
+        Args:
+            item_specs: the specs to specify the raster items directly. If None, the
+                raster items are automatically detected from the files in the source
+                directory.
+        """
+        self.item_specs = item_specs
+
+    def list_items(self, src_dir: UPath) -> list[Item]:
         """Extract a list of Items from the source directory.
 
         Args:
-            config: the configuration of the layer.
             src_dir: the source directory.
         """
-        item_specs: list[RasterItemSpec] = []
+        item_specs: list[RasterItemSpec]
+
         # See if user has provided the item specs directly.
-        if "item_specs" in config.data_source.config_dict:
-            for spec_dict in config.data_source.config_dict["item_specs"]:
-                spec = RasterItemSpec.from_config(src_dir, spec_dict)
-                item_specs.append(spec)
+        if self.item_specs is not None:
+            item_specs = self.item_specs
         else:
             # Otherwise we need to list files and assume each one is separate.
             # And we'll need to autodetect the bands later.
+            item_specs = []
             file_paths = src_dir.glob("**/*.*")
             for path in file_paths:
-                spec = RasterItemSpec(fnames=[path], bands=None)
+                # Ignore JSON files.
+                if path.name.endswith(".json"):
+                    continue
+
+                # Ignore temporary files that may be created by open_atomic.
+                # The suffix should be like "X.tif.tmp.1234".
+                parts = path.name.split(".")
+                if len(parts) >= 4 and parts[-2] == "tmp" and parts[-1].isdigit():
+                    continue
+
+                spec = RasterItemSpec(
+                    fnames=[get_relative_suffix(src_dir, path)], bands=None
+                )
                 item_specs.append(spec)
 
         items: list[Item] = []
         for spec in item_specs:
             # Get geometry from the first raster file.
             # We assume files are readable with rasterio.
-            with spec.fnames[0].open("rb") as f:
-                with rasterio.open(f) as src:
-                    crs = src.crs
-                    left = src.transform.c
-                    top = src.transform.f
-                    # Resolutions in projection units per pixel.
-                    x_resolution = src.transform.a
-                    y_resolution = src.transform.e
-                    start = (int(left / x_resolution), int(top / y_resolution))
-                    shp = shapely.box(
-                        start[0], start[1], start[0] + src.width, start[1] + src.height
-                    )
-                    projection = Projection(crs, x_resolution, y_resolution)
-                    geometry = STGeometry(projection, shp, None)
+            fname = join_upath(src_dir, spec.fnames[0])
+            with open_rasterio_upath_reader(fname) as src:
+                crs = src.crs
+                left = src.transform.c
+                top = src.transform.f
+                # Resolutions in projection units per pixel.
+                x_resolution = src.transform.a
+                y_resolution = src.transform.e
+                start = (int(left / x_resolution), int(top / y_resolution))
+                shp = shapely.box(
+                    start[0], start[1], start[0] + src.width, start[1] + src.height
+                )
+                projection = Projection(crs, x_resolution, y_resolution)
+                geometry = STGeometry(projection, shp, None)
+
+            if geometry.is_too_large():
+                geometry = get_global_geometry(time_range=None)
+                logger.warning(
+                    "Global geometry detected: this geometry will be matched against all "
+                    "windows in the rslearn dataset. When using settings like "
+                    "max_matches=1 and space_mode=MOSAIC, this may cause windows outside "
+                    "the geometry’s valid bounds to be materialized from the global raster "
+                    "instead of a more appropriate source. Consider using COMPOSITE mode, "
+                    "or increasing max_matches if this behavior is unintended."
+                )
 
             if spec.name:
                 item_name = spec.name
             else:
-                item_name = spec.fnames[0].name.split(".")[0]
+                item_name = fname.name.split(".")[0]
 
-            items.append(RasterItem(item_name, geometry, spec))
+            logger.debug(
+                "RasterImporter.list_items: got bounds of %s: %s", item_name, geometry
+            )
+            items.append(RasterItem(item_name, geometry, str(src_dir), spec))
 
+        logger.debug("RasterImporter.list_items: discovered %d items", len(items))
         return items
 
     def ingest_item(
         self,
-        config: LayerConfig,
-        tile_store: TileStore,
+        tile_store: TileStoreWithLayer,
         item: Item,
         cur_geometries: list[STGeometry],
-    ):
+    ) -> None:
         """Ingest the specified local file item.
 
         Args:
-            config: the configuration of the layer.
-            tile_store: the TileStore to ingest the data into.
-            item: the Item to ingest
+            tile_store: the tile store to ingest the data into.
+            item: the RasterItem to ingest
             cur_geometries: the geometries where the item is needed.
         """
         assert isinstance(item, RasterItem)
         for file_idx, fname in enumerate(item.spec.fnames):
-            with fname.open("rb") as f:
-                with rasterio.open(f) as src:
-                    if item.spec.bands:
-                        bands = item.spec.bands[file_idx]
-                    else:
-                        bands = [f"B{band_idx+1}" for band_idx in range(src.count)]
-                    cur_tile_store = PrefixedTileStore(tile_store, ("_".join(bands),))
-                    needed_projections = get_needed_projections(
-                        cur_tile_store, bands, config.band_sets, cur_geometries
-                    )
-                    if not needed_projections:
-                        return
-
-                    for projection in needed_projections:
-                        ingest_raster(
-                            tile_store=cur_tile_store,
-                            raster=src,
-                            projection=projection,
-                            time_range=item.geometry.time_range,
-                            layer_config=config,
-                        )
+            fname_upath = join_upath(UPath(item.src_dir), fname)
+            with open_rasterio_upath_reader(fname_upath) as src:
+                if item.spec.bands:
+                    bands = item.spec.bands[file_idx]
+                else:
+                    bands = [f"B{band_idx + 1}" for band_idx in range(src.count)]
+
+            if tile_store.is_raster_ready(item.name, bands):
+                continue
+            tile_store.write_raster_file(item.name, bands, fname_upath)
 
 
-@Importers.register("vector")
 class VectorImporter(Importer):
     """An Importer for vector data."""
 
-    def list_items(self, config: LayerConfig, src_dir: UPath) -> list[Item]:
+    # We need some buffer around GeoJSON bounds in case it just contains one point.
+    item_buffer_epsilon = 1e-4
+
+    def list_items(self, src_dir: UPath) -> list[Item]:
         """Extract a list of Items from the source directory.
 
         Args:
-            config: the configuration of the layer.
             src_dir: the source directory.
         """
         file_paths = src_dir.glob("**/*.*")
         items: list[Item] = []
 
         for path in file_paths:
+            # Ignore JSON files.
+            if path.name.endswith(".json"):
+                continue
+
             # Get the bounds of the features in the vector file, which we assume fiona can
             # read.
             # For shapefile, to open it we need to copy all the aux files.
             aux_files: list[UPath] = []
-            if path.name.split(".")[-1] == "shp":
+            if path.name.endswith(".shp"):
                 prefix = ".".join(path.name.split(".")[:-1])
                 for ext in SHAPEFILE_AUX_EXTENSIONS:
                     aux_files.append(path.parent / (prefix + ext))
@@ -299,45 +324,50 @@ class VectorImporter(Importer):
                             bounds[2] = max(bounds[2], cur_bounds[2])
                             bounds[3] = max(bounds[3], cur_bounds[3])
 
+                    # Normal GeoJSON should have coordinates in CRS coordinates, i.e. it
+                    # should be 1 projection unit/pixel.
                     projection = Projection(crs, 1, 1)
-                    geometry = STGeometry(projection, shapely.box(*bounds), None)
+                    geometry = STGeometry(
+                        projection,
+                        shapely.box(*bounds).buffer(self.item_buffer_epsilon),
+                        None,
+                    )
 
+                    # There can be problems with GeoJSON files that have large spatial
+                    # coverage, since the bounds may not re-project correctly to match
+                    # windows that are using projections with limited validity.
+                    # We check if there is a large spatial coverage here, and mark the
+                    # item's geometry as having global coverage if so.
+                    if geometry.is_too_large():
+                        geometry = get_global_geometry(time_range=None)
+
+            logger.debug(
+                "VectorImporter.list_items: got bounds of %s: %s", path, geometry
+            )
             items.append(
                 VectorItem(path.name.split(".")[0], geometry, path.absolute().as_uri())
             )
 
+        logger.debug("VectorImporter.list_items: discovered %d items", len(items))
         return items
 
     def ingest_item(
         self,
-        config: LayerConfig,
-        tile_store: TileStore,
+        tile_store: TileStoreWithLayer,
         item: Item,
         cur_geometries: list[STGeometry],
-    ):
+    ) -> None:
         """Ingest the specified local file item.
 
         Args:
-            config: the configuration of the layer.
             tile_store: the TileStore to ingest the data into.
             item: the Item to ingest
             cur_geometries: the geometries where the item is needed.
         """
-        assert isinstance(config, VectorLayerConfig)
-
-        needed_projections = set()
-        for geometry in cur_geometries:
-            projection, _ = config.get_final_projection_and_bounds(
-                geometry.projection, None
-            )
-            ts_layer = tile_store.get_layer((str(projection),))
-            if ts_layer and ts_layer.get_metadata().properties.get("completed"):
-                continue
-            needed_projections.add(projection)
-
-        if not needed_projections:
+        if tile_store.is_vector_ready(item.name):
             return
 
+        assert isinstance(item, VectorItem)
         path = UPath(item.path_uri)
 
         aux_files: list[UPath] = []
@@ -347,14 +377,18 @@ class VectorImporter(Importer):
                 aux_files.append(path.parent / (prefix + ext))
 
         # TODO: move converting fiona file to list[Feature] to utility function.
-        # TODO: don't assume WGS-84 projection here.
         with get_upath_local(path, extra_paths=aux_files) as local_fname:
             with fiona.open(local_fname) as src:
+                crs = CRS.from_wkt(src.crs.to_wkt())
+                # Normal GeoJSON should have coordinates in CRS coordinates, i.e. it
+                # should be 1 projection unit/pixel.
+                projection = Projection(crs, 1, 1)
+
                 features = []
                 for feat in src:
                     features.append(
                         Feature.from_geojson(
-                            WGS84_PROJECTION,
+                            projection,
                             {
                                 "type": "Feature",
                                 "geometry": dict(feat.geometry),
@@ -363,37 +397,72 @@ class VectorImporter(Importer):
                         )
                     )
 
-                for projection in needed_projections:
-                    cur_features = [feat.to_projection(projection) for feat in features]
-                    layer = tile_store.create_layer(
-                        (str(projection),),
-                        LayerMetadata(projection, None, {}),
-                    )
-                    layer.write_vector(cur_features)
-                    layer.set_property("completed", True)
+                tile_store.write_vector(item.name, features)
 
 
 class LocalFiles(DataSource):
     """A data source for ingesting data from local files."""
 
-    def __init__(self, config: LayerConfig, src_dir: UPath) -> None:
+    def __init__(
+        self,
+        src_dir: str,
+        raster_item_specs: list[RasterItemSpec] | None = None,
+        layer_type: LayerType | None = None,
+        context: DataSourceContext = DataSourceContext(),
+    ) -> None:
         """Initialize a new LocalFiles instance.
 
         Args:
-            config: configuration for this layer.
             src_dir: source directory to ingest
+            raster_item_specs: the specs to specify the raster items directly. If None,
+                the raster items are automatically detected from the files in the
+                source directory.
+            layer_type: the layer type. It only needs to be set if the layer_config is
+                missing from the context.
+            context: the data source context. The layer config must be in the context.
         """
-        self.config = config
-        self.src_dir = src_dir
+        if context.ds_path is not None:
+            self.src_dir = join_upath(context.ds_path, src_dir)
+        else:
+            self.src_dir = UPath(src_dir)
 
-        self.importer: Importer = Importers[config.layer_type.value]
-        self.items = self.importer.list_items(config, src_dir)
+        # Determine layer type.
+        if context.layer_config is not None:
+            self.layer_type = context.layer_config.type
+        elif layer_type is not None:
+            self.layer_type = layer_type
+        else:
+            raise ValueError(
+                "layer type must be specified if the layer config is not in the context"
+            )
 
-    @staticmethod
-    def from_config(config: LayerConfig, ds_path: UPath) -> "LocalFiles":
-        """Creates a new LocalFiles instance from a configuration dictionary."""
-        d = config.data_source.config_dict
-        return LocalFiles(config=config, src_dir=join_upath(ds_path, d["src_dir"]))
+        self.importer: Importer
+        if self.layer_type == LayerType.RASTER:
+            self.importer = RasterImporter(item_specs=raster_item_specs)
+        elif self.layer_type == LayerType.VECTOR:
+            self.importer = VectorImporter()
+        else:
+            raise ValueError(f"unknown layer type {self.layer_type}")
+
+    @functools.cache
+    def list_items(self) -> list[Item]:
+        """Lists items from the source directory while maintaining a cache file."""
+        cache_fname = self.src_dir / "summary.json"
+        if not cache_fname.exists():
+            logger.debug("cache at %s does not exist, listing items", cache_fname)
+            items = self.importer.list_items(self.src_dir)
+            serialized_items = [item.serialize() for item in items]
+            with cache_fname.open("w") as f:
+                json.dump(serialized_items, f)
+            return items
+
+        logger.debug("loading item list from cache at %s", cache_fname)
+        with cache_fname.open() as f:
+            serialized_items = json.load(f)
+        return [
+            self.deserialize_item(serialized_item)
+            for serialized_item in serialized_items
+        ]
 
     def get_items(
         self, geometries: list[STGeometry], query_config: QueryConfig
@@ -410,7 +479,7 @@ class LocalFiles(DataSource):
         groups = []
         for geometry in geometries:
             cur_items = []
-            for item in self.items:
+            for item in self.list_items():
                 if not item.geometry.intersects(geometry):
                     continue
                 cur_items.append(item)
@@ -421,17 +490,18 @@ class LocalFiles(DataSource):
             groups.append(cur_groups)
         return groups
 
-    def deserialize_item(self, serialized_item: Any) -> Item:
+    def deserialize_item(self, serialized_item: Any) -> RasterItem | VectorItem:
         """Deserializes an item from JSON-decoded data."""
-        assert isinstance(serialized_item, dict)
-        if self.config.layer_type == LayerType.RASTER:
+        if self.layer_type == LayerType.RASTER:
             return RasterItem.deserialize(serialized_item)
-        elif self.config.layer_type == LayerType.VECTOR:
+        elif self.layer_type == LayerType.VECTOR:
             return VectorItem.deserialize(serialized_item)
+        else:
+            raise ValueError(f"Unknown layer type: {self.layer_type}")
 
     def ingest(
         self,
-        tile_store: TileStore,
+        tile_store: TileStoreWithLayer,
         items: list[Item],
         geometries: list[list[STGeometry]],
     ) -> None:
@@ -443,5 +513,4 @@ class LocalFiles(DataSource):
             geometries: a list of geometries needed for each item
         """
         for item, cur_geometries in zip(items, geometries):
-            cur_tile_store = PrefixedTileStore(tile_store, (item.name,))
-            self.importer.ingest_item(self.config, cur_tile_store, item, cur_geometries)
+            self.importer.ingest_item(tile_store, item, cur_geometries)