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

rslearn/data_sources/openstreetmap.py

@@ -1,4 +1,4 @@
-"""Data source for raster data on public Cloud Storage buckets."""
+"""Data source for OpenStreetMap vector features."""
 
 import json
 import shutil
@@ -7,17 +7,21 @@ from enum import Enum
 from typing import Any
 
 import osmium
+import osmium.osm.types
 import shapely
 from upath import UPath
 
-from rslearn.config import LayerConfig, QueryConfig, VectorLayerConfig
+from rslearn.config import QueryConfig
 from rslearn.const import WGS84_PROJECTION
-from rslearn.data_sources import DataSource, Item
+from rslearn.data_sources import DataSource, DataSourceContext, Item
 from rslearn.data_sources.utils import match_candidate_items_to_window
-from rslearn.tile_stores import LayerMetadata, TileStore
+from rslearn.log_utils import get_logger
+from rslearn.tile_stores import TileStoreWithLayer
 from rslearn.utils import Feature, GridIndex, STGeometry
 from rslearn.utils.fsspec import get_upath_local, join_upath
 
+logger = get_logger(__name__)
+
 
 class FeatureType(Enum):
     """OpenStreetMap feature type."""
@@ -36,7 +40,7 @@ class Filter:
         tag_conditions: dict[str, list[str]] | None = None,
         tag_properties: dict[str, str] | None = None,
         to_geometry: str | None = None,
-    ):
+    ) -> None:
         """Create a new Filter instance.
 
         Args:
@@ -54,27 +58,6 @@ class Filter:
         self.tag_properties = tag_properties
         self.to_geometry = to_geometry
 
-    @staticmethod
-    def from_config(d: dict[str, Any]) -> "Filter":
-        """Creates a Filter from a config dict.
-
-        Args:
-            d: the config dict
-
-        Returns:
-            the Filter object
-        """
-        kwargs = {}
-        if "feature_types" in d:
-            kwargs["feature_types"] = [FeatureType(el) for el in d["feature_types"]]
-        if "tag_conditions" in d:
-            kwargs["tag_conditions"] = d["tag_conditions"]
-        if "tag_properties" in d:
-            kwargs["tag_properties"] = d["tag_properties"]
-        if "to_geometry" in d:
-            kwargs["to_geometry"] = d["to_geometry"]
-        return Filter(**kwargs)
-
     def match_tags(self, tags: dict[str, str]) -> bool:
         """Returns whether this filter matches based on the tags."""
         if not self.tag_conditions:
@@ -104,12 +87,12 @@ class Filter:
 class BoundsHandler(osmium.SimpleHandler):
     """An osmium handler for computing the bounds of an input file."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Initialize a new BoundsHandler."""
         osmium.SimpleHandler.__init__(self)
-        self.bounds = (180, 90, -180, -90)
+        self.bounds: tuple[float, float, float, float] = (180, 90, -180, -90)
 
-    def node(self, n):
+    def node(self, n: osmium.osm.types.Node) -> None:
         """Handle nodes and update the computed bounds."""
         lon = n.location.lon
         lat = n.location.lat
@@ -130,7 +113,7 @@ class OsmHandler(osmium.SimpleHandler):
         geometries: list[STGeometry],
         grid_size: float = 0.03,
         padding: float = 0.03,
