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

rslearn/data_sources/planetary_computer.py

@@ -1,6 +1,5 @@
 """Data on Planetary Computer."""
 
-import json
 import os
 import tempfile
 import xml.etree.ElementTree as ET
@@ -10,18 +9,14 @@ from typing import Any
 import affine
 import numpy.typing as npt
 import planetary_computer
-import pystac
-import pystac_client
 import rasterio
 import requests
-import shapely
 from rasterio.enums import Resampling
 from upath import UPath
 
-from rslearn.config import LayerConfig, QueryConfig
-from rslearn.const import WGS84_PROJECTION
-from rslearn.data_sources import DataSource, DataSourceContext, Item
-from rslearn.data_sources.utils import match_candidate_items_to_window
+from rslearn.config import LayerConfig
+from rslearn.data_sources import DataSourceContext
+from rslearn.data_sources.stac import SourceItem, StacDataSource
 from rslearn.dataset import Window
 from rslearn.dataset.materialize import RasterMaterializer
 from rslearn.log_utils import get_logger
@@ -35,38 +30,7 @@ from .copernicus import get_harmonize_callback
 logger = get_logger(__name__)
 
 
-class PlanetaryComputerItem(Item):
-    """An item in the PlanetaryComputer data source."""
-
-    def __init__(self, name: str, geometry: STGeometry, asset_urls: dict[str, str]):
-        """Creates a new PlanetaryComputerItem.
-
-        Args:
-            name: unique name of the item
-            geometry: the spatial and temporal extent of the item
-            asset_urls: map from asset key to the unsigned asset URL.
-        """
-        super().__init__(name, geometry)
-        self.asset_urls = asset_urls
-
-    def serialize(self) -> dict[str, Any]:
-        """Serializes the item to a JSON-encodable dictionary."""
-        d = super().serialize()
-        d["asset_urls"] = self.asset_urls
-        return d
-
-    @staticmethod
-    def deserialize(d: dict[str, Any]) -> "PlanetaryComputerItem":
-        """Deserializes an item from a JSON-decoded dictionary."""
-        item = super(PlanetaryComputerItem, PlanetaryComputerItem).deserialize(d)
-        return PlanetaryComputerItem(
-            name=item.name,
-            geometry=item.geometry,
-            asset_urls=d["asset_urls"],
-        )
-
-
-class PlanetaryComputer(DataSource, TileStore):
+class PlanetaryComputer(StacDataSource, TileStore):
     """Modality-agnostic data source for data on Microsoft Planetary Computer.
 
     If there is a subclass available for a modality, it is recommended to use the
@@ -83,10 +47,6 @@ class PlanetaryComputer(DataSource, TileStore):
 
     STAC_ENDPOINT = "https://planetarycomputer.microsoft.com/api/stac/v1"
 
-    # Default threshold for recreating the STAC client to prevent memory leaks
-    # from the pystac Catalog's resolved objects cache growing unbounded
-    DEFAULT_MAX_ITEMS_PER_CLIENT = 1000
-
     def __init__(
         self,
         collection_name: str,
@@ -97,7 +57,6 @@ class PlanetaryComputer(DataSource, TileStore):
         timeout: timedelta = timedelta(seconds=10),
         skip_items_missing_assets: bool = False,
         cache_dir: str | None = None,
-        max_items_per_client: int | None = None,
         context: DataSourceContext = DataSourceContext(),
     ):
         """Initialize a new PlanetaryComputer instance.
@@ -115,228 +74,40 @@ class PlanetaryComputer(DataSource, TileStore):
             cache_dir: optional directory to cache items by name, including asset URLs.
                 If not set, there will be no cache and instead STAC requests will be
                 needed each time.
