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

rslearn/data_sources/usda_cdl.py

@@ -0,0 +1,195 @@
+"""Crop type data from the USDA Cropland Data Layer."""
+
+import os
+import tempfile
+import zipfile
+from datetime import UTC, datetime, timedelta
+from typing import Any
+
+import requests
+import requests.auth
+import shapely
+from upath import UPath
+
+from rslearn.config import QueryConfig
+from rslearn.const import WGS84_PROJECTION
+from rslearn.data_sources import DataSource, DataSourceContext, Item
+from rslearn.data_sources.utils import match_candidate_items_to_window
+from rslearn.log_utils import get_logger
+from rslearn.tile_stores import TileStoreWithLayer
+from rslearn.utils.geometry import STGeometry
+
+logger = get_logger(__name__)
+
+
+class CDL(DataSource):
+    """Data source for crop type data from the USDA Cropland Data Layer.
+
+    See https://www.nass.usda.gov/Research_and_Science/Cropland/SARS1a.php for details
+    about the data.
+
+    There is one GeoTIFF item per year from 2008. Each GeoTIFF spans the entire
+    continental US, and has a single band.
+    """
+
+    BASE_URL = (
+        "https://www.nass.usda.gov/Research_and_Science/Cropland/Release/datasets/"
+    )
+    ZIP_FILENAMES = {
+        2024: "2024_30m_cdls.zip",
+        2023: "2023_30m_cdls.zip",
+        2022: "2022_30m_cdls.zip",
+        2021: "2021_30m_cdls.zip",
+        2020: "2020_30m_cdls.zip",
+        2019: "2019_30m_cdls.zip",
+        2018: "2018_30m_cdls.zip",
+        2017: "2017_30m_cdls.zip",
+        2016: "2016_30m_cdls.zip",
+        2015: "2015_30m_cdls.zip",
+        2014: "2014_30m_cdls.zip",
+        2013: "2013_30m_cdls.zip",
+        2012: "2012_30m_cdls.zip",
+        2011: "2011_30m_cdls.zip",
+        2010: "2010_30m_cdls.zip",
+        2009: "2009_30m_cdls.zip",
+        2008: "2008_30m_cdls.zip",
+    }
+
+    # The bounds of each GeoTIFF in WGS84 coordinates, based on the 2023 map.
+    BOUNDS = shapely.box(-127.9, 23.0, -65.3, 48.3)
+
+    def __init__(
+        self,
+        timeout: timedelta = timedelta(seconds=10),
+        context: DataSourceContext = DataSourceContext(),
+    ):
+        """Initialize a new CDL instance.
+
+        Args:
+            timeout: timeout for requests.
+            context: the data source context.
+        """
+        self.timeout = timeout
+
+        # Get the band name from the layer config, which should have a single band set
+        # with a single band. If the layer config is not available in the context, we
+        # default to "cdl".
+        if context.layer_config is not None:
+            if len(context.layer_config.band_sets) != 1:
+                raise ValueError("expected a single band set")
+            if len(context.layer_config.band_sets[0].bands) != 1:
+                raise ValueError("expected band set to have a single band")
+            self.band_name = context.layer_config.band_sets[0].bands[0]
+        else:
+            self.band_name = "cdl"
+
+    def get_item_by_name(self, name: str) -> Item:
+        """Gets an item by name.
+
+        Args:
+            name: the name of the item to get. For CDL, the item name is the filename
+                of the zip file containing the per-year GeoTIFF.
+
+        Returns:
+            the Item object
+        """
+        year = int(name[0:4])
+        geometry = STGeometry(
+            WGS84_PROJECTION,
+            self.BOUNDS,
+            (
+                datetime(year, 1, 1, tzinfo=UTC),
+                datetime(year + 1, 1, 1, tzinfo=UTC),
+            ),
+        )
+        return Item(name, geometry)
+
+    def get_items(
+        self, geometries: list[STGeometry], query_config: QueryConfig
+    ) -> list[list[list[Item]]]:
+        """Get a list of items in the data source intersecting the given geometries.
+
+        Args:
+            geometries: the spatiotemporal geometries
+            query_config: the query configuration
+
+        Returns:
+            List of groups of items that should be retrieved for each geometry.
+        """
+        # First enumerate all items.
+        # Then we simply pass this to match_candidate_items_to_window.
+        items: list[Item] = []
+        for year, fname in self.ZIP_FILENAMES.items():
+            geometry = STGeometry(
+                WGS84_PROJECTION,
+                self.BOUNDS,
+                (
+                    datetime(year, 1, 1, tzinfo=UTC),
+                    datetime(year + 1, 1, 1, tzinfo=UTC),
+                ),
+            )
+            items.append(Item(fname, geometry))
+
+        groups = []
+        for geometry in geometries:
+            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 Item.deserialize(serialized_item)
+
+    def ingest(
+        self,
+        tile_store: TileStoreWithLayer,
+        items: list[Item],
+        geometries: list[list[STGeometry]],
+    ) -> None:
+        """Ingest items into the given tile store.
+
+        Args:
+            tile_store: the tile store to ingest into
+            items: the items to ingest
+            geometries: a list of geometries needed for each item
+        """
+        for item in items:
+            if tile_store.is_raster_ready(item.name, [self.band_name]):
+                continue
+
+            # Download the zip file.
+            url = self.BASE_URL + item.name
+            logger.debug(f"Downloading CDL GeoTIFF from {url}")
+            response = requests.get(
+                url, stream=True, timeout=self.timeout.total_seconds()
+            )
+            response.raise_for_status()
+
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                # Store it in temporary directory.
+                zip_fname = os.path.join(tmp_dir, "data.zip")
+                with open(zip_fname, "wb") as f:
+                    for chunk in response.iter_content(chunk_size=8192):
+                        f.write(chunk)
+
+                # Extract the .tif file.
+                logger.debug(f"Extracting GeoTIFF from {item.name}")
+                with zipfile.ZipFile(zip_fname) as zip_f:
+                    candidate_member_names = [
+                        member_name
+                        for member_name in zip_f.namelist()
+                        if member_name.endswith(".tif")
+                    ]
+                    if len(candidate_member_names) != 1:
+                        raise ValueError(
+                            f"expected CDL zip to have one .tif file but got {candidate_member_names}"
+                        )
+                    local_fname = zip_f.extract(candidate_member_names[0], path=tmp_dir)
+
+                # Now we can ingest it.
+                logger.debug(f"Ingesting data for {item.name}")
+                tile_store.write_raster_file(
+                    item.name, [self.band_name], UPath(local_fname)
+                )

