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

rslearn/data_sources/__init__.py

@@ -17,6 +17,7 @@ from .data_source import (
     ItemLookupDataSource,
     RetrieveItemDataSource,
 )
+from .direct_materialize_data_source import DirectMaterializeDataSource
 
 __all__ = (
     "DataSource",
@@ -24,5 +25,6 @@ __all__ = (
     "Item",
     "ItemLookupDataSource",
     "RetrieveItemDataSource",
+    "DirectMaterializeDataSource",
     "data_source_from_config",
 )

rslearn/data_sources/aws_landsat.py

@@ -9,32 +9,28 @@ import urllib.request
 import zipfile
 from collections.abc import Generator
 from datetime import datetime
-from typing import Any, BinaryIO
+from typing import BinaryIO
 
-import affine
 import boto3
 import dateutil.parser
 import fiona
 import fiona.transform
-import numpy.typing as npt
-import rasterio
 import shapely
 import shapely.geometry
 import tqdm
-from rasterio.enums import Resampling
 from upath import UPath
 
 import rslearn.data_sources.utils
-from rslearn.config import LayerConfig
 from rslearn.const import SHAPEFILE_AUX_EXTENSIONS, WGS84_PROJECTION
-from rslearn.dataset import Window
-from rslearn.dataset.materialize import RasterMaterializer
-from rslearn.tile_stores import TileStore, TileStoreWithLayer
+from rslearn.data_sources.direct_materialize_data_source import (
+    DirectMaterializeDataSource,
+)
+from rslearn.tile_stores import TileStoreWithLayer
 from rslearn.utils.fsspec import get_upath_local, join_upath, open_atomic
-from rslearn.utils.geometry import PixelBounds, Projection, STGeometry
+from rslearn.utils.geometry import STGeometry
 from rslearn.utils.grid_index import GridIndex
 
-from .data_source import DataSource, DataSourceContext, Item, QueryConfig
+from .data_source import DataSourceContext, Item, QueryConfig
 
 WRS2_GRID_SIZE = 1.0
 
@@ -79,7 +75,7 @@ class LandsatOliTirsItem(Item):
         )
 
 
-class LandsatOliTirs(DataSource, TileStore):
+class LandsatOliTirs(DirectMaterializeDataSource[LandsatOliTirsItem]):
     """A data source for Landsat 8/9 OLI-TIRS imagery on AWS.
 
     Specifically, uses the usgs-landsat S3 bucket maintained by USGS. The data includes
@@ -91,7 +87,7 @@ class LandsatOliTirs(DataSource, TileStore):
 
     bucket_name = "usgs-landsat"
     bucket_prefix = "collection02/level-1/standard/oli-tirs"
-    bands = ["B1", "B2", "B3", "B4", "B5", "B6", "B7", "B8", "B9", "B10", "B11"]
+    BANDS = ["B1", "B2", "B3", "B4", "B5", "B6", "B7", "B8", "B9", "B10", "B11"]
 
     wrs2_url = "https://d9-wret.s3.us-west-2.amazonaws.com/assets/palladium/production/s3fs-public/atoms/files/WRS2_descending_0.zip"  # noqa
     """URL to download shapefile specifying polygon of each (path, row)."""
@@ -110,6 +106,10 @@ class LandsatOliTirs(DataSource, TileStore):
                 SpaceMode.WITHIN.
             context: the data source context.
         """
+        # Each band is a separate single-band asset.
+        asset_bands = {band: [band] for band in self.BANDS}
+        super().__init__(asset_bands=asset_bands)
+
         # If context is provided, we join the directory with the dataset path,
         # otherwise we treat it directly as UPath.
         if context.ds_path is not None:
@@ -342,16 +342,43 @@ class LandsatOliTirs(DataSource, TileStore):
                 return item
         raise ValueError(f"item {name} not found")
 
