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

rslearn/const.py

@@ -1,23 +1,17 @@
 """Constants."""
 
-from rasterio.crs import CRS
-
-from rslearn.utils import PixelBounds, Projection
-
-WGS84_EPSG = 4326
-"""The EPSG code for WGS-84."""
-
-WGS84_PROJECTION = Projection(CRS.from_epsg(WGS84_EPSG), 1, 1)
-"""The Projection for WGS-84 assuming 1 degree per pixel.
-
-This can be used to create STGeometry with shapes in longitude/latitude coordinates.
-"""
-
-WGS84_BOUNDS: PixelBounds = (-180, -90, 180, 90)
-"""The bounds of the WGS-84 projection."""
+from rslearn.utils.geometry import WGS84_BOUNDS, WGS84_EPSG, WGS84_PROJECTION
 
 TILE_SIZE = 512
 """Default tile size. TODO: remove this or move it elsewhere."""
 
 SHAPEFILE_AUX_EXTENSIONS = [".cpg", ".dbf", ".prj", ".sbn", ".sbx", ".shx", ".txt"]
 """Extensions of potential auxiliary files to .shp file."""
+
+__all__ = (
+    "WGS84_PROJECTION",
+    "WGS84_EPSG",
+    "WGS84_BOUNDS",
+    "TILE_SIZE",
+    "SHAPEFILE_AUX_EXTENSIONS",
+)

rslearn/data_sources/__init__.py

@@ -10,32 +10,17 @@ Each source supports operations to lookup items that match with spatiotemporal
 geometries, and ingest those items.
 """
 
-import importlib
-
-from upath import UPath
-
-from rslearn.config import LayerConfig
-
-from .data_source import DataSource, Item, ItemLookupDataSource, RetrieveItemDataSource
-
-
-def data_source_from_config(config: LayerConfig, ds_path: UPath) -> DataSource:
-    """Loads a data source from config dict.
-
-    Args:
-        config: the LayerConfig containing this data source.
-        ds_path: the dataset root directory.
-    """
-    name = config.data_source.name
-    module_name = ".".join(name.split(".")[:-1])
-    class_name = name.split(".")[-1]
-    module = importlib.import_module(module_name)
-    class_ = getattr(module, class_name)
-    return class_.from_config(config, ds_path)
-
+from .data_source import (
+    DataSource,
+    DataSourceContext,
+    Item,
+    ItemLookupDataSource,
+    RetrieveItemDataSource,
+)
 
 __all__ = (
     "DataSource",
+    "DataSourceContext",
     "Item",
     "ItemLookupDataSource",
     "RetrieveItemDataSource",

rslearn/data_sources/aws_landsat.py

@@ -2,33 +2,41 @@
 
 import io
 import json
+import os
 import shutil
+import tempfile
 import urllib.request
 import zipfile
 from collections.abc import Generator
-from datetime import timedelta
+from datetime import datetime
 from typing import Any, BinaryIO
 
+import affine
 import boto3
 import dateutil.parser
 import fiona
 import fiona.transform
-import pytimeparse
+import numpy.typing as npt
 import rasterio
 import shapely
+import shapely.geometry
 import tqdm
+from rasterio.enums import Resampling
 from upath import UPath
 
 import rslearn.data_sources.utils
-import rslearn.utils.mgrs
-from rslearn.config import LayerConfig, RasterLayerConfig
+from rslearn.config import LayerConfig
 from rslearn.const import SHAPEFILE_AUX_EXTENSIONS, WGS84_PROJECTION
-from rslearn.tile_stores import PrefixedTileStore, TileStore
-from rslearn.utils import STGeometry
+from rslearn.dataset import Window
+from rslearn.dataset.materialize import RasterMaterializer
+from rslearn.tile_stores import TileStore, TileStoreWithLayer
 from rslearn.utils.fsspec import get_upath_local, join_upath, open_atomic
+from rslearn.utils.geometry import PixelBounds, Projection, STGeometry
+from rslearn.utils.grid_index import GridIndex
 
-from .data_source import DataSource, Item, QueryConfig
-from .raster_source import get_needed_projections, ingest_raster
+from .data_source import DataSource, DataSourceContext, Item, QueryConfig
+
+WRS2_GRID_SIZE = 1.0
 
 
 class LandsatOliTirsItem(Item):
@@ -36,7 +44,7 @@ class LandsatOliTirsItem(Item):
 
     def __init__(
         self, name: str, geometry: STGeometry, blob_path: str, cloud_cover: float
-    ):
+    ) -> None:
         """Creates a new LandsatOliTirsItem.
 
         Args:
@@ -58,7 +66,7 @@ class LandsatOliTirsItem(Item):
         return d
 
     @staticmethod
-    def deserialize(d: dict) -> Item:
+    def deserialize(d: dict) -> "LandsatOliTirsItem":
         """Deserializes an item from a JSON-decoded dictionary."""
         if "name" not in d:
             d["name"] = d["blob_path"].split("/")[-1].split(".tif")[0]
@@ -71,7 +79,7 @@ class LandsatOliTirsItem(Item):
         )
 
 
-class LandsatOliTirs(DataSource):
+class LandsatOliTirs(DataSource, TileStore):
     """A data source for Landsat 8/9 OLI-TIRS imagery on AWS.
 
     Specifically, uses the usgs-landsat S3 bucket maintained by USGS. The data includes
@@ -90,53 +98,37 @@ class LandsatOliTirs(DataSource):
 
     def __init__(
         self,
-        config: LayerConfig,
-        metadata_cache_dir: UPath,
-        max_time_delta: timedelta = timedelta(days=30),
+        metadata_cache_dir: str,
         sort_by: str | None = None,
+        context: DataSourceContext = DataSourceContext(),
     ) -> None:
         """Initialize a new LandsatOliTirs instance.
 
         Args:
-            config: configuration of this layer
-            metadata_cache_dir: directory to cache product metadata files.
-            max_time_delta: maximum time before a query start time or after a
-                query end time to look for products. This is required due to the large
-                number of available products, and defaults to 30 days.
+            metadata_cache_dir: directory to cache produtc metadata files.
             sort_by: can be "cloud_cover", default arbitrary order; only has effect for
                 SpaceMode.WITHIN.
+            context: the data source context.
         """
-        self.config = config
-        self.metadata_cache_dir = metadata_cache_dir
-        self.max_time_delta = max_time_delta
+        # If context is provided, we join the directory with the dataset path,
+        # otherwise we treat it directly as UPath.
+        if context.ds_path is not None:
+            self.metadata_cache_dir = join_upath(context.ds_path, metadata_cache_dir)
+        else:
+            self.metadata_cache_dir = UPath(metadata_cache_dir)
+
         self.sort_by = sort_by
 
+        self.client = boto3.client("s3")
         self.bucket = boto3.resource("s3").Bucket(self.bucket_name)
-
         self.metadata_cache_dir.mkdir(parents=True, exist_ok=True)
 
-    @staticmethod
-    def from_config(config: LayerConfig, ds_path: UPath) -> "LandsatOliTirs":
-        """Creates a new LandsatOliTirs instance from a configuration dictionary."""
-        assert isinstance(config, RasterLayerConfig)
-        d = config.data_source.config_dict
-        kwargs = dict(
-            config=config,
-            metadata_cache_dir=join_upath(ds_path, d["metadata_cache_dir"]),
-        )
-        if "max_time_delta" in d:
-            kwargs["max_time_delta"] = timedelta(
-                seconds=pytimeparse.parse(d["max_time_delta"])
-            )
-        if "sort_by" in d:
-            kwargs["sort_by"] = d["sort_by"]
-
-        return LandsatOliTirs(**kwargs)
+        self.wrs2_index: GridIndex | None = None
 
     def _read_products(
         self, needed_year_pathrows: set[tuple[int, str, str]]
     ) -> Generator[LandsatOliTirsItem, None, None]:
-        """Read _MTL.json files and yield relevant LandsatOliTirsItems.
+        """Read _stac.json files and yield relevant LandsatOliTirsItems.
 
         Args:
             needed_year_pathrows: set of (year, path, row) where we need to search for
@@ -155,7 +147,10 @@ class LandsatOliTirs(DataSource):
                 for obj in self.bucket.objects.filter(
                     Prefix=prefix, RequestPayer="requester"
                 ):
-                    if not obj.key.endswith("_MTL.json"):
+                    # Only read the _stac.json files.
+                    # Previously we used _MTL.json but those files don't have the full
+                    # geometry of the Landsat scene, only the bounding box.
+                    if not obj.key.endswith("_stac.json"):
                         continue
                     # Load JSON data.
                     buf = io.BytesIO()
@@ -163,33 +158,32 @@ class LandsatOliTirs(DataSource):
                         obj.key, buf, ExtraArgs={"RequestPayer": "requester"}
                     )
                     buf.seek(0)
-                    product = json.load(buf)
-                    metadata = product["LANDSAT_METADATA_FILE"]
-                    image_attributes = metadata["IMAGE_ATTRIBUTES"]
-                    projection_attributes = metadata["PROJECTION_ATTRIBUTES"]
+                    stac_data = json.load(buf)
 
                     # Get polygon coordinates.
-                    coordinates = []
-                    for corner_id in ["UL", "UR", "LR", "LL"]:
-                        lon = projection_attributes[f"CORNER_{corner_id}_LON_PRODUCT"]
-                        lat = projection_attributes[f"CORNER_{corner_id}_LAT_PRODUCT"]
-                        coordinates.append((lon, lat))
+                    shp = shapely.geometry.shape(stac_data["geometry"])
 
                     # Get datetime.
-                    date_str = image_attributes["DATE_ACQUIRED"]
-                    time_str = image_attributes["SCENE_CENTER_TIME"]
-                    ts = dateutil.parser.isoparse(date_str + "T" + time_str)
-
-                    blob_path = obj.key.split("MTL.json")[0]
-                    geometry = STGeometry(
-                        WGS84_PROJECTION, shapely.Polygon(coordinates), [ts, ts]
-                    )
+                    ts = dateutil.parser.isoparse(stac_data["properties"]["datetime"])
+
+                    blob_path = obj.key.split("stac.json")[0]
+                    time_range: tuple[datetime, datetime] = (ts, ts)
+                    geometry = STGeometry(WGS84_PROJECTION, shp, time_range)
+                    cloud_cover: float
+                    if "eo:cloud_cover" in stac_data["properties"]:
+                        cloud_cover = stac_data["properties"]["eo:cloud_cover"]
+                    elif "landsat:cloud_cover_land" in stac_data["properties"]:
+                        cloud_cover = stac_data["properties"][
+                            "landsat:cloud_cover_land"
+                        ]
+                    else:
+                        cloud_cover = -1
                     items.append(
                         LandsatOliTirsItem(
-                            name=metadata["PRODUCT_CONTENTS"]["LANDSAT_PRODUCT_ID"],
+                            name=stac_data["id"],
                             geometry=geometry,
                             blob_path=blob_path,
-                            cloud_cover=image_attributes["CLOUD_COVER"],
+                            cloud_cover=cloud_cover,
                         )
                     )
 
@@ -205,7 +199,7 @@ class LandsatOliTirs(DataSource):
 
             yield from items
 
-    def get_wrs2_polygons(self) -> list[tuple[shapely.Geometry, str, str]]:
+    def _get_wrs2_polygons(self) -> list[tuple[shapely.Geometry, str, str]]:
         """Get polygons for each (path, row) in the WRS2 grid.
 
         Returns:
@@ -216,6 +210,7 @@ class LandsatOliTirs(DataSource):
         if not shp_fname.exists():
             # Download and extract zip to cache dir.
             zip_fname = self.metadata_cache_dir / f"{prefix}.zip"
+            print(f"Downloading {self.wrs2_url} to {zip_fname}")
             with urllib.request.urlopen(self.wrs2_url) as response:
                 with zip_fname.open("wb") as f:
                     shutil.copyfileobj(response, f)
@@ -257,9 +252,22 @@ class LandsatOliTirs(DataSource):
                     polygons.append((shp, path, row))
                 return polygons
 
+    def _get_wrs2_index(self) -> GridIndex:
+        """Get a grid index over the WRS2 polygons."""
+        if self.wrs2_index is not None:
+            return self.wrs2_index
+
+        # Index doesn't exist so we need to build it.
+        # We cache it with the object since it takes a bit of time to create it.
+        polygons = self._get_wrs2_polygons()
+        self.wrs2_index = GridIndex(WRS2_GRID_SIZE)
+        for polygon, path, row in polygons:
+            self.wrs2_index.insert(polygon.bounds, (polygon, path, row))
+        return self.wrs2_index
+
     def get_items(
         self, geometries: list[STGeometry], query_config: QueryConfig
-    ) -> list[list[list[Item]]]:
+    ) -> list[list[list[LandsatOliTirsItem]]]:
         """Get a list of items in the data source intersecting the given geometries.
 
         Args:
@@ -269,7 +277,7 @@ class LandsatOliTirs(DataSource):
         Returns:
             List of groups of items that should be retrieved for each geometry.
         """
-        wrs2_polygons = self.get_wrs2_polygons()
+        wrs2_index = self._get_wrs2_index()
         needed_year_pathrows = set()
         wgs84_geometries = [
             geometry.to_projection(WGS84_PROJECTION) for geometry in geometries
@@ -280,13 +288,13 @@ class LandsatOliTirs(DataSource):
                     "Landsat on AWS requires geometry time ranges to be set"
                 )
             cur_pathrows = set()
-            for polygon, path, row in wrs2_polygons:
+            for polygon, path, row in wrs2_index.query(wgs84_geometry.shp.bounds):
                 if wgs84_geometry.shp.intersects(polygon):
                     cur_pathrows.add((path, row))
             for path, row in cur_pathrows:
                 for year in range(
-                    (wgs84_geometry.time_range[0] - self.max_time_delta).year,
-                    (wgs84_geometry.time_range[1] + self.max_time_delta).year + 1,
+                    wgs84_geometry.time_range[0].year,
+                    wgs84_geometry.time_range[1].year + 1,
                 ):
                     needed_year_pathrows.add((year, path, row))
 
@@ -301,18 +309,22 @@ class LandsatOliTirs(DataSource):
                 cur_items.append(item)
 
             if self.sort_by == "cloud_cover":
-                items.sort(key=lambda item: item.cloud_cover)
+                cur_items.sort(
+                    key=lambda item: item.cloud_cover if item.cloud_cover >= 0 else 100
+                )
             elif self.sort_by is not None:
                 raise ValueError(f"invalid sort_by setting ({self.sort_by})")
 
-            cur_groups = rslearn.data_sources.utils.match_candidate_items_to_window(
-                geometry, cur_items, query_config
+            cur_groups: list[list[LandsatOliTirsItem]] = (
+                rslearn.data_sources.utils.match_candidate_items_to_window(
+                    geometry, cur_items, query_config
+                )
             )
             groups.append(cur_groups)
 
         return groups
 
-    def get_item_by_name(self, name: str) -> Item:
+    def get_item_by_name(self, name: str) -> LandsatOliTirsItem:
         """Gets an item by name."""
         # Product name is like LC08_L1TP_046027_20230715_20230724_02_T1.
         # We want to use _read_products so we need to extract:
@@ -330,12 +342,14 @@ class LandsatOliTirs(DataSource):
                 return item
         raise ValueError(f"item {name} not found")
 
-    def deserialize_item(self, serialized_item: Any) -> Item:
+    def deserialize_item(self, serialized_item: Any) -> LandsatOliTirsItem:
         """Deserializes an item from JSON-decoded data."""
         assert isinstance(serialized_item, dict)
         return LandsatOliTirsItem.deserialize(serialized_item)
 
-    def retrieve_item(self, item: Item) -> Generator[tuple[str, BinaryIO], None, None]:
+    def retrieve_item(
+        self, item: LandsatOliTirsItem
+    ) -> Generator[tuple[str, BinaryIO], None, None]:
         """Retrieves the rasters corresponding to an item as file streams."""
         for band in self.bands:
             buf = io.BytesIO()
@@ -350,8 +364,8 @@ class LandsatOliTirs(DataSource):
 
     def ingest(
         self,
-        tile_store: TileStore,
-        items: list[Item],
+        tile_store: TileStoreWithLayer,
+        items: list[LandsatOliTirsItem],
         geometries: list[list[STGeometry]],
     ) -> None:
         """Ingest items into the given tile store.
@@ -364,28 +378,158 @@ class LandsatOliTirs(DataSource):
         for item, cur_geometries in zip(items, geometries):
             for band in self.bands:
                 band_names = [band]
-                cur_tile_store = PrefixedTileStore(
-                    tile_store, (item.name, "_".join(band_names))
-                )
-                needed_projections = get_needed_projections(
-                    cur_tile_store, band_names, self.config.band_sets, cur_geometries
-                )
-                if not needed_projections:
+                if tile_store.is_raster_ready(item.name, band_names):
                     continue
 
-                buf = io.BytesIO()
-                self.bucket.download_fileobj(
-                    item.blob_path + f"{band}.TIF",
-                    buf,
-                    ExtraArgs={"RequestPayer": "requester"},
-                )
-                buf.seek(0)
-                with rasterio.open(buf) as raster:
-                    for projection in needed_projections:
-                        ingest_raster(
-                            tile_store=cur_tile_store,
-                            raster=raster,
-                            projection=projection,
-                            time_range=item.geometry.time_range,
-                            layer_config=self.config,
-                        )
+                with tempfile.TemporaryDirectory() as tmp_dir:
+                    fname = os.path.join(tmp_dir, f"{band}.tif")
+                    self.bucket.download_file(
+                        item.blob_path + f"{band}.TIF",
+                        fname,
+                        ExtraArgs={"RequestPayer": "requester"},
+                    )
+                    tile_store.write_raster_file(item.name, band_names, UPath(fname))
+
+    # The functions below are to emulate TileStore functionality so we can easily
+    # support materialization directly from the COGs.
+    def is_raster_ready(
+        self, layer_name: str, item_name: str, bands: list[str]
+    ) -> bool:
+        """Checks if this raster has been written to the store.
+
+        Args:
+            layer_name: the layer name or alias.
+            item_name: 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 access it on AWS bucket.
+        return True
+
+    def get_raster_bands(self, layer_name: str, item_name: str) -> 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_name: 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 [[band] for band in self.bands]
+
+    def get_raster_bounds(
+        self, layer_name: str, item_name: str, 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_name: 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.
+        """
+        item = self.get_item_by_name(item_name)
+        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 read_raster(
+        self,
+        layer_name: str,
+        item_name: str,
+        bands: list[str],
+        projection: Projection,
+        bounds: PixelBounds,
+        resampling: Resampling = Resampling.bilinear,
+    ) -> npt.NDArray[Any]:
+        """Read raster data from the store.
+
+        Args:
+            layer_name: the layer name or alias.
+            item_name: 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
+        """
+        # Landsat assets have single band per asset.
+        assert len(bands) == 1
+        band = bands[0]
+
+        # Get the item since it has the blob path.
+        item = self.get_item_by_name(item_name)
+
+        # Create pre-signed URL for rasterio access.
+        # We do this because accessing via URL is much faster since rasterio can use
+        # the URL directly.
+        blob_key = item.blob_path + f"{band}.TIF"
+        url = self.client.generate_presigned_url(
+            "get_object",
+            Params={
+                "Bucket": self.bucket_name,
+                "Key": blob_key,
+                "RequestPayer": "requester",
+            },
+        )
+
+        # Construct the transform to use for the warped dataset.
+        wanted_transform = affine.Affine(
+            projection.x_resolution,
+            0,
+            bounds[0] * projection.x_resolution,
+            0,
+            projection.y_resolution,
+            bounds[1] * projection.y_resolution,
+        )
+
+        with rasterio.open(url) as src:
+            with rasterio.vrt.WarpedVRT(
+                src,
+                crs=projection.crs,
+                transform=wanted_transform,
+                width=bounds[2] - bounds[0],
+                height=bounds[3] - bounds[1],
+                resampling=resampling,
+            ) as vrt:
+                return vrt.read()
+
+    def materialize(
+        self,
+        window: Window,
+        item_groups: list[list[LandsatOliTirsItem]],
+        layer_name: str,
+        layer_cfg: LayerConfig,
+    ) -> 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
+        """
+        RasterMaterializer().materialize(
+            TileStoreWithLayer(self, layer_name),
+            window,
+            layer_name,
+            layer_cfg,
+            item_groups,
+        )