rslearn.data_sources.google_earth_engine

google_earth_engine

Data source for raster or vector data in local files.

NoValidPixelsException

Bases: Exception

Exception when GEE API reports that export failed due to no valid pixels.

Source code in rslearn/data_sources/google_earth_engine.py

class NoValidPixelsException(Exception):
    """Exception when GEE API reports that export failed due to no valid pixels."""

    # Expected GEE error_message when the task fails.
    GEE_MESSAGE = "No valid (un-masked) pixels in export region."

ExportException

Bases: Exception

GEE API export error.

Source code in rslearn/data_sources/google_earth_engine.py

class ExportException(Exception):
    """GEE API export error."""

    pass

MergedRaster dataclass

Result of merging one or more GEE-exported raster files.

Source code in rslearn/data_sources/google_earth_engine.py

@dataclass
class MergedRaster:
    """Result of merging one or more GEE-exported raster files."""

    array: npt.NDArray[np.generic]
    projection: Projection
    bounds: PixelBounds
    nodata_value: int | float | None

GEE

Bases: DataSource, TileStore

A data source for ingesting images from Google Earth Engine.

Source code in rslearn/data_sources/google_earth_engine.py

class GEE(DataSource, TileStore):
    """A data source for ingesting images from Google Earth Engine."""

    def __init__(
        self,
        collection_name: str,
        gcs_bucket_name: str,
        index_cache_dir: str,
        service_account_name: str,
        service_account_credentials: str,
        bands: list[str] | None = None,
        filters: list[tuple[str, Any]] | None = None,
        dtype: DType | None = None,
        context: DataSourceContext = DataSourceContext(),
    ) -> None:
        """Initialize a new GEE instance.

        Args:
            collection_name: the Earth Engine ImageCollection to ingest images from
            gcs_bucket_name: the Cloud Storage bucket to export GEE images to
            index_cache_dir: cache directory to store rtree index
            service_account_name: name of the service account to use for authentication
            service_account_credentials: service account credentials filename
            bands: the list of bands to ingest, in case the layer config is not present
                in the context.
            filters: optional list of tuples (property_name, property_value) to filter
                images (using ee.Filter.eq)
            dtype: optional desired array data type. If the data obtained from GEE does
                not match this type, then it is converted.
            context: the data source context.
        """
        self.collection_name = collection_name
        self.gcs_bucket_name = gcs_bucket_name
        self.filters = filters
        self.dtype = dtype

        # Get index cache dir depending on dataset path.
        if context.ds_path is not None:
            self.index_cache_dir = join_upath(context.ds_path, index_cache_dir)
        else:
            self.index_cache_dir = UPath(index_cache_dir)

        # Get bands we need to export.
        if context.layer_config is not None:
            self.bands = [
                band
                for band_set in context.layer_config.band_sets
                for band in band_set.bands
            ]
        elif bands is not None:
            self.bands = bands
        else:
            raise ValueError(
                "bands must be specified if layer_config is not present in the context"
            )

        self.bucket = storage.Client().bucket(self.gcs_bucket_name)

        credentials = ee.ServiceAccountCredentials(
            service_account_name, service_account_credentials
        )
        ee.Initialize(credentials)

        self._cached_nodata_value: int | float | None = None

        self.index_cache_dir.mkdir(parents=True, exist_ok=True)
        self.rtree_index = get_cached_rtree(self.index_cache_dir, self._build_index)

    def get_collection(self) -> ee.ImageCollection:
        """Returns the Earth Engine image collection for this data source."""
        image_collection = ee.ImageCollection(self.collection_name)
        if self.filters is None:
            return image_collection

        for k, v in self.filters:
            cur_filter = ee.Filter.eq(k, v)
            image_collection = image_collection.filter(cur_filter)
        return image_collection

    def _build_index(self, rtree_index: RtreeIndex) -> None:
        csv_blob = self.bucket.blob(f"{self.collection_name}/index.csv")

        if not csv_blob.exists():
            # Export feature collection of image metadata to GCS.
            def image_to_feature(image: ee.Image) -> ee.Feature:
                geometry = image.geometry().transform(proj="EPSG:4326", maxError=0.001)
                return ee.Feature(geometry, {"time": image.date().format()})

            fc = self.get_collection().map(image_to_feature)
            task = ee.batch.Export.table.toCloudStorage(
                collection=fc,
                description="rslearn GEE index export task",
                bucket=self.gcs_bucket_name,
                fileNamePrefix=f"{self.collection_name}/index",
                fileFormat="CSV",
            )
            task.start()
            logger.info(
                "Started task to export GEE index for image collection %s",
                self.collection_name,
            )
            while True:
                time.sleep(10)
                status_dict = task.status()
                logger.debug(
                    "Waiting for export task to complete, current status is %s",
                    status_dict,
                )
                if status_dict["state"] in ["UNSUBMITTED", "READY", "RUNNING"]:
                    continue
                elif status_dict["state"] != "COMPLETED":
                    raise ValueError(
                        f"got unexpected GEE task state {status_dict['state']}"
                    )
                break

        # Read the CSV and add rows into the rtree index.
        with csv_blob.open() as f:
            reader = csv.DictReader(f)
            for row in tqdm.tqdm(reader, desc="Building index"):
                shp = shapely.geometry.shape(json.loads(row[".geo"]))
                if "E" in row["time"]:
                    unix_time = float(row["time"]) / 1000
                    ts = datetime.fromtimestamp(unix_time, tz=UTC)
                else:
                    ts = datetime.fromisoformat(row["time"]).replace(tzinfo=UTC)
                geometry = STGeometry(WGS84_PROJECTION, shp, (ts, ts))
                item = Item(row["system:index"], geometry)
                rtree_index.insert(shp.bounds, json.dumps(item.serialize()))

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

        Args:
            name: the name of the item to get

        Returns:
            the item object
        """
        filtered = self.get_collection().filter(ee.Filter.eq("system:index", name))
        image = filtered.first()
        shp = shapely.geometry.shape(
            image.geometry().transform(proj="EPSG:4326", maxError=0.001).getInfo()
        )
        ts = datetime.fromisoformat(image.date().format().getInfo()).replace(tzinfo=UTC)
        geometry = STGeometry(WGS84_PROJECTION, shp, (ts, ts))
        return Item(name, geometry)

    def get_items(
        self, geometries: list[STGeometry], query_config: QueryConfig
    ) -> list[list[MatchedItemGroup[Item]]]:
        """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.
        """
        wgs84_geometries = [
            geometry.to_projection(WGS84_PROJECTION) for geometry in geometries
        ]

        groups = []
        for geometry in wgs84_geometries:
            cur_items = []
            encoded_items = self.rtree_index.query(geometry.shp.bounds)
            for encoded_item in encoded_items:
                item = Item.deserialize(json.loads(encoded_item))
                if not item.geometry.shp.intersects(geometry.shp):
                    continue
                cur_items.append(item)

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

            cur_groups = rslearn.data_sources.utils.match_candidate_items_to_window(
                geometry, cur_items, query_config
            )
            groups.append(cur_groups)

        return groups

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

    def item_to_image(self, item: Item) -> ee.image.Image:
        """Get the Image corresponding to the Item.

        This function is separated so it can be overriden if subclasses want to add
        modifications to the image.
        """
        filtered = self.get_collection().filter(ee.Filter.eq("system:index", item.name))
        image = filtered.first()
        image = image.select(self.bands)
        return image

    def export_item(
        self,
        item: Item,
        blob_prefix: str,
        projection_and_bounds: tuple[Projection, PixelBounds] | None = None,
    ) -> None:
        """Export the item to the specified folder.

        Args:
            item: the item to export.
            blob_prefix: the prefix (folder) to use.
            projection_and_bounds: optionally use this projection and bounds instead of
                the extent of the image.
        """
        image = self.item_to_image(item)
        projection = image.select(self.bands[0]).projection().getInfo()
        logger.info("Starting task to retrieve image %s", item.name)

        extent_kwargs: dict[str, Any]
        if projection_and_bounds is not None:
            projection, bounds = projection_and_bounds
            transform = get_transform_from_projection_and_bounds(projection, bounds)
            width = bounds[2] - bounds[0]
            height = bounds[3] - bounds[1]
            extent_kwargs = dict(
                crs=str(projection.crs),
                crsTransform=[
                    transform.a,
                    transform.b,
                    transform.c,
                    transform.d,
                    transform.e,
                    transform.f,
                ],
                dimensions=f"{width}x{height}",
            )
        else:
            # Use the native projection of the image.
            # We pass scale instead of crsTransform since some images have positive y
            # resolution which means they are upside down and rasterio cannot merge
            # them.
            extent_kwargs = dict(
                crs=projection["crs"],
                scale=projection["transform"][0],
            )

        task = ee.batch.Export.image.toCloudStorage(
            image=image,
            description=item.name,
            bucket=self.gcs_bucket_name,
            fileNamePrefix=blob_prefix,
            maxPixels=10000000000,
            fileFormat="GeoTIFF",
            skipEmptyTiles=True,
            **extent_kwargs,
        )
        task.start()
        while True:
            time.sleep(10)
            status_dict = task.status()
            if status_dict["state"] in ["UNSUBMITTED", "READY", "RUNNING"]:
                continue
            if status_dict["state"] == "COMPLETED":
                break
            if status_dict["state"] != "FAILED":
                raise ValueError(
                    f"got unexpected GEE task state {status_dict['state']}"
                )
            # The task failed. We see if it is an okay failure case or if we need to
            # raise exception.
            if status_dict["error_message"] == NoValidPixelsException.GEE_MESSAGE:
                raise NoValidPixelsException()
            raise ExportException(f"GEE task failed: {status_dict['error_message']}")

    def _merge_rasters(
        self,
        blobs: list[storage.Blob],
        crs_bounds: tuple[float, float, float, float] | None = None,
        res: float | None = None,
    ) -> MergedRaster:
        """Merge multiple rasters split up during export by GEE.

        GEE can produce multiple rasters if it determines the file size exceeds its
        internal limit. So in this case we stitch them back together.

        Args:
            blobs: the list of GCS blobs where the rasters were written.
            crs_bounds: generate merged output under this bounds, in CRS coordinates
                (not pixel units).
            res: generate merged output under this resolution.

        Returns:
            a MergedRaster with the stitched array, projection, bounds, and nodata.
        """
        with tempfile.TemporaryDirectory() as tmp_dir_name:
            rasterio_datasets = []
            for blob in blobs:
                local_fname = os.path.join(tmp_dir_name, blob.name.split("/")[-1])
                blob.download_to_filename(local_fname)
                src = rasterio.open(local_fname)
                rasterio_datasets.append(src)

            nodata_vals: list[int | float] = []
            for ds in rasterio_datasets:
                if ds.nodata is not None:
                    nodata_vals.append(ds.nodata)
            nodata_value = unique_nodata_value(nodata_vals)

            merge_kwargs: dict[str, Any] = dict(
                sources=rasterio_datasets,
                bounds=crs_bounds,
                res=res,
            )
            if self.dtype:
                merge_kwargs["dtype"] = self.dtype.value
            array, transform = rasterio.merge.merge(**merge_kwargs)
            projection, bounds = get_raster_projection_and_bounds_from_transform(
                rasterio_datasets[0].crs,
                transform,
                array.shape[2],
                array.shape[1],
            )

            for ds in rasterio_datasets:
                ds.close()

        return MergedRaster(
            array=array,
            projection=projection,
            bounds=bounds,
            nodata_value=nodata_value,
        )

    def ingest(
        self,
        tile_store: TileStoreWithLayer,
        items: list[Item],
        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 in items:
            if tile_store.is_raster_ready(item, self.bands):
                continue

            # Export the item to GCS.
            blob_prefix = f"{self.collection_name}/{item.name}.{os.getpid()}/"
            self.export_item(item, blob_prefix)

            # See what files the export produced.
            # If there are multiple, then we merge them into one file since that's the
            # simplest way to handle it.
            blobs = list(self.bucket.list_blobs(prefix=blob_prefix))

            with tempfile.TemporaryDirectory() as tmp_dir_name:
                if len(blobs) == 1:
                    local_fname = os.path.join(
                        tmp_dir_name, blobs[0].name.split("/")[-1]
                    )
                    blobs[0].download_to_filename(local_fname)
                    tile_store.write_raster_file(
                        item,
                        self.bands,
                        UPath(local_fname),
                        time_range=item.geometry.time_range,
                    )

                else:
                    merged = self._merge_rasters(blobs)
                    tile_store.write_raster(
                        item,
                        self.bands,
                        merged.projection,
                        merged.bounds,
                        RasterArray(
                            chw_array=merged.array,
                            time_range=item.geometry.time_range,
                            metadata=RasterMetadata(nodata_value=merged.nodata_value),
                        ),
                    )

            for blob in blobs:
                blob.delete()

    def is_raster_ready(self, layer_name: str, item: Item, bands: list[str]) -> bool:
        """Checks if this raster has been written to the store.

        Args:
            layer_name: the layer name or alias.
            item: 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 Planetary Computer.
        return True

    def get_raster_bands(self, layer_name: str, item: Item) -> 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: 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 [self.bands]

    def get_raster_bounds(
        self, layer_name: str, item: Item, 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: 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.
        """
        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 get_raster_metadata(
        self, layer_name: str, item: Item, bands: list[str]
    ) -> RasterMetadata:
        """Export and read the raster header to obtain nodata.

        The nodata value is cached after the first successful export that
        yields a non-None value, so subsequent calls skip the export entirely.
        """
        if self._cached_nodata_value is not None:
            return RasterMetadata(nodata_value=self._cached_nodata_value)

        blob_prefix = f"{self.collection_name}/{item.name}.{os.getpid()}/"
        try:
            self.export_item(item, blob_prefix)
        except NoValidPixelsException:
            return RasterMetadata()

        blobs = list(self.bucket.list_blobs(prefix=blob_prefix))
        buf = io.BytesIO()
        blobs[0].download_to_file(buf)
        buf.seek(0)
        with rasterio.open(buf) as src:
            nodata_value = src.nodata
        if nodata_value is not None:
            self._cached_nodata_value = nodata_value
        return RasterMetadata(nodata_value=nodata_value)

    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.

        Args:
            layer_name: the layer name or alias.
            item: 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
        """
        # Extract the requested extent and export to GCS.
        bounds_str = f"{bounds[0]}_{bounds[1]}_{bounds[2]}_{bounds[3]}"
        blob_prefix = f"{self.collection_name}/{item.name}.{bounds_str}.{os.getpid()}/"

        try:
            self.export_item(
                item, blob_prefix, projection_and_bounds=(projection, bounds)
            )
        except NoValidPixelsException:
            # No valid pixels means the result should be empty.
            logger.info(
                f"No valid pixels in item {item.name} with projection={projection}, bounds={bounds}, returning empty image"
            )
            return RasterArray(
                chw_array=np.zeros(
                    (len(bands), bounds[3] - bounds[1], bounds[2] - bounds[0]),
                    dtype=np.float32,
                ),
                time_range=item.geometry.time_range,
            )

        wanted_transform = get_transform_from_projection_and_bounds(projection, bounds)
        crs_bounds = (
            bounds[0] * projection.x_resolution,
            bounds[3] * projection.y_resolution,
            bounds[2] * projection.x_resolution,
            bounds[1] * projection.y_resolution,
        )

        blobs = list(self.bucket.list_blobs(prefix=blob_prefix))

        if len(blobs) == 1:
            # With a single output, we can simply read it with vrt.
            buf = io.BytesIO()
            blobs[0].download_to_file(buf)
            buf.seek(0)
            with rasterio.open(buf) as src:
                nodata_value = src.nodata
                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 RasterArray(
                        chw_array=vrt.read(),
                        time_range=item.geometry.time_range,
                        metadata=RasterMetadata(nodata_value=nodata_value),
                    )

        else:
            # With multiple outputs, we need to merge them together.
            # We can set the bounds in CRS coordinates when we do the merging.
            if projection.x_resolution != -projection.y_resolution:
                raise NotImplementedError(
                    "Only projection with x_res=-y_res is supported for GEE direct materialization"
                )
            merged = self._merge_rasters(
                blobs, crs_bounds=crs_bounds, res=projection.x_resolution
            )
            metadata = RasterMetadata(nodata_value=merged.nodata_value)

            # We copy the array if its bounds don't match exactly.
            if merged.bounds == bounds:
                return RasterArray(
                    chw_array=merged.array,
                    time_range=item.geometry.time_range,
                    metadata=metadata,
                )
            dst_array = np.zeros(
                (
                    merged.array.shape[0],
                    bounds[3] - bounds[1],
                    bounds[2] - bounds[0],
                ),
                dtype=merged.array.dtype,
            )
            copy_spatial_array(merged.array, dst_array, merged.bounds[0:2], bounds[0:2])
            return RasterArray(
                chw_array=dst_array,
                time_range=item.geometry.time_range,
                metadata=metadata,
            )

    def materialize(
        self,
        window: Window,
        item_groups: list[list[Item]],
        layer_name: str,
        layer_cfg: LayerConfig,
        group_time_ranges: list[tuple[datetime, datetime] | None] | None = None,
    ) -> 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
            group_time_ranges: optional request time range for each item group
        """
        RasterMaterializer().materialize(
            TileStoreWithLayer(self, layer_name),
            window,
            layer_name,
            layer_cfg,
            item_groups,
            group_time_ranges=group_time_ranges,
        )

get_collection

get_collection() -> ImageCollection

Returns the Earth Engine image collection for this data source.

Source code in rslearn/data_sources/google_earth_engine.py

def get_collection(self) -> ee.ImageCollection:
    """Returns the Earth Engine image collection for this data source."""
    image_collection = ee.ImageCollection(self.collection_name)
    if self.filters is None:
        return image_collection

    for k, v in self.filters:
        cur_filter = ee.Filter.eq(k, v)
        image_collection = image_collection.filter(cur_filter)
    return image_collection

get_item_by_name

get_item_by_name(name: str) -> Item

Gets an item by name.

Parameters:

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

Returns:

Type	Description
Item	the item object

Source code in rslearn/data_sources/google_earth_engine.py

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

    Args:
        name: the name of the item to get

    Returns:
        the item object
    """
    filtered = self.get_collection().filter(ee.Filter.eq("system:index", name))
    image = filtered.first()
    shp = shapely.geometry.shape(
        image.geometry().transform(proj="EPSG:4326", maxError=0.001).getInfo()
    )
    ts = datetime.fromisoformat(image.date().format().getInfo()).replace(tzinfo=UTC)
    geometry = STGeometry(WGS84_PROJECTION, shp, (ts, ts))
    return Item(name, geometry)

get_items

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

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

Source code in rslearn/data_sources/google_earth_engine.py

def get_items(
    self, geometries: list[STGeometry], query_config: QueryConfig
) -> list[list[MatchedItemGroup[Item]]]:
    """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.
    """
    wgs84_geometries = [
        geometry.to_projection(WGS84_PROJECTION) for geometry in geometries
    ]

    groups = []
    for geometry in wgs84_geometries:
        cur_items = []
        encoded_items = self.rtree_index.query(geometry.shp.bounds)
        for encoded_item in encoded_items:
            item = Item.deserialize(json.loads(encoded_item))
            if not item.geometry.shp.intersects(geometry.shp):
                continue
            cur_items.append(item)

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

        cur_groups = rslearn.data_sources.utils.match_candidate_items_to_window(
            geometry, cur_items, query_config
        )
        groups.append(cur_groups)

    return groups

deserialize_item

deserialize_item(serialized_item: dict) -> Item

Deserializes an item from JSON-decoded data.

Source code in rslearn/data_sources/google_earth_engine.py

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

item_to_image

item_to_image(item: Item) -> Image

Get the Image corresponding to the Item.

This function is separated so it can be overriden if subclasses want to add modifications to the image.

Source code in rslearn/data_sources/google_earth_engine.py

def item_to_image(self, item: Item) -> ee.image.Image:
    """Get the Image corresponding to the Item.

    This function is separated so it can be overriden if subclasses want to add
    modifications to the image.
    """
    filtered = self.get_collection().filter(ee.Filter.eq("system:index", item.name))
    image = filtered.first()
    image = image.select(self.bands)
    return image

export_item

export_item(item: Item, blob_prefix: str, projection_and_bounds: tuple[Projection, PixelBounds] | None = None) -> None

Export the item to the specified folder.

Parameters:

Name	Type	Description	Default
item	Item	the item to export.	required
blob_prefix	str	the prefix (folder) to use.	required
projection_and_bounds	tuple[Projection, PixelBounds] | None	optionally use this projection and bounds instead of the extent of the image.	None

Source code in rslearn/data_sources/google_earth_engine.py

def export_item(
    self,
    item: Item,
    blob_prefix: str,
    projection_and_bounds: tuple[Projection, PixelBounds] | None = None,
) -> None:
    """Export the item to the specified folder.

    Args:
        item: the item to export.
        blob_prefix: the prefix (folder) to use.
        projection_and_bounds: optionally use this projection and bounds instead of
            the extent of the image.
    """
    image = self.item_to_image(item)
    projection = image.select(self.bands[0]).projection().getInfo()
    logger.info("Starting task to retrieve image %s", item.name)

    extent_kwargs: dict[str, Any]
    if projection_and_bounds is not None:
        projection, bounds = projection_and_bounds
        transform = get_transform_from_projection_and_bounds(projection, bounds)
        width = bounds[2] - bounds[0]
        height = bounds[3] - bounds[1]
        extent_kwargs = dict(
            crs=str(projection.crs),
            crsTransform=[
                transform.a,
                transform.b,
                transform.c,
                transform.d,
                transform.e,
                transform.f,
            ],
            dimensions=f"{width}x{height}",
        )
    else:
        # Use the native projection of the image.
        # We pass scale instead of crsTransform since some images have positive y
        # resolution which means they are upside down and rasterio cannot merge
        # them.
        extent_kwargs = dict(
            crs=projection["crs"],
            scale=projection["transform"][0],
        )

    task = ee.batch.Export.image.toCloudStorage(
        image=image,
        description=item.name,
        bucket=self.gcs_bucket_name,
        fileNamePrefix=blob_prefix,
        maxPixels=10000000000,
        fileFormat="GeoTIFF",
        skipEmptyTiles=True,
        **extent_kwargs,
    )
    task.start()
    while True:
        time.sleep(10)
        status_dict = task.status()
        if status_dict["state"] in ["UNSUBMITTED", "READY", "RUNNING"]:
            continue
        if status_dict["state"] == "COMPLETED":
            break
        if status_dict["state"] != "FAILED":
            raise ValueError(
                f"got unexpected GEE task state {status_dict['state']}"
            )
        # The task failed. We see if it is an okay failure case or if we need to
        # raise exception.
        if status_dict["error_message"] == NoValidPixelsException.GEE_MESSAGE:
            raise NoValidPixelsException()
        raise ExportException(f"GEE task failed: {status_dict['error_message']}")

ingest

ingest(tile_store: TileStoreWithLayer, items: list[Item], 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[Item]	the items to ingest	required
geometries	list[list[STGeometry]]	a list of geometries needed for each item	required

Source code in rslearn/data_sources/google_earth_engine.py

def ingest(
    self,
    tile_store: TileStoreWithLayer,
    items: list[Item],
    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 in items:
        if tile_store.is_raster_ready(item, self.bands):
            continue

        # Export the item to GCS.
        blob_prefix = f"{self.collection_name}/{item.name}.{os.getpid()}/"
        self.export_item(item, blob_prefix)

        # See what files the export produced.
        # If there are multiple, then we merge them into one file since that's the
        # simplest way to handle it.
        blobs = list(self.bucket.list_blobs(prefix=blob_prefix))

        with tempfile.TemporaryDirectory() as tmp_dir_name:
            if len(blobs) == 1:
                local_fname = os.path.join(
                    tmp_dir_name, blobs[0].name.split("/")[-1]
                )
                blobs[0].download_to_filename(local_fname)
                tile_store.write_raster_file(
                    item,
                    self.bands,
                    UPath(local_fname),
                    time_range=item.geometry.time_range,
                )

            else:
                merged = self._merge_rasters(blobs)
                tile_store.write_raster(
                    item,
                    self.bands,
                    merged.projection,
                    merged.bounds,
                    RasterArray(
                        chw_array=merged.array,
                        time_range=item.geometry.time_range,
                        metadata=RasterMetadata(nodata_value=merged.nodata_value),
                    ),
                )

        for blob in blobs:
            blob.delete()

is_raster_ready

is_raster_ready(layer_name: str, item: Item, bands: list[str]) -> bool

Checks if this raster has been written to the store.

Parameters:

Name	Type	Description	Default
layer_name	str	the layer name or alias.	required
item	Item	the item.	required
bands	list[str]	the list of bands identifying which specific raster to read.	required

Returns:

Type	Description
bool	whether there is a raster in the store matching the source, item, and bands.

Source code in rslearn/data_sources/google_earth_engine.py

def is_raster_ready(self, layer_name: str, item: Item, bands: list[str]) -> bool:
    """Checks if this raster has been written to the store.

    Args:
        layer_name: the layer name or alias.
        item: 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 Planetary Computer.
    return True

get_raster_bands

get_raster_bands(layer_name: str, item: Item) -> list[list[str]]

Get the sets of bands that have been stored for the specified item.

Parameters:

Name	Type	Description	Default
layer_name	str	the layer name or alias.	required
item	Item	the item.	required

Returns:

Type	Description
list[list[str]]	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.

Source code in rslearn/data_sources/google_earth_engine.py

def get_raster_bands(self, layer_name: str, item: Item) -> 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: 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 [self.bands]

get_raster_bounds

get_raster_bounds(layer_name: str, item: Item, bands: list[str], projection: Projection) -> PixelBounds

Get the bounds of the raster in the specified projection.

Parameters:

Name	Type	Description	Default
layer_name	str	the layer name or alias.	required
item	Item	the item to check.	required
bands	list[str]	the list of bands identifying which specific raster to read. These bands must match the bands of a stored raster.	required
projection	Projection	the projection to get the raster's bounds in.	required

Returns:

Type	Description
PixelBounds	the bounds of the raster in the projection.

Source code in rslearn/data_sources/google_earth_engine.py

def get_raster_bounds(
    self, layer_name: str, item: Item, 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: 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.
    """
    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]),
    )

