rslearn.tile_stores

tile_stores

Tile stores that store ingested raster and vector data before materialization.

DefaultTileStore

Bases: TileStore

Default TileStore implementation.

It stores raster and vector data under the provided UPath.

Raster data is always stored as a geo-referenced image, while vector data can use any provided VectorFormat. This is because for raster data we support reading in an arbitrary projection, but this is not supported in GeotiffRasterFormat, so we directly use rasterio to access the file.

Source code in rslearn/tile_stores/default.py

class DefaultTileStore(TileStore):
    """Default TileStore implementation.

    It stores raster and vector data under the provided UPath.

    Raster data is always stored as a geo-referenced image, while vector data can use
    any provided VectorFormat. This is because for raster data we support reading in an
    arbitrary projection, but this is not supported in GeotiffRasterFormat, so we
    directly use rasterio to access the file.
    """

    def __init__(
        self,
        path_suffix: str = "tiles",
        convert_rasters_to_cogs: bool = True,
        tile_size: int = 256,
        geotiff_options: dict[str, Any] = {},
        vector_format: VectorFormat = GeojsonVectorFormat(),
    ):
        """Create a new DefaultTileStore.

        Args:
            path_suffix: the path suffix to store files under, which is joined with
                the dataset path if it does not contain a protocol string. See
                rslearn.utils.fsspec.join_upath.
            convert_rasters_to_cogs: whether to re-encode all raster files to tiled
                GeoTIFFs. Re-encoding will also handle rasters that specify ground
                control points instead of a CRS and transform; for such rasters, always
                use convert_rasters_to_cogs since materialization cannot handle those
                rasters.
            tile_size: if converting to COGs, the tile size to use.
            geotiff_options: other options to pass to rasterio.open (for writes).
            vector_format: format to use for storing vector data.
        """
        self.path_suffix = path_suffix
        self.convert_rasters_to_cogs = convert_rasters_to_cogs
        self.tile_size = tile_size
        self.geotiff_options = geotiff_options
        self.vector_format = vector_format

        self.path: UPath | None = None

    @override
    def set_dataset_path(self, ds_path: UPath) -> None:
        self.path = join_upath(ds_path, self.path_suffix)

    def _get_raster_dir(
        self, layer_name: str, item_name: str, bands: list[str], write: bool = False
    ) -> UPath:
        """Get the directory where the specified raster is stored.

        Args:
            layer_name: the name of the dataset layer.
            item_name: the name of the item from the data source.
            bands: list of band names that are expected to be stored together.
            write: whether to create the directory and write the bands to a file inside
                the directory.

        Returns:
            the UPath directory where the raster should be stored.
        """
        assert self.path is not None
        dir_name = self.path / layer_name / item_name / get_bandset_dirname(bands)

        if write:
            dir_name.mkdir(parents=True, exist_ok=True)
            with (dir_name / BANDS_FNAME).open("w") as f:
                json.dump(bands, f)

        return dir_name

    def _get_raster_fname(
        self, layer_name: str, item_name: str, bands: list[str]
    ) -> UPath:
        """Get the filename of the specified raster.

        Args:
            layer_name: the name of the dataset layer.
            item_name: the name of the item from the data source.
            bands: list of band names that are expected to be stored together.

        Returns:
            the UPath filename of the raster, which should be readable by rasterio.

        Raises:
            ValueError: if no file is found.
        """
        raster_dir = self._get_raster_dir(layer_name, item_name, bands)
        for fname in iter_nonhidden_files(raster_dir):
            # Ignore completed sentinel files, bands files, as well as temporary files created by
            # open_atomic (in case this tile store is on local filesystem).
            if fname.name == COMPLETED_FNAME:
                continue
            if fname.name == BANDS_FNAME:
                continue
            if fname.name == METADATA_FNAME:
                continue
            if ".tmp." in fname.name:
                continue
            return fname
        raise ValueError(f"no raster found in {raster_dir}")

    @override
    def is_raster_ready(self, layer_name: str, item: Item, bands: list[str]) -> bool:
        raster_dir = self._get_raster_dir(layer_name, item.name, bands)
        return (raster_dir / COMPLETED_FNAME).exists()

    @override
    def get_raster_bands(self, layer_name: str, item: Item) -> list[list[str]]:
        assert isinstance(self.path, UPath)
        item_dir = self.path / layer_name / item.name
        if not item_dir.exists():
            return []

        bands: list[list[str]] = []
        for raster_dir in iter_nonhidden_subdirs(item_dir):
            if not (raster_dir / BANDS_FNAME).exists():
                # This is likely a legacy directory where the bands are only encoded in
                # the directory name, so we have to rely on that.
                parts = raster_dir.name.split("_")
                bands.append(parts)
                continue

            # We use the BANDS_FNAME here -- although it is slower to read the file, it
            # is more reliable since sometimes the directory name is a hash of the
            # bands in case there are too many bands (filename too long) or some bands
            # contain the underscore character.
            with (raster_dir / BANDS_FNAME).open() as f:
                bands.append(json.load(f))

        return bands

    @override
    def get_raster_bounds(
        self, layer_name: str, item: Item, bands: list[str], projection: Projection
    ) -> PixelBounds:
        raster_fname = self._get_raster_fname(layer_name, item.name, bands)

        with open_rasterio_upath_reader(raster_fname) as src:
            with rasterio.vrt.WarpedVRT(src, crs=projection.crs) as vrt:
                bounds = (
                    vrt.bounds[0] / projection.x_resolution,
                    vrt.bounds[1] / projection.y_resolution,
                    vrt.bounds[2] / projection.x_resolution,
                    vrt.bounds[3] / projection.y_resolution,
                )
                return (
                    math.floor(min(bounds[0], bounds[2])),
                    math.floor(min(bounds[1], bounds[3])),
                    math.ceil(max(bounds[0], bounds[2])),
                    math.ceil(max(bounds[1], bounds[3])),
                )

    @override
    def get_raster_metadata(
        self, layer_name: str, item: Item, bands: list[str]
    ) -> RasterMetadata:
        raster_fname = self._get_raster_fname(layer_name, item.name, bands)
        with open_rasterio_upath_reader(raster_fname) as src:
            return RasterMetadata(nodata_value=src.nodata)

    @override
    def read_raster(
        self,
        layer_name: str,
        item: Item,
        bands: list[str],
        projection: Projection,
        bounds: PixelBounds,
        resampling: Resampling = Resampling.bilinear,
    ) -> RasterArray:
        raster_fname = self._get_raster_fname(layer_name, item.name, bands)
        return GeotiffRasterFormat().decode_raster(
            path=raster_fname.parent,
            fname=raster_fname.name,
            projection=projection,
            bounds=bounds,
            resampling=resampling,
        )

    @override
    def write_raster(
        self,
        layer_name: str,
        item: Item,
        bands: list[str],
        projection: Projection,
        bounds: PixelBounds,
        raster: RasterArray,
    ) -> None:
        raster_dir = self._get_raster_dir(layer_name, item.name, bands, write=True)
        raster_format = GeotiffRasterFormat(geotiff_options=self.geotiff_options)
        raster_format.encode_raster(raster_dir, projection, bounds, raster)
        (raster_dir / COMPLETED_FNAME).touch()

    @override
    def write_raster_file(
        self,
        layer_name: str,
        item: Item,
        bands: list[str],
        fname: UPath,
        time_range: tuple[datetime, datetime] | None = None,
    ) -> None:
        raster_dir = self._get_raster_dir(layer_name, item.name, bands, write=True)
        raster_dir.mkdir(parents=True, exist_ok=True)

        if self.convert_rasters_to_cogs:
            with open_rasterio_upath_reader(fname) as src:
                nodata = src.nodata

                # If raster specifies ground control points, use WarpedVRT to get it in
                # an appropriate projection.
                # Previously we used rasterio.transform.from_gcps(gcps) but I think the
                # problem is that it computes one transform for the entire raster but
                # the raster might actually need warping.
                gcps, gcp_crs = src.gcps
                if src.crs is None and len(gcps) > 0:
                    # Use the first ground control point to pick a UTM/UPS projection.
                    first_gcp_orig = STGeometry(
                        Projection(gcp_crs, 1, 1),
                        shapely.Point(gcps[0].x, gcps[0].y),
                        None,
                    )
                    first_gcp_wgs84 = first_gcp_orig.to_projection(WGS84_PROJECTION)
                    crs = get_utm_ups_crs(first_gcp_wgs84.shp.x, first_gcp_wgs84.shp.y)
                    with rasterio.vrt.WarpedVRT(
                        src, crs=crs, resampling=Resampling.cubic
                    ) as vrt:
                        array = vrt.read()
                        transform = vrt.transform

                else:
                    array = src.read()
                    crs = src.crs
                    transform = src.transform

            output_profile = {
                "driver": "GTiff",
                "compress": "lzw",
                "width": array.shape[2],
                "height": array.shape[1],
                "count": array.shape[0],
                "dtype": array.dtype.name,
                "crs": crs,
                "transform": transform,
                "BIGTIFF": "IF_SAFER",
                "tiled": True,
                "blockxsize": self.tile_size,
                "blockysize": self.tile_size,
            }
            if nodata is not None:
                output_profile["nodata"] = nodata

            output_profile.update(self.geotiff_options)

            with open_rasterio_upath_writer(
                raster_dir / "geotiff.tif", **output_profile
            ) as dst:
                dst.write(array)

        else:
            # Just copy the file directly.
            dst_fname = raster_dir / fname.name
            with fname.open("rb") as src:
                with open_atomic(dst_fname, "wb") as dst:
                    shutil.copyfileobj(src, dst)

        if time_range is not None:
            # Add GeotiffRasterFormat metadata file so that GeotiffRasterFormat will
            # see it upon read and use it to obtain the time range information for the
            # raster.
            GeotiffRasterFormat.encode_metadata(
                raster_dir,
                GeotiffRasterMetadata(
                    num_channels=len(bands),
                    num_timesteps=1,
                    timestamps=[(time_range[0], time_range[1])],
                ),
            )

        (raster_dir / COMPLETED_FNAME).touch()

    def _get_vector_dir(self, layer_name: str, item_name: str) -> UPath:
        assert self.path is not None
        return self.path / layer_name / item_name

    @override
    def is_vector_ready(self, layer_name: str, item: Item) -> bool:
        vector_dir = self._get_vector_dir(layer_name, item.name)
        return (vector_dir / COMPLETED_FNAME).exists()

    @override
    def read_vector(
        self,
        layer_name: str,
        item: Item,
        projection: Projection,
        bounds: PixelBounds,
    ) -> list[Feature]:
        vector_dir = self._get_vector_dir(layer_name, item.name)
        return self.vector_format.decode_vector(vector_dir, projection, bounds)

    @override
    def write_vector(
        self, layer_name: str, item: Item, features: list[Feature]
    ) -> None:
        vector_dir = self._get_vector_dir(layer_name, item.name)
        vector_dir.mkdir(parents=True, exist_ok=True)
        self.vector_format.encode_vector(vector_dir, features)
        (vector_dir / COMPLETED_FNAME).touch()

