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

rslearn/data_sources/local_files.py

@@ -1,6 +1,6 @@
 """Data source for raster or vector data in local files."""
 
-from typing import Any
+from typing import Any, Generic, TypeVar
 
 import fiona
 import rasterio
@@ -11,8 +11,8 @@ 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.config import LayerConfig, LayerType, RasterLayerConfig, VectorLayerConfig
+from rslearn.const import SHAPEFILE_AUX_EXTENSIONS
 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
@@ -22,11 +22,15 @@ from .raster_source import get_needed_projections, ingest_raster
 
 Importers = ClassRegistry()
 
+ItemType = TypeVar("ItemType", bound=Item)
+LayerConfigType = TypeVar("LayerConfigType", bound=LayerConfig)
+ImporterType = TypeVar("ImporterType", bound="Importer")
 
-class Importer:
-    """An abstract class for importing data from local files."""
 
-    def list_items(self, config: LayerConfig, src_dir: UPath) -> list[Item]:
+class Importer(Generic[ItemType, LayerConfigType]):
+    """An abstract base class for importing data from local files."""
+
+    def list_items(self, config: LayerConfigType, src_dir: UPath) -> list[ItemType]:
         """Extract a list of Items from the source directory.
 
         Args:
@@ -37,11 +41,11 @@ class Importer:
 
     def ingest_item(
         self,
-        config: LayerConfig,
+        config: LayerConfigType,
         tile_store: TileStore,
-        item: Item,
+        item: ItemType,
         cur_geometries: list[STGeometry],
-    ):
+    ) -> None:
         """Ingest the specified local file item.
 
         Args:
@@ -131,7 +135,7 @@ class RasterItem(Item):
         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)
         spec = RasterItemSpec.deserialize(d["spec"])
@@ -159,7 +163,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(
@@ -171,7 +175,7 @@ class VectorItem(Item):
 class RasterImporter(Importer):
     """An Importer for raster data."""
 
-    def list_items(self, config: LayerConfig, src_dir: UPath) -> list[Item]:
+    def list_items(self, config: LayerConfig, src_dir: UPath) -> list[RasterItem]:
         """Extract a list of Items from the source directory.
 
         Args:
@@ -179,6 +183,8 @@ class RasterImporter(Importer):
             src_dir: the source directory.
         """
         item_specs: list[RasterItemSpec] = []
+        if config.data_source is None:
+            raise ValueError("RasterImporter requires a data source config")
         # 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"]:
@@ -192,7 +198,7 @@ class RasterImporter(Importer):
                 spec = RasterItemSpec(fnames=[path], bands=None)
                 item_specs.append(spec)
 
-        items: list[Item] = []
+        items = []
         for spec in item_specs:
             # Get geometry from the first raster file.
             # We assume files are readable with rasterio.
@@ -222,20 +228,19 @@ class RasterImporter(Importer):
 
     def ingest_item(
         self,
-        config: LayerConfig,
+        config: RasterLayerConfig,
         tile_store: TileStore,
-        item: Item,
+        item: RasterItem,
         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
+            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:
@@ -264,7 +269,10 @@ class RasterImporter(Importer):
 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, config: LayerConfig, src_dir: UPath) -> list[VectorItem]:
         """Extract a list of Items from the source directory.
 
         Args:
@@ -272,7 +280,7 @@ class VectorImporter(Importer):
             src_dir: the source directory.
         """
         file_paths = src_dir.glob("**/*.*")
-        items: list[Item] = []
+        items: list[VectorItem] = []
 
         for path in file_paths:
             # Get the bounds of the features in the vector file, which we assume fiona can
@@ -299,8 +307,14 @@ 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,
+                    )
 
             items.append(
                 VectorItem(path.name.split(".")[0], geometry, path.absolute().as_uri())
@@ -310,11 +324,11 @@ class VectorImporter(Importer):
 
     def ingest_item(
         self,
-        config: LayerConfig,
+        config: VectorLayerConfig,
         tile_store: TileStore,
-        item: Item,
+        item: VectorItem,
         cur_geometries: list[STGeometry],
-    ):
+    ) -> None:
         """Ingest the specified local file item.
 
         Args:
@@ -323,7 +337,8 @@ class VectorImporter(Importer):
             item: the Item to ingest
             cur_geometries: the geometries where the item is needed.
         """
-        assert isinstance(config, VectorLayerConfig)
+        if not isinstance(config, VectorLayerConfig):
+            raise ValueError("VectorImporter requires a VectorLayerConfig")
 
         needed_projections = set()
         for geometry in cur_geometries:
@@ -347,14 +362,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),
@@ -376,7 +395,11 @@ class VectorImporter(Importer):
 class LocalFiles(DataSource):
     """A data source for ingesting data from local files."""
 