-    def deserialize_item(self, serialized_item: Any) -> LandsatOliTirsItem:
+    # --- DirectMaterializeDataSource implementation ---
+
+    def get_asset_url(self, item_name: str, asset_key: str) -> str:
+        """Get the presigned 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 (the band name).
+
+        Returns:
+            the presigned URL to read the asset from.
+        """
+        # Get the item since it has the blob path.
+        item = self.get_item_by_name(item_name)
+
+        # For Landsat, the asset_key is the band name (e.g., "B1", "B2", etc.).
+        blob_key = item.blob_path + f"{asset_key}.TIF"
+        return self.client.generate_presigned_url(
+            "get_object",
+            Params={
+                "Bucket": self.bucket_name,
+                "Key": blob_key,
+                "RequestPayer": "requester",
+            },
+        )
+
+    # --- DataSource implementation ---
+
+    def deserialize_item(self, serialized_item: dict) -> LandsatOliTirsItem:
         """Deserializes an item from JSON-decoded data."""
-        assert isinstance(serialized_item, dict)
         return LandsatOliTirsItem.deserialize(serialized_item)
 
     def retrieve_item(
         self, item: LandsatOliTirsItem
     ) -> Generator[tuple[str, BinaryIO], None, None]:
         """Retrieves the rasters corresponding to an item as file streams."""