TileStore

An abstract class for a tile store.

A tile store supports operations to read and write raster and vector data.

Source code in rslearn/tile_stores/tile_store.py

class TileStore:
    """An abstract class for a tile store.

    A tile store supports operations to read and write raster and vector data.
    """

    def set_dataset_path(self, ds_path: UPath) -> None:
        """Set the dataset path.

        This is in case the TileStore wants to use the ds_path to help determine where
        to store data.

        Args:
            ds_path: the dataset path that this TileStore is a part of.
        """
        pass

    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.
        """
        raise NotImplementedError

    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.
        """
        raise NotImplementedError

    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.
        """
        raise NotImplementedError

    def get_raster_metadata(
        self, layer_name: str, item: Item, bands: list[str]
    ) -> RasterMetadata:
        """Get metadata for a stored raster without reading pixel data.

        Args:
            layer_name: the layer name or alias.
            item: the item.
            bands: the list of bands identifying which specific raster to read.

        Returns:
            a RasterMetadata instance (may have None fields if the store
            does not track a given piece of metadata).
        """
        raise NotImplementedError

    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
        """
        raise NotImplementedError

    def write_raster(
        self,
        layer_name: str,
        item: Item,
        bands: list[str],
        projection: Projection,
        bounds: PixelBounds,
        raster: RasterArray,
    ) -> None:
        """Write raster data to the store.

        Args:
            layer_name: the layer name or alias.
            item: the item to write.
            bands: the list of bands in the array.
            projection: the projection of the array.
            bounds: the bounds of the array.
            raster: the raster data.
        """
        raise NotImplementedError

    def write_raster_file(
        self,
        layer_name: str,
        item: Item,
        bands: list[str],
        fname: UPath,
        time_range: tuple[datetime, datetime] | None = None,
    ) -> None:
        """Write raster data to the store.

        Args:
            layer_name: the layer name or alias.
            item: the item to write.
            bands: the list of bands in the array.
            fname: the raster file.
            time_range: optional time range for the raster.
        """
        raise NotImplementedError

    def is_vector_ready(self, layer_name: str, item: Item) -> bool:
        """Checks if this vector item has been written to the store.

        Args:
            layer_name: the layer name or alias.
            item: the item.

        Returns:
            whether the vector data from the item has been stored.
        """
        raise NotImplementedError

    def read_vector(
        self,
        layer_name: str,
        item: Item,
        projection: Projection,
        bounds: PixelBounds,
    ) -> list[Feature]:
        """Read vector data from the store.

        Args:
            layer_name: the layer name or alias.
            item: the item to read.
            projection: the projection to read in.
            bounds: the bounds within which to read.

        Returns:
            the vector data
        """
        raise NotImplementedError

    def write_vector(
        self, layer_name: str, item: Item, features: list[Feature]
    ) -> None:
        """Write vector data to the store.

        Args:
            layer_name: the layer name or alias.
            item: the item to write.
            features: the vector data.
        """
        raise NotImplementedError

