rslearn.data_sources.aws_landsat

aws_landsat

Data source for raster data in Registry of Open Data on AWS.

LandsatOliTirsItem

Bases: Item

An item in the LandsatOliTirs data source.

Source code in rslearn/data_sources/aws_landsat.py

class LandsatOliTirsItem(Item):
    """An item in the LandsatOliTirs data source."""

    def __init__(
        self, name: str, geometry: STGeometry, blob_path: str, cloud_cover: float
    ) -> None:
        """Creates a new LandsatOliTirsItem.

        Args:
            name: unique name of the item
            geometry: the spatial and temporal extent of the item
            blob_path: path in bucket, e.g.,
                collection02/level-1/standard/oli-tirs/2024/032/028/LC09_L1GT_032028_20240214_20240214_02_T2/LC09_L1GT_032028_20240214_20240214_02_T2_
            cloud_cover: the scene's cloud cover (between 0 and 100)
        """
        super().__init__(name, geometry)
        self.blob_path = blob_path
        self.cloud_cover = cloud_cover

    def serialize(self) -> dict:
        """Serializes the item to a JSON-encodable dictionary."""
        d = super().serialize()
        d["blob_path"] = self.blob_path
        d["cloud_cover"] = self.cloud_cover
        return d

    @staticmethod
    def deserialize(d: dict) -> "LandsatOliTirsItem":
        """Deserializes an item from a JSON-decoded dictionary."""
        if "name" not in d:
            d["name"] = d["blob_path"].split("/")[-1].split(".tif")[0]
        item = super(LandsatOliTirsItem, LandsatOliTirsItem).deserialize(d)
        return LandsatOliTirsItem(
            name=item.name,
            geometry=item.geometry,
            blob_path=d["blob_path"],
            cloud_cover=d["cloud_cover"],
        )

serialize

serialize() -> dict

Serializes the item to a JSON-encodable dictionary.

Source code in rslearn/data_sources/aws_landsat.py

def serialize(self) -> dict:
    """Serializes the item to a JSON-encodable dictionary."""
    d = super().serialize()
    d["blob_path"] = self.blob_path
    d["cloud_cover"] = self.cloud_cover
    return d

deserialize staticmethod

deserialize(d: dict) -> LandsatOliTirsItem

Deserializes an item from a JSON-decoded dictionary.

Source code in rslearn/data_sources/aws_landsat.py

@staticmethod
def deserialize(d: dict) -> "LandsatOliTirsItem":
    """Deserializes an item from a JSON-decoded dictionary."""
    if "name" not in d:
        d["name"] = d["blob_path"].split("/")[-1].split(".tif")[0]
    item = super(LandsatOliTirsItem, LandsatOliTirsItem).deserialize(d)
    return LandsatOliTirsItem(
        name=item.name,
        geometry=item.geometry,
        blob_path=d["blob_path"],
        cloud_cover=d["cloud_cover"],
    )

LandsatOliTirs

Bases: DirectMaterializeDataSource[LandsatOliTirsItem], ItemLookupDataSource[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 Tier 1/2 scenes but does not seem to include Real-Time scenes.

See https://aws.amazon.com/marketplace/pp/prodview-ivr4jeq6flk7u for details about the bucket.

Source code in rslearn/data_sources/aws_landsat.py

class LandsatOliTirs(
    DirectMaterializeDataSource[LandsatOliTirsItem],
    ItemLookupDataSource[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
    Tier 1/2 scenes but does not seem to include Real-Time scenes.

    See https://aws.amazon.com/marketplace/pp/prodview-ivr4jeq6flk7u for details about
    the bucket.
    """

    bucket_name = "usgs-landsat"
    bucket_prefix = "collection02/level-1/standard/oli-tirs"
    BANDS = ["B1", "B2", "B3", "B4", "B5", "B6", "B7", "B8", "B9", "B10", "B11"]

    def __init__(
        self,
        metadata_cache_dir: str,
        sort_by: str | None = None,
        context: DataSourceContext = DataSourceContext(),
    ) -> None:
        """Initialize a new LandsatOliTirs instance.

        Args:
            metadata_cache_dir: directory to cache produtc metadata files.
            sort_by: can be "cloud_cover", default arbitrary order; only has effect for
                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:
            self.metadata_cache_dir = join_upath(context.ds_path, metadata_cache_dir)
        else:
            self.metadata_cache_dir = UPath(metadata_cache_dir)

        self.sort_by = sort_by

        self.client = boto3.client("s3")
        self.bucket = boto3.resource("s3").Bucket(self.bucket_name)
        self.metadata_cache_dir.mkdir(parents=True, exist_ok=True)

        self.wrs2_index: GridIndex | None = None

    def _read_products(
        self, needed_year_pathrows: set[tuple[int, str, str]]
    ) -> Generator[LandsatOliTirsItem, None, None]:
        """Read _stac.json files and yield relevant LandsatOliTirsItems.

        Args:
            needed_year_pathrows: set of (year, path, row) where we need to search for
                images.
        """
        year_pathrows = sorted(needed_year_pathrows)
        total = len(year_pathrows)
        for cur_idx, (year, path, row) in enumerate(year_pathrows, start=1):
            logger.debug(
                "Reading product infos (%s) year=%s path=%s row=%s",
                f"{cur_idx}/{total}",
                year,
                path,
                row,
            )
            assert len(path) == 3
            assert len(row) == 3
            local_fname = self.metadata_cache_dir / f"{year}_{path}_{row}.json"

            if not local_fname.exists():
                prefix = f"{self.bucket_prefix}/{year}/{path}/{row}/"
                items = []
                for obj in self.bucket.objects.filter(
                    Prefix=prefix, RequestPayer="requester"
                ):
                    # Only read the _stac.json files.
                    # Previously we used _MTL.json but those files don't have the full
                    # geometry of the Landsat scene, only the bounding box.
                    if not obj.key.endswith("_stac.json"):
                        continue
                    # Load JSON data.
                    buf = io.BytesIO()
                    self.bucket.download_fileobj(
                        obj.key, buf, ExtraArgs={"RequestPayer": "requester"}
                    )
                    buf.seek(0)
                    stac_data = json.load(buf)

                    # Get polygon coordinates.
                    shp = shapely.geometry.shape(stac_data["geometry"])

                    # Get datetime.
                    ts = dateutil.parser.isoparse(stac_data["properties"]["datetime"])

                    blob_path = obj.key.split("stac.json")[0]
                    time_range: tuple[datetime, datetime] = (ts, ts)
                    geometry = STGeometry(WGS84_PROJECTION, shp, time_range)
                    cloud_cover: float
                    if "eo:cloud_cover" in stac_data["properties"]:
                        cloud_cover = stac_data["properties"]["eo:cloud_cover"]
                    elif "landsat:cloud_cover_land" in stac_data["properties"]:
                        cloud_cover = stac_data["properties"][
                            "landsat:cloud_cover_land"
                        ]
                    else:
                        cloud_cover = -1
                    items.append(
                        LandsatOliTirsItem(
                            name=stac_data["id"],
                            geometry=geometry,
                            blob_path=blob_path,
                            cloud_cover=cloud_cover,
                        )
                    )

                with open_atomic(local_fname, "w") as f:
                    json.dump([item.serialize() for item in items], f)

            else:
                with local_fname.open() as f:
                    items = [
                        LandsatOliTirsItem.deserialize(item_dict)
                        for item_dict in json.load(f)
                    ]

            # OLI-only products (LO08_*/LO09_*) lack thermal bands B10/B11;
            # only yield OLI-TIRS scenes (LC08_*/LC09_*).
            yield from (
                item for item in items if not item.name.startswith(("LO08_", "LO09_"))
            )

    def _get_wrs2_index(self) -> GridIndex:
        """Get a grid index over the WRS2 polygons."""
        if self.wrs2_index is not None:
            return self.wrs2_index

        # Index doesn't exist so we need to build it.
        # We cache it with the object since it takes a bit of time to create it.
        self.wrs2_index = build_wrs2_grid_index(self.metadata_cache_dir)
        return self.wrs2_index

    def get_items(
        self, geometries: list[STGeometry], query_config: QueryConfig
    ) -> list[list[MatchedItemGroup[LandsatOliTirsItem]]]:
        """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.
        """
        wrs2_index = self._get_wrs2_index()
        needed_year_pathrows = set()
        wgs84_geometries = [geometry.to_wgs84() for geometry in geometries]
        for wgs84_geometry in wgs84_geometries:
            if wgs84_geometry.time_range is None:
                raise ValueError(
                    "Landsat on AWS requires geometry time ranges to be set"
                )
            cur_pathrows = get_pathrows_for_geometry(wrs2_index, wgs84_geometry)
            for path, row in cur_pathrows:
                for year in range(
                    wgs84_geometry.time_range[0].year,
                    wgs84_geometry.time_range[1].year + 1,
                ):
                    needed_year_pathrows.add((year, path, row))

        items = list(self._read_products(needed_year_pathrows))

        groups = []
        for geometry, wgs84_geometry in zip(geometries, wgs84_geometries):
            cur_items = []
            for item in items:
                if not wgs84_geometry.shp.intersects(item.geometry.shp):
                    continue
                cur_items.append(item)

            if self.sort_by == "cloud_cover":
                cur_items.sort(
                    key=lambda item: item.cloud_cover if item.cloud_cover >= 0 else 100
                )
            elif self.sort_by is not None:
                raise ValueError(f"invalid sort_by setting ({self.sort_by})")

            cur_groups: list[MatchedItemGroup[LandsatOliTirsItem]] = (
                rslearn.data_sources.utils.match_candidate_items_to_window(
                    geometry, cur_items, query_config
                )
            )
            groups.append(cur_groups)

        return groups

    def get_item_by_name(self, name: str) -> LandsatOliTirsItem:
        """Gets an item by name."""
        # Product name is like LC08_L1TP_046027_20230715_20230724_02_T1.
        # We want to use _read_products so we need to extract:
        # - year: 2023
        # - path: 046
        # - row: 027
        parts = name.split("_")
        assert len(parts[2]) == 6
        assert len(parts[3]) == 8
        year = int(parts[3][0:4])
        path = parts[2][0:3]
        row = parts[2][3:6]
        for item in self._read_products({(year, path, row)}):
            if item.name == name:
                return item
        raise ValueError(f"item {name} not found")

    # --- DirectMaterializeDataSource implementation ---

    def get_asset_url(self, item: LandsatOliTirsItem, asset_key: str) -> str:
        """Get the presigned URL to read the asset for the given item and asset key.

        Args:
            item: the item.
            asset_key: the key identifying which asset to get (the band name).

        Returns:
            the presigned URL to read the asset from.
        """
        # 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."""
        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:
            buf = io.BytesIO()
            self.bucket.download_fileobj(
                item.blob_path + f"{band}.TIF",
                buf,
                ExtraArgs={"RequestPayer": "requester"},
            )
            buf.seek(0)
            fname = item.blob_path.split("/")[-1] + f"_{band}.TIF"
            yield (fname, buf)

    def ingest(
        self,
        tile_store: TileStoreWithLayer,
        items: list[LandsatOliTirsItem],
        geometries: list[list[STGeometry]],
    ) -> None:
        """Ingest items into the given tile store.

        Args:
            tile_store: the tile store to ingest into
            items: the items to ingest
            geometries: a list of geometries needed for each item
        """
        for item, cur_geometries in zip(items, geometries):
            for band in self.BANDS:
                band_names = [band]
                if tile_store.is_raster_ready(item, band_names):
                    continue

                with tempfile.TemporaryDirectory() as tmp_dir:
                    fname = os.path.join(tmp_dir, f"{band}.tif")
                    self.bucket.download_file(
                        item.blob_path + f"{band}.TIF",
                        fname,
                        ExtraArgs={"RequestPayer": "requester"},
                    )
                    tile_store.write_raster_file(
                        item,
                        band_names,
                        UPath(fname),
                        time_range=item.geometry.time_range,
                    )

get_items

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

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[LandsatOliTirsItem]]]	List of groups of items that should be retrieved for each geometry.

Source code in rslearn/data_sources/aws_landsat.py

def get_items(
    self, geometries: list[STGeometry], query_config: QueryConfig
) -> list[list[MatchedItemGroup[LandsatOliTirsItem]]]:
    """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.
    """
    wrs2_index = self._get_wrs2_index()
    needed_year_pathrows = set()
    wgs84_geometries = [geometry.to_wgs84() for geometry in geometries]
    for wgs84_geometry in wgs84_geometries:
        if wgs84_geometry.time_range is None:
            raise ValueError(
                "Landsat on AWS requires geometry time ranges to be set"
            )
        cur_pathrows = get_pathrows_for_geometry(wrs2_index, wgs84_geometry)
        for path, row in cur_pathrows:
            for year in range(
                wgs84_geometry.time_range[0].year,
                wgs84_geometry.time_range[1].year + 1,
            ):
                needed_year_pathrows.add((year, path, row))

    items = list(self._read_products(needed_year_pathrows))

    groups = []
    for geometry, wgs84_geometry in zip(geometries, wgs84_geometries):
        cur_items = []
        for item in items:
            if not wgs84_geometry.shp.intersects(item.geometry.shp):
                continue
            cur_items.append(item)

        if self.sort_by == "cloud_cover":
            cur_items.sort(
                key=lambda item: item.cloud_cover if item.cloud_cover >= 0 else 100
            )
        elif self.sort_by is not None:
            raise ValueError(f"invalid sort_by setting ({self.sort_by})")

        cur_groups: list[MatchedItemGroup[LandsatOliTirsItem]] = (
            rslearn.data_sources.utils.match_candidate_items_to_window(
                geometry, cur_items, query_config
            )
        )
        groups.append(cur_groups)

    return groups

get_item_by_name

get_item_by_name(name: str) -> LandsatOliTirsItem

Gets an item by name.

Source code in rslearn/data_sources/aws_landsat.py

def get_item_by_name(self, name: str) -> LandsatOliTirsItem:
    """Gets an item by name."""
    # Product name is like LC08_L1TP_046027_20230715_20230724_02_T1.
    # We want to use _read_products so we need to extract:
    # - year: 2023
    # - path: 046
    # - row: 027
    parts = name.split("_")
    assert len(parts[2]) == 6
    assert len(parts[3]) == 8
    year = int(parts[3][0:4])
    path = parts[2][0:3]
    row = parts[2][3:6]
    for item in self._read_products({(year, path, row)}):
        if item.name == name:
            return item
    raise ValueError(f"item {name} not found")

get_asset_url

get_asset_url(item: LandsatOliTirsItem, asset_key: str) -> str

Get the presigned URL to read the asset for the given item and asset key.

Parameters:

Name	Type	Description	Default
item	LandsatOliTirsItem	the item.	required
asset_key	str	the key identifying which asset to get (the band name).	required

Returns:

Type	Description
str	the presigned URL to read the asset from.

Source code in rslearn/data_sources/aws_landsat.py

def get_asset_url(self, item: LandsatOliTirsItem, asset_key: str) -> str:
    """Get the presigned URL to read the asset for the given item and asset key.

    Args:
        item: the item.
        asset_key: the key identifying which asset to get (the band name).

    Returns:
        the presigned URL to read the asset from.
    """
    # 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",
        },
    )

deserialize_item

deserialize_item(serialized_item: dict) -> LandsatOliTirsItem

Deserializes an item from JSON-decoded data.

Source code in rslearn/data_sources/aws_landsat.py

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

retrieve_item

retrieve_item(item: LandsatOliTirsItem) -> Generator[tuple[str, BinaryIO], None, None]

Retrieves the rasters corresponding to an item as file streams.

Source code in rslearn/data_sources/aws_landsat.py

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:
        buf = io.BytesIO()
        self.bucket.download_fileobj(
            item.blob_path + f"{band}.TIF",
            buf,
            ExtraArgs={"RequestPayer": "requester"},
        )
        buf.seek(0)
        fname = item.blob_path.split("/")[-1] + f"_{band}.TIF"
        yield (fname, buf)

ingest

ingest(tile_store: TileStoreWithLayer, items: list[LandsatOliTirsItem], geometries: list[list[STGeometry]]) -> None

Ingest items into the given tile store.

Parameters:

Name	Type	Description	Default
tile_store	TileStoreWithLayer	the tile store to ingest into	required
items	list[LandsatOliTirsItem]	the items to ingest	required
geometries	list[list[STGeometry]]	a list of geometries needed for each item	required

Source code in rslearn/data_sources/aws_landsat.py

def ingest(
    self,
    tile_store: TileStoreWithLayer,
    items: list[LandsatOliTirsItem],
    geometries: list[list[STGeometry]],
) -> None:
    """Ingest items into the given tile store.

    Args:
        tile_store: the tile store to ingest into
        items: the items to ingest
        geometries: a list of geometries needed for each item
    """
    for item, cur_geometries in zip(items, geometries):
        for band in self.BANDS:
            band_names = [band]
            if tile_store.is_raster_ready(item, band_names):
                continue

            with tempfile.TemporaryDirectory() as tmp_dir:
                fname = os.path.join(tmp_dir, f"{band}.tif")
                self.bucket.download_file(
                    item.blob_path + f"{band}.TIF",
                    fname,
                    ExtraArgs={"RequestPayer": "requester"},
                )
                tile_store.write_raster_file(
                    item,
                    band_names,
                    UPath(fname),
                    time_range=item.geometry.time_range,
                )