get_raster_metadata

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

Export and read the raster header to obtain nodata.

The nodata value is cached after the first successful export that yields a non-None value, so subsequent calls skip the export entirely.

Source code in rslearn/data_sources/google_earth_engine.py

def get_raster_metadata(
    self, layer_name: str, item: Item, bands: list[str]
) -> RasterMetadata:
    """Export and read the raster header to obtain nodata.

    The nodata value is cached after the first successful export that
    yields a non-None value, so subsequent calls skip the export entirely.
    """
    if self._cached_nodata_value is not None:
        return RasterMetadata(nodata_value=self._cached_nodata_value)

    blob_prefix = f"{self.collection_name}/{item.name}.{os.getpid()}/"
    try:
        self.export_item(item, blob_prefix)
    except NoValidPixelsException:
        return RasterMetadata()

    blobs = list(self.bucket.list_blobs(prefix=blob_prefix))
    buf = io.BytesIO()
    blobs[0].download_to_file(buf)
    buf.seek(0)
    with rasterio.open(buf) as src:
        nodata_value = src.nodata
    if nodata_value is not None:
        self._cached_nodata_value = nodata_value
    return RasterMetadata(nodata_value=nodata_value)

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.

Parameters:

Name	Type	Description	Default
layer_name	str	the layer name or alias.	required
item	Item	the item to read.	required
bands	list[str]	the list of bands identifying which specific raster to read. These bands must match the bands of a stored raster.	required
projection	Projection	the projection to read in.	required
bounds	PixelBounds	the bounds to read.	required
resampling	Resampling	the resampling method to use in case reprojection is needed.	bilinear