set_dataset_path

set_dataset_path(ds_path: UPath) -> None

Set the dataset path.

This is in case the TileStore wants to use the ds_path to help determine where to store data.

Parameters:

Name	Type	Description	Default
ds_path	UPath	the dataset path that this TileStore is a part of.	required

Source code in rslearn/tile_stores/tile_store.py

def set_dataset_path(self, ds_path: UPath) -> None:
    """Set the dataset path.

    This is in case the TileStore wants to use the ds_path to help determine where
    to store data.

    Args:
        ds_path: the dataset path that this TileStore is a part of.
    """
    pass

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/tile_stores/tile_store.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.
    """
    raise NotImplementedError

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/tile_stores/tile_store.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.
    """
    raise NotImplementedError

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/tile_stores/tile_store.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.
    """
    raise NotImplementedError

get_raster_metadata

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

Get metadata for a stored raster without reading pixel data.

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
RasterMetadata	a RasterMetadata instance (may have None fields if the store
RasterMetadata	does not track a given piece of metadata).

Source code in rslearn/tile_stores/tile_store.py

def get_raster_metadata(
    self, layer_name: str, item: Item, bands: list[str]
) -> RasterMetadata:
    """Get metadata for a stored raster without reading pixel data.

    Args:
        layer_name: the layer name or alias.
        item: the item.
        bands: the list of bands identifying which specific raster to read.

    Returns:
        a RasterMetadata instance (may have None fields if the store
        does not track a given piece of metadata).
    """
    raise NotImplementedError

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/tile_stores/tile_store.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
    """
    raise NotImplementedError

