rslearn 0.0.26__py3-none-any.whl → 0.0.27__py3-none-any.whl

rslearn/data_sources/direct_materialize_data_source.py

@@ -0,0 +1,336 @@
+"""Base class for data sources that support direct materialization via TileStore."""
+
+from abc import abstractmethod
+from collections.abc import Callable
+from typing import Any, Generic
+
+import affine
+import numpy.typing as npt
+import rasterio
+import rasterio.vrt
+from rasterio.enums import Resampling
+
+from rslearn.config import LayerConfig
+from rslearn.data_sources.data_source import DataSource, ItemType
+from rslearn.dataset import Window
+from rslearn.dataset.materialize import RasterMaterializer
+from rslearn.tile_stores import TileStore, TileStoreWithLayer
+from rslearn.utils.geometry import PixelBounds, Projection
+
+
+class DirectMaterializeDataSource(DataSource[ItemType], TileStore, Generic[ItemType]):
+    """Base class for data sources that support direct materialization via TileStore.
+
+    This class provides common TileStore functionality for data sources that can read
+    raster data on-demand from remote sources (like cloud buckets or APIs) without
+    first ingesting into a local tile store.
+
+    Subclasses must implement:
+        - get_asset_url(): Get the URL for an asset given item name and bands
+        - get_item_by_name(): Get an item by its name
+
+    Subclasses may optionally override:
+        - get_raster_bands(): By default, we assume that items have all assets. If
+            items may have a subset of assets, override get_raster_bands to return
+            the sets of bands available for that item.
+        - get_read_callback(): Returns a callback to transform the raster array,
+            for post-processing like Sentinel-2 harmonization.
+    """
+
+    def __init__(self, asset_bands: dict[str, list[str]]):
+        """Initialize the DirectMaterializeDataSource.
+
+        Args:
+            asset_bands: mapping from asset key to the list of band names in that asset.
+        """
+        self.asset_bands = asset_bands
+
+    def _get_asset_key_by_bands(self, bands: list[str]) -> str:
+        """Get the asset key based on the band names.
+
+        Args:
+            bands: list of band names to look up.
+
+        Returns:
+            the asset key that provides those bands.
+
+        Raises:
+            ValueError: if no asset provides those bands.
+        """
+        for asset_key, asset_bands in self.asset_bands.items():
+            if bands == asset_bands:
+                return asset_key
+        raise ValueError(f"no known asset with bands {bands}")
+
+    # --- Methods that subclasses must implement ---
+
+    @abstractmethod
+    def get_asset_url(self, item_name: str, asset_key: str) -> str:
+        """Get the URL to read the asset for the given item and asset key.
+
+        Args:
+            item_name: the name of the item.
+            asset_key: the key identifying which asset to get.
+
+        Returns:
+            the URL to read the asset from (must be readable by rasterio).
+        """
+        raise NotImplementedError
+
+    def get_item_by_name(self, name: str) -> ItemType:
+        """Get an item by its name.
+
+        Subclasses must implement this method, either directly or by inheriting from
+        a class that provides it (e.g., StacDataSource).
+
+        Args:
+            name: the name of the item to get.
+
+        Returns:
+            the item object.
+        """
+        raise NotImplementedError
+
+    # --- Optional hooks for subclasses ---
+
+    def get_read_callback(
+        self, item_name: str, asset_key: str
+    ) -> Callable[[npt.NDArray[Any]], npt.NDArray[Any]] | None:
+        """Return a callback to post-process raster data (e.g., harmonization).
+
+        Subclasses can override this to apply transformations to the raw raster data
+        after reading, such as harmonization for Sentinel-2 data.
+
+        Args:
+            item_name: the name of the item being read.
+            asset_key: the key identifying which asset is being read.
+
+        Returns:
+            A callback function that takes an array and returns a modified array,
+            or None if no post-processing is needed.
+        """
+        return None
+
+    # --- TileStore implementation ---
+
+    def is_raster_ready(
+        self, layer_name: str, item_name: str, bands: list[str]
+    ) -> bool:
+        """Checks if this raster has been written to the store.
+
+        For remote-backed tile stores, this always returns True since data is
+        read on-demand from the remote source.
+
+        Args:
+            layer_name: the layer name or alias.
+            item_name: the item.
+            bands: the list of bands identifying which specific raster to read.
+
+        Returns:
+            True, since data is always available from the remote source.
+        """
+        return True
+
+    def get_raster_bands(self, layer_name: str, item_name: str) -> list[list[str]]:
+        """Get the sets of bands that have been stored for the specified item.
+
+        By default, returns all band sets from the asset_bands configuration.
+        Subclasses can override this if not all items have all assets.
+
+        Args:
+            layer_name: the layer name or alias.
+            item_name: the item.
+
+        Returns:
+            a list of lists of bands available for this item.
+        """
+        return list(self.asset_bands.values())
+
+    def get_raster_bounds(
+        self, layer_name: str, item_name: str, bands: list[str], projection: Projection
+    ) -> PixelBounds:
+        """Get the bounds of the raster in the specified projection.
+
+        Args:
+            layer_name: the layer name or alias.
+            item_name: the item to check.
+            bands: the list of bands identifying which specific raster to read.
+            projection: the projection to get the raster's bounds in.
+
+        Returns:
+            the bounds of the raster in the projection.
+        """
+        item = self.get_item_by_name(item_name)
+        geom = item.geometry.to_projection(projection)
+        return (
+            int(geom.shp.bounds[0]),
+            int(geom.shp.bounds[1]),
+            int(geom.shp.bounds[2]),
+            int(geom.shp.bounds[3]),
+        )
+
+    def _read_raster_from_url(
+        self,
+        url: str,
+        projection: Projection,
+        bounds: PixelBounds,
+        resampling: Resampling,
+    ) -> npt.NDArray[Any]:
+        """Read raster data from a URL with reprojection.
+
+        This is the common logic for reading raster data from a URL and reprojecting
+        it to the target projection and bounds using rasterio's WarpedVRT.
+
+        Args:
+            url: the URL to read from (must be readable by rasterio).
+            projection: the projection to read in.
+            bounds: the bounds to read.
+            resampling: the resampling method to use.
+
+        Returns:
+            the raster data as a numpy array.
+        """
+        # Construct the transform to use for the warped dataset.
+        wanted_transform = affine.Affine(
+            projection.x_resolution,
+            0,
+            bounds[0] * projection.x_resolution,
+            0,
+            projection.y_resolution,
+            bounds[1] * projection.y_resolution,
+        )
+
+        with rasterio.open(url) as src:
+            with rasterio.vrt.WarpedVRT(
+                src,
+                crs=projection.crs,
+                transform=wanted_transform,
+                width=bounds[2] - bounds[0],
+                height=bounds[3] - bounds[1],
+                resampling=resampling,
+            ) as vrt:
+                return vrt.read()
+
+    def read_raster(
+        self,
+        layer_name: str,
+        item_name: str,
+        bands: list[str],
+        projection: Projection,
+        bounds: PixelBounds,
+        resampling: Resampling = Resampling.bilinear,
+    ) -> npt.NDArray[Any]:
+        """Read raster data from the store.
+
+        Args:
+            layer_name: the layer name or alias.
+            item_name: the item to read.
+            bands: the list of bands identifying which specific raster to read.
+            projection: the projection to read in.
+            bounds: the bounds to read.
+            resampling: the resampling method to use in case reprojection is needed.
+
+        Returns:
+            the raster data as a numpy array.
+        """
+        # Get the asset key for the requested bands
+        asset_key = self._get_asset_key_by_bands(bands)
+
+        # Get the asset URL from the subclass
+        asset_url = self.get_asset_url(item_name, asset_key)
+
+        # Read the raster data
+        raw_data = self._read_raster_from_url(asset_url, projection, bounds, resampling)
+
+        # Apply any post-processing callback
+        callback = self.get_read_callback(item_name, asset_key)
+        if callback is not None:
+            raw_data = callback(raw_data)
+
+        return raw_data
+
+    def materialize(
+        self,
+        window: Window,
+        item_groups: list[list[ItemType]],
+        layer_name: str,
+        layer_cfg: LayerConfig,
+    ) -> None:
+        """Materialize data for the window.
+
+        Args:
+            window: the window to materialize.
+            item_groups: the items from get_items.
+            layer_name: the name of this layer.
+            layer_cfg: the config of this layer.
+        """
+        RasterMaterializer().materialize(
+            TileStoreWithLayer(self, layer_name),
+            window,
+            layer_name,
+            layer_cfg,
+            item_groups,
+        )
+
+    # --- TileStore methods that are not supported ---
+
+    def write_raster(
+        self,
+        layer_name: str,
+        item_name: str,
+        bands: list[str],
+        projection: Projection,
+        bounds: PixelBounds,
+        array: npt.NDArray[Any],
+    ) -> None:
+        """Write raster data to the store.
+
+        This is not supported for remote-backed tile stores.
+        """
+        raise NotImplementedError(
+            "DirectMaterializeDataSource does not support writing raster data"
+        )
+
+    def write_raster_file(
+        self, layer_name: str, item_name: str, bands: list[str], fname: Any
+    ) -> None:
+        """Write raster data to the store.
+
+        This is not supported for remote-backed tile stores.
+        """
+        raise NotImplementedError(
+            "DirectMaterializeDataSource does not support writing raster files"
+        )
+
+    def is_vector_ready(self, layer_name: str, item_name: str) -> bool:
+        """Checks if this vector item has been written to the store.
+
+        This is not supported for remote-backed tile stores.
+        """
+        raise NotImplementedError(
+            "DirectMaterializeDataSource does not support vector operations"
+        )
+
+    def read_vector(
+        self,
+        layer_name: str,
+        item_name: str,
+        projection: Projection,
+        bounds: PixelBounds,
+    ) -> Any:
+        """Read vector data from the store.
+
+        This is not supported for remote-backed tile stores.
+        """
+        raise NotImplementedError(
+            "DirectMaterializeDataSource does not support vector operations"
+        )
+
+    def write_vector(self, layer_name: str, item_name: str, features: Any) -> None:
+        """Write vector data to the store.
+
+        This is not supported for remote-backed tile stores.
+        """
+        raise NotImplementedError(
+            "DirectMaterializeDataSource does not support vector operations"
+        )

