rslearn.data_sources.stac

stac

A partial data source implementation providing get_items using a STAC API.

SourceItem

Bases: Item

An item in the StacDataSource data source.

Source code in rslearn/data_sources/stac.py

class SourceItem(Item):
    """An item in the StacDataSource data source."""

    def __init__(
        self,
        name: str,
        geometry: STGeometry,
        asset_urls: dict[str, str],
        properties: dict[str, str],
    ):
        """Creates a new SourceItem.

        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.
            properties: properties requested by the data source implementation.
        """
        super().__init__(name, geometry)
        self.asset_urls = asset_urls
        self.properties = properties

    def serialize(self) -> dict[str, Any]:
        """Serializes the item to a JSON-encodable dictionary."""
        d = super().serialize()
        d["asset_urls"] = self.asset_urls
        d["properties"] = self.properties
        return d

    @staticmethod
    def deserialize(d: dict[str, Any]) -> "SourceItem":
        """Deserializes an item from a JSON-decoded dictionary."""
        item = super(SourceItem, SourceItem).deserialize(d)
        return SourceItem(
            name=item.name,
            geometry=item.geometry,
            asset_urls=d["asset_urls"],
            properties=d["properties"],
        )

serialize

serialize() -> dict[str, Any]

Serializes the item to a JSON-encodable dictionary.

Source code in rslearn/data_sources/stac.py

def serialize(self) -> dict[str, Any]:
    """Serializes the item to a JSON-encodable dictionary."""
    d = super().serialize()
    d["asset_urls"] = self.asset_urls
    d["properties"] = self.properties
    return d

deserialize staticmethod

deserialize(d: dict[str, Any]) -> SourceItem

Deserializes an item from a JSON-decoded dictionary.

Source code in rslearn/data_sources/stac.py

@staticmethod
def deserialize(d: dict[str, Any]) -> "SourceItem":
    """Deserializes an item from a JSON-decoded dictionary."""
    item = super(SourceItem, SourceItem).deserialize(d)
    return SourceItem(
        name=item.name,
        geometry=item.geometry,
        asset_urls=d["asset_urls"],
        properties=d["properties"],
    )

StacDataSource

Bases: ItemLookupDataSource[SourceItem]

A partial data source implementing get_items using a STAC API.

This is a helper class that full implementations can extend to not have to worry about the get_items and get_item_by_name implementation.

Source code in rslearn/data_sources/stac.py