-    def __init__(self, config: LayerConfig, src_dir: UPath) -> None:
+    def __init__(
+        self,
+        config: LayerConfig,
+        src_dir: UPath,
+    ) -> None:
         """Initialize a new LocalFiles instance.
 
         Args:
@@ -384,14 +407,17 @@ class LocalFiles(DataSource):
             src_dir: source directory to ingest
         """
         self.config = config
+
+        self.importer = Importers[config.layer_type.value]
         self.src_dir = src_dir
 
-        self.importer: Importer = Importers[config.layer_type.value]
-        self.items = self.importer.list_items(config, src_dir)
+        self.items = self.importer.list_items(self.config, src_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 LocalFiles(config=config, src_dir=join_upath(ds_path, d["src_dir"]))
 
@@ -421,13 +447,14 @@ 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:
             return RasterItem.deserialize(serialized_item)
         elif self.config.layer_type == LayerType.VECTOR:
             return VectorItem.deserialize(serialized_item)
+        else:
+            raise ValueError(f"Unknown layer type: {self.config.layer_type}")
 
     def ingest(
         self,

rslearn/data_sources/openstreetmap.py

@@ -7,10 +7,11 @@ 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, VectorLayerConfig
 from rslearn.const import WGS84_PROJECTION
 from rslearn.data_sources import DataSource, Item
 from rslearn.data_sources.utils import match_candidate_items_to_window
@@ -36,7 +37,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:
@@ -64,7 +65,7 @@ class Filter:
         Returns:
             the Filter object
         """
-        kwargs = {}
+        kwargs: dict[str, Any] = {}
         if "feature_types" in d:
             kwargs["feature_types"] = [FeatureType(el) for el in d["feature_types"]]
         if "tag_conditions" in d:
@@ -104,12 +105,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 +131,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 +164,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 +194,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 +202,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 +236,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 +268,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 +367,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
@@ -420,9 +421,10 @@ class OpenStreetMap(DataSource):
         self.pbf_bounds = self._get_pbf_bounds()
 
     @staticmethod
-    def from_config(config: LayerConfig, ds_path: UPath) -> "OpenStreetMap":
+    def from_config(config: VectorLayerConfig, ds_path: UPath) -> "OpenStreetMap":
         """Creates a new OpenStreetMap instance from a configuration dictionary."""
-        assert isinstance(config, VectorLayerConfig)
+        if config.data_source is None:
+            raise ValueError("data_source is required")
         d = config.data_source.config_dict
         categories = {
             category_name: Filter.from_config(filter_config_dict)
@@ -437,7 +439,9 @@ class OpenStreetMap(DataSource):
             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:
@@ -458,7 +462,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 +491,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],
+        items: list[OsmItem],
         geometries: list[list[STGeometry]],
     ) -> None:
         """Ingest items into the given tile store.

rslearn/data_sources/planet.py

@@ -6,6 +6,7 @@ import pathlib
 import shutil
 import tempfile
 from datetime import datetime
+from pathlib import Path
 from typing import Any
 
 import planet
@@ -14,7 +15,7 @@ import shapely
 from fsspec.implementations.local import LocalFileSystem
 from upath import UPath
 
-from rslearn.config import LayerConfig, QueryConfig, RasterLayerConfig
+from rslearn.config import QueryConfig, RasterLayerConfig
 from rslearn.const import WGS84_PROJECTION
 from rslearn.data_sources import DataSource, Item
 from rslearn.data_sources.utils import match_candidate_items_to_window
@@ -33,7 +34,7 @@ class Planet(DataSource):
 
     def __init__(
         self,
-        config: LayerConfig,
+        config: RasterLayerConfig,
         item_type_id: str,
         cache_dir: UPath | None = None,
         asset_type_id: str = "ortho_analytic_sr",
@@ -73,9 +74,10 @@ class Planet(DataSource):
         self.bands = bands
 
     @staticmethod
-    def from_config(config: LayerConfig, ds_path: UPath) -> "Planet":
+    def from_config(config: RasterLayerConfig, ds_path: UPath) -> "Planet":
         """Creates a new Planet instance from a configuration dictionary."""
-        assert isinstance(config, RasterLayerConfig)
+        if config.data_source is None:
+            raise ValueError("data_source is required")
         d = config.data_source.config_dict
         kwargs = dict(
             config=config,
@@ -101,11 +103,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]),
             ]
@@ -265,7 +266,7 @@ class Planet(DataSource):
                 if not needed_projections:
                     continue
 
-                asset_path = asyncio.run(self._download_asset(item, tmp_dir))
+                asset_path = asyncio.run(self._download_asset(item, Path(tmp_dir)))
                 with asset_path.open("rb") as f:
                     with rasterio.open(f) as raster:
                         for projection in needed_projections: