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

rslearn/data_sources/planet_basemap.py

@@ -0,0 +1,303 @@
+"""Data source for Planet Labs Basemaps API."""
+
+import os
+import tempfile
+from datetime import datetime
+from typing import Any
+
+import rasterio
+import requests
+import shapely
+from upath import UPath
+
+from rslearn.config import LayerConfig, QueryConfig, RasterLayerConfig
+from rslearn.const import WGS84_PROJECTION
+from rslearn.data_sources import DataSource, Item
+from rslearn.data_sources.utils import match_candidate_items_to_window
+from rslearn.log_utils import get_logger
+from rslearn.tile_stores import PrefixedTileStore, TileStore
+from rslearn.utils import STGeometry
+
+from .raster_source import get_needed_projections, ingest_raster
+
+logger = get_logger(__name__)
+
+
+class PlanetItem(Item):
+    """An item referencing a particular mosaic and quad in Basemaps API."""
+
+    def __init__(self, name: str, geometry: STGeometry, mosaic_id: str, quad_id: str):
+        """Create a new PlanetItem.
+
+        Args:
+            name: the item name (combination of mosaic and quad ID).
+            geometry: the geometry associated with this quad.
+            mosaic_id: the mosaic ID in API
+            quad_id: the quad ID in API
+        """
+        super().__init__(name, geometry)
+        self.mosaic_id = mosaic_id
+        self.quad_id = quad_id
+
+    def serialize(self) -> dict:
+        """Serializes the item to a JSON-encodable dictionary."""
+        d = super().serialize()
+        d["mosaic_id"] = self.mosaic_id
+        d["quad_id"] = self.quad_id
+        return d
+
+    @staticmethod
+    def deserialize(d: dict) -> Item:
+        """Deserializes an item from a JSON-decoded dictionary."""
+        item = super(PlanetItem, PlanetItem).deserialize(d)
+        return PlanetItem(
+            name=item.name,
+            geometry=item.geometry,
+            mosaic_id=d["mosaic_id"],
+            quad_id=d["quad_id"],
+        )
+
+
+class ApiError(Exception):
+    """An error from Planet Labs API."""
+
+    pass
+
+
+class PlanetBasemap(DataSource):
+    """A data source for Planet Labs Basemaps API."""
+
+    api_url = "https://api.planet.com/basemaps/v1/"
+
+    def __init__(
+        self,
+        config: RasterLayerConfig,
+        series_id: str,
+        bands: list[str],
+        api_key: str | None = None,
+    ):
+        """Initialize a new Planet instance.
+
+        Args:
+            config: the LayerConfig of the layer containing this data source
+            series_id: the series of mosaics to use.
+            bands: list of band names to use.
+            api_key: optional Planet API key (it can also be provided via PL_API_KEY
+                environmnet variable).
+        """
+        self.config = config
+        self.bands = bands
+
+        self.session = requests.Session()
+        if api_key is None:
+            api_key = os.environ["PL_API_KEY"]
+        self.session.auth = (api_key, "")
+
+        # List mosaics.
+        self.mosaics = {}
+        for mosaic_dict in self._api_get_paginate(
+            path=f"series/{series_id}/mosaics", list_key="mosaics"
+        ):
+            shp = shapely.box(*mosaic_dict["bbox"])
+            time_range = (
+                datetime.fromisoformat(mosaic_dict["first_acquired"]),
+                datetime.fromisoformat(mosaic_dict["last_acquired"]),
+            )
+            geom = STGeometry(WGS84_PROJECTION, shp, time_range)
+            self.mosaics[mosaic_dict["id"]] = geom
+
+    @staticmethod
+    def from_config(config: LayerConfig, ds_path: UPath) -> "PlanetBasemap":
+        """Creates a new PlanetBasemap instance from a configuration dictionary."""
+        assert isinstance(config, RasterLayerConfig)
+        if config.data_source is None:
+            raise ValueError("data_source is required")
+        d = config.data_source.config_dict
+        kwargs = dict(
+            config=config,
+            series_id=d["series_id"],
+            bands=d["bands"],
+        )
+        optional_keys = [
+            "api_key",
+        ]
+        for optional_key in optional_keys:
+            if optional_key in d:
+                kwargs[optional_key] = d[optional_key]
+        return PlanetBasemap(**kwargs)
+
+    def _api_get(
+        self,
+        path: str | None = None,
+        url: str | None = None,
+        query_args: dict[str, str] | None = None,
+    ) -> list[Any] | dict[str, Any]:
+        """Perform a GET request on the API.
+
+        Args:
+            path: the path to GET, like "series".
+            url: the full URL to GET. Only one of path or url should be set.
+            query_args: optional params to include with the request.
+
+        Returns:
+            the JSON response data.
+
+        Raises:
+            ApiError: if the API returned an error response.
+        """
+        if path is None and url is None:
+            raise ValueError("Only one of path or url should be set")
+        if query_args:
+            kwargs = dict(params=query_args)
+        else:
+            kwargs = {}
+
+        if path:
+            url = self.api_url + path
+        if url is None:
+            raise ValueError("url is required")
+        response = self.session.get(url, **kwargs)  # type: ignore
+
+        if response.status_code != 200:
+            raise ApiError(
+                f"{url}: got status code {response.status_code}: {response.text}"
+            )
+        return response.json()
+
+    def _api_get_paginate(
+        self, path: str, list_key: str, query_args: dict[str, str] | None = None
+    ) -> list:
+        """Get all items in a paginated response.
+
+        Args:
+            path: the path to GET.
+            list_key: the key in the response containing the list that should be
+                concatenated across all available pages.
+            query_args: optional params to include with the requests.
+
+        Returns:
+            the concatenated list of items.
+
+        Raises:
+            ApiError if the API returned an error response.
+        """
+        next_url = self.api_url + path
+        items = []
+        while True:
+            json_data = self._api_get(url=next_url, query_args=query_args)
+            if not isinstance(json_data, dict):
+                logger.warning(f"Expected dict, got {type(json_data)}")
+                continue
+            items += json_data[list_key]
+
+            if "_next" in json_data["_links"]:
+                next_url = json_data["_links"]["_next"]
+            else:
+                return items
+
+    def get_items(
+        self, geometries: list[STGeometry], query_config: QueryConfig
+    ) -> list[list[list[PlanetItem]]]:
+        """Get a list of items in the data source intersecting the given geometries.
+
+        Args:
+            geometries: the spatiotemporal geometries
+            query_config: the query configuration
+
+        Returns:
+            List of groups of items that should be retrieved for each geometry.
+        """
+        groups = []
+        for geometry in geometries:
+            geom_bbox = geometry.to_projection(WGS84_PROJECTION).shp.bounds
+            geom_bbox_str = ",".join([str(value) for value in geom_bbox])
+
+            # Find the relevant mosaics that the geometry intersects.
+            # For each relevant mosaic, identify the intersecting quads.
+            items = []
+            for mosaic_id, mosaic_geom in self.mosaics.items():
+                if not geometry.intersects(mosaic_geom):
+                    continue
+                logger.info(f"found mosaic {mosaic_geom} for geom {geometry}")
+                # List all quads that intersect the current geometry's
+                # longitude/latitude bbox in this mosaic.
+                for quad_dict in self._api_get_paginate(
+                    path=f"mosaics/{mosaic_id}/quads",
+                    list_key="items",
+                    query_args={"bbox": geom_bbox_str},
+                ):
+                    logger.info(f"found quad {quad_dict}")
+                    shp = shapely.box(*quad_dict["bbox"])
+                    geom = STGeometry(WGS84_PROJECTION, shp, mosaic_geom.time_range)
+                    quad_id = quad_dict["id"]
+                    items.append(
+                        PlanetItem(f"{mosaic_id}_{quad_id}", geom, mosaic_id, quad_id)
+                    )
+            logger.info(f"found {len(items)} items for geom {geometry}")
+            cur_groups = match_candidate_items_to_window(geometry, items, query_config)
+            groups.append(cur_groups)
+
+        return groups
+
+    def deserialize_item(self, serialized_item: Any) -> Item:
+        """Deserializes an item from JSON-decoded data."""
+        assert isinstance(serialized_item, dict)
+        return PlanetItem.deserialize(serialized_item)
+
+    def ingest(
+        self,
+        tile_store: TileStore,
+        items: list[Item],
+        geometries: list[list[STGeometry]],
+    ) -> None:
+        """Ingest items into the given tile store.
+
+        Args:
+            tile_store: the tile store to ingest into
+            items: the items to ingest
+            geometries: a list of geometries needed for each item
+        """
+        for item, cur_geometries in zip(items, geometries):
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                band_names = self.bands
+                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:
+                    continue
+
+                assert isinstance(item, PlanetItem)
+                download_url = (
+                    self.api_url + f"mosaics/{item.mosaic_id}/quads/{item.quad_id}/full"
+                )
+                response = self.session.get(
+                    download_url, allow_redirects=True, stream=True
+                )
+                if response.status_code != 200:
+                    # # temporary skip for now
+                    # logger.error(
+                    #     f"{download_url}: got status code {response.status_code}: {response.text}"
+                    # )
+                    # continue
+                    raise ApiError(
+                        f"{download_url}: got status code {response.status_code}: {response.text}"
+                    )
+
+                with tempfile.TemporaryDirectory() as tmp_dir:
+                    local_fname = os.path.join(tmp_dir, "temp.tif")
+                    with open(local_fname, "wb") as f:
+                        for chunk in response.iter_content(chunk_size=8192):
+                            f.write(chunk)
+
+                    with rasterio.open(local_fname) 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,
+                            )