rslearn/data_sources/earthdaily.py

@@ -6,27 +6,24 @@ import tempfile
 from datetime import timedelta
 from typing import Any, Literal
 
-import affine
-import numpy.typing as npt
 import pystac
 import pystac_client
-import rasterio
 import requests
 import shapely
 from earthdaily import EDSClient, EDSConfig
-from rasterio.enums import Resampling
 from upath import UPath
 
-from rslearn.config import LayerConfig, QueryConfig
+from rslearn.config import QueryConfig
 from rslearn.const import WGS84_PROJECTION
-from rslearn.data_sources import DataSource, DataSourceContext, Item
+from rslearn.data_sources import DataSourceContext, Item
+from rslearn.data_sources.direct_materialize_data_source import (
+    DirectMaterializeDataSource,
+)
 from rslearn.data_sources.utils import match_candidate_items_to_window
-from rslearn.dataset import Window
-from rslearn.dataset.materialize import RasterMaterializer
 from rslearn.log_utils import get_logger
-from rslearn.tile_stores import TileStore, TileStoreWithLayer
+from rslearn.tile_stores import TileStoreWithLayer
 from rslearn.utils.fsspec import join_upath
-from rslearn.utils.geometry import PixelBounds, Projection, STGeometry
+from rslearn.utils.geometry import STGeometry
 
 logger = get_logger(__name__)
 