-        for band in self.bands:
+        for band in self.BANDS:
             buf = io.BytesIO()
             self.bucket.download_fileobj(
                 item.blob_path + f"{band}.TIF",
@@ -376,7 +403,7 @@ class LandsatOliTirs(DataSource, TileStore):
             geometries: a list of geometries needed for each item
         """
         for item, cur_geometries in zip(items, geometries):
-            for band in self.bands:
+            for band in self.BANDS:
                 band_names = [band]
                 if tile_store.is_raster_ready(item.name, band_names):
                     continue
@@ -389,147 +416,3 @@ class LandsatOliTirs(DataSource, TileStore):
                         ExtraArgs={"RequestPayer": "requester"},
                     )
                     tile_store.write_raster_file(item.name, band_names, UPath(fname))
-
-    # The functions below are to emulate TileStore functionality so we can easily
-    # support materialization directly from the COGs.
-    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 access it on AWS bucket.
-        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.
-
-        Returns:
-            a list of lists of bands that are in the tile store (with one raster
-                stored corresponding to each inner list). If no rasters are ready for
-                this item, returns empty list.
-        """
-        return [[band] for band in self.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
-        """
-        # Landsat assets have single band per asset.
-        assert len(bands) == 1
-        band = bands[0]
-
-        # Get the item since it has the blob path.
-        item = self.get_item_by_name(item_name)
-
-        # Create pre-signed URL for rasterio access.
-        # We do this because accessing via URL is much faster since rasterio can use
-        # the URL directly.
-        blob_key = item.blob_path + f"{band}.TIF"
-        url = self.client.generate_presigned_url(
-            "get_object",
-            Params={
-                "Bucket": self.bucket_name,
-                "Key": blob_key,
-                "RequestPayer": "requester",
-            },
-        )
-
-        # 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 materialize(
-        self,
-        window: Window,
-        item_groups: list[list[LandsatOliTirsItem]],
-        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,
-        )

rslearn/data_sources/aws_open_data.py

@@ -318,9 +318,8 @@ class Naip(DataSource):
             groups.append(cur_groups)
         return groups
 
-    def deserialize_item(self, serialized_item: Any) -> NaipItem:
+    def deserialize_item(self, serialized_item: dict) -> NaipItem:
         """Deserializes an item from JSON-decoded data."""
-        assert isinstance(serialized_item, dict)
         return NaipItem.deserialize(serialized_item)
 
     def ingest(
@@ -639,9 +638,8 @@ class Sentinel2(
                 return item
         raise ValueError(f"item {name} not found")
 
-    def deserialize_item(self, serialized_item: Any) -> Sentinel2Item:
+    def deserialize_item(self, serialized_item: dict) -> Sentinel2Item:
         """Deserializes an item from JSON-decoded data."""
-        assert isinstance(serialized_item, dict)
         return Sentinel2Item.deserialize(serialized_item)
 
     def retrieve_item(

rslearn/data_sources/aws_sentinel1.py

@@ -2,7 +2,6 @@
 
 import os
 import tempfile
-from typing import Any
 
 import boto3
 from upath import UPath
@@ -78,9 +77,8 @@ class Sentinel1(DataSource, TileStore):
         """Gets an item by name."""
         return self.sentinel1.get_item_by_name(name)
 
-    def deserialize_item(self, serialized_item: Any) -> CopernicusItem:
+    def deserialize_item(self, serialized_item: dict) -> CopernicusItem:
         """Deserializes an item from JSON-decoded data."""
-        assert isinstance(serialized_item, dict)
         return CopernicusItem.deserialize(serialized_item)
 
     def ingest(

rslearn/data_sources/aws_sentinel2_element84.py

@@ -6,23 +6,20 @@ from collections.abc import Callable
 from datetime import timedelta
 from typing import Any
 
-import affine
 import numpy as np
 import numpy.typing as npt
 import rasterio
 import requests
-from rasterio.enums import Resampling
 from upath import UPath
 
-from rslearn.config import LayerConfig
+from rslearn.data_sources.direct_materialize_data_source import (
+    DirectMaterializeDataSource,
+)
 from rslearn.data_sources.stac import SourceItem, StacDataSource
-from rslearn.dataset import Window
-from rslearn.dataset.manage import RasterMaterializer
 from rslearn.log_utils import get_logger
-from rslearn.tile_stores import TileStore, TileStoreWithLayer
-from rslearn.utils import Projection, STGeometry
+from rslearn.tile_stores import TileStoreWithLayer
+from rslearn.utils import STGeometry
 from rslearn.utils.fsspec import join_upath
-from rslearn.utils.geometry import PixelBounds
 from rslearn.utils.raster_format import get_raster_projection_and_bounds
 
 from .data_source import (
@@ -32,7 +29,7 @@ from .data_source import (
 logger = get_logger(__name__)
 
 
-class Sentinel2(StacDataSource, TileStore):
+class Sentinel2(DirectMaterializeDataSource[SourceItem], StacDataSource):
     """A data source for Sentinel-2 L2A imagery on AWS from s3://sentinel-cogs.
 
     The S3 bucket has COGs so this data source supports direct materialization. It also
@@ -97,31 +94,36 @@ class Sentinel2(StacDataSource, TileStore):
             cache_upath.mkdir(parents=True, exist_ok=True)
 
         # Determine which assets we need based on the bands in the layer config.
-        self.asset_bands: dict[str, list[str]]
+        asset_bands: dict[str, list[str]]
         if context.layer_config is not None:
-            self.asset_bands = {}
+            asset_bands = {}
             for asset_key, band_names in self.ASSET_BANDS.items():
                 # See if the bands provided by this asset intersect with the bands in
                 # at least one configured band set.
                 for band_set in context.layer_config.band_sets:
                     if not set(band_set.bands).intersection(set(band_names)):
                         continue
-                    self.asset_bands[asset_key] = band_names
+                    asset_bands[asset_key] = band_names
                     break
         elif assets is not None:
-            self.asset_bands = {
+            asset_bands = {
                 asset_key: self.ASSET_BANDS[asset_key] for asset_key in assets
             }
         else:
-            self.asset_bands = self.ASSET_BANDS
+            asset_bands = dict(self.ASSET_BANDS)
+
+        # Initialize DirectMaterializeDataSource with asset_bands
+        DirectMaterializeDataSource.__init__(self, asset_bands=asset_bands)
 
-        super().__init__(
+        # Initialize StacDataSource
+        StacDataSource.__init__(
+            self,
             endpoint=self.STAC_ENDPOINT,
             collection_name=self.COLLECTION_NAME,
             query=query,
             sort_by=sort_by,
             sort_ascending=sort_ascending,
-            required_assets=list(self.asset_bands.keys()),
+            required_assets=list(asset_bands.keys()),
             cache_dir=cache_upath,
             properties_to_record=[self.HARMONIZE_PROPERTY_NAME],
         )
@@ -129,6 +131,42 @@ class Sentinel2(StacDataSource, TileStore):
         self.harmonize = harmonize
         self.timeout = timeout
 
+    # --- 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_read_callback(
+        self, item_name: str, asset_key: str
+    ) -> Callable[[npt.NDArray[Any]], npt.NDArray[Any]] | None:
+        """Return a callback to harmonize Sentinel-2 data if needed.
+
+        Args:
+            item_name: the name of the item being read.
+            asset_key: the key identifying which asset is being read.
+
+        Returns:
+            A callback function for harmonization, or None if not needed.
+        """
+        # Visual bands do not need harmonization.
+        if not self.harmonize or asset_key == "visual":
+            return None
+
+        item = self.get_item_by_name(item_name)
+        return self._get_harmonize_callback(item)
+
+    # --- Harmonization helpers ---
+
     def _get_harmonize_callback(
         self, item: SourceItem
     ) -> Callable[[npt.NDArray], npt.NDArray] | None:
@@ -223,152 +261,3 @@ class Sentinel2(StacDataSource, 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 underlying API.
-        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.
-
-        Returns:
-            a list of lists of bands that are in the tile store (with one raster
-                stored corresponding to each inner list). If no rasters are ready for
-                this item, returns empty list.
-        """
-        return list(self.asset_bands.values())
-
-    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 known asset 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,
-        )
-
-        # Read from the raster under the specified projection/bounds.
-        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:
-                raw_data = vrt.read()
-
-        # We can return the data now if harmonization is not needed.
-        if not self.harmonize or bands == self.ASSET_BANDS["visual"]:
-            return raw_data
-
-        # Otherwise we apply the harmonize_callback.
-        item = self.get_item_by_name(item_name)
-        harmonize_callback = self._get_harmonize_callback(item)
-
-        if harmonize_callback is None:
-            return raw_data
-
-        array = harmonize_callback(raw_data)
-        return array
-
-    def materialize(
-        self,
-        window: Window,
-        item_groups: list[list[SourceItem]],
-        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,
-        )

rslearn/data_sources/climate_data_store.py

@@ -3,7 +3,6 @@
 import os
 import tempfile
 from datetime import UTC, datetime
-from typing import Any
 
 import cdsapi
 import netCDF4
@@ -160,9 +159,8 @@ class ERA5Land(DataSource):
 
         return all_groups
 
-    def deserialize_item(self, serialized_item: Any) -> Item:
+    def deserialize_item(self, serialized_item: dict) -> Item:
         """Deserializes an item from JSON-decoded data."""
-        assert isinstance(serialized_item, dict)
         return Item.deserialize(serialized_item)
 
     def _convert_nc_to_tif(self, nc_path: UPath, tif_path: UPath) -> None:

rslearn/data_sources/copernicus.py

@@ -353,9 +353,8 @@ class Copernicus(DataSource):
                 self.username = os.environ["COPERNICUS_USERNAME"]
                 self.password = os.environ["COPERNICUS_PASSWORD"]
 
-    def deserialize_item(self, serialized_item: Any) -> CopernicusItem:
+    def deserialize_item(self, serialized_item: dict) -> CopernicusItem:
         """Deserializes an item from JSON-decoded data."""
-        assert isinstance(serialized_item, dict)
         return CopernicusItem.deserialize(serialized_item)
 
     def _get(self, path: str) -> dict[str, Any]:

rslearn/data_sources/data_source.py

@@ -76,7 +76,7 @@ class DataSource(Generic[ItemType]):
         """
         raise NotImplementedError
 
-    def deserialize_item(self, serialized_item: Any) -> ItemType:
+    def deserialize_item(self, serialized_item: dict) -> ItemType:
         """Deserializes an item from JSON-decoded data."""
         raise NotImplementedError