-    ):
+    ) -> None:
         """Initialize a new OsmHandler.
 
         Args:
@@ -163,12 +146,12 @@ class OsmHandler(osmium.SimpleHandler):
             )
             self.grid_index.insert(bounds, 1)
 
-        self.cached_nodes = {}
-        self.cached_ways = {}
+        self.cached_nodes: dict = {}
+        self.cached_ways: dict = {}
 
-        self.features = []
+        self.features: list[Feature] = []
 
-    def node(self, n):
+    def node(self, n: osmium.osm.types.Node) -> None:
         """Handle nodes."""
         # Check if node is relevant to our geometries.
         lon = n.location.lon
@@ -193,7 +176,7 @@ class OsmHandler(osmium.SimpleHandler):
             )
             self.features.append(feat)
 
-    def _get_way_coords(self, node_ids):
+    def _get_way_coords(self, node_ids: list[int]) -> list:
         coords = []
         for id in node_ids:
             if id not in self.cached_nodes:
@@ -201,7 +184,7 @@ class OsmHandler(osmium.SimpleHandler):
             coords.append(self.cached_nodes[id])
         return coords
 
-    def way(self, w):
+    def way(self, w: osmium.osm.types.Way) -> None:
         """Handle ways."""
         # Collect nodes, skip if too few.
         node_ids = [member.ref for member in w.nodes]
@@ -235,7 +218,7 @@ class OsmHandler(osmium.SimpleHandler):
             )
             self.features.append(feat)
 
-    def match_relation(self, r):
+    def match_relation(self, r: osmium.osm.types.Relation) -> None:
         """Handle relations."""
         # Collect ways and distinguish exterior vs holes, skip if none found.
         exterior_ways = []
@@ -267,7 +250,7 @@ class OsmHandler(osmium.SimpleHandler):
         # Merge the ways in case some exterior/interior polygons are split into
         # multiple ways.
         # And convert them from node IDs to coordinates.
-        def get_polygons(ways):
+        def get_polygons(ways: list) -> list:
             polygons: list[list[int]] = []
             for way in ways:
                 # Attempt to match the way to an existing polygon.
@@ -366,13 +349,13 @@ class OsmItem(Item):
         return d
 
     @staticmethod
-    def deserialize(d: dict) -> Item:
+    def deserialize(d: dict) -> "OsmItem":
         """Deserializes an item from a JSON-decoded dictionary."""
         item = super(OsmItem, OsmItem).deserialize(d)
         return OsmItem(name=item.name, geometry=item.geometry, path_uri=d["path_uri"])
 
 
-class OpenStreetMap(DataSource):
+class OpenStreetMap(DataSource[OsmItem]):
     """A data source for OpenStreetMap data from PBF file.
 
     An existing local PBF file can be used, or if the provided path doesn't exist, then
@@ -386,12 +369,12 @@ class OpenStreetMap(DataSource):
 
     def __init__(
         self,
-        config: VectorLayerConfig,
-        pbf_fnames: list[UPath],
-        bounds_fname: UPath,
+        pbf_fnames: list[str],
+        bounds_fname: str,
         categories: dict[str, Filter],
+        context: DataSourceContext = DataSourceContext(),
     ):
-        """Initialize a new Sentinel2 instance.
+        """Initialize a new OpenStreetMap instance.
 
         Args:
             config: the configuration of this layer.
@@ -401,14 +384,21 @@ class OpenStreetMap(DataSource):
             bounds_fname: filename where the bounds of the PBF are cached.
             categories: dictionary of (category name, filter). Features that match the
                 filter will be emitted under the corresponding category.
+            context: the data source context.
         """
-        self.config = config
-        self.pbf_fnames = pbf_fnames
-        self.bounds_fname = bounds_fname
         self.categories = categories
 
+        if context.ds_path is not None:
+            self.pbf_fnames = [
+                join_upath(context.ds_path, pbf_fname) for pbf_fname in pbf_fnames
+            ]
+            self.bounds_fname = join_upath(context.ds_path, bounds_fname)
+        else:
+            self.pbf_fnames = [UPath(pbf_fname) for pbf_fname in pbf_fnames]
+            self.bounds_fname = UPath(bounds_fname)
+
         if len(self.pbf_fnames) == 1 and not self.pbf_fnames[0].exists():
-            print(
+            logger.info(
                 "Downloading planet.osm.pbf from "
                 + f"{self.planet_pbf_url} to {self.pbf_fnames[0]}"
             )
@@ -419,29 +409,13 @@ class OpenStreetMap(DataSource):
         # Detect bounds of each pbf file if needed.
         self.pbf_bounds = self._get_pbf_bounds()
 
-    @staticmethod
-    def from_config(config: LayerConfig, ds_path: UPath) -> "OpenStreetMap":
-        """Creates a new OpenStreetMap instance from a configuration dictionary."""
-        assert isinstance(config, VectorLayerConfig)
-        d = config.data_source.config_dict
-        categories = {
-            category_name: Filter.from_config(filter_config_dict)
-            for category_name, filter_config_dict in d["categories"].items()
-        }
-        pbf_fnames = [join_upath(ds_path, pbf_fname) for pbf_fname in d["pbf_fnames"]]
-        bounds_fname = join_upath(ds_path, d["bounds_fname"])
-        return OpenStreetMap(
-            config=config,
-            pbf_fnames=pbf_fnames,
-            bounds_fname=bounds_fname,
-            categories=categories,
-        )
-
-    def _get_pbf_bounds(self):
+    def _get_pbf_bounds(self) -> list[tuple[float, float, float, float]]:
+        # Determine WGS84 bounds of each PBF file by processing them through
+        # BoundsHandler.
         if not self.bounds_fname.exists():
             pbf_bounds = []
             for pbf_fname in self.pbf_fnames:
-                print(f"detecting bounds of {pbf_fname}")
+                logger.info(f"detecting bounds of {pbf_fname}")
                 handler = BoundsHandler()
                 with get_upath_local(pbf_fname) as local_fname:
                     handler.apply_file(local_fname)
@@ -458,7 +432,7 @@ class OpenStreetMap(DataSource):
 
     def get_items(
         self, geometries: list[STGeometry], query_config: QueryConfig
-    ) -> list[list[list[Item]]]:
+    ) -> list[list[list[OsmItem]]]:
         """Get a list of items in the data source intersecting the given geometries.
 
         Args:
@@ -487,14 +461,14 @@ class OpenStreetMap(DataSource):
             groups.append(cur_groups)
         return groups
 
-    def deserialize_item(self, serialized_item: Any) -> Item:
+    def deserialize_item(self, serialized_item: Any) -> OsmItem:
         """Deserializes an item from JSON-decoded data."""
         return OsmItem.deserialize(serialized_item)
 
     def ingest(
         self,
-        tile_store: TileStore,
-        items: list[Item],
+        tile_store: TileStoreWithLayer,
+        items: list[OsmItem],
         geometries: list[list[STGeometry]],
     ) -> None:
         """Ingest items into the given tile store.
@@ -504,10 +478,11 @@ class OpenStreetMap(DataSource):
             items: the items to ingest
             geometries: a list of geometries needed for each item
         """
-        item_names = [item.name for item in items]
-        item_names.sort()
         for cur_item, cur_geometries in zip(items, geometries):
-            print(
+            if tile_store.is_vector_ready(cur_item.name):
+                continue
+
+            logger.info(
                 f"ingesting osm item {cur_item.name} "
                 + f"with {len(cur_geometries)} geometries"
             )
@@ -515,17 +490,4 @@ class OpenStreetMap(DataSource):
             with get_upath_local(UPath(cur_item.path_uri)) as local_fname:
                 handler.apply_file(local_fname)
 
-            projections = set()
-            for geometry in cur_geometries:
-                projection, _ = self.config.get_final_projection_and_bounds(
-                    geometry.projection, None
-                )
-                projections.add(projection)
-
-            for projection in projections:
-                features = [feat.to_projection(projection) for feat in handler.features]
-                layer = tile_store.create_layer(
-                    (cur_item.name, str(projection)),
-                    LayerMetadata(projection, None, {}),
-                )
-                layer.write_vector(features)
+            tile_store.write_vector(cur_item.name, handler.features)

rslearn/data_sources/planet.py

@@ -6,24 +6,22 @@ import pathlib
 import shutil
 import tempfile
 from datetime import datetime
+from pathlib import Path
 from typing import Any
 
 import planet
-import rasterio
 import shapely
 from fsspec.implementations.local import LocalFileSystem
 from upath import UPath
 
-from rslearn.config import LayerConfig, QueryConfig, RasterLayerConfig
+from rslearn.config import QueryConfig
 from rslearn.const import WGS84_PROJECTION
-from rslearn.data_sources import DataSource, Item
+from rslearn.data_sources import DataSource, DataSourceContext, Item
 from rslearn.data_sources.utils import match_candidate_items_to_window
-from rslearn.tile_stores import PrefixedTileStore, TileStore
+from rslearn.tile_stores import TileStoreWithLayer
 from rslearn.utils import STGeometry
 from rslearn.utils.fsspec import join_upath
 
-from .raster_source import get_needed_projections, ingest_raster
-
 
 class Planet(DataSource):
     """A data source for Planet Labs API.
@@ -33,19 +31,18 @@ class Planet(DataSource):
 
     def __init__(
         self,
-        config: LayerConfig,
         item_type_id: str,
-        cache_dir: UPath | None = None,
+        cache_dir: str | None = None,
         asset_type_id: str = "ortho_analytic_sr",
         range_filters: dict[str, dict[str, Any]] = {},
         use_permission_filter: bool = True,
         sort_by: str | None = None,
         bands: list[str] = ["b01", "b02", "b03", "b04"],
+        context: DataSourceContext = DataSourceContext(),
     ):
         """Initialize a new Planet instance.
 
         Args:
-            config: the LayerConfig of the layer containing this data source
             item_type_id: the item type ID, like "PSScene" or "SkySatCollect".
             cache_dir: where to store downloaded assets, or None to just store it in
                 temporary directory before putting into tile store.
@@ -62,38 +59,22 @@ class Planet(DataSource):
                 "-clear_percent" or "cloud_cover" (if it starts with minus sign then we
                 sort descending.)
             bands: what to call the bands in the asset.
+            context: the data source context.
         """
-        self.config = config
         self.item_type_id = item_type_id
-        self.cache_dir = cache_dir
         self.asset_type_id = asset_type_id
         self.range_filters = range_filters
         self.use_permission_filter = use_permission_filter
         self.sort_by = sort_by
         self.bands = bands
 
-    @staticmethod
-    def from_config(config: LayerConfig, ds_path: UPath) -> "Planet":
-        """Creates a new Planet instance from a configuration dictionary."""
-        assert isinstance(config, RasterLayerConfig)
-        d = config.data_source.config_dict
-        kwargs = dict(
-            config=config,
-            item_type_id=d["item_type_id"],
-        )
-        optional_keys = [
-            "asset_type_id",
-            "range_filters",
-            "use_permission_filter",
-            "sort_by",
-            "bands",
-        ]
-        for optional_key in optional_keys:
-            if optional_key in d:
-                kwargs[optional_key] = d[optional_key]
-        if "cache_dir" in d:
-            kwargs["cache_dir"] = join_upath(ds_path, d["cache_dir"])
-        return Planet(**kwargs)
+        if cache_dir is None:
+            self.cache_dir = None
+        else:
+            if context.ds_path is not None:
+                self.cache_dir = join_upath(context.ds_path, cache_dir)
+            else:
+                self.cache_dir = UPath(cache_dir)
 
     async def _search_items(self, geometry: STGeometry) -> list[dict[str, Any]]:
         wgs84_geometry = geometry.to_projection(WGS84_PROJECTION)
@@ -101,11 +82,10 @@ class Planet(DataSource):
 
         async with planet.Session() as session:
             client = session.client("data")
-
+            gte = geometry.time_range[0] if geometry.time_range is not None else None
+            lte = geometry.time_range[1] if geometry.time_range is not None else None
             filter_list = [
-                planet.data_filter.date_range_filter(
-                    "acquired", gte=geometry.time_range[0], lte=geometry.time_range[1]
-                ),
+                planet.data_filter.date_range_filter("acquired", gte=gte, lte=lte),
                 planet.data_filter.geometry_filter(geojson_data),
                 planet.data_filter.asset_filter([self.asset_type_id]),
             ]
@@ -242,7 +222,7 @@ class Planet(DataSource):
 
     def ingest(
         self,
-        tile_store: TileStore,
+        tile_store: TileStoreWithLayer,
         items: list[Item],
         geometries: list[list[STGeometry]],
     ) -> None:
@@ -253,26 +233,10 @@ class Planet(DataSource):
             items: the items to ingest
             geometries: a list of geometries needed for each item
         """
-        for item, cur_geometries in zip(items, geometries):
+        for item in items:
+            if tile_store.is_raster_ready(item.name, self.bands):
+                continue
+
             with tempfile.TemporaryDirectory() as tmp_dir:
-                band_names = self.bands
-                cur_tile_store = PrefixedTileStore(
-                    tile_store, (item.name, "_".join(band_names))
-                )
-                needed_projections = get_needed_projections(
-                    cur_tile_store, band_names, self.config.band_sets, cur_geometries
-                )
-                if not needed_projections:
-                    continue
-
-                asset_path = asyncio.run(self._download_asset(item, tmp_dir))
-                with asset_path.open("rb") as f:
-                    with rasterio.open(f) as raster:
-                        for projection in needed_projections:
-                            ingest_raster(
-                                tile_store=cur_tile_store,
-                                raster=raster,
-                                projection=projection,
-                                time_range=item.geometry.time_range,
-                                layer_config=self.config,
-                            )
+                asset_path = asyncio.run(self._download_asset(item, Path(tmp_dir)))
+                tile_store.write_raster_file(item.name, self.bands, asset_path)