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

rslearn/data_sources/earthdata_srtm.py

@@ -0,0 +1,282 @@
+"""Elevation data from the Shuttle Radar Topography Mission via NASA Earthdata."""
+
+import math
+import os
+import tempfile
+import zipfile
+from datetime import timedelta
+from typing import Any
+
+import requests
+import requests.auth
+import shapely
+from upath import UPath
+
+from rslearn.config import QueryConfig, SpaceMode
+from rslearn.const import WGS84_PROJECTION
+from rslearn.data_sources import DataSource, DataSourceContext, Item
+from rslearn.log_utils import get_logger
+from rslearn.tile_stores import TileStoreWithLayer
+from rslearn.utils.geometry import STGeometry
+
+logger = get_logger(__name__)
+
+
+class SRTM(DataSource):
+    """Data source for SRTM elevation data using NASA Earthdata credentials.
+
+    See https://e4ftl01.cr.usgs.gov/MEASURES/SRTMGL1.003/ and
+    https://dwtkns.com/srtm30m/ for details about the data.
+
+    The data is split into 1x1-degree tiles, where the filename ends with e.g.
+    S28W055.SRTMGL1.hgt.zip (so only the first seven characters change).
+
+    These URLs can only be accessed with a NASA Earthdata username and password.
+
+    The zip file contains a single hgt file which can be read by rasterio. It has a
+    single 16-bit signed integer band indicating the elevation.
+
+    Items from this data source do not come with a time range. The band name will match
+    that specified in the band set, which should have a single band.
+    """
+
+    BASE_URL = "https://e4ftl01.cr.usgs.gov/MEASURES/SRTMGL1.003/2000.02.11/"
+    FILENAME_SUFFIX = ".SRTMGL1.hgt.zip"
+
+    def __init__(
+        self,
+        username: str | None = None,
+        password: str | None = None,
+        timeout: timedelta = timedelta(seconds=10),
+        context: DataSourceContext = DataSourceContext(),
+    ):
+        """Initialize a new SRTM instance.
+
+        Args:
+            username: NASA Earthdata account username. If not set, it is read from the
+                NASA_EARTHDATA_USERNAME environment variable.
+            password: NASA Earthdata account password. If not set, it is read from the
+                NASA_EARTHDATA_PASSWORD environment variable.
+            timeout: timeout for requests.
+            context: the data source context.
+        """
+        # Get band name from context if possible, falling back to "srtm".
+        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 = "srtm"
+
+        self.timeout = timeout
+
+        if username is None:
+            username = os.environ["NASA_EARTHDATA_USERNAME"]
+        self.username = username
+
+        if password is None:
+            password = os.environ["NASA_EARTHDATA_PASSWORD"]
+        self.password = password
+
+        self.session = requests.session()
+
+    def get_item_by_name(self, name: str) -> Item:
+        """Gets an item by name.
+
+        Args:
+            name: the name of the item to get. For SRTM, the item name is the filename
+                of the zip file containing the hgt file.
+
+        Returns:
+            the Item object
+        """
+        if not name.endswith(self.FILENAME_SUFFIX):
+            raise ValueError(
+                f"expected item name to end with {self.FILENAME_SUFFIX}, but got {name}"
+            )
+        # Parse the first seven characters, e.g. S28W055.
+        # We do this to reconstruct the geometry of the item.
+        lat_sign = name[0]
+        lat_degrees = int(name[1:3])
+        lon_sign = name[4]
+        lon_degrees = int(name[5:8])
+
+        if lat_sign == "N":
+            lat_min = lat_degrees
+        elif lat_sign == "S":
+            lat_min = -lat_degrees
+        else:
+            raise ValueError(f"invalid item name {name}")
+
+        if lon_sign == "E":
+            lon_min = lon_degrees
+        elif lon_sign == "W":
+            lon_min = -lon_degrees
+        else:
+            raise ValueError(f"invalid item name {name}")
+
+        geometry = STGeometry(
+            WGS84_PROJECTION,
+            shapely.box(lon_min, lat_min, lon_min + 1, lat_min + 1),
+            None,
+        )
+        return Item(name, geometry)
+
+    def _lon_lat_to_item(self, lon_min: int, lat_min: int) -> Item:
+        """Get an item based on the 1x1 longitude/latitude grid.
+
+        Args:
+            lon_min: the starting longitude integer of the grid cell.
+            lat_min: the starting latitude integer of the grid cell.
+
+        Returns:
+            the Item object.
+        """
+        # Construct the filename for this grid cell.
+        # The item name is just the filename.
+        if lon_min < 0:
+            lon_part = f"W{-lon_min:03d}"
+        else:
+            lon_part = f"E{lon_min:03d}"
+        if lat_min < 0:
+            lat_part = f"S{-lat_min:02d}"
+        else:
+            lat_part = f"N{lat_min:02d}"
+        fname = lat_part + lon_part + self.FILENAME_SUFFIX
+
+        # We also need the geometry for the item.
+        geometry = STGeometry(
+            WGS84_PROJECTION,
+            shapely.box(lon_min, lat_min, lon_min + 1, lat_min + 1),
+            None,
+        )
+
+        return Item(fname, 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.
+        """
+        # We only support mosaic here, other query modes don't really make sense.
+        if query_config.space_mode != SpaceMode.MOSAIC or query_config.max_matches != 1:
+            raise ValueError(
+                "expected mosaic with max_matches=1 for the query configuration"
+            )
+
+        groups = []
+        for geometry in geometries:
+            # We iterate over each 1x1 cell that this geometry intersects and include
+            # the corresponing item in this item group.
+            # Since it is a mosaic with one match, there will just be one item group
+            # for each item.
+            wgs84_geometry = geometry.to_projection(WGS84_PROJECTION)
+            shp_bounds = wgs84_geometry.shp.bounds
+            cell_bounds = (
+                math.floor(shp_bounds[0]),
+                math.floor(shp_bounds[1]),
+                math.ceil(shp_bounds[2]),
+                math.ceil(shp_bounds[3]),
+            )
+            # lon_min/lat_min are the lower range of each cell.
+            items = []
+            for lon_min in range(cell_bounds[0], cell_bounds[2]):
+                for lat_min in range(cell_bounds[1], cell_bounds[3]):
+                    items.append(self._lon_lat_to_item(lon_min, lat_min))
+
+            logger.debug(f"Got {len(items)} items (grid cells) for geometry")
+            groups.append([items])
+
+        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 item.
+            # We first attempt to access it directly, which works if we have already
+            # authenticated. If not, we get redirected to a login endpoint where we
+            # need to use basic authentication; the endpoint will redirect us back to
+            # the original URL.
+            url = self.BASE_URL + item.name
+            logger.debug(f"Downloading SRTM data for {item.name} from {url}")
+
+            # Try to access directly.
+            response = self.session.get(
+                url,
+                stream=True,
+                timeout=self.timeout.total_seconds(),
+                allow_redirects=False,
+            )
+
+            if response.status_code == 302:
+                # Encountered redirect, so set response to actually access the redirect
+                # URL. This time we follow redirects since it will take us back to the
+                # original URL.
+                redirect_url = response.headers["Location"]
+                logger.debug(f"Following redirect to {redirect_url}")
+                auth = requests.auth.HTTPBasicAuth(self.username, self.password)
+                response = self.session.get(
+                    redirect_url,
+                    stream=True,
+                    timeout=self.timeout.total_seconds(),
+                    auth=auth,
+                )
+
+            if response.status_code == 404:
+                # Some grid cells don't exist so this isn't a big issue.
+                logger.warning(
+                    f"Skipping item {item.name} because there is no data at that cell"
+                )
+                continue
+            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 .hgt file.
+                logger.debug(f"Extracting data for {item.name}")
+                with zipfile.ZipFile(zip_fname) as zip_f:
+                    member_names = zip_f.namelist()
+                    if len(member_names) != 1:
+                        raise ValueError(
+                            f"expected SRTM zip to have one member but got {member_names}"
+                        )
+                    local_fname = zip_f.extract(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/eurocrops.py

@@ -0,0 +1,242 @@
+"""Data source for vector EuroCrops crop type data."""
+
+import glob
+import os
+import tempfile
+import zipfile
+from datetime import UTC, datetime, timedelta
+from typing import Any
+
+import fiona
+import requests
+from rasterio.crs import CRS
+
+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.feature import Feature
+from rslearn.utils.geometry import Projection, STGeometry, get_global_geometry
+
+logger = get_logger(__name__)
+
+
+class EuroCropsItem(Item):
+    """An item in the EuroCrops data source.
+
+    For simplicity, we have just one item per year, so each item combines all of the
+    country-level files for that year.
+    """
+
+    def __init__(self, name: str, geometry: STGeometry, zip_fnames: list[str]):
+        """Creates a new EuroCropsItem.
+
+        Args:
+            name: unique name of the item. It is just the year that this item
+                corresponds to.
+            geometry: the spatial and temporal extent of the item
+            zip_fnames: the filenames of the zip files that contain country-level crop
+                type data for this year.
+        """
+        super().__init__(name, geometry)
+        self.zip_fnames = zip_fnames
+
+    def serialize(self) -> dict:
+        """Serializes the item to a JSON-encodable dictionary."""
+        d = super().serialize()
+        d["zip_fnames"] = self.zip_fnames
+        return d
+
+    @staticmethod
+    def deserialize(d: dict) -> "EuroCropsItem":
+        """Deserializes an item from a JSON-decoded dictionary."""
+        item = super(EuroCropsItem, EuroCropsItem).deserialize(d)
+        return EuroCropsItem(
+            name=item.name, geometry=item.geometry, zip_fnames=d["zip_fnames"]
+        )
+
+
+class EuroCrops(DataSource[EuroCropsItem]):
+    """A data source for EuroCrops vector data (v11).
+
+    See https://zenodo.org/records/14094196 for details.
+
+    While the source data is split into country-level files, this data source uses one
+    item per year for simplicity. So each item corresponds to all of the country-level
+    files for that year.
+
+    Note that the RO_ny.zip file is not used.
+    """
+
+    BASE_URL = "https://zenodo.org/records/14094196/files/"
+    FILENAMES_BY_YEAR = {
+        2018: [
+            "FR_2018.zip",
+        ],
+        2019: [
+            "DK_2019.zip",
+        ],
+        2020: [
+            "ES_NA_2020.zip",
+            "FI_2020.zip",
+            "HR_2020.zip",
+            "NL_2020.zip",
+        ],
+        2021: [
+            "AT_2021.zip",
+            "BE_VLG_2021.zip",
+            "BE_WAL_2021.zip",
+            "EE_2021.zip",
+            "LT_2021.zip",
+            "LV_2021.zip",
+            "PT_2021.zip",
+            "SE_2021.zip",
+            "SI_2021.zip",
+            "SK_2021.zip",
+        ],
+        2023: [
+            "CZ_2023.zip",
+            "DE_BB_2023.zip",
+            "DE_LS_2021.zip",
+            "DE_NRW_2021.zip",
+            "ES_2023.zip",
+            "IE_2023.zip",
+        ],
+    }
+    TIMEOUT = timedelta(seconds=10)
+
+    def __init__(self, context: DataSourceContext = DataSourceContext()):
+        """Create a new EuroCrops."""
+        pass
+
+    def _get_all_items(self) -> list[EuroCropsItem]:
+        """Get a list of all available items in the data source."""
+        items: list[EuroCropsItem] = []
+        for year, fnames in self.FILENAMES_BY_YEAR.items():
+            items.append(
+                EuroCropsItem(
+                    str(year),
+                    get_global_geometry(
+                        time_range=(
+                            datetime(year, 1, 1, tzinfo=UTC),
+                            datetime(year + 1, 1, 1, tzinfo=UTC),
+                        ),
+                    ),
+                    fnames,
+                )
+            )
+        return items
+
+    def get_items(
+        self, geometries: list[STGeometry], query_config: QueryConfig
+    ) -> list[list[list[EuroCropsItem]]]:
+        """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.
+        """
+        wgs84_geometries = [
+            geometry.to_projection(WGS84_PROJECTION) for geometry in geometries
+        ]
+        all_items = self._get_all_items()
+        groups = []
+        for geometry in wgs84_geometries:
+            cur_groups = match_candidate_items_to_window(
+                geometry, all_items, query_config
+            )
+            groups.append(cur_groups)
+        return groups
+
+    def deserialize_item(self, serialized_item: Any) -> EuroCropsItem:
+        """Deserializes an item from JSON-decoded data."""
+        return EuroCropsItem.deserialize(serialized_item)
+
+    def _extract_features(self, fname: str) -> list[Feature]:
+        """Download the given zip file, extract shapefile, and return list of features."""
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            # Download the zip file.
+            url = self.BASE_URL + fname
+            logger.debug(f"Downloading zip file from {url}")
+            response = requests.get(
+                url,
+                stream=True,
+                timeout=self.TIMEOUT.total_seconds(),
+                allow_redirects=False,
+            )
+            response.raise_for_status()
+            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 all of the files and look for shapefile filename.
+            logger.debug(f"Extracting zip file {fname}")
+            with zipfile.ZipFile(zip_fname) as zip_f:
+                zip_f.extractall(path=tmp_dir)
+
+            # The shapefiles or geopackage files can appear at any level in the hierarchy.
+            # Most zip files contain one but some contain multiple (one per region).
+            shp_fnames = glob.glob(
+                "**/*.shp", root_dir=tmp_dir, recursive=True
+            ) + glob.glob("**/*.gpkg", root_dir=tmp_dir, recursive=True)
+            if len(shp_fnames) == 0:
+                tmp_dir_fnames = os.listdir(tmp_dir)
+                raise ValueError(
+                    f"expected {fname} to contain .shp file but none found (matches={shp_fnames}, ls={tmp_dir_fnames})"
+                )
+
+            # Load the features from the shapefile(s).
+            features = []
+            for shp_fname in shp_fnames:
+                logger.debug(f"Loading feature list from {shp_fname}")
+                with fiona.open(os.path.join(tmp_dir, shp_fname)) as src:
+                    crs = CRS.from_wkt(src.crs.to_wkt())
+                    # Normal GeoJSON should have coordinates in CRS coordinates, i.e. it
+                    # should be 1 projection unit/pixel.
+                    projection = Projection(crs, 1, 1)
+
+                    for feat in src:
+                        features.append(
+                            Feature.from_geojson(
+                                projection,
+                                {
+                                    "type": "Feature",
+                                    "geometry": dict(feat.geometry),
+                                    "properties": dict(feat.properties),
+                                },
+                            )
+                        )
+
+            return features
+
+    def ingest(
+        self,
+        tile_store: TileStoreWithLayer,
+        items: list[EuroCropsItem],
+        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_vector_ready(item.name):
+                continue
+
+            # Get features across all shapefiles.
+            features: list[Feature] = []
+            for fname in item.zip_fnames:
+                logger.debug(f"Getting features from {fname} for item {item.name}")
+                features.extend(self._extract_features(fname))
+
+            logger.debug(f"Writing features for {item.name} to the tile store")
+            tile_store.write_vector(item.name, features)