rslearn/data_sources/usgs_landsat.py

@@ -1,29 +1,30 @@
-"""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
+import os
 import shutil
+import tempfile
 import time
 import uuid
 from collections.abc import Generator
-from datetime import datetime, timedelta, timezone
+from datetime import UTC, datetime, timedelta
 from typing import Any, BinaryIO
 
-import pytimeparse
-import rasterio
 import requests
 import shapely
 from upath import UPath
 
-from rslearn.config import LayerConfig, QueryConfig, RasterLayerConfig
+from rslearn.config import QueryConfig
 from rslearn.const import WGS84_PROJECTION
-from rslearn.data_sources import DataSource, Item
+from rslearn.data_sources import DataSource, DataSourceContext, Item
 from rslearn.data_sources.utils import match_candidate_items_to_window
-from rslearn.tile_stores import PrefixedTileStore, TileStore
+from rslearn.tile_stores import TileStoreWithLayer
 from rslearn.utils import STGeometry
 
-from .raster_source import get_needed_projections, ingest_raster
-
 
 class APIException(Exception):
     """Exception raised for M2M API errors."""
@@ -37,17 +38,46 @@ class M2MAPIClient:
     api_url = "https://m2m.cr.usgs.gov/api/api/json/stable/"
     pagination_size = 1000
 
-    def __init__(self, username, password):
+    def __init__(
+        self,
+        username: str,
+        password: str | None = None,
+        token: str | None = None,
+        timeout: timedelta = timedelta(seconds=120),
+    ) -> None:
         """Initialize a new M2MAPIClient.
 
         Args:
             username: the EROS username
             password: the EROS password
