rslearn.data_sources.aws_google_satellite_embedding_v1

aws_google_satellite_embedding_v1

Data source for Google Satellite Embedding V1 dataset on AWS Open Data.

GoogleSatelliteEmbeddingV1Item

Bases: Item

An item in the GoogleSatelliteEmbeddingV1 data source.

Source code in rslearn/data_sources/aws_google_satellite_embedding_v1.py

class GoogleSatelliteEmbeddingV1Item(Item):
    """An item in the GoogleSatelliteEmbeddingV1 data source."""

    def __init__(
        self,
        name: str,
        geometry: STGeometry,
        s3_path: str,
    ) -> None:
        """Creates a new GoogleSatelliteEmbeddingV1Item.

        Args:
            name: unique name of the item (the filename without extension)
            geometry: the spatial and temporal extent of the item
            s3_path: full S3 path to the TIFF file
        """
        super().__init__(name, geometry)
        self.s3_path = s3_path

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

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

serialize

serialize() -> dict

Serializes the item to a JSON-encodable dictionary.

Source code in rslearn/data_sources/aws_google_satellite_embedding_v1.py

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

deserialize staticmethod

deserialize(d: dict) -> GoogleSatelliteEmbeddingV1Item

Deserializes an item from a JSON-decoded dictionary.

Source code in rslearn/data_sources/aws_google_satellite_embedding_v1.py

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

GoogleSatelliteEmbeddingV1

Bases: DirectMaterializeDataSource[GoogleSatelliteEmbeddingV1Item], ItemLookupDataSource[GoogleSatelliteEmbeddingV1Item]

Data source for Google Satellite Embedding V1 on AWS Open Data.

It consists of annual satellite embeddings at 10m resolution with 64 bands (A00-A63). The data is stored as Cloud-Optimized GeoTIFFs organized by year and UTM zone. Each file covers 8192x8192 pixels.

Available years: 2018-2024.

See https://registry.opendata.aws/aef-source/ for details.

Source code in rslearn/data_sources/aws_google_satellite_embedding_v1.py

class GoogleSatelliteEmbeddingV1(
    DirectMaterializeDataSource[GoogleSatelliteEmbeddingV1Item],
    ItemLookupDataSource[GoogleSatelliteEmbeddingV1Item],
):
    """Data source for Google Satellite Embedding V1 on AWS Open Data.

    It consists of annual satellite embeddings at 10m resolution with 64 bands
    (A00-A63). The data is stored as Cloud-Optimized GeoTIFFs organized by year and UTM
    zone. Each file covers 8192x8192 pixels.

    Available years: 2018-2024.

    See https://registry.opendata.aws/aef-source/ for details.
    """

    def __init__(
        self,
        metadata_cache_dir: str,
        apply_dequantization: bool = True,
        context: DataSourceContext = DataSourceContext(),
    ) -> None:
        """Initialize a new GoogleSatelliteEmbeddingV1 instance.

        Args:
            metadata_cache_dir: directory to cache the index file.
            apply_dequantization: whether to apply de-quantization to convert
                int8 values to float32. The raw data is quantized int8; the
                de-quantization maps values to [-1, 1] using the formula:
                ((values / 127.5) ** 2) * sign(values). The raw data has nodata value
                -128 while with dequantization the nodata value is -1.0. See
                https://source.coop/tge-labs/aef for details.
            context: the data source context.
        """
        # We have a single asset containing all 64 bands. Here "image" is an arbitrary
        # name, since DirectMaterializeDataSource requires an asset name.
        super().__init__(asset_bands={"image": BANDS})

        self.apply_dequantization = apply_dequantization

        # Set up cache directory
        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.metadata_cache_dir.mkdir(parents=True, exist_ok=True)

        # S3 client with anonymous access (only used for downloading index)
        self.s3_client = boto3.client(
            "s3",
            config=Config(signature_version=UNSIGNED),
            region_name="us-west-2",
        )

        # Lazy-loaded grid index
        self._grid_index: GridIndex | None = None
        self._items_by_name: dict[str, GoogleSatelliteEmbeddingV1Item] | None = None

    def _read_index_csv(self) -> pd.DataFrame:
        """Read the index CSV, downloading from S3 if not cached.

        Returns:
            DataFrame with WKT, path, and year columns.
        """
        cache_file = self.metadata_cache_dir / "aef_index.csv"
        if not cache_file.exists():
            response = self.s3_client.get_object(Bucket=BUCKET_NAME, Key=INDEX_KEY)
            content = response["Body"].read()
            with cache_file.open("wb") as f:
                f.write(content)

        return pd.read_csv(cache_file)

    def _load_index(
        self,
    ) -> tuple[GridIndex, dict[str, GoogleSatelliteEmbeddingV1Item]]:
        """Load the index file and build spatial index.

        Returns:
            Tuple of (grid_index, items_by_name dict).
        """
        if self._grid_index is not None and self._items_by_name is not None:
            return self._grid_index, self._items_by_name

        df = self._read_index_csv()

        grid_index = GridIndex(GRID_SIZE)
        items_by_name: dict[str, GoogleSatelliteEmbeddingV1Item] = {}

        for _, row in df.iterrows():
            shp = shapely.wkt.loads(row["WKT"])

            year = int(row["year"])
            time_range = (
                datetime(year, 1, 1, tzinfo=UTC),
                datetime(year, 12, 31, 23, 59, 59, tzinfo=UTC),
            )

            s3_path = row["path"]
            name = s3_path.split("/")[-1].replace(".tiff", "")

            geometry = STGeometry(WGS84_PROJECTION, shp, time_range)
            item = GoogleSatelliteEmbeddingV1Item(
                name=name,
                geometry=geometry,
                s3_path=s3_path,
            )

            grid_index.insert(shp.bounds, item)
            items_by_name[name] = item

        self._grid_index = grid_index
        self._items_by_name = items_by_name
        return grid_index, items_by_name

    # --- DataSource implementation ---

    def get_items(
        self, geometries: list[STGeometry], query_config: QueryConfig
    ) -> list[list[MatchedItemGroup[GoogleSatelliteEmbeddingV1Item]]]:
        """Get a list of items in the data source intersecting the given geometries."""
        grid_index, _ = self._load_index()

        wgs84_geometries = [
            geometry.to_projection(WGS84_PROJECTION) for geometry in geometries
        ]

        groups = []
        for geometry, wgs84_geometry in zip(geometries, wgs84_geometries):
            cur_items = []
            for item in grid_index.query(wgs84_geometry.shp.bounds):
                if not wgs84_geometry.shp.intersects(item.geometry.shp):
                    continue
                # Check time range if specified
                if wgs84_geometry.time_range is not None:
                    item_start, item_end = item.geometry.time_range
                    query_start, query_end = wgs84_geometry.time_range
                    if item_end < query_start or item_start > query_end:
                        continue
                cur_items.append(item)

            cur_items.sort(key=lambda item: item.geometry.time_range[0])

            cur_groups: list[MatchedItemGroup[GoogleSatelliteEmbeddingV1Item]] = (
                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) -> GoogleSatelliteEmbeddingV1Item:
        """Gets an item by name."""
        _, items_by_name = self._load_index()
        if name not in items_by_name:
            raise ValueError(f"item {name} not found")
        return items_by_name[name]

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

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

        Note: Each file is 2-3GB so this can be slow. Direct materialization via
        read_raster or materialize is recommended for most use cases.

        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 in items:
            if tile_store.is_raster_ready(item, BANDS):
                continue

            # Download the TIFF file directly to disk
            key = item.s3_path.replace(f"s3://{BUCKET_NAME}/", "")
            with tempfile.TemporaryDirectory() as tmp_dir:
                local_path = os.path.join(tmp_dir, f"{item.name}.tiff")
                self.s3_client.download_file(BUCKET_NAME, key, local_path)

                if self.apply_dequantization:
                    with rasterio.open(local_path) as src:
                        array = src.read()
                        projection, bounds = get_raster_projection_and_bounds(src)
                    array = self._dequantize(array)
                    raster_metadata = RasterMetadata(nodata_value=-1.0)
                    tile_store.write_raster(
                        item,
                        BANDS,
                        projection,
                        bounds,
                        RasterArray(
                            chw_array=array,
                            time_range=item.geometry.time_range,
                            metadata=raster_metadata,
                        ),
                    )
                else:
                    tile_store.write_raster_file(
                        item,
                        BANDS,
                        UPath(local_path),
                        time_range=item.geometry.time_range,
                    )

    # --- DirectMaterializeDataSource implementation ---

    def get_raster_metadata(
        self, layer_name: str, item: Item, bands: list[str]
    ) -> RasterMetadata:
        """Return metadata reflecting the effective nodata value."""
        if self.apply_dequantization:
            return RasterMetadata(nodata_value=-1.0)
        return RasterMetadata(nodata_value=-128)

    def get_asset_url(
        self, item: GoogleSatelliteEmbeddingV1Item, asset_key: str
    ) -> str:
        """Get the HTTP URL to read the asset."""
        # Convert s3://bucket/path to HTTP URL
        key = item.s3_path.replace(f"s3://{BUCKET_NAME}/", "")
        return f"{HTTP_URL_BASE}/{key}"

    def _dequantize(self, data: npt.NDArray[Any]) -> npt.NDArray[np.float32]:
        """Apply de-quantization to convert int8 values to float32.

        The raw data is quantized int8; this maps values to [-1, 1] using the
        formula: ((values / 127.5) ** 2) * sign(values). NODATA (-128) is mapped
        to -1.0. See https://source.coop/tge-labs/aef for details.

        Args:
            data: the raw int8 raster array.

        Returns:
            the de-quantized float32 array.
        """
        nodata_mask = data == -128
        float_data = data.astype(np.float32)
        result = ((float_data / 127.5) ** 2) * np.sign(float_data)
        result[nodata_mask] = -1.0
        return result

    def get_read_callback(
        self, item: GoogleSatelliteEmbeddingV1Item, asset_key: str
    ) -> Callable[[npt.NDArray[Any]], npt.NDArray[Any]] | None:
        """Return a callback to apply de-quantization if enabled."""
        if not self.apply_dequantization:
            return None
        return self._dequantize

    def read_raster(
        self,
        layer_name: str,
        item: Item,
        bands: list[str],
        projection: Projection,
        bounds: PixelBounds,
        resampling: Resampling = Resampling.bilinear,
    ) -> RasterArray:
        """Read raster data from the store.

        Overrides base class to handle band selection (the base class reads all bands).
        """
        if not isinstance(item, GoogleSatelliteEmbeddingV1Item):
            raise TypeError(
                f"expected GoogleSatelliteEmbeddingV1Item, got {type(item)}"
            )
        asset_url = self.get_asset_url(item, "image")

        # Determine which band indices to read (1-indexed for rasterio)
        if bands == BANDS:
            band_indices = list(range(1, 65))
        else:
            band_indices = [BANDS.index(b) + 1 for b in bands]

        # Construct the transform for the requested bounds
        wanted_transform = rasterio.transform.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:
                data = vrt.read(indexes=band_indices)

        callback = self.get_read_callback(item, "image")
        if callback is not None:
            data = callback(data)

        if self.apply_dequantization:
            raster_metadata = RasterMetadata(nodata_value=-1.0)
        else:
            raster_metadata = RasterMetadata(nodata_value=-128)

        return RasterArray(
            chw_array=data,
            time_range=item.geometry.time_range,
            metadata=raster_metadata,
        )

get_items

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

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

Source code in rslearn/data_sources/aws_google_satellite_embedding_v1.py

def get_items(
    self, geometries: list[STGeometry], query_config: QueryConfig
) -> list[list[MatchedItemGroup[GoogleSatelliteEmbeddingV1Item]]]:
    """Get a list of items in the data source intersecting the given geometries."""
    grid_index, _ = self._load_index()

    wgs84_geometries = [
        geometry.to_projection(WGS84_PROJECTION) for geometry in geometries
    ]

    groups = []
    for geometry, wgs84_geometry in zip(geometries, wgs84_geometries):
        cur_items = []
        for item in grid_index.query(wgs84_geometry.shp.bounds):
            if not wgs84_geometry.shp.intersects(item.geometry.shp):
                continue
            # Check time range if specified
            if wgs84_geometry.time_range is not None:
                item_start, item_end = item.geometry.time_range
                query_start, query_end = wgs84_geometry.time_range
                if item_end < query_start or item_start > query_end:
                    continue
            cur_items.append(item)

        cur_items.sort(key=lambda item: item.geometry.time_range[0])

        cur_groups: list[MatchedItemGroup[GoogleSatelliteEmbeddingV1Item]] = (
            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) -> GoogleSatelliteEmbeddingV1Item

Gets an item by name.

Source code in rslearn/data_sources/aws_google_satellite_embedding_v1.py

def get_item_by_name(self, name: str) -> GoogleSatelliteEmbeddingV1Item:
    """Gets an item by name."""
    _, items_by_name = self._load_index()
    if name not in items_by_name:
        raise ValueError(f"item {name} not found")
    return items_by_name[name]

deserialize_item

deserialize_item(serialized_item: dict) -> GoogleSatelliteEmbeddingV1Item

Deserializes an item from JSON-decoded data.

Source code in rslearn/data_sources/aws_google_satellite_embedding_v1.py

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

ingest

ingest(tile_store: Any, items: list[GoogleSatelliteEmbeddingV1Item], geometries: list[list[STGeometry]]) -> None

Ingest items into the given tile store.

Note: Each file is 2-3GB so this can be slow. Direct materialization via read_raster or materialize is recommended for most use cases.

Parameters:

Name	Type	Description	Default
tile_store	Any	the tile store to ingest into	required
items	list[GoogleSatelliteEmbeddingV1Item]	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_google_satellite_embedding_v1.py

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

    Note: Each file is 2-3GB so this can be slow. Direct materialization via
    read_raster or materialize is recommended for most use cases.

    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 in items:
        if tile_store.is_raster_ready(item, BANDS):
            continue

        # Download the TIFF file directly to disk
        key = item.s3_path.replace(f"s3://{BUCKET_NAME}/", "")
        with tempfile.TemporaryDirectory() as tmp_dir:
            local_path = os.path.join(tmp_dir, f"{item.name}.tiff")
            self.s3_client.download_file(BUCKET_NAME, key, local_path)

            if self.apply_dequantization:
                with rasterio.open(local_path) as src:
                    array = src.read()
                    projection, bounds = get_raster_projection_and_bounds(src)
                array = self._dequantize(array)
                raster_metadata = RasterMetadata(nodata_value=-1.0)
                tile_store.write_raster(
                    item,
                    BANDS,
                    projection,
                    bounds,
                    RasterArray(
                        chw_array=array,
                        time_range=item.geometry.time_range,
                        metadata=raster_metadata,
                    ),
                )
            else:
                tile_store.write_raster_file(
                    item,
                    BANDS,
                    UPath(local_path),
                    time_range=item.geometry.time_range,
                )

get_raster_metadata

get_raster_metadata(layer_name: str, item: Item, bands: list[str]) -> RasterMetadata

Return metadata reflecting the effective nodata value.

Source code in rslearn/data_sources/aws_google_satellite_embedding_v1.py

def get_raster_metadata(
    self, layer_name: str, item: Item, bands: list[str]
) -> RasterMetadata:
    """Return metadata reflecting the effective nodata value."""
    if self.apply_dequantization:
        return RasterMetadata(nodata_value=-1.0)
    return RasterMetadata(nodata_value=-128)

get_asset_url

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

Get the HTTP URL to read the asset.

Source code in rslearn/data_sources/aws_google_satellite_embedding_v1.py

def get_asset_url(
    self, item: GoogleSatelliteEmbeddingV1Item, asset_key: str
) -> str:
    """Get the HTTP URL to read the asset."""
    # Convert s3://bucket/path to HTTP URL
    key = item.s3_path.replace(f"s3://{BUCKET_NAME}/", "")
    return f"{HTTP_URL_BASE}/{key}"

get_read_callback

get_read_callback(item: GoogleSatelliteEmbeddingV1Item, asset_key: str) -> Callable[[NDArray[Any]], NDArray[Any]] | None

Return a callback to apply de-quantization if enabled.

Source code in rslearn/data_sources/aws_google_satellite_embedding_v1.py

def get_read_callback(
    self, item: GoogleSatelliteEmbeddingV1Item, asset_key: str
) -> Callable[[npt.NDArray[Any]], npt.NDArray[Any]] | None:
    """Return a callback to apply de-quantization if enabled."""
    if not self.apply_dequantization:
        return None
    return self._dequantize

read_raster

read_raster(layer_name: str, item: Item, bands: list[str], projection: Projection, bounds: PixelBounds, resampling: Resampling = bilinear) -> RasterArray

Read raster data from the store.

Overrides base class to handle band selection (the base class reads all bands).

Source code in rslearn/data_sources/aws_google_satellite_embedding_v1.py

def read_raster(
    self,
    layer_name: str,
    item: Item,
    bands: list[str],
    projection: Projection,
    bounds: PixelBounds,
    resampling: Resampling = Resampling.bilinear,
) -> RasterArray:
    """Read raster data from the store.

    Overrides base class to handle band selection (the base class reads all bands).
    """
    if not isinstance(item, GoogleSatelliteEmbeddingV1Item):
        raise TypeError(
            f"expected GoogleSatelliteEmbeddingV1Item, got {type(item)}"
        )
    asset_url = self.get_asset_url(item, "image")

    # Determine which band indices to read (1-indexed for rasterio)
    if bands == BANDS:
        band_indices = list(range(1, 65))
    else:
        band_indices = [BANDS.index(b) + 1 for b in bands]

    # Construct the transform for the requested bounds
    wanted_transform = rasterio.transform.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:
            data = vrt.read(indexes=band_indices)

    callback = self.get_read_callback(item, "image")
    if callback is not None:
        data = callback(data)

    if self.apply_dequantization:
        raster_metadata = RasterMetadata(nodata_value=-1.0)
    else:
        raster_metadata = RasterMetadata(nodata_value=-128)

    return RasterArray(
        chw_array=data,
        time_range=item.geometry.time_range,
        metadata=raster_metadata,
    )