-            max_items_per_client: number of STAC items to process before recreating
-                the client to prevent memory leaks from the resolved objects cache.
-                Defaults to DEFAULT_MAX_ITEMS_PER_CLIENT.
             context: the data source context.
         """
-        self.collection_name = collection_name
-        self.asset_bands = asset_bands
-        self.query = query
-        self.sort_by = sort_by
-        self.sort_ascending = sort_ascending
-        self.timeout = timeout
-        self.skip_items_missing_assets = skip_items_missing_assets
-        self.max_items_per_client = (
-            max_items_per_client or self.DEFAULT_MAX_ITEMS_PER_CLIENT
-        )
-
+        # Determine the cache_dir to use.
+        cache_upath: UPath | None = None
         if cache_dir is not None:
             if context.ds_path is not None:
-                self.cache_dir = join_upath(context.ds_path, cache_dir)
+                cache_upath = join_upath(context.ds_path, cache_dir)
             else:
-                self.cache_dir = UPath(cache_dir)
-
-            self.cache_dir.mkdir(parents=True, exist_ok=True)
-        else:
-            self.cache_dir = None
-
-        self.client: pystac_client.Client | None = None
-        self._client_item_count = 0
+                cache_upath = UPath(cache_dir)
 
-    def _load_client(
-        self,
-    ) -> pystac_client.Client:
-        """Lazily load pystac client.
+            cache_upath.mkdir(parents=True, exist_ok=True)
 
-        We don't load it when creating the data source because it takes time and caller
-        may not be calling get_items. Additionally, loading it during the get_items
-        call enables leveraging the retry loop functionality in
-        prepare_dataset_windows.
+        # We pass required_assets to StacDataSource of skip_items_missing_assets is set.
+        required_assets: list[str] | None = None
+        if skip_items_missing_assets:
+            required_assets = list(asset_bands.keys())
 