write_raster

write_raster(layer_name: str, item: Item, bands: list[str], projection: Projection, bounds: PixelBounds, raster: RasterArray) -> None

Write raster data to the store.

Parameters:

Name	Type	Description	Default
layer_name	str	the layer name or alias.	required
item	Item	the item to write.	required
bands	list[str]	the list of bands in the array.	required
projection	Projection	the projection of the array.	required
bounds	PixelBounds	the bounds of the array.	required
raster	RasterArray	the raster data.	required

Source code in rslearn/tile_stores/tile_store.py

def write_raster(
    self,
    layer_name: str,
    item: Item,
    bands: list[str],
    projection: Projection,
    bounds: PixelBounds,
    raster: RasterArray,
) -> None:
    """Write raster data to the store.

    Args:
        layer_name: the layer name or alias.
        item: the item to write.
        bands: the list of bands in the array.
        projection: the projection of the array.
        bounds: the bounds of the array.
        raster: the raster data.
    """
    raise NotImplementedError

write_raster_file

write_raster_file(layer_name: str, item: Item, bands: list[str], fname: UPath, time_range: tuple[datetime, datetime] | None = None) -> None

Write raster data to the store.

Parameters:

Name	Type	Description	Default
layer_name	str	the layer name or alias.	required
item	Item	the item to write.	required
bands	list[str]	the list of bands in the array.	required
fname	UPath	the raster file.	required
time_range	tuple[datetime, datetime] | None	optional time range for the raster.	None

Source code in rslearn/tile_stores/tile_store.py

def write_raster_file(
    self,
    layer_name: str,
    item: Item,
    bands: list[str],
    fname: UPath,
    time_range: tuple[datetime, datetime] | None = None,
) -> None:
    """Write raster data to the store.

    Args:
        layer_name: the layer name or alias.
        item: the item to write.
        bands: the list of bands in the array.
        fname: the raster file.
        time_range: optional time range for the raster.
    """
    raise NotImplementedError