@@ -62,7 +59,7 @@ class EarthDailyItem(Item):
         )
 
 
-class EarthDaily(DataSource, TileStore):
+class EarthDaily(DirectMaterializeDataSource[EarthDailyItem]):
     """A data source for EarthDaily data.
 
     This requires the following environment variables to be set:
@@ -111,8 +108,9 @@ class EarthDaily(DataSource, TileStore):
                 services "legacy" and "internal" are not supported.
             context: the data source context.
         """
+        super().__init__(asset_bands=asset_bands)
+
         self.collection_name = collection_name
-        self.asset_bands = asset_bands
         self.query = query
         self.sort_by = sort_by
         self.sort_ascending = sort_ascending
@@ -221,6 +219,47 @@ class EarthDaily(DataSource, TileStore):
 
         return item
 
+    # --- DirectMaterializeDataSource implementation ---
+
+    def get_asset_url(self, item_name: str, asset_key: str) -> str:
+        """Get the URL to read the asset for the given item and asset key.
+
+        Args:
+            item_name: the name of the item.
+            asset_key: the key identifying which asset to get.
+
+        Returns:
+            the URL to read the asset from.
+        """
+        item = self.get_item_by_name(item_name)
+        return item.asset_urls[asset_key]
+
+    def get_raster_bands(self, layer_name: str, item_name: str) -> list[list[str]]:
+        """Get the sets of bands that have been stored for the specified item.
+
+        Args:
+            layer_name: the layer name or alias.
+            item_name: the item.
+
+        Returns:
+            a list of lists of bands available for this item.
+        """
+        if self.skip_items_missing_assets:
+            # In this case we can assume that the item has all of the assets.
+            return list(self.asset_bands.values())
+
+        # Otherwise we have to lookup the STAC item to see which assets it has.
+        # Here we use get_item_by_name since it handles caching.
+        item = self.get_item_by_name(item_name)
+        all_bands = []
+        for asset_key, band_names in self.asset_bands.items():
+            if asset_key not in item.asset_urls:
+                continue
+            all_bands.append(band_names)
+        return all_bands
+
+    # --- DataSource implementation ---
+
     def get_items(
         self, geometries: list[STGeometry], query_config: QueryConfig
     ) -> list[list[list[EarthDailyItem]]]:
@@ -285,9 +324,8 @@ class EarthDaily(DataSource, TileStore):
 
         return groups
 
-    def deserialize_item(self, serialized_item: Any) -> EarthDailyItem:
+    def deserialize_item(self, serialized_item: dict) -> EarthDailyItem:
         """Deserializes an item from JSON-decoded data."""