rslearn/data_sources/raster_source.py

@@ -15,10 +15,13 @@ from rasterio.crs import CRS
 from rslearn.config import BandSetConfig, RasterFormatConfig, RasterLayerConfig
 from rslearn.const import TILE_SIZE
 from rslearn.dataset import Window
+from rslearn.log_utils import get_logger
 from rslearn.tile_stores import LayerMetadata, TileStore
 from rslearn.utils import Projection, STGeometry
 from rslearn.utils.raster_format import load_raster_format
 
+logger = get_logger(__name__)
+
 
 class ArrayWithTransform:
     """Stores an array along with the transform associated with the array."""
@@ -70,7 +73,7 @@ class ArrayWithTransform:
         """
         return self.array
 
-    def close(self):
+    def close(self) -> None:
         """This is to mimic the rasterio.DatasetReader API.
 
         The close function is a no-op.
@@ -144,23 +147,25 @@ def get_needed_projections(
         list of Projection objects for which the item has not been ingested yet
     """
     # Identify which band set configs are relevant to this raster.
-    raster_bands = set(raster_bands)
-    relevant_band_sets = []
+    raster_bands_set = set(raster_bands)
+    relevant_band_set_list = []
     for band_set in band_sets:
         is_match = False
+        if band_set.bands is None:
+            continue
         for band in band_set.bands:
-            if band not in raster_bands:
+            if band not in raster_bands_set:
                 continue
             is_match = True
             break
         if not is_match:
             continue
-        relevant_band_sets.append(band_set)
+        relevant_band_set_list.append(band_set)
 
     all_projections = {geometry.projection for geometry in geometries}
     needed_projections = []
     for projection in all_projections:
-        for band_set in relevant_band_sets:
+        for band_set in relevant_band_set_list:
             final_projection, _ = band_set.get_final_projection_and_bounds(
                 projection, None
             )
@@ -216,16 +221,17 @@ def ingest_raster(
     else:
         # Compute the suggested target transform.
         # rasterio negates the y resolution itself so here we have to negate it.
+        raster_bounds: rasterio.coords.BoundingBox = raster.bounds
         (dst_transform, dst_width, dst_height) = (
             rasterio.warp.calculate_default_transform(
                 # Source info.
                 src_crs=raster.crs,
                 width=raster.width,
                 height=raster.height,
-                left=raster.bounds[0],
-                bottom=raster.bounds[1],
-                right=raster.bounds[2],
-                top=raster.bounds[3],
+                left=raster_bounds.left,
+                bottom=raster_bounds.bottom,
+                right=raster_bounds.right,
+                top=raster_bounds.top,
                 # Destination info.
                 dst_crs=projection.crs,
                 resolution=(projection.x_resolution, -projection.y_resolution),
@@ -258,7 +264,7 @@ def materialize_raster(
     window: Window,
     layer_name: str,
     band_cfg: BandSetConfig,
-):
+) -> None:
     """Materialize a given raster for a window.
 
     Currently it is only supported for materializing one band set.
@@ -272,7 +278,8 @@ def materialize_raster(
     window_projection, window_bounds = band_cfg.get_final_projection_and_bounds(
         window.projection, window.bounds
     )
-
+    if window_bounds is None:
+        raise ValueError(f"No windowbounds specified for {layer_name}")
     # Re-project to just extract the window.
     array = raster.read()
     window_width = window_bounds[2] - window_bounds[0]
@@ -297,7 +304,10 @@ def materialize_raster(
         dst_transform=dst_transform,
         resampling=rasterio.enums.Resampling.bilinear,
     )
-
+    if band_cfg.bands is None or band_cfg.format is None:
+        raise ValueError(
+            f"No bands or format specified for {layer_name} materialization"
+        )
     # Write the array to layer directory.
     layer_dir = window.path / "layers" / layer_name
     out_dir = layer_dir / "_".join(band_cfg.bands)

rslearn/data_sources/usgs_landsat.py

@@ -1,4 +1,7 @@
-"""Data source for Landsat data from USGS M2M API."""
+"""Data source for Landsat data from USGS M2M API.
+
+# TODO: Handle the requests in a helper function for none checking
+"""
 
 import io
 import json
@@ -15,7 +18,7 @@ import requests
 import shapely
 from upath import UPath
 
-from rslearn.config import LayerConfig, QueryConfig, RasterLayerConfig
+from rslearn.config import QueryConfig, RasterLayerConfig
 from rslearn.const import WGS84_PROJECTION
 from rslearn.data_sources import DataSource, Item
 from rslearn.data_sources.utils import match_candidate_items_to_window
@@ -36,8 +39,9 @@ class M2MAPIClient:
 
     api_url = "https://m2m.cr.usgs.gov/api/api/json/stable/"
     pagination_size = 1000
+    TIMEOUT = 1000000  # Set very high to start
 
-    def __init__(self, username, password):
+    def __init__(self, username: str, password: str) -> None:
         """Initialize a new M2MAPIClient.
 
         Args:
@@ -47,7 +51,9 @@ class M2MAPIClient:
         self.username = username
         self.password = password
         json_data = json.dumps({"username": self.username, "password": self.password})
-        response = requests.post(self.api_url + "login", data=json_data)
+        response = requests.post(
+            self.api_url + "login", data=json_data, timeout=self.TIMEOUT
+        )
         response.raise_for_status()
         self.auth_token = response.json()["data"]
 
@@ -67,24 +73,26 @@ class M2MAPIClient:
             self.api_url + endpoint,
             headers={"X-Auth-Token": self.auth_token},
             data=json.dumps(data),
+            timeout=self.TIMEOUT,
         )
         response.raise_for_status()
         if response.text:
-            data = response.json()
-            if data["errorMessage"]:
-                raise APIException(data["errorMessage"])
-            return data
+            response_dict = response.json()
+
+            if response_dict["errorMessage"]:
+                raise APIException(response_dict["errorMessage"])
+            return response_dict
         return None
 
-    def close(self):
+    def close(self) -> None:
         """Logout from the API."""
         self.request("logout")
 
-    def __enter__(self):
+    def __enter__(self) -> "M2MAPIClient":
         """Enter function to provide with semantics."""
         return self
 
-    def __exit__(self):
+    def __exit__(self) -> None:
         """Exit function to provide with semantics.
 
         Logs out the API.
@@ -100,7 +108,10 @@ class M2MAPIClient:
         Returns:
             list of filter objects
         """
-        return self.request("dataset-filters", {"datasetName": dataset_name})["data"]
+        response_dict = self.request("dataset-filters", {"datasetName": dataset_name})
+        if response_dict is None:
+            raise APIException("No response from API")
+        return response_dict["data"]
 
     def scene_search(
         self,
@@ -119,7 +130,7 @@ class M2MAPIClient:
             bbox: optional spatial filter
             metadata_filter: optional metadata filter dict
         """
-        base_data = {"datasetName": dataset_name, "sceneFilter": {}}
+        base_data: dict[str, Any] = {"datasetName": dataset_name, "sceneFilter": {}}
         if acquisition_time_range:
             base_data["sceneFilter"]["acquisitionFilter"] = {
                 "start": acquisition_time_range[0].isoformat(),
@@ -146,7 +157,10 @@ class M2MAPIClient:
             cur_data = base_data.copy()
             cur_data["startingNumber"] = starting_number
             cur_data["maxResults"] = self.pagination_size
-            data = self.request("scene-search", cur_data)["data"]
+            response_dict = self.request("scene-search", cur_data)
+            if response_dict is None:
+                raise APIException("No response from API")
+            data = response_dict["data"]
             results.extend(data["results"])
             if data["recordsReturned"] < self.pagination_size:
                 break
@@ -164,14 +178,17 @@ class M2MAPIClient:
         Returns:
             full scene metadata
         """
-        return self.request(
+        response_dict = self.request(
             "scene-metadata",
             {
                 "datasetName": dataset_name,
                 "entityId": entity_id,
                 "metadataType": "full",
             },
-        )["data"]
+        )
+        if response_dict is None:
+            raise APIException("No response from API")
+        return response_dict["data"]
 
     def get_downloadable_products(
         self, dataset_name: str, entity_id: str
@@ -186,7 +203,10 @@ class M2MAPIClient:
             list of downloadable products
         """
         data = {"datasetName": dataset_name, "entityIds": [entity_id]}
-        return self.request("download-options", data)["data"]
+        response_dict = self.request("download-options", data)
+        if response_dict is None:
+            raise APIException("No response from API")
+        return response_dict["data"]
 
     def get_download_url(self, entity_id: str, product_id: str) -> str:
         """Get the download URL for a given product.
@@ -204,9 +224,15 @@ class M2MAPIClient:
                 {"label": label, "entityId": entity_id, "productId": product_id}
             ]
         }
-        response = self.request("download-request", data)["data"]
+        response_dict = self.request("download-request", data)
+        if response_dict is None:
+            raise APIException("No response from API")
+        response = response_dict["data"]
         while True:
-            response = self.request("download-retrieve", {"label": label})["data"]
+            response_dict = self.request("download-retrieve", {"label": label})
+            if response_dict is None:
+                raise APIException("No response from API")
+            response = response_dict["data"]
             if len(response["available"]) > 0:
                 return response["available"][0]["url"]
             if len(response["requested"]) == 0:
@@ -264,7 +290,7 @@ class LandsatOliTirs(DataSource):
 
     def __init__(
         self,
-        config: LayerConfig,
+        config: RasterLayerConfig,
         username: str,
         password: str,
         max_time_delta: timedelta = timedelta(days=30),
@@ -289,9 +315,10 @@ class LandsatOliTirs(DataSource):
         self.client = M2MAPIClient(username, password)
 
     @staticmethod
-    def from_config(config: LayerConfig, ds_path: UPath) -> "LandsatOliTirs":
+    def from_config(config: RasterLayerConfig, ds_path: UPath) -> "LandsatOliTirs":
         """Creates a new LandsatOliTirs instance from a configuration dictionary."""
-        assert isinstance(config, RasterLayerConfig)
+        if config.data_source is None:
+            raise ValueError("data_source is required")
         d = config.data_source.config_dict
         if "max_time_delta" in d:
             max_time_delta = timedelta(seconds=pytimeparse.parse(d["max_time_delta"]))
@@ -328,7 +355,7 @@ class LandsatOliTirs(DataSource):
 
     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:
@@ -400,7 +427,7 @@ class LandsatOliTirs(DataSource):
         assert isinstance(serialized_item, dict)
         return LandsatOliTirsItem.deserialize(serialized_item)
 
-    def _get_download_urls(self, item: Item) -> dict[str, str]:
+    def _get_download_urls(self, item: Item) -> dict[str, tuple[str, str]]:
         """Gets the download URLs for each band.
 
         Args:
@@ -438,7 +465,7 @@ class LandsatOliTirs(DataSource):
         download_urls = self._get_download_urls(item)
         for _, (display_id, download_url) in download_urls.items():
             buf = io.BytesIO()
-            with requests.get(download_url, stream=True) as r:
+            with requests.get(download_url, stream=True, timeout=self.TIMEOUT) as r:
                 r.raise_for_status()
                 shutil.copyfileobj(r.raw, buf)
             buf.seek(0)
@@ -447,7 +474,7 @@ class LandsatOliTirs(DataSource):
     def ingest(
         self,
         tile_store: TileStore,
-        items: list[Item],
+        items: list[LandsatOliTirsItem],
         geometries: list[list[STGeometry]],
     ) -> None:
         """Ingest items into the given tile store.
@@ -471,7 +498,9 @@ class LandsatOliTirs(DataSource):
                     continue
 
                 buf = io.BytesIO()
-                with requests.get(download_urls[band][1], stream=True) as r:
+                with requests.get(
+                    download_urls[band][1], stream=True, timeout=self.TIMEOUT
+                ) as r:
                     r.raise_for_status()
                     shutil.copyfileobj(r.raw, buf)
                 buf.seek(0)

rslearn/data_sources/utils.py

@@ -1,5 +1,7 @@
 """Utilities shared by data sources."""
 
+from typing import TypeVar
+
 from rslearn.config import QueryConfig, SpaceMode, TimeMode
 from rslearn.data_sources import Item
 from rslearn.utils import STGeometry, shp_intersects
@@ -11,10 +13,12 @@ MOSAIC_REMAINDER_EPSILON = 0.01
 """Fraction of original geometry area below which mosaic is considered to contain the
 entire geometry."""
 
+ItemType = TypeVar("ItemType", bound=Item)
+
 
 def match_candidate_items_to_window(
-    geometry: STGeometry, items: list[Item], query_config: QueryConfig
-) -> list[list[Item]]:
+    geometry: STGeometry, items: list[ItemType], query_config: QueryConfig
+) -> list[list[ItemType]]:
     """Match candidate items to a window based on the query configuration.
 
     Candidate items should be collected that intersect with the window's spatial
@@ -45,17 +49,20 @@ def match_candidate_items_to_window(
                 items = [
                     item
                     for item in items
-                    if not item.time_range or item.range[1] <= geometry.time_range[0]
+                    if not item.geometry.time_range
+                    or item.geometry.time_range[1] <= geometry.time_range[0]
                 ]
             elif query_config.time_mode == TimeMode.AFTER:
                 items = [
                     item
                     for item in items
-                    if not item.time_range
-                    or item.time_range[0] >= geometry.time_range[1]
+                    if not item.geometry.time_range
+                    or item.geometry.time_range[0] >= geometry.time_range[1]
                 ]
             items.sort(
-                key=lambda item: geometry.distance_to_time_range(item.time_range)
+                key=lambda item: geometry.distance_to_time_range(
+                    item.geometry.time_range
+                )
             )
 
     # Now apply space mode.

rslearn/data_sources/vector_source.py

@@ -0,0 +1 @@
+"""Placeholder for a vector data source."""