is_vector_ready

is_vector_ready(layer_name: str, item: Item) -> bool

Checks if this vector item 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

Returns:

Type	Description
bool	whether the vector data from the item has been stored.

Source code in rslearn/tile_stores/tile_store.py

def is_vector_ready(self, layer_name: str, item: Item) -> bool:
    """Checks if this vector item has been written to the store.

    Args:
        layer_name: the layer name or alias.
        item: the item.

    Returns:
        whether the vector data from the item has been stored.
    """
    raise NotImplementedError

read_vector

read_vector(layer_name: str, item: Item, projection: Projection, bounds: PixelBounds) -> list[Feature]

Read vector 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
projection	Projection	the projection to read in.	required
bounds	PixelBounds	the bounds within which to read.	required

Returns:

Type	Description
list[Feature]	the vector data

Source code in rslearn/tile_stores/tile_store.py

def read_vector(
    self,
    layer_name: str,
    item: Item,
    projection: Projection,
    bounds: PixelBounds,
) -> list[Feature]:
    """Read vector data from the store.

    Args:
        layer_name: the layer name or alias.
        item: the item to read.
        projection: the projection to read in.
        bounds: the bounds within which to read.

    Returns:
        the vector data
    """
    raise NotImplementedError

write_vector

write_vector(layer_name: str, item: Item, features: list[Feature]) -> None

Write vector data to the store.

Parameters:

Name	Type	Description	Default
layer_name	str	the layer name or alias.	required
item	Item	the item to write.	required
features	list[Feature]	the vector data.	required

Source code in rslearn/tile_stores/tile_store.py

def write_vector(
    self, layer_name: str, item: Item, features: list[Feature]
) -> None:
    """Write vector data to the store.

    Args:
        layer_name: the layer name or alias.
        item: the item to write.
        features: the vector data.
    """
    raise NotImplementedError

TileStoreWithLayer

Convenience class to access TileStore in the context of a layer.

Source code in rslearn/tile_stores/tile_store.py

class TileStoreWithLayer:
    """Convenience class to access TileStore in the context of a layer."""

    def __init__(self, tile_store: TileStore, layer_name: str):
        """Create a new TileStoreWithLayer.

        Args:
            tile_store: underlying TileStore.
            layer_name: the layer name.
        """
        self.tile_store = tile_store
        self.layer_name = layer_name

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

        Args:
            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.
        """
        return self.tile_store.is_raster_ready(self.layer_name, item, bands)

    def get_raster_bands(self, item: Item) -> list[list[str]]:
        """Get the sets of bands that have been stored for the specified item.

        Args:
            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.tile_store.get_raster_bands(self.layer_name, item)

    def get_raster_bounds(
        self, item: Item, bands: list[str], projection: Projection
    ) -> PixelBounds:
        """Get the bounds of the raster in the specified projection.

        Args:
            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.
        """
        return self.tile_store.get_raster_bounds(
            self.layer_name, item, bands, projection
        )

    def get_raster_metadata(self, item: Item, bands: list[str]) -> RasterMetadata:
        """Get metadata for a stored raster without reading pixel data.

        Args:
            item: the item.
            bands: the list of bands identifying which specific raster to read.

        Returns:
            a RasterMetadata instance.
        """
        return self.tile_store.get_raster_metadata(self.layer_name, item, bands)

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

        Args:
            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
        """
        return self.tile_store.read_raster(
            self.layer_name, item, bands, projection, bounds, resampling
        )

    def write_raster(
        self,
        item: Item,
        bands: list[str],
        projection: Projection,
        bounds: PixelBounds,
        raster: RasterArray,
    ) -> None:
        """Write raster data to the store.

        Args:
            item: the item to write.
            bands: the list of bands in the array.
            projection: the projection of the array.
            bounds: the bounds of the array.
            raster: the raster data.
        """
        self.tile_store.write_raster(
            self.layer_name, item, bands, projection, bounds, raster
        )

    def write_raster_file(
        self,
        item: Item,
        bands: list[str],
        fname: UPath,
        time_range: tuple[datetime, datetime] | None = None,
    ) -> None:
        """Write raster data to the store.

        Args:
            item: the item to write.
            bands: the list of bands in the array.
            fname: the raster file.
            time_range: optional time range for the raster.
        """
        self.tile_store.write_raster_file(
            self.layer_name, item, bands, fname, time_range=time_range
        )

    def is_vector_ready(self, item: Item) -> bool:
        """Checks if this vector item has been written to the store.

        Args:
            item: the item.

        Returns:
            whether the vector data from the item has been stored.
        """
        return self.tile_store.is_vector_ready(self.layer_name, item)

    def read_vector(
        self, item: Item, projection: Projection, bounds: PixelBounds
    ) -> list[Feature]:
        """Read vector data from the store.

        Args:
            item: the item to read.
            projection: the projection to read in.
            bounds: the bounds within which to read.

        Returns:
            the vector data
        """
        return self.tile_store.read_vector(self.layer_name, item, projection, bounds)

    def write_vector(self, item: Item, features: list[Feature]) -> None:
        """Write vector data to the store.

        Args:
            item: the item to write.
            features: the vector data.
        """
        self.tile_store.write_vector(self.layer_name, item, features)

is_raster_ready

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

Checks if this raster has been written to the store.

Parameters:

Name	Type	Description	Default
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/tile_stores/tile_store.py

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

    Args:
        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.
    """
    return self.tile_store.is_raster_ready(self.layer_name, item, bands)

get_raster_bands

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

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

Parameters:

Name	Type	Description	Default
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/tile_stores/tile_store.py

def get_raster_bands(self, item: Item) -> list[list[str]]:
    """Get the sets of bands that have been stored for the specified item.

    Args:
        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.tile_store.get_raster_bands(self.layer_name, item)

get_raster_bounds

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

Get the bounds of the raster in the specified projection.

Parameters:

Name	Type	Description	Default
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/tile_stores/tile_store.py

def get_raster_bounds(
    self, item: Item, bands: list[str], projection: Projection
) -> PixelBounds:
    """Get the bounds of the raster in the specified projection.

    Args:
        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.
    """
    return self.tile_store.get_raster_bounds(
        self.layer_name, item, bands, projection
    )

get_raster_metadata

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

Get metadata for a stored raster without reading pixel data.

Parameters:

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

Returns:

Type	Description
RasterMetadata	a RasterMetadata instance.

Source code in rslearn/tile_stores/tile_store.py

def get_raster_metadata(self, item: Item, bands: list[str]) -> RasterMetadata:
    """Get metadata for a stored raster without reading pixel data.

    Args:
        item: the item.
        bands: the list of bands identifying which specific raster to read.

    Returns:
        a RasterMetadata instance.
    """
    return self.tile_store.get_raster_metadata(self.layer_name, item, bands)

read_raster

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

Read raster data from the store.

Parameters:

Name	Type	Description	Default
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/tile_stores/tile_store.py

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

    Args:
        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
    """
    return self.tile_store.read_raster(
        self.layer_name, item, bands, projection, bounds, resampling
    )

write_raster

write_raster(item: Item, bands: list[str], projection: Projection, bounds: PixelBounds, raster: RasterArray) -> None

Write raster data to the store.

Parameters:

Name	Type	Description	Default
item	Item	the item to write.	required
bands	list[str]	the list of bands in the array.	required
projection	Projection	the projection of the array.	required
bounds	PixelBounds	the bounds of the array.	required
raster	RasterArray	the raster data.	required

Source code in rslearn/tile_stores/tile_store.py

def write_raster(
    self,
    item: Item,
    bands: list[str],
    projection: Projection,
    bounds: PixelBounds,
    raster: RasterArray,
) -> None:
    """Write raster data to the store.

    Args:
        item: the item to write.
        bands: the list of bands in the array.
        projection: the projection of the array.
        bounds: the bounds of the array.
        raster: the raster data.
    """
    self.tile_store.write_raster(
        self.layer_name, item, bands, projection, bounds, raster
    )

write_raster_file

write_raster_file(item: Item, bands: list[str], fname: UPath, time_range: tuple[datetime, datetime] | None = None) -> None

Write raster data to the store.

Parameters:

Name	Type	Description	Default
item	Item	the item to write.	required
bands	list[str]	the list of bands in the array.	required
fname	UPath	the raster file.	required
time_range	tuple[datetime, datetime] | None	optional time range for the raster.	None

Source code in rslearn/tile_stores/tile_store.py