+            token: the application token. One of password or token must be specified.
+            timeout: timeout for requests.
         """
         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)
+        self.timeout = timeout
+
+        if password is not None and token is not None:
+            raise ValueError("only one of password or token can be specified")
+
+        if password is not None:
+            json_data = json.dumps({"username": self.username, "password": password})
+            response = requests.post(
+                self.api_url + "login",
+                data=json_data,
+                timeout=self.timeout.total_seconds(),
+            )
+
+        elif token is not None:
+            json_data = json.dumps({"username": username, "token": token})
+            response = requests.post(
+                self.api_url + "login-token",
+                data=json_data,
+                timeout=self.timeout.total_seconds(),
+            )
+
+        else:
+            raise ValueError("one of password or token must be specified")
+
         response.raise_for_status()
         self.auth_token = response.json()["data"]
 
@@ -67,24 +97,26 @@ class M2MAPIClient:
             self.api_url + endpoint,
             headers={"X-Auth-Token": self.auth_token},
             data=json.dumps(data),
+            timeout=self.timeout.total_seconds(),
         )
         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 +132,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 +154,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 +181,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 +202,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 +227,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 +248,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,45 +314,29 @@ class LandsatOliTirs(DataSource):
 
     def __init__(
         self,
-        config: LayerConfig,
         username: str,
-        password: str,
-        max_time_delta: timedelta = timedelta(days=30),
         sort_by: str | None = None,
+        password: str | None = None,
+        token: str | None = None,
+        timeout: timedelta = timedelta(seconds=10),
+        context: DataSourceContext = DataSourceContext(),
     ):
         """Initialize a new LandsatOliTirs instance.
 
         Args:
-            config: the LayerConfig of the layer containing this data source
             username: EROS username
-            password: EROS password
-            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.
             sort_by: can be "cloud_cover", default arbitrary order; only has effect for
                 SpaceMode.WITHIN.
+            password: EROS password (see M2MAPIClient).
+            token: EROS application token (see M2MAPIClient).
+            timeout: timeout for requests.
+            context: the data source context.
         """
-        self.config = config
-        self.max_time_delta = max_time_delta
         self.sort_by = sort_by
+        self.timeout = timeout
 
-        self.client = M2MAPIClient(username, password)
-
-    @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
-        if "max_time_delta" in d:
-            max_time_delta = timedelta(seconds=pytimeparse.parse(d["max_time_delta"]))
-        else:
-            max_time_delta = timedelta(days=30)
-        return LandsatOliTirs(
-            config=config,
-            username=d["username"],
-            password=d["password"],
-            max_time_delta=max_time_delta,
-            sort_by=d.get("sort_by"),
+        self.client = M2MAPIClient(
+            username, password=password, token=token, timeout=timeout
         )
 
     def _scene_metadata_to_item(self, result: dict[str, Any]) -> LandsatOliTirsItem:
@@ -317,7 +351,7 @@ class LandsatOliTirs(DataSource):
             ts = datetime.strptime(metadata_dict["Start Time"], "%Y-%m-%d %H:%M:%S.%f")
         else:
             ts = datetime.strptime(metadata_dict["Start Time"], "%Y-%m-%d %H:%M:%S")
-        ts = ts.replace(tzinfo=timezone.utc)
+        ts = ts.replace(tzinfo=UTC)
 
         return LandsatOliTirsItem(
             name=result["displayId"],
@@ -328,7 +362,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 +434,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 +472,9 @@ 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.total_seconds()
+            ) as r:
                 r.raise_for_status()
                 shutil.copyfileobj(r.raw, buf)
             buf.seek(0)
@@ -446,8 +482,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.
@@ -457,30 +493,26 @@ class LandsatOliTirs(DataSource):
             items: the items to ingest
             geometries: a list of geometries needed for each item
         """
-        for item, cur_geometries in zip(items, geometries):
+        for item in items:
             download_urls = self._get_download_urls(item)
+
             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()
-                with requests.get(download_urls[band][1], stream=True) as r:
-                    r.raise_for_status()
-                    shutil.copyfileobj(r.raw, buf)
-                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:
+                    local_filename = os.path.join(tmp_dir, "data.tif")
+
+                    with requests.get(
+                        download_urls[band][1],
+                        stream=True,
+                        timeout=self.timeout.total_seconds(),
+                    ) as r:
+                        r.raise_for_status()
+                        with open(local_filename, "wb") as f:
+                            shutil.copyfileobj(r.raw, f)
+
+                    tile_store.write_raster_file(
+                        item.name, band_names, UPath(local_filename)
+                    )