rslearn 0.0.1__py3-none-any.whl → 0.0.21__py3-none-any.whl

rslearn/utils/raster_format.py

@@ -1,5 +1,6 @@
 """Abstract RasterFormat class."""
 
+import hashlib
 import json
 from typing import Any, BinaryIO
 
@@ -7,16 +8,134 @@ import affine
 import numpy as np
 import numpy.typing as npt
 import rasterio
-from class_registry import ClassRegistry
 from PIL import Image
+from rasterio.crs import CRS
+from rasterio.enums import Resampling
 from upath import UPath
 
-from rslearn.config import RasterFormatConfig
 from rslearn.const import TILE_SIZE
+from rslearn.log_utils import get_logger
+from rslearn.utils.fsspec import open_rasterio_upath_reader, open_rasterio_upath_writer
 
 from .geometry import PixelBounds, Projection
 
-RasterFormats = ClassRegistry()
+logger = get_logger(__name__)
+
+
+def get_bandset_dirname(bands: list[str]) -> str:
+    """Get the directory name that should be used to store the given group of bands."""
+    # We try to use a human-readable name with underscore as the delimiter, but if that
+    # isn't straightforward then we use hash instead.
+    if any(["_" in band for band in bands]):
+        # In this case we hash the JSON representation of the bands.
+        return hashlib.sha256(json.dumps(bands).encode()).hexdigest()
+    dirname = "_".join(bands)
+    if len(dirname) > 64:
+        # Previously we simply joined the bands, but this can result in directory name
+        # that is too long. In this case, now we use hash instead.
+        # We use a different code path here where we hash the initial directory name
+        # instead of the JSON, for historical reasons (to maintain backwards
+        # compatibility).
+        dirname = hashlib.sha256(dirname.encode()).hexdigest()
+    return dirname
+
+
+def get_raster_projection_and_bounds_from_transform(
+    crs: CRS, transform: affine.Affine, width: int, height: int
+) -> tuple[Projection, PixelBounds]:
+    """Determine Projection and bounds from the specified CRS and transform.
+
+    Args:
+        crs: the coordinate reference system.
+        transform: corresponding affine transform matrix.
+        width: the array width
+        height: the array height
+
+    Returns:
+        a tuple (projection, bounds).
+    """
+    x_resolution = transform.a
+    y_resolution = transform.e
+    projection = Projection(crs, x_resolution, y_resolution)
+    offset = (
+        int(round(transform.c / x_resolution)),
+        int(round(transform.f / y_resolution)),
+    )
+    bounds = (offset[0], offset[1], offset[0] + width, offset[1] + height)
+    return (projection, bounds)
+
+
+def get_raster_projection_and_bounds(
+    raster: rasterio.DatasetReader,
+) -> tuple[Projection, PixelBounds]:
+    """Determine the Projection and bounds of the specified raster.
+
+    Args:
+        raster: the raster dataset opened with rasterio.
+
+    Returns:
+        a tuple (projection, bounds).
+    """
+    return get_raster_projection_and_bounds_from_transform(
+        raster.crs, raster.transform, raster.width, raster.height
+    )
+
+
+def get_transform_from_projection_and_bounds(
+    projection: Projection, bounds: PixelBounds
+) -> affine.Affine:
+    """Get the affine transform that corresponds to the given projection and bounds.
+
+    Args:
+        projection: the projection. Only the resolutions are used.
+        bounds: the bounding box. Only the top-left corner is used.
+    """
+    return affine.Affine(
+        projection.x_resolution,
+        0,
+        bounds[0] * projection.x_resolution,
+        0,
+        projection.y_resolution,
+        bounds[1] * projection.y_resolution,
+    )
+
+
+def adjust_projection_and_bounds_for_array(
+    projection: Projection, bounds: PixelBounds, array: npt.NDArray
+) -> tuple[Projection, PixelBounds]:
+    """Adjust the projection and bounds to correspond to the resolution of the array.
+
+    The returned projection and bounds cover the same spatial extent as the inputs, but
+    are updated so that the width and height match that of the array.
+
+    Args:
+        projection: the original projection.
+        bounds: the original bounds.
+        array: the CHW array for which to compute an updated projection and bounds. The
+            returned bounds will have the same width and height as this array.
+
+    Returns:
+        a tuple of adjusted (projection, bounds)
+    """
+    if array.shape[2] == (bounds[2] - bounds[0]) and array.shape[1] == (
+        bounds[3] - bounds[1]
+    ):
+        return (projection, bounds)
+
+    x_factor = array.shape[2] / (bounds[2] - bounds[0])
+    y_factor = array.shape[1] / (bounds[3] - bounds[1])
+    adjusted_projection = Projection(
+        projection.crs,
+        projection.x_resolution / x_factor,
+        projection.y_resolution / y_factor,
+    )
+    adjusted_bounds = (
+        round(bounds[0] * x_factor),
+        round(bounds[1] * y_factor),
+        round(bounds[0] * x_factor) + array.shape[2],
+        round(bounds[1] * y_factor) + array.shape[1],
+    )
+    return (adjusted_projection, adjusted_bounds)
 
 
 class RasterFormat:
@@ -44,21 +163,39 @@ class RasterFormat:
         raise NotImplementedError
 
     def decode_raster(
-        self, path: UPath, bounds: PixelBounds
-    ) -> npt.NDArray[Any] | None:
+        self,
+        path: UPath,
+        projection: Projection,
+        bounds: PixelBounds,
+        resampling: Resampling = Resampling.bilinear,
+    ) -> npt.NDArray[Any]:
         """Decodes raster data.
 
         Args:
             path: the directory to read from
-            bounds: the bounds of the raster to read
+            projection: the projection to read the raster in.
+            bounds: the bounds to read in the given projection.
+            resampling: resampling method to use in case resampling is needed.
+
+        Returns:
+            the raster data
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def from_config(name: str, config: dict[str, Any]) -> "RasterFormat":
+        """Create a RasterFormat from a config dict.
+
+        Args:
+            name: the name of this format
+            config: the config dict
 
         Returns:
-            the raster data, or None if no image content is found
+            the RasterFormat instance
         """
         raise NotImplementedError
 
 
-@RasterFormats.register("image_tile")
 class ImageTileRasterFormat(RasterFormat):
     """A RasterFormat that stores data in image tiles corresponding to grid cells.
 
@@ -152,6 +289,19 @@ class ImageTileRasterFormat(RasterFormat):
             bounds: the bounds of the raster data in the projection
             array: the raster data (must be CHW)
         """
+        # Write metadata about the projection that we are writing under.
+        # We also save dtype and number of bands so we can return correct shape when
+        # there are no intersecting tiles.
+        with (path / "metadata.json").open("w") as f:
+            json.dump(
+                {
+                    "projection": projection.serialize(),
+                    "dtype": array.dtype.name,
+                    "num_bands": array.shape[0],
+                },
+                f,
+            )
+
         start_tile = (bounds[0] // self.tile_size, bounds[1] // self.tile_size)
         end_tile = (bounds[2] // self.tile_size + 1, bounds[3] // self.tile_size + 1)
         extension = self.get_extension()
@@ -190,17 +340,34 @@ class ImageTileRasterFormat(RasterFormat):
                     self.encode_tile(f, projection, cur_bounds, cur_array)
 
     def decode_raster(
-        self, path: UPath, bounds: PixelBounds
-    ) -> npt.NDArray[Any] | None:
+        self,
+        path: UPath,
+        projection: Projection,
+        bounds: PixelBounds,
+        resampling: Resampling = Resampling.bilinear,
+    ) -> npt.NDArray[Any]:
         """Decodes raster data.
 
         Args:
             path: the directory to read from
-            bounds: the bounds of the raster to read
+            projection: the projection to read the raster in.
+            bounds: the bounds to read in the given projection.
+            resampling: resampling method to use in case resampling is needed.
 
         Returns:
-            the raster data, or None if no image content is found
+            the raster data
         """
+        # Verify that the source data has the same projection as the requested one.
+        # ImageTileRasterFormat currently does not support re-projecting.
+        with (path / "metadata.json").open() as f:
+            image_metadata = json.load(f)
+        source_data_projection = Projection.deserialize(image_metadata["projection"])
+        if source_data_projection != projection:
+            raise NotImplementedError(
+                "not implemented to re-project source data "
+                + f"(source projection {source_data_projection} does not match requested projection {projection})"
+            )
+
         extension = self.get_extension()
 
         # Load tiles one at a time.
@@ -209,7 +376,12 @@ class ImageTileRasterFormat(RasterFormat):
             (bounds[2] - 1) // self.tile_size + 1,
             (bounds[3] - 1) // self.tile_size + 1,
         )
-        dst = None
+        dst_shape = (
+            image_metadata["num_bands"],
+            bounds[3] - bounds[1],
+            bounds[2] - bounds[0],
+        )
+        dst = np.zeros(dst_shape, dtype=image_metadata["dtype"])
         for col in range(start_tile[0], end_tile[0]):
             for row in range(start_tile[1], end_tile[1]):
                 fname = path / f"{col}_{row}.{extension}"
@@ -272,13 +444,17 @@ class ImageTileRasterFormat(RasterFormat):
         )
 
 
-@RasterFormats.register("geotiff")
 class GeotiffRasterFormat(RasterFormat):
     """A raster format that uses one big, tiled GeoTIFF with small block size."""
 
     fname = "geotiff.tif"
 
-    def __init__(self, block_size: int = TILE_SIZE, always_enable_tiling: bool = False):
+    def __init__(
+        self,
+        block_size: int = TILE_SIZE,
+        always_enable_tiling: bool = False,
+        geotiff_options: dict[str, Any] = {},
+    ):
         """Initializes a GeotiffRasterFormat.
 
         Args:
@@ -287,9 +463,11 @@ class GeotiffRasterFormat(RasterFormat):
                 GeoTIFFs. The default is False so that tiling is only used if the size
                 of the GeoTIFF exceeds the block_size on either dimension. If True,
                 then tiling is always enabled (cloud-optimized GeoTIFF).
+            geotiff_options: other options to pass to rasterio.open (for writes).
         """
         self.block_size = block_size
         self.always_enable_tiling = always_enable_tiling
+        self.geotiff_options = geotiff_options
 
     def encode_raster(
         self,
@@ -297,6 +475,7 @@ class GeotiffRasterFormat(RasterFormat):
         projection: Projection,
         bounds: PixelBounds,
         array: npt.NDArray[Any],
+        fname: str | None = None,
     ) -> None:
         """Encodes raster data.
 
@@ -305,7 +484,11 @@ class GeotiffRasterFormat(RasterFormat):
             projection: the projection of the raster data
             bounds: the bounds of the raster data in the projection
             array: the raster data
+            fname: override the filename to save as
         """
+        if fname is None:
+            fname = self.fname
+
         crs = projection.crs
         transform = affine.Affine(
             projection.x_resolution,
@@ -338,76 +521,48 @@ class GeotiffRasterFormat(RasterFormat):
             profile["blockxsize"] = self.block_size
             profile["blockysize"] = self.block_size
 
+        profile.update(self.geotiff_options)
+
         path.mkdir(parents=True, exist_ok=True)
-        with (path / self.fname).open("wb") as f:
-            with rasterio.open(f, "w", **profile) as dst:
-                dst.write(array)
+        logger.debug(f"Writing geotiff to {path / fname}")
+        with open_rasterio_upath_writer(path / fname, **profile) as dst:
+            dst.write(array)
 
     def decode_raster(
-        self, path: UPath, bounds: PixelBounds
-    ) -> npt.NDArray[Any] | None:
+        self,
+        path: UPath,
+        projection: Projection,
+        bounds: PixelBounds,
+        resampling: Resampling = Resampling.bilinear,
+        fname: str | None = None,
+    ) -> npt.NDArray[Any]:
         """Decodes raster data.
 
         Args:
             path: the directory to read from
-            bounds: the bounds of the raster to read
+            projection: the projection to read the raster in.
+            bounds: the bounds to read in the given projection.
+            resampling: resampling method to use in case resampling is needed.
+            fname: override the filename to read from
 
         Returns:
-            the raster data, or None if no image content is found
+            the raster data
         """
-        with (path / self.fname).open("rb") as f:
-            with rasterio.open(f) as src:
-                transform = src.transform
-                x_resolution = transform.a
-                y_resolution = transform.e
-                offset = (
-                    int(transform.c / x_resolution),
-                    int(transform.f / y_resolution),
-                )
-                # bounds is in global pixel coordinates.
-                # We first convert that to pixels relative to top-left of the raster.
-                relative_bounds = [
-                    bounds[0] - offset[0],
-                    bounds[1] - offset[1],
-                    bounds[2] - offset[0],
-                    bounds[3] - offset[1],
-                ]
-                if (
-                    relative_bounds[2] < 0
-                    or relative_bounds[3] < 0
-                    or relative_bounds[0] >= src.width
-                    or relative_bounds[1] >= src.height
-                ):
-                    return None
-                # Now get the actual pixels we will read, which must be contained in
-                # the GeoTIFF.
-                # Padding is (before_x, before_y, after_x, after_y) and will be used to
-                # pad the output back to the originally requested bounds.
-                padding = [0, 0, 0, 0]
-                if relative_bounds[0] < 0:
-                    padding[0] = -relative_bounds[0]
-                    relative_bounds[0] = 0
-                if relative_bounds[1] < 0:
-                    padding[1] = -relative_bounds[1]
-                    relative_bounds[1] = 0
-                if relative_bounds[2] > src.width:
-                    padding[2] = relative_bounds[2] - src.width
-                    relative_bounds[2] = src.width
-                if relative_bounds[3] > src.height:
-                    padding[3] = relative_bounds[3] - src.height
-                    relative_bounds[3] = src.height
-
-                window = rasterio.windows.Window(
-                    relative_bounds[0],
-                    relative_bounds[1],
-                    relative_bounds[2] - relative_bounds[0],
-                    relative_bounds[3] - relative_bounds[1],
-                )
-                array = src.read(window=window)
-                array = np.pad(
-                    array, ((0, 0), (padding[1], padding[3]), (padding[0], padding[2]))
-                )
-                return array
+        if fname is None:
+            fname = self.fname
+
+        # Construct the transform to use for the warped dataset.
+        wanted_transform = get_transform_from_projection_and_bounds(projection, bounds)
+        with open_rasterio_upath_reader(path / fname) as src:
+            with rasterio.vrt.WarpedVRT(
+                src,
+                crs=projection.crs,
+                transform=wanted_transform,
+                width=bounds[2] - bounds[0],
+                height=bounds[3] - bounds[1],
+                resampling=resampling,
+            ) as vrt:
+                return vrt.read()
 
     def get_raster_bounds(self, path: UPath) -> PixelBounds:
         """Returns the bounds of the stored raster.
@@ -418,21 +573,9 @@ class GeotiffRasterFormat(RasterFormat):
         Returns:
             the PixelBounds of the raster
         """
-        with (path / self.fname).open("rb") as f:
-            with rasterio.open(f) as src:
-                transform = src.transform
-                x_resolution = transform.a
-                y_resolution = transform.e
-                offset = (
-                    int(transform.c / x_resolution),
-                    int(transform.f / y_resolution),
-                )
-                return (
-                    offset[0],
-                    offset[1],
-                    offset[0] + src.width,
-                    offset[1] + src.height,
-                )
+        with open_rasterio_upath_reader(path / self.fname) as src:
+            _, bounds = get_raster_projection_and_bounds(src)
+            return bounds
 
     @staticmethod
     def from_config(name: str, config: dict[str, Any]) -> "GeotiffRasterFormat":
@@ -450,10 +593,11 @@ class GeotiffRasterFormat(RasterFormat):
             kwargs["block_size"] = config["block_size"]
         if "always_enable_tiling" in config:
             kwargs["always_enable_tiling"] = config["always_enable_tiling"]
+        if "geotiff_options" in config:
+            kwargs["geotiff_options"] = config["geotiff_options"]
         return GeotiffRasterFormat(**kwargs)
 
 
-@RasterFormats.register("single_image")
 class SingleImageRasterFormat(RasterFormat):
     """A raster format that produces a single image called image.png/jpg.
 
@@ -503,35 +647,57 @@ class SingleImageRasterFormat(RasterFormat):
             if array.shape[2] == 1:
                 array = array[:, :, 0]
             Image.fromarray(array).save(f, format=self.format.upper())
+
+        # Since the image file doesn't include the georeferencing, we store it in an
+        # auxiliary metadata file.
         with (path / "metadata.json").open("w") as f:
             json.dump(
                 {
+                    "projection": projection.serialize(),
                     "bounds": bounds,
                 },
                 f,
             )
 
     def decode_raster(
-        self, path: UPath, bounds: PixelBounds
-    ) -> npt.NDArray[Any] | None:
+        self,
+        path: UPath,
+        projection: Projection,
+        bounds: PixelBounds,
+        resampling: Resampling = Resampling.bilinear,
+    ) -> npt.NDArray[Any]:
         """Decodes raster data.
 
         Args:
             path: the directory to read from
-            bounds: the bounds of the raster to read
+            projection: the projection to read the raster in.
+            bounds: the bounds to read in the given projection.
+            resampling: resampling method to use in case resampling is needed.
 
         Returns:
-            the raster data, or None if no image content is found
+            the raster data
         """
-        image_fname = path / ("image." + self.get_extension())
+        # Try to get the bounds of the saved image from the metadata file.
+        # In old versions, the file may be missing the projection key.
         metadata_fname = path / "metadata.json"
-        if metadata_fname.exists():
-            with metadata_fname.open() as f:
-                image_bounds = json.load(f)["bounds"]
-        else:
-            # Backwards compatibility -- assume that requested bounds matches the window bounds.
-            image_bounds = bounds
+        with metadata_fname.open() as f:
+            image_metadata = json.load(f)
+
+        image_bounds = image_metadata["bounds"]
 
+        # If the projection key is set, verify that it matches the requested projection
+        # since SingleImageRasterFormat currently does not support re-projecting.
+        if "projection" in image_metadata:
+            source_data_projection = Projection.deserialize(
+                image_metadata["projection"]
+            )
+            if projection != source_data_projection:
+                raise NotImplementedError(
+                    "not implemented to re-project source data "
+                    + f"(source projection {source_data_projection} does not match requested projection {projection})"
+                )
+
+        image_fname = path / ("image." + self.get_extension())
         with image_fname.open("rb") as f:
             array = np.array(Image.open(f, formats=[self.format.upper()]))
 
@@ -583,17 +749,3 @@ class SingleImageRasterFormat(RasterFormat):
         if "format" in config:
             kwargs["format"] = config["format"]
         return SingleImageRasterFormat(**kwargs)
-
-
-def load_raster_format(config: RasterFormatConfig) -> RasterFormat:
-    """Loads a RasterFormat from a RasterFormatConfig.
-
-    Args:
-        config: the RasterFormatConfig configuration object specifying the
-            RasterFormat.
-
-    Returns:
-        the loaded RasterFormat implementation
-    """
-    cls = RasterFormats.get_class(config.name)
-    return cls.from_config(config.name, config.config_dict)

rslearn/utils/rtree_index.py

@@ -1,7 +1,9 @@
 """RtreeIndex spatial index implementation."""
 
+import hashlib
 import os
 import shutil
+import tempfile
 from collections.abc import Callable
 from typing import Any
 
@@ -9,8 +11,11 @@ import fsspec
 from rtree import index
 from upath import UPath
 
+from rslearn.log_utils import get_logger
 from rslearn.utils.spatial_index import SpatialIndex
 
+logger = get_logger(__name__)
+
 
 class RtreeIndex(SpatialIndex):
     """An index of spatiotemporal geometries backed by an rtree index.
@@ -18,7 +23,7 @@ class RtreeIndex(SpatialIndex):
     Both in-memory and on-disk options are supported.
     """
 
-    def __init__(self, fname: str | None = None):
+    def __init__(self, fname: str | None = None) -> None:
         """Initialize a new RtreeIndex.
 
         If fname is set, the index is persisted on disk, otherwise it is in-memory.
@@ -50,6 +55,7 @@ class RtreeIndex(SpatialIndex):
         self.counter += 1
         self.index.insert(id=self.counter, coordinates=box, obj=data)
 
+    # TODO: Make a named tuple for all the bounding box stuff
     def query(self, box: tuple[float, float, float, float]) -> list[Any]:
         """Query the index for objects intersecting a box.
 
@@ -63,20 +69,51 @@ class RtreeIndex(SpatialIndex):
         return [r.object for r in results]
 
 
+def delete_partially_created_local_files(fname: str) -> None:
+    """Delete partially created .dat and .idx files."""
+    extensions = [".dat", ".idx"]
+    for ext in extensions:
+        cur_fname = fname + ext
+        if os.path.exists(cur_fname):
+            os.unlink(cur_fname)
+
+
+def _get_tmp_dir_for_cached_rtree(cache_dir: UPath) -> str:
+    """Get a local temporary directory to store the rtree from the specified cache_dir.
+
+    This function is deterministic so the same cache_dir will always yield the same
+    local temporary directory.
+
+    Note that the directory is not cleaned up after the program exits, so the rtree
+    will remain there. This is because this function may be called from multiple worker
+    processes but the index should be reused across workers.
+
+    Args:
+        cache_dir: the non-local directory where the rtree files are stored.
+
+    Returns:
+        the temporary local directory to copy the cached rtree to.
+    """
+    cache_id = hashlib.sha256(str(cache_dir).encode()).hexdigest()
+    tmp_dir = os.path.join(
+        tempfile.gettempdir(), "rslearn_cache", "rtree_index", cache_id
+    )
+    os.makedirs(tmp_dir, exist_ok=True)
+    return tmp_dir
+
+
 def get_cached_rtree(
-    cache_dir: UPath, tmp_dir: str, build_fn: Callable[[RtreeIndex], None]
+    cache_dir: UPath, build_fn: Callable[[RtreeIndex], None]
 ) -> RtreeIndex:
     """Returns an RtreeIndex cached in cache_dir, creating it if needed.
 
     The .dat and .idx files are cached in cache_dir. Since RtreeIndex expects local
-    filesystem, it is copied to a local temporary directory if needed. If the index
-    doesn't exist yet, then it is created using build_fn.
+    filesystem, it is copied to a local temporary directory if needed (it is not needed
+    if the cache_dir is already on local filesystem). If the index doesn't exist yet,
+    then it is created using build_fn.
 
     Args:
         cache_dir: directory to cache the index files.
-        tmp_dir: temporary local directory to use in case cache_dir is on a remote
-            filesystem. The caller is responsible for cleaning this up when they don't
-            need the index anymore.
         build_fn: function to build the index in case it doesn't exist yet.
 
     Returns:
@@ -95,14 +132,13 @@ def get_cached_rtree(
         if is_local_cache:
             local_fname = (cache_dir / "rtree_index").path
         else:
+            tmp_dir = _get_tmp_dir_for_cached_rtree(cache_dir)
             local_fname = os.path.join(tmp_dir, "rtree_index")
+        delete_partially_created_local_files(local_fname)
 
-        # Delete any local files that might be partially created.
-        for ext in extensions:
-            cur_fname = local_fname + ext
-            if os.path.exists(cur_fname):
-                os.unlink(cur_fname)
-
+        logger.info(
+            "building rtree index at %s to be cached at %s", local_fname, cache_dir
+        )
         rtree_index = RtreeIndex(local_fname)
         build_fn(rtree_index)
         del rtree_index
@@ -115,6 +151,7 @@ def get_cached_rtree(
 
         # Create the completed file to indicate index is ready in cache.
         completed_fname.touch()
+        logger.info("rtree index is built and ready")
 
     else:
         # Initialize the index from the cached version.
@@ -122,10 +159,20 @@ def get_cached_rtree(
         if is_local_cache:
             local_fname = (cache_dir / "rtree_index").path
         else:
+            tmp_dir = _get_tmp_dir_for_cached_rtree(cache_dir)
             local_fname = os.path.join(tmp_dir, "rtree_index")
-            for ext in extensions:
-                with (cache_dir / f"rtree_index{ext}").open("rb") as src:
-                    with open(local_fname + ext, "wb") as dst:
-                        shutil.copyfileobj(src, dst)
+
+            if not os.path.exists(local_fname + extensions[-1]):
+                logger.info(
+                    "copying rtree index from non-local cache at %s to local temporary directory %s",
+                    cache_dir,
+                    local_fname,
+                )
+                for ext in extensions:
+                    with (cache_dir / f"rtree_index{ext}").open("rb") as src:
+                        with open(local_fname + ext + ".tmp", "wb") as dst:
+                            shutil.copyfileobj(src, dst)
+                    os.rename(local_fname + ext + ".tmp", local_fname + ext)
+                logger.info("rtree index is ready")
 
     return RtreeIndex(local_fname)