def write_raster_file(
    self,
    item: Item,
    bands: list[str],
    fname: UPath,
    time_range: tuple[datetime, datetime] | None = None,
) -> None:
    """Write raster data to the store.

    Args:
        item: the item to write.
        bands: the list of bands in the array.
        fname: the raster file.
        time_range: optional time range for the raster.
    """
    self.tile_store.write_raster_file(
        self.layer_name, item, bands, fname, time_range=time_range
    )

is_vector_ready

is_vector_ready(item: Item) -> bool

Checks if this vector item has been written to the store.

Parameters:

Name	Type	Description	Default
item	Item	the item.	required

Returns:

Type	Description
bool	whether the vector data from the item has been stored.

Source code in rslearn/tile_stores/tile_store.py

def is_vector_ready(self, item: Item) -> bool:
    """Checks if this vector item has been written to the store.

    Args:
        item: the item.

    Returns:
        whether the vector data from the item has been stored.
    """
    return self.tile_store.is_vector_ready(self.layer_name, item)

read_vector

read_vector(item: Item, projection: Projection, bounds: PixelBounds) -> list[Feature]

Read vector data from the store.

Parameters:

Name	Type	Description	Default
item	Item	the item to read.	required
projection	Projection	the projection to read in.	required
bounds	PixelBounds	the bounds within which to read.	required

Returns:

Type	Description
list[Feature]	the vector data

Source code in rslearn/tile_stores/tile_store.py

def read_vector(
    self, item: Item, projection: Projection, bounds: PixelBounds
) -> list[Feature]:
    """Read vector data from the store.

    Args:
        item: the item to read.
        projection: the projection to read in.
        bounds: the bounds within which to read.

    Returns:
        the vector data
    """
    return self.tile_store.read_vector(self.layer_name, item, projection, bounds)

write_vector

write_vector(item: Item, features: list[Feature]) -> None

Write vector data to the store.

Parameters:

Name	Type	Description	Default
item	Item	the item to write.	required
features	list[Feature]	the vector data.	required

Source code in rslearn/tile_stores/tile_store.py

def write_vector(self, item: Item, features: list[Feature]) -> None:
    """Write vector data to the store.

    Args:
        item: the item to write.
        features: the vector data.
    """
    self.tile_store.write_vector(self.layer_name, item, features)

load_tile_store

load_tile_store(config: dict[str, Any], ds_path: UPath) -> TileStore

Load a tile store from a configuration.

Parameters:

Name	Type	Description	Default
config	dict[str, Any]	the tile store configuration.	required
ds_path	UPath	the dataset root path.	required

Returns:

Type	Description
TileStore	the TileStore

Source code in rslearn/tile_stores/__init__.py

def load_tile_store(config: dict[str, Any], ds_path: UPath) -> TileStore:
    """Load a tile store from a configuration.

    Args:
        config: the tile store configuration.
        ds_path: the dataset root path.

    Returns:
        the TileStore
    """
    init_jsonargparse()
    parser = jsonargparse.ArgumentParser()
    parser.add_argument("--tile_store", type=TileStore)
    cfg = parser.parse_object({"tile_store": config})
    tile_store = parser.instantiate_classes(cfg).tile_store
    tile_store.set_dataset_path(ds_path)
    return tile_store

get_tile_store_with_layer

get_tile_store_with_layer(tile_store: TileStore, layer_name: str, layer_cfg: LayerConfig) -> TileStoreWithLayer

Get the TileStoreWithLayer for the specified layer.

Uses alias of the layer if it is set, otherwise just the layer name.

Parameters:

Name	Type	Description	Default
tile_store	TileStore	the tile store.	required
layer_name	str	the layer name.	required
layer_cfg	LayerConfig	the layer configuration which can specify an alias.	required

Returns:

Type	Description
TileStoreWithLayer	corresponding TileStoreWithLayer

Source code in rslearn/tile_stores/__init__.py

def get_tile_store_with_layer(
    tile_store: TileStore, layer_name: str, layer_cfg: LayerConfig
) -> TileStoreWithLayer:
    """Get the TileStoreWithLayer for the specified layer.

    Uses alias of the layer if it is set, otherwise just the layer name.

    Args:
        tile_store: the tile store.
        layer_name: the layer name.
        layer_cfg: the layer configuration which can specify an alias.

    Returns:
        corresponding TileStoreWithLayer
    """
    if layer_cfg.alias is not None:
        return TileStoreWithLayer(tile_store, layer_cfg.alias)
    return TileStoreWithLayer(tile_store, layer_name)