-        Note: We periodically recreate the client to prevent memory leaks from the
-        pystac Catalog's resolved objects cache, which grows unbounded as STAC items
-        are deserialized and cached. The cache cannot be cleared or disabled.
-        """
-        if self.client is None:
-            logger.info("Creating initial STAC client")
-            self.client = pystac_client.Client.open(self.STAC_ENDPOINT)
-            return self.client
-
-        if self._client_item_count < self.max_items_per_client:
-            return self.client
-
-        # Recreate client to clear the resolved objects cache
-        current_client = self.client
-        logger.debug(
-            "Recreating STAC client after processing %d items (threshold: %d)",
-            self._client_item_count,
-            self.max_items_per_client,
+        super().__init__(
+            endpoint=self.STAC_ENDPOINT,
+            collection_name=collection_name,
+            query=query,
+            sort_by=sort_by,
+            sort_ascending=sort_ascending,
+            required_assets=required_assets,
+            cache_dir=cache_upath,
         )
-        client_root = current_client.get_root()
-        client_root.clear_links()
-        client_root.clear_items()
-        client_root.clear_children()
-        self._client_item_count = 0
-        self.client = pystac_client.Client.open(self.STAC_ENDPOINT)
-        return self.client
-
-    def _stac_item_to_item(self, stac_item: pystac.Item) -> PlanetaryComputerItem:
-        shp = shapely.geometry.shape(stac_item.geometry)
-
-        # Get time range.
-        metadata = stac_item.common_metadata
-        if metadata.start_datetime is not None and metadata.end_datetime is not None:
-            time_range = (
-                metadata.start_datetime,
-                metadata.end_datetime,
-            )
-        elif stac_item.datetime is not None:
-            time_range = (stac_item.datetime, stac_item.datetime)
-        else:
-            raise ValueError(
-                f"item {stac_item.id} unexpectedly missing start_datetime, end_datetime, and datetime"
-            )
-
-        geom = STGeometry(WGS84_PROJECTION, shp, time_range)
-        asset_urls = {
-            asset_key: asset_obj.href
-            for asset_key, asset_obj in stac_item.assets.items()
-        }
-        return PlanetaryComputerItem(stac_item.id, geom, asset_urls)
-
-    def get_item_by_name(self, name: str) -> PlanetaryComputerItem:
-        """Gets an item by name.
-
-        Args:
-            name: the name of the item to get
-
-        Returns:
-            the item object
-        """
-        # If cache_dir is set, we cache the item. First here we check if it is already
-        # in the cache.
-        cache_fname: UPath | None = None
-        if self.cache_dir:
-            cache_fname = self.cache_dir / f"{name}.json"
-        if cache_fname is not None and cache_fname.exists():
-            with cache_fname.open() as f:
-                return PlanetaryComputerItem.deserialize(json.load(f))
-
-        # No cache or not in cache, so we need to make the STAC request.
-        logger.debug("Getting STAC item {name}")
-        client = self._load_client()
-
-        search_result = client.search(ids=[name], collections=[self.collection_name])
-        stac_items = list(search_result.items())
-
-        if not stac_items:
-            raise ValueError(
-                f"Item {name} not found in collection {self.collection_name}"
-            )
-        if len(stac_items) > 1:
-            raise ValueError(
-                f"Multiple items found for ID {name} in collection {self.collection_name}"
-            )
-
-        stac_item = stac_items[0]
-        item = self._stac_item_to_item(stac_item)
-
-        # Track items processed for client recreation threshold (after deserialization)
-        self._client_item_count += 1
-
-        # Finally we cache it if cache_dir is set.
-        if cache_fname is not None:
-            with cache_fname.open("w") as f:
-                json.dump(item.serialize(), f)
-
-        return item
-
-    def get_items(
-        self, geometries: list[STGeometry], query_config: QueryConfig
-    ) -> list[list[list[PlanetaryComputerItem]]]:
-        """Get a list of items in the data source intersecting the given geometries.
-
-        Args:
-            geometries: the spatiotemporal geometries
-            query_config: the query configuration
-
-        Returns:
-            List of groups of items that should be retrieved for each geometry.
-        """
-        client = self._load_client()
-
-        groups = []
-        for geometry in geometries:
-            # Get potentially relevant items from the collection by performing one search
-            # for each requested geometry.
-            wgs84_geometry = geometry.to_projection(WGS84_PROJECTION)
-            logger.debug("performing STAC search for geometry %s", wgs84_geometry)
-            result = client.search(
-                collections=[self.collection_name],
-                intersects=shapely.to_geojson(wgs84_geometry.shp),
-                datetime=wgs84_geometry.time_range,
-                query=self.query,
-            )
-            stac_items = [item for item in result.items()]
-            # Track items processed for client recreation threshold (after deserialization)
-            self._client_item_count += len(stac_items)
-            logger.debug("STAC search yielded %d items", len(stac_items))
-
-            if self.skip_items_missing_assets:
-                # Filter out items that are missing any of the assets in self.asset_bands.
-                good_stac_items = []
-                for stac_item in stac_items:
-                    good = True
-                    for asset_key in self.asset_bands.keys():
-                        if asset_key in stac_item.assets:
-                            continue
-                        good = False
-                        break
-                    if good:
-                        good_stac_items.append(stac_item)
-                logger.debug(
-                    "skip_items_missing_assets filter from %d to %d items",
-                    len(stac_items),
-                    len(good_stac_items),
-                )
-                stac_items = good_stac_items
-
-            if self.sort_by is not None:
-                stac_items.sort(
-                    key=lambda stac_item: stac_item.properties[self.sort_by],
-                    reverse=not self.sort_ascending,
-                )
-
-            candidate_items = [
-                self._stac_item_to_item(stac_item) for stac_item in stac_items
-            ]
-
-            # Since we made the STAC request, might as well save these to the cache.
-            if self.cache_dir is not None:
-                for item in candidate_items:
-                    cache_fname = self.cache_dir / f"{item.name}.json"
-                    if cache_fname.exists():
-                        continue
-                    with cache_fname.open("w") as f:
-                        json.dump(item.serialize(), f)
-
-            cur_groups = match_candidate_items_to_window(
-                geometry, candidate_items, query_config
-            )
-            groups.append(cur_groups)
-
-        return groups
-
-    def deserialize_item(self, serialized_item: Any) -> PlanetaryComputerItem:
-        """Deserializes an item from JSON-decoded data."""
-        assert isinstance(serialized_item, dict)
-        return PlanetaryComputerItem.deserialize(serialized_item)
+        self.asset_bands = asset_bands
+        self.timeout = timeout
+        self.skip_items_missing_assets = skip_items_missing_assets
 
     def ingest(
         self,
         tile_store: TileStoreWithLayer,
-        items: list[PlanetaryComputerItem],
+        items: list[SourceItem],
         geometries: list[list[STGeometry]],
     ) -> None:
         """Ingest items into the given tile store.
@@ -512,7 +283,7 @@ class PlanetaryComputer(DataSource, TileStore):
     def materialize(
         self,
         window: Window,
-        item_groups: list[list[Item]],
+        item_groups: list[list[SourceItem]],
         layer_name: str,
         layer_cfg: LayerConfig,
     ) -> None:
@@ -601,7 +372,7 @@ class Sentinel2(PlanetaryComputer):
             **kwargs,
         )
 
-    def _get_product_xml(self, item: PlanetaryComputerItem) -> ET.Element:
+    def _get_product_xml(self, item: SourceItem) -> ET.Element:
         asset_url = planetary_computer.sign(item.asset_urls["product-metadata"])
         response = requests.get(asset_url, timeout=self.timeout.total_seconds())
         response.raise_for_status()
@@ -610,7 +381,7 @@ class Sentinel2(PlanetaryComputer):
     def ingest(
         self,
         tile_store: TileStoreWithLayer,
-        items: list[PlanetaryComputerItem],
+        items: list[SourceItem],
         geometries: list[list[STGeometry]],
     ) -> None:
         """Ingest items into the given tile store.
@@ -796,3 +567,53 @@ class Naip(PlanetaryComputer):
             context=context,
             **kwargs,
         )
+
+
+class CopDemGlo30(PlanetaryComputer):
+    """A data source for Copernicus DEM GLO-30 (30m) on Microsoft Planetary Computer.
+
+    See https://planetarycomputer.microsoft.com/dataset/cop-dem-glo-30.
+    """
+
+    COLLECTION_NAME = "cop-dem-glo-30"
+    DATA_ASSET = "data"
+
+    def __init__(
+        self,
+        band_name: str = "DEM",
+        context: DataSourceContext = DataSourceContext(),
+        **kwargs: Any,
+    ):
+        """Initialize a new CopDemGlo30 instance.
+
+        Args:
+            band_name: band name to use if the layer config is missing from the
+                context.
+            context: the data source context.
+            kwargs: additional arguments to pass to PlanetaryComputer.
+        """
+        if context.layer_config is not None:
+            if len(context.layer_config.band_sets) != 1:
+                raise ValueError("expected a single band set")
+            if len(context.layer_config.band_sets[0].bands) != 1:
+                raise ValueError("expected band set to have a single band")
+            band_name = context.layer_config.band_sets[0].bands[0]
+
+        super().__init__(
+            collection_name=self.COLLECTION_NAME,
+            asset_bands={self.DATA_ASSET: [band_name]},
+            # Skip since all items should have the same asset(s).
+            skip_items_missing_assets=True,
+            context=context,
+            **kwargs,
+        )
+
+    def _stac_item_to_item(self, stac_item: Any) -> SourceItem:
+        # Copernicus DEM is static; ignore item timestamps so it matches any window.
+        item = super()._stac_item_to_item(stac_item)
+        item.geometry = STGeometry(item.geometry.projection, item.geometry.shp, None)
+        return item
+
+    def _get_search_time_range(self, geometry: STGeometry) -> None:
+        # Copernicus DEM is static; do not filter STAC searches by time.
+        return None