-        assert isinstance(serialized_item, dict)
         return EarthDailyItem.deserialize(serialized_item)
 
     def ingest(
@@ -341,144 +379,3 @@ class EarthDaily(DataSource, TileStore):
                     item.name,
                     asset_key,
                 )
-
-    def is_raster_ready(
-        self, layer_name: str, item_name: str, bands: list[str]
-    ) -> bool:
-        """Checks if this raster has been written to the store.
-
-        Args:
-            layer_name: the layer name or alias.
-            item_name: the item.
-            bands: the list of bands identifying which specific raster to read.
-
-        Returns:
-            whether there is a raster in the store matching the source, item, and
-                bands.
-        """
-        # Always ready since we wrap accesses to EarthDaily.
-        return True
-
-    def get_raster_bands(self, layer_name: str, item_name: str) -> list[list[str]]:
-        """Get the sets of bands that have been stored for the specified item.
-
-        Args:
-            layer_name: the layer name or alias.
-            item_name: the item.
-        """
-        if self.skip_items_missing_assets:
-            # In this case we can assume that the item has all of the assets.
-            return list(self.asset_bands.values())
-
-        # Otherwise we have to lookup the STAC item to see which assets it has.
-        # Here we use get_item_by_name since it handles caching.
-        item = self.get_item_by_name(item_name)
-        all_bands = []
-        for asset_key, band_names in self.asset_bands.items():
-            if asset_key not in item.asset_urls:
-                continue
-            all_bands.append(band_names)
-        return all_bands
-
-    def _get_asset_by_band(self, bands: list[str]) -> str:
-        """Get the name of the asset based on the band names."""
-        for asset_key, asset_bands in self.asset_bands.items():
-            if bands == asset_bands:
-                return asset_key
-
-        raise ValueError(f"no raster with bands {bands}")
-
-    def get_raster_bounds(
-        self, layer_name: str, item_name: str, bands: list[str], projection: Projection
-    ) -> PixelBounds:
-        """Get the bounds of the raster in the specified projection.
-
-        Args:
-            layer_name: the layer name or alias.
-            item_name: the item to check.
-            bands: the list of bands identifying which specific raster to read. These
-                bands must match the bands of a stored raster.
-            projection: the projection to get the raster's bounds in.
-
-        Returns:
-            the bounds of the raster in the projection.
-        """
-        item = self.get_item_by_name(item_name)
-        geom = item.geometry.to_projection(projection)
-        return (
-            int(geom.shp.bounds[0]),
-            int(geom.shp.bounds[1]),
-            int(geom.shp.bounds[2]),
-            int(geom.shp.bounds[3]),
-        )
-
-    def read_raster(
-        self,
-        layer_name: str,
-        item_name: str,
-        bands: list[str],
-        projection: Projection,
-        bounds: PixelBounds,
-        resampling: Resampling = Resampling.bilinear,
-    ) -> npt.NDArray[Any]:
-        """Read raster data from the store.
-
-        Args:
-            layer_name: the layer name or alias.
-            item_name: the item to read.
-            bands: the list of bands identifying which specific raster to read. These
-                bands must match the bands of a stored raster.
-            projection: the projection to read in.
-            bounds: the bounds to read.
-            resampling: the resampling method to use in case reprojection is needed.
-
-        Returns:
-            the raster data
-        """
-        asset_key = self._get_asset_by_band(bands)
-        item = self.get_item_by_name(item_name)
-        asset_url = item.asset_urls[asset_key]
-
-        # Construct the transform to use for the warped dataset.
-        wanted_transform = affine.Affine(
-            projection.x_resolution,
-            0,
-            bounds[0] * projection.x_resolution,
-            0,
-            projection.y_resolution,
-            bounds[1] * projection.y_resolution,
-        )
-
-        with rasterio.open(asset_url) as src:
-            with rasterio.vrt.WarpedVRT(
-                src,
-                crs=projection.crs,
-                transform=wanted_transform,
-                width=bounds[2] - bounds[0],
-                height=bounds[3] - bounds[1],
-                resampling=resampling,
-            ) as vrt:
-                return vrt.read()
-
-    def materialize(
-        self,
-        window: Window,
-        item_groups: list[list[Item]],
-        layer_name: str,
-        layer_cfg: LayerConfig,
-    ) -> None:
-        """Materialize data for the window.
-
-        Args:
-            window: the window to materialize
-            item_groups: the items from get_items
-            layer_name: the name of this layer
-            layer_cfg: the config of this layer
-        """
-        RasterMaterializer().materialize(
-            TileStoreWithLayer(self, layer_name),
-            window,
-            layer_name,
-            layer_cfg,
-            item_groups,
-        )