Returns:

Type	Description
RasterArray	the raster data

Source code in rslearn/data_sources/google_earth_engine.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.

    Args:
        layer_name: the layer name or alias.
        item: 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
    """
    # Extract the requested extent and export to GCS.
    bounds_str = f"{bounds[0]}_{bounds[1]}_{bounds[2]}_{bounds[3]}"
    blob_prefix = f"{self.collection_name}/{item.name}.{bounds_str}.{os.getpid()}/"

    try:
        self.export_item(
            item, blob_prefix, projection_and_bounds=(projection, bounds)
        )
    except NoValidPixelsException:
        # No valid pixels means the result should be empty.
        logger.info(
            f"No valid pixels in item {item.name} with projection={projection}, bounds={bounds}, returning empty image"
        )
        return RasterArray(
            chw_array=np.zeros(
                (len(bands), bounds[3] - bounds[1], bounds[2] - bounds[0]),
                dtype=np.float32,
            ),
            time_range=item.geometry.time_range,
        )

    wanted_transform = get_transform_from_projection_and_bounds(projection, bounds)
    crs_bounds = (
        bounds[0] * projection.x_resolution,
        bounds[3] * projection.y_resolution,
        bounds[2] * projection.x_resolution,
        bounds[1] * projection.y_resolution,
    )

    blobs = list(self.bucket.list_blobs(prefix=blob_prefix))

    if len(blobs) == 1:
        # With a single output, we can simply read it with vrt.
        buf = io.BytesIO()
        blobs[0].download_to_file(buf)
        buf.seek(0)
        with rasterio.open(buf) as src:
            nodata_value = src.nodata
            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 RasterArray(
                    chw_array=vrt.read(),
                    time_range=item.geometry.time_range,
                    metadata=RasterMetadata(nodata_value=nodata_value),
                )

    else:
        # With multiple outputs, we need to merge them together.
        # We can set the bounds in CRS coordinates when we do the merging.
        if projection.x_resolution != -projection.y_resolution:
            raise NotImplementedError(
                "Only projection with x_res=-y_res is supported for GEE direct materialization"
            )
        merged = self._merge_rasters(
            blobs, crs_bounds=crs_bounds, res=projection.x_resolution
        )
        metadata = RasterMetadata(nodata_value=merged.nodata_value)

        # We copy the array if its bounds don't match exactly.
        if merged.bounds == bounds:
            return RasterArray(
                chw_array=merged.array,
                time_range=item.geometry.time_range,
                metadata=metadata,
            )
        dst_array = np.zeros(
            (
                merged.array.shape[0],
                bounds[3] - bounds[1],
                bounds[2] - bounds[0],
            ),
            dtype=merged.array.dtype,
        )
        copy_spatial_array(merged.array, dst_array, merged.bounds[0:2], bounds[0:2])
        return RasterArray(
            chw_array=dst_array,
            time_range=item.geometry.time_range,
            metadata=metadata,
        )

materialize

materialize(window: Window, item_groups: list[list[Item]], layer_name: str, layer_cfg: LayerConfig, group_time_ranges: list[tuple[datetime, datetime] | None] | None = None) -> None

Materialize data for the window.

Parameters:

Name	Type	Description	Default
window	Window	the window to materialize	required
item_groups	list[list[Item]]	the items from get_items	required
layer_name	str	the name of this layer	required
layer_cfg	LayerConfig	the config of this layer	required
group_time_ranges	list[tuple[datetime, datetime] | None] | None	optional request time range for each item group	None

Source code in rslearn/data_sources/google_earth_engine.py

def materialize(
    self,
    window: Window,
    item_groups: list[list[Item]],
    layer_name: str,
    layer_cfg: LayerConfig,
    group_time_ranges: list[tuple[datetime, datetime] | None] | None = None,
) -> 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
        group_time_ranges: optional request time range for each item group
    """
    RasterMaterializer().materialize(
        TileStoreWithLayer(self, layer_name),
        window,
        layer_name,
        layer_cfg,
        item_groups,
        group_time_ranges=group_time_ranges,
    )