class StacDataSource(ItemLookupDataSource[SourceItem]):
    """A partial data source implementing get_items using a STAC API.

    This is a helper class that full implementations can extend to not have to worry
    about the get_items and get_item_by_name implementation.
    """

    def __init__(
        self,
        endpoint: str,
        collection_name: str | list[str],
        query: dict[str, Any] | None = None,
        sort_by: str | None = None,
        sort_ascending: bool = True,
        required_assets: list[str] | None = None,
        limit: int = 100,
        properties_to_record: list[str] = [],
    ):
        """Create a new StacDataSource.

        Args:
            endpoint: the STAC endpoint to use.
            collection_name: the STAC collection name or names.
            query: optional STAC query dict to include in searches, e.g. {"eo:cloud_cover": {"lt": 50}}.
            sort_by: sort results by this STAC property.
            sort_ascending: if sort_by is set, sort in ascending order (default).
                Otherwise sort in descending order.
            required_assets: if set, we ignore items that do not have all of these
                asset keys.
            limit: limit to pass to search queries.
            properties_to_record: if these properties on the STAC item exist, they are
                are retained in the SourceItem when we initialize it.
        """
        self.client = StacClient(endpoint)
        if isinstance(collection_name, str):
            self.collection_names = [collection_name]
        else:
            self.collection_names = list(collection_name)
        self.query = query
        self.sort_by = sort_by
        self.sort_ascending = sort_ascending
        self.required_assets = required_assets
        self.limit = limit
        self.properties_to_record = properties_to_record

    def _collection_description(self) -> str:
        """Return a human-readable description of the configured collection filter."""
        if len(self.collection_names) == 1:
            return f"collection {self.collection_names[0]}"
        return f"collections {self.collection_names}"

    def _stac_item_to_item(self, stac_item: StacItem) -> SourceItem:
        # Make sure geometry, time range, and assets are set.
        if stac_item.geometry is None:
            raise ValueError("got unexpected item with no geometry")
        if stac_item.time_range is None:
            raise ValueError("got unexpected item with no time range")
        if stac_item.assets is None:
            raise ValueError("got unexpected item with no assets")

        shp = shapely.geometry.shape(stac_item.geometry)
        geom = STGeometry(WGS84_PROJECTION, shp, stac_item.time_range)
        asset_urls = {
            asset_key: asset_obj.href
            for asset_key, asset_obj in stac_item.assets.items()
        }

        # Keep any properties requested by the data source implementation.
        properties = {}
        for prop_name in self.properties_to_record:
            if prop_name not in stac_item.properties:
                continue
            properties[prop_name] = stac_item.properties[prop_name]

        return SourceItem(stac_item.id, geom, asset_urls, properties)

    def _get_search_time_range(
        self, geometry: STGeometry
    ) -> datetime | tuple[datetime, datetime] | None:
        """Get time range to include in STAC API search.

        By default, we filter STAC searches to the window's time range. Subclasses can
        override this to disable time filtering for "static" datasets.

        Args:
            geometry: the geometry we are searching for.

        Returns:
            the time range (or timestamp) to pass to the STAC search, or None to avoid
                temporal filtering in the search request.
        """
        # Note: StacClient.search accepts either a datetime or a (start, end) tuple.
        return geometry.time_range

    def _get_sort_key(self, stac_item: StacItem) -> Any:
        """Return the key to use when sorting STAC items."""
        assert self.sort_by is not None
        return stac_item.properties[self.sort_by]

    def _get_search_intersects_shape(self, shp: shapely.Geometry) -> shapely.Geometry:
        """Return the geometry to use for STAC intersects searches."""
        return shp

    def _should_include_item(self, item: SourceItem) -> bool:
        """Return whether the converted item should be included in search results."""
        return True

    def get_item_by_name(self, name: str) -> SourceItem:
        """Gets an item by name.

        Args:
            name: the name of the item to get

        Returns:
            the item object
        """
        logger.debug(f"Getting STAC item {name}")
        stac_items = self.client.search(ids=[name], collections=self.collection_names)

        if len(stac_items) == 0:
            raise ValueError(
                f"Item {name} not found in {self._collection_description()}"
            )
        if len(stac_items) > 1:
            raise ValueError(
                f"Multiple items found for ID {name} in {self._collection_description()}"
            )

        stac_item = stac_items[0]
        return self._stac_item_to_item(stac_item)

    def get_items(
        self, geometries: list[STGeometry], query_config: QueryConfig
    ) -> list[list[MatchedItemGroup[SourceItem]]]:
        """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.
        """
        groups = []
        for geometry in geometries:
            # Get potentially relevant items from the collection by performing one search
            # for each requested geometry.
            wgs84_geometry = geometry.to_wgs84()
            logger.debug("performing STAC search for geometry %s", wgs84_geometry)
            search_time_range = self._get_search_time_range(wgs84_geometry)
            stac_items = self.client.search(
                collections=self.collection_names,
                intersects=json.loads(
                    shapely.to_geojson(
                        self._get_search_intersects_shape(wgs84_geometry.shp)
                    )
                ),
                date_time=search_time_range,
                query=self.query,
                limit=self.limit,
            )
            logger.debug("STAC search yielded %d items", len(stac_items))

            if self.required_assets is not None:
                # Filter out items that are missing any of the assets in self.asset_bands.
                good_stac_items = []
                for stac_item in stac_items:
                    if stac_item.assets is None:
                        raise ValueError(f"got STAC item {stac_item.id} with no assets")

                    good = True
                    for asset_key in self.required_assets:
                        if asset_key in stac_item.assets:
                            continue
                        good = False
                        break
                    if good:
                        good_stac_items.append(stac_item)
                logger.debug(
                    "required_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=self._get_sort_key,
                    reverse=not self.sort_ascending,
                )

            candidate_items = []
            for stac_item in stac_items:
                candidate_item = self._stac_item_to_item(stac_item)
                if not self._should_include_item(candidate_item):
                    continue
                candidate_items.append(candidate_item)

            cur_groups = match_candidate_items_to_window(
                geometry, candidate_items, query_config
            )
            groups.append(cur_groups)

        return groups

    def deserialize_item(self, serialized_item: dict) -> SourceItem:
        """Deserializes an item from JSON-decoded data."""
        return SourceItem.deserialize(serialized_item)

get_item_by_name

get_item_by_name(name: str) -> SourceItem

Gets an item by name.

Parameters:

Name	Type	Description	Default
name	str	the name of the item to get	required

Returns:

Type	Description
SourceItem	the item object

Source code in rslearn/data_sources/stac.py

def get_item_by_name(self, name: str) -> SourceItem:
    """Gets an item by name.

    Args:
        name: the name of the item to get

    Returns:
        the item object
    """
    logger.debug(f"Getting STAC item {name}")
    stac_items = self.client.search(ids=[name], collections=self.collection_names)

    if len(stac_items) == 0:
        raise ValueError(
            f"Item {name} not found in {self._collection_description()}"
        )
    if len(stac_items) > 1:
        raise ValueError(
            f"Multiple items found for ID {name} in {self._collection_description()}"
        )

    stac_item = stac_items[0]
    return self._stac_item_to_item(stac_item)

get_items

get_items(geometries: list[STGeometry], query_config: QueryConfig) -> list[list[MatchedItemGroup[SourceItem]]]

Get a list of items in the data source intersecting the given geometries.

Parameters:

Name	Type	Description	Default
geometries	list[STGeometry]	the spatiotemporal geometries	required
query_config	QueryConfig	the query configuration	required

Returns:

Type	Description
list[list[MatchedItemGroup[SourceItem]]]	List of groups of items that should be retrieved for each geometry.

Source code in rslearn/data_sources/stac.py

def get_items(
    self, geometries: list[STGeometry], query_config: QueryConfig
) -> list[list[MatchedItemGroup[SourceItem]]]:
    """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.
    """
    groups = []
    for geometry in geometries:
        # Get potentially relevant items from the collection by performing one search
        # for each requested geometry.
        wgs84_geometry = geometry.to_wgs84()
        logger.debug("performing STAC search for geometry %s", wgs84_geometry)
        search_time_range = self._get_search_time_range(wgs84_geometry)
        stac_items = self.client.search(
            collections=self.collection_names,
            intersects=json.loads(
                shapely.to_geojson(
                    self._get_search_intersects_shape(wgs84_geometry.shp)
                )
            ),
            date_time=search_time_range,
            query=self.query,
            limit=self.limit,
        )
        logger.debug("STAC search yielded %d items", len(stac_items))

        if self.required_assets is not None:
            # Filter out items that are missing any of the assets in self.asset_bands.
            good_stac_items = []
            for stac_item in stac_items:
                if stac_item.assets is None:
                    raise ValueError(f"got STAC item {stac_item.id} with no assets")

                good = True
                for asset_key in self.required_assets:
                    if asset_key in stac_item.assets:
                        continue
                    good = False
                    break
                if good:
                    good_stac_items.append(stac_item)
            logger.debug(
                "required_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=self._get_sort_key,
                reverse=not self.sort_ascending,
            )

        candidate_items = []
        for stac_item in stac_items:
            candidate_item = self._stac_item_to_item(stac_item)
            if not self._should_include_item(candidate_item):
                continue
            candidate_items.append(candidate_item)

        cur_groups = match_candidate_items_to_window(
            geometry, candidate_items, query_config
        )
        groups.append(cur_groups)

    return groups

deserialize_item

deserialize_item(serialized_item: dict) -> SourceItem

Deserializes an item from JSON-decoded data.

Source code in rslearn/data_sources/stac.py

def deserialize_item(self, serialized_item: dict) -> SourceItem:
    """Deserializes an item from JSON-decoded data."""
    return SourceItem.deserialize(serialized_item)

AxisAlignedStacDataSource

Bases: StacDataSource

STAC data source variant for APIs requiring axis-aligned search geometry.

Source code in rslearn/data_sources/stac.py

class AxisAlignedStacDataSource(StacDataSource):
    """STAC data source variant for APIs requiring axis-aligned search geometry."""

    def _get_search_intersects_shape(self, shp: shapely.Geometry) -> shapely.Geometry:
        """Return an axis-aligned geometry to use for STAC intersects searches."""
        return _get_axis_aligned_search_intersects_shape(shp)