GoogleSatelliteEmbeddings

Bases: GEE

GEE data source for the Google Satellite Embeddings.

See here for details: https://developers.google.com/earth-engine/datasets/catalog/GOOGLE_SATELLITE_EMBEDDING_V1_ANNUAL

Source code in rslearn/data_sources/google_earth_engine.py

class GoogleSatelliteEmbeddings(GEE):
    """GEE data source for the Google Satellite Embeddings.

    See here for details:
    https://developers.google.com/earth-engine/datasets/catalog/GOOGLE_SATELLITE_EMBEDDING_V1_ANNUAL
    """

    COLLECTION_NAME = "GOOGLE/SATELLITE_EMBEDDING/V1/ANNUAL"

    def __init__(
        self,
        gcs_bucket_name: str,
        index_cache_dir: str,
        service_account_name: str,
        service_account_credentials: str,
        context: DataSourceContext = DataSourceContext(),
    ):
        """Create a new GoogleSatelliteEmbeddings. See GEE for the arguments."""
        super().__init__(
            bands=[f"A{idx:02d}" for idx in range(64)],
            collection_name=self.COLLECTION_NAME,
            gcs_bucket_name=gcs_bucket_name,
            index_cache_dir=index_cache_dir,
            service_account_name=service_account_name,
            service_account_credentials=service_account_credentials,
            context=context,
        )

    # Override to add conversion to uint16.
    def item_to_image(self, item: Item) -> ee.image.Image:
        """Get the Image corresponding to the Item."""
        filtered = self.get_collection().filter(ee.Filter.eq("system:index", item.name))
        image = filtered.first()
        image = image.select(self.bands)
        return image.multiply(8192).add(8192).toUint16()

item_to_image

item_to_image(item: Item) -> Image

Get the Image corresponding to the Item.

Source code in rslearn/data_sources/google_earth_engine.py

def item_to_image(self, item: Item) -> ee.image.Image:
    """Get the Image corresponding to the Item."""
    filtered = self.get_collection().filter(ee.Filter.eq("system:index", item.name))
    image = filtered.first()
    image = image.select(self.bands)
    return image.multiply(8192).add(8192).toUint16()