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

rslearn/data_sources/soilgrids.py

@@ -0,0 +1,331 @@
+"""Data source for SoilGrids via the `soilgrids` Python package.
+
+This source is intended to be used with `ingest: false` (direct materialization),
+since data is fetched on-demand per window.
+"""
+
+from __future__ import annotations
+
+import tempfile
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+import rasterio
+import rasterio.warp
+import shapely
+from rasterio.crs import CRS
+from rasterio.enums import Resampling
+from upath import UPath
+
+from rslearn.config import LayerConfig, QueryConfig
+from rslearn.dataset import Window
+from rslearn.dataset.materialize import RasterMaterializer
+from rslearn.tile_stores import TileStore, TileStoreWithLayer
+from rslearn.utils import PixelBounds, Projection, STGeometry
+from rslearn.utils.geometry import get_global_geometry
+from rslearn.utils.raster_format import get_transform_from_projection_and_bounds
+
+from .data_source import DataSource, DataSourceContext, Item
+from .utils import match_candidate_items_to_window
+
+SOILGRIDS_NODATA_VALUE = -32768.0
+"""Default nodata value used by SoilGrids GeoTIFF responses (GEOTIFF_INT16)."""
+
+
+def _crs_to_rasterio(crs: str) -> CRS:
+    """Best-effort conversion of CRS strings used by `soilgrids` to rasterio CRS."""
+    try:
+        return CRS.from_string(crs)
+    except Exception:
+        # Fallback: if rasterio can't parse the string but it contains an EPSG code,
+        # extract the trailing integer and build a CRS from it.
+        parts = [p for p in crs.replace(":", " ").split() if p.isdigit()]
+        if parts:
+            return CRS.from_epsg(int(parts[-1]))
+        raise
+
+
+def _crs_to_soilgrids_urn(crs: str) -> str:
+    """Convert common CRS spellings to the URN form expected by `soilgrids`.
+
+    The `soilgrids` package compares CRS strings against the supported CRS URNs from
+    OWSLib (e.g. "urn:ogc:def:crs:EPSG::3857"). This helper allows users to specify
+    simpler forms like "EPSG:3857" while still working.
+    """
+    s = crs.strip()
+
+    # If already an EPSG URN, canonicalize to the form soilgrids expects.
+    if s.lower().startswith("urn:ogc:def:crs:") and "epsg" in s.lower():
+        parts = [p for p in s.replace(":", " ").split() if p.isdigit()]
+        if parts:
+            return f"urn:ogc:def:crs:EPSG::{parts[-1]}"
+        return s
+
+    # Accept "EPSG:3857", "epsg:3857", or other strings containing an EPSG code.
+    if "epsg" in s.lower():
+        parts = [p for p in s.replace(":", " ").split() if p.isdigit()]
+        if parts:
+            return f"urn:ogc:def:crs:EPSG::{parts[-1]}"
+
+    return s
+
+
+class SoilGrids(DataSource, TileStore):
+    """Access SoilGrids coverages as an rslearn raster data source."""
+
+    def __init__(
+        self,
+        service_id: str,
+        coverage_id: str,
+        crs: str = "EPSG:3857",
+        width: int | None = None,
+        height: int | None = None,
+        resx: float | None = None,
+        resy: float | None = None,
+        response_crs: str | None = None,
+        band_names: list[str] = ["B1"],
+        context: DataSourceContext = DataSourceContext(),
+    ):
+        """Create a new SoilGrids data source.
+
+        Args:
+            service_id: SoilGrids map service id (e.g., "clay", "phh2o").
+            coverage_id: coverage id within the service (e.g., "clay_0-5cm_mean").
+            crs: request CRS string passed through to `soilgrids.SoilGrids`, typically
+                a URN like "urn:ogc:def:crs:EPSG::4326" or "urn:ogc:def:crs:EPSG::152160".
+            width: optional WCS WIDTH parameter. Required by SoilGrids WCS when CRS is
+                EPSG:4326.
+            height: optional WCS HEIGHT parameter.
+            resx: optional WCS RESX parameter (projection units / pixel).
+            resy: optional WCS RESY parameter (projection units / pixel).
+            response_crs: optional response CRS (defaults to `crs`).
+            band_names: band names exposed to rslearn. For a single coverage, this
+                should have length 1.
+            context: rslearn data source context.
+        """
+        if len(band_names) != 1:
+            raise ValueError("SoilGrids currently supports only single-band coverages")
+        if (width is None) != (height is None):
+            raise ValueError("width and height must be specified together")
+        if (resx is None) != (resy is None):
+            raise ValueError("resx and resy must be specified together")
+        if width is not None and resx is not None:
+            raise ValueError("specify either width/height or resx/resy, not both")
+
+        self.service_id = service_id
+        self.coverage_id = coverage_id
+        self.crs = crs
+        self.width = width
+        self.height = height
+        self.resx = resx
+        self.resy = resy
+        self.response_crs = response_crs
+        self.band_names = band_names
+
+        # Represent the coverage as a single item that matches all windows.
+        item_name = f"{self.service_id}:{self.coverage_id}"
+        self._items = [Item(item_name, get_global_geometry(time_range=None))]
+
+    def get_items(
+        self, geometries: list[STGeometry], query_config: QueryConfig
+    ) -> list[list[list[Item]]]:
+        """Get item groups matching each requested geometry."""
+        groups = []
+        for geometry in geometries:
+            cur_groups = match_candidate_items_to_window(
+                geometry, self._items, query_config
+            )
+            groups.append(cur_groups)
+        return groups
+
+    def deserialize_item(self, serialized_item: Any) -> Item:
+        """Deserialize an item from JSON-decoded data."""
+        return Item.deserialize(serialized_item)
+
+    def ingest(
+        self,
+        tile_store: TileStoreWithLayer,
+        items: list[Item],
+        geometries: list[list[STGeometry]],
+    ) -> None:
+        """Ingest is not supported (direct materialization only)."""
+        raise NotImplementedError(
+            "SoilGrids is intended for direct materialization; set data_source.ingest=false."
+        )
+
+    def is_raster_ready(
+        self, layer_name: str, item_name: str, bands: list[str]
+    ) -> bool:
+        """Return whether the requested raster is ready (always true for direct reads)."""
+        return True
+
+    def get_raster_bands(self, layer_name: str, item_name: str) -> list[list[str]]:
+        """Return the band sets available for this coverage."""
+        return [self.band_names]
+
+    def get_raster_bounds(
+        self, layer_name: str, item_name: str, bands: list[str], projection: Projection
+    ) -> PixelBounds:
+        """Return (approximate) bounds for this raster in the requested projection."""
+        # We don't know bounds without an extra metadata request; treat as "very large"
+        # so materialization always attempts reads for windows.
+        return (-(10**9), -(10**9), 10**9, 10**9)
+
+    def _download_geotiff(
+        self,
+        west: float,
+        south: float,
+        east: float,
+        north: float,
+        output: str,
+        width: int | None,
+        height: int | None,
+        resx: float | None,
+        resy: float | None,
+    ) -> None:
+        from soilgrids import SoilGrids as SoilGridsClient
+
+        client = SoilGridsClient()
+        kwargs: dict[str, Any] = dict(
+            service_id=self.service_id,
+            coverage_id=self.coverage_id,
+            crs=_crs_to_soilgrids_urn(self.crs),
+            west=west,
+            south=south,
+            east=east,
+            north=north,
+            output=output,
+        )
+        if width is not None and height is not None:
+            kwargs["width"] = width
+            kwargs["height"] = height
+        elif resx is not None and resy is not None:
+            kwargs["resx"] = resx
+            kwargs["resy"] = resy
+
+        if self.response_crs is not None:
+            kwargs["response_crs"] = _crs_to_soilgrids_urn(self.response_crs)
+
+        client.get_coverage_data(**kwargs)
+
+    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 and reproject a SoilGrids coverage subset into the requested grid."""
+        if bands != self.band_names:
+            raise ValueError(
+                f"expected request for bands {self.band_names} but got {bands}"
+            )
+
+        # Compute bounding box in CRS coordinates for the request.
+        request_crs = _crs_to_rasterio(self.crs)
+        request_projection = Projection(request_crs, 1.0, 1.0)
+        request_geom = STGeometry(projection, shapely.box(*bounds), None).to_projection(
+            request_projection
+        )
+        west, south, east, north = request_geom.shp.bounds
+
+        # Determine output grid for the WCS request.
+        #
+        # If the user explicitly configured an output grid (width/height or resx/resy),
+        # we respect it.
+        #
+        # Otherwise, default to requesting at ~250 m resolution in the request CRS
+        # (when it is projected), and then reprojecting to the window grid.
+        #
+        # For EPSG:4326 requests, SoilGrids WCS requires WIDTH/HEIGHT, so we default
+        # to matching the window pixel size.
+        window_width = bounds[2] - bounds[0]
+        window_height = bounds[3] - bounds[1]
+
+        out_width = self.width
+        out_height = self.height
+        out_resx = self.resx
+        out_resy = self.resy
+
+        if request_crs.to_epsg() == 4326 and out_width is None:
+            # Required by the SoilGrids WCS for EPSG:4326; resx/resy is not accepted.
+            out_width = window_width
+            out_height = window_height
+            out_resx = None
+            out_resy = None
+        elif out_width is None and out_resx is None:
+            # Default to native-ish SoilGrids resolution (~250 m) in projected CRSs.
+            out_resx = 250.0
+            out_resy = 250.0
+
+        with tempfile.TemporaryDirectory(prefix="rslearn_soilgrids_") as tmpdir:
+            output_path = str(UPath(tmpdir) / "coverage.tif")
+            self._download_geotiff(
+                west=west,
+                south=south,
+                east=east,
+                north=north,
+                output=output_path,
+                width=out_width,
+                height=out_height,
+                resx=out_resx,
+                resy=out_resy,
+            )
+
+            with rasterio.open(output_path) as src:
+                src_array = src.read(1).astype(np.float32)
+                src_nodata = src.nodata
+                scale = float(src.scales[0]) if src.scales else 1.0
+                offset = float(src.offsets[0]) if src.offsets else 0.0
+
+                if src_nodata is not None:
+                    valid_mask = src_array != float(src_nodata)
+                    src_array[valid_mask] = src_array[valid_mask] * scale + offset
+                    dst_nodata = float(src_nodata)
+                    src_nodata_val = dst_nodata
+                else:
+                    src_array = src_array * scale + offset
+                    dst_nodata = SOILGRIDS_NODATA_VALUE
+                    src_nodata_val = None
+
+                src_chw = src_array[None, :, :]
+                dst = np.full(
+                    (1, bounds[3] - bounds[1], bounds[2] - bounds[0]),
+                    dst_nodata,
+                    dtype=np.float32,
+                )
+                dst_transform = get_transform_from_projection_and_bounds(
+                    projection, bounds
+                )
+
+                rasterio.warp.reproject(
+                    source=src_chw,
+                    src_crs=src.crs,
+                    src_transform=src.transform,
+                    src_nodata=src_nodata_val,
+                    destination=dst,
+                    dst_crs=projection.crs,
+                    dst_transform=dst_transform,
+                    dst_nodata=dst_nodata,
+                    resampling=resampling,
+                )
+                return dst
+
+    def materialize(
+        self,
+        window: Window,
+        item_groups: list[list[Item]],
+        layer_name: str,
+        layer_cfg: LayerConfig,
+    ) -> None:
+        """Materialize a window by reading from SoilGrids on-demand."""
+        RasterMaterializer().materialize(
+            TileStoreWithLayer(self, layer_name),
+            window,
+            layer_name,
+            layer_cfg,
+            item_groups,
+        )

rslearn/data_sources/stac.py

@@ -0,0 +1,275 @@
+"""A partial data source implementation providing get_items using a STAC API."""
+
+import json
+from datetime import datetime
+from typing import Any
+
+import shapely
+from upath import UPath
+
+from rslearn.config import QueryConfig
+from rslearn.const import WGS84_PROJECTION
+from rslearn.data_sources.data_source import Item, ItemLookupDataSource
+from rslearn.data_sources.utils import match_candidate_items_to_window
+from rslearn.log_utils import get_logger
+from rslearn.utils.geometry import STGeometry
+from rslearn.utils.stac import StacClient, StacItem
+
+logger = get_logger(__name__)
+
+
+class SourceItem(Item):
+    """An item in the StacDataSource data source."""
+
+    def __init__(
+        self,
+        name: str,
+        geometry: STGeometry,
+        asset_urls: dict[str, str],
+        properties: dict[str, str],
+    ):
+        """Creates a new SourceItem.
+
+        Args:
+            name: unique name of the item
+            geometry: the spatial and temporal extent of the item
+            asset_urls: map from asset key to the unsigned asset URL.
+            properties: properties requested by the data source implementation.
+        """
+        super().__init__(name, geometry)
+        self.asset_urls = asset_urls
+        self.properties = properties
+
+    def serialize(self) -> dict[str, Any]:
+        """Serializes the item to a JSON-encodable dictionary."""
+        d = super().serialize()
+        d["asset_urls"] = self.asset_urls
+        d["properties"] = self.properties
+        return d
+
+    @staticmethod
+    def deserialize(d: dict[str, Any]) -> "SourceItem":
+        """Deserializes an item from a JSON-decoded dictionary."""
+        item = super(SourceItem, SourceItem).deserialize(d)
+        return SourceItem(
+            name=item.name,
+            geometry=item.geometry,
+            asset_urls=d["asset_urls"],
+            properties=d["properties"],
+        )
+
+
+class StacDataSource(ItemLookupDataSource[SourceItem]):
+    """A partial data source implementing get_items using a STAC API.
+
+    This is a helper class that full implementations can extend to not have to worry
+    about the get_items and get_item_by_name implementation.
+    """
+
+    def __init__(
+        self,
+        endpoint: str,
+        collection_name: str,
+        query: dict[str, Any] | None = None,
+        sort_by: str | None = None,
+        sort_ascending: bool = True,
+        required_assets: list[str] | None = None,
+        cache_dir: UPath | None = None,
+        limit: int = 100,
+        properties_to_record: list[str] = [],
+    ):
+        """Create a new StacDataSource.
+
+        Args:
+            endpoint: the STAC endpoint to use.
+            collection_name: the STAC collection name.
+            query: optional STAC query dict to include in searches, e.g. {"eo:cloud_cover": {"lt": 50}}.
+            sort_by: sort results by this STAC property.
+            sort_ascending: if sort_by is set, sort in ascending order (default).
+                Otherwise sort in descending order.
+            required_assets: if set, we ignore items that do not have all of these
+                asset keys.
+            cache_dir: optional cache directory to cache items. This is recommended if
+                allowing direct materialization from the data source, since it will
+                likely be necessary to make lots of get_item_by_name calls during
+                materialization. TODO: give direct materialization access to the Item
+                object.
+            limit: limit to pass to search queries.
+            properties_to_record: if these properties on the STAC item exist, they are
+                are retained in the SourceItem when we initialize it.
+        """
+        self.client = StacClient(endpoint)
+        self.collection_name = collection_name
+        self.query = query
+        self.sort_by = sort_by
+        self.sort_ascending = sort_ascending
+        self.required_assets = required_assets
+        self.cache_dir = cache_dir
+        self.limit = limit
+        self.properties_to_record = properties_to_record
+
+    def _stac_item_to_item(self, stac_item: StacItem) -> SourceItem:
+        # Make sure geometry, time range, and assets are set.
+        if stac_item.geometry is None:
+            raise ValueError("got unexpected item with no geometry")
+        if stac_item.time_range is None:
+            raise ValueError("got unexpected item with no time range")
+        if stac_item.assets is None:
+            raise ValueError("got unexpected item with no assets")
+
+        shp = shapely.geometry.shape(stac_item.geometry)
+        geom = STGeometry(WGS84_PROJECTION, shp, stac_item.time_range)
+        asset_urls = {
+            asset_key: asset_obj.href
+            for asset_key, asset_obj in stac_item.assets.items()
+        }
+
+        # Keep any properties requested by the data source implementation.
+        properties = {}
+        for prop_name in self.properties_to_record:
+            if prop_name not in stac_item.properties:
+                continue
+            properties[prop_name] = stac_item.properties[prop_name]
+
+        return SourceItem(stac_item.id, geom, asset_urls, properties)
+
+    def _get_search_time_range(
+        self, geometry: STGeometry
+    ) -> datetime | tuple[datetime, datetime] | None:
+        """Get time range to include in STAC API search.
+
+        By default, we filter STAC searches to the window's time range. Subclasses can
+        override this to disable time filtering for "static" datasets.
+
+        Args:
+            geometry: the geometry we are searching for.
+
+        Returns:
+            the time range (or timestamp) to pass to the STAC search, or None to avoid
+                temporal filtering in the search request.
+        """
+        # Note: StacClient.search accepts either a datetime or a (start, end) tuple.
+        return geometry.time_range
+
+    def get_item_by_name(self, name: str) -> SourceItem:
+        """Gets an item by name.
+
+        Args:
+            name: the name of the item to get
+
+        Returns:
+            the item object
+        """
+        # If cache_dir is set, we cache the item. First here we check if it is already
+        # in the cache.
+        cache_fname: UPath | None = None
+        if self.cache_dir:
+            cache_fname = self.cache_dir / f"{name}.json"
+        if cache_fname is not None and cache_fname.exists():
+            with cache_fname.open() as f:
+                return SourceItem.deserialize(json.load(f))
+
+        # No cache or not in cache, so we need to make the STAC request.
+        logger.debug(f"Getting STAC item {name}")
+        stac_items = self.client.search(ids=[name], collections=[self.collection_name])
+
+        if len(stac_items) == 0:
+            raise ValueError(
+                f"Item {name} not found in collection {self.collection_name}"
+            )
+        if len(stac_items) > 1:
+            raise ValueError(
+                f"Multiple items found for ID {name} in collection {self.collection_name}"
+            )
+
+        stac_item = stac_items[0]
+        item = self._stac_item_to_item(stac_item)
+
+        # Finally we cache it if cache_dir is set.
+        if cache_fname is not None:
+            with cache_fname.open("w") as f:
+                json.dump(item.serialize(), f)
+
+        return item
+
+    def get_items(
+        self, geometries: list[STGeometry], query_config: QueryConfig
+    ) -> list[list[list[SourceItem]]]:
+        """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:
+            # Get potentially relevant items from the collection by performing one search
+            # for each requested geometry.
+            wgs84_geometry = geometry.to_projection(WGS84_PROJECTION)
+            logger.debug("performing STAC search for geometry %s", wgs84_geometry)
+            search_time_range = self._get_search_time_range(wgs84_geometry)
+            stac_items = self.client.search(
+                collections=[self.collection_name],
+                intersects=json.loads(shapely.to_geojson(wgs84_geometry.shp)),
+                date_time=search_time_range,
+                query=self.query,
+                limit=self.limit,
+            )
+            logger.debug("STAC search yielded %d items", len(stac_items))
+
+            if self.required_assets is not None:
+                # Filter out items that are missing any of the assets in self.asset_bands.
+                good_stac_items = []
+                for stac_item in stac_items:
+                    if stac_item.assets is None:
+                        raise ValueError(f"got STAC item {stac_item.id} with no assets")
+
+                    good = True
+                    for asset_key in self.required_assets:
+                        if asset_key in stac_item.assets:
+                            continue
+                        good = False
+                        break
+                    if good:
+                        good_stac_items.append(stac_item)
+                logger.debug(
+                    "required_assets filter from %d to %d items",
+                    len(stac_items),
+                    len(good_stac_items),
+                )
+                stac_items = good_stac_items
+
+            if self.sort_by is not None:
+                sort_by = self.sort_by
+                stac_items.sort(
+                    key=lambda stac_item: stac_item.properties[sort_by],
+                    reverse=not self.sort_ascending,
+                )
+
+            candidate_items = [
+                self._stac_item_to_item(stac_item) for stac_item in stac_items
+            ]
+
+            # Since we made the STAC request, might as well save these to the cache.
+            if self.cache_dir is not None:
+                for item in candidate_items:
+                    cache_fname = self.cache_dir / f"{item.name}.json"
+                    if cache_fname.exists():
+                        continue
+                    with cache_fname.open("w") as f:
+                        json.dump(item.serialize(), f)
+
+            cur_groups = match_candidate_items_to_window(
+                geometry, candidate_items, query_config
+            )
+            groups.append(cur_groups)
+
+        return groups
+
+    def deserialize_item(self, serialized_item: Any) -> SourceItem:
+        """Deserializes an item from JSON-decoded data."""
+        assert isinstance(serialized_item, dict)
+        return SourceItem.deserialize(serialized_item)

rslearn/main.py

@@ -2,6 +2,7 @@
 
 import argparse
 import multiprocessing
+import os
 import random
 import sys
 import time
@@ -45,6 +46,7 @@ handler_registry = {}
 ItemType = TypeVar("ItemType", bound="Item")
 
 MULTIPROCESSING_CONTEXT = "forkserver"
+MP_CONTEXT_ENV_VAR = "RSLEARN_MULTIPROCESSING_CONTEXT"
 
 
 def register_handler(category: Any, command: str) -> Callable:
@@ -837,7 +839,8 @@ def model_predict() -> None:
 def main() -> None:
     """CLI entrypoint."""
     try:
-        multiprocessing.set_start_method(MULTIPROCESSING_CONTEXT)
+        mp_context = os.environ.get(MP_CONTEXT_ENV_VAR, MULTIPROCESSING_CONTEXT)
+        multiprocessing.set_start_method(mp_context)
     except RuntimeError as e:
         logger.error(
             f"Multiprocessing context already set to {multiprocessing.get_context()}: "

rslearn/models/attention_pooling.py

@@ -150,8 +150,11 @@ class AttentionPool(IntermediateComponent):
             D // self.num_heads
         )
         attn_weights = F.softmax(attn_scores, dim=-1)
-        x = torch.matmul(attn_weights, v)  # [B, head, 1, D_head]
-        return x.reshape(B, D, H, W)
+        x = torch.matmul(attn_weights, v)  # [B*H*W, num_heads, 1, D_head]
+        x = x.squeeze(-2)  # [B*H*W, num_heads, D_head]
+        return rearrange(
+            x, "(b h w) nh dh -> b (nh dh) h w", b=B, h=H, w=W
+        )  # [B, D, H, W]
 
     def forward(self, intermediates: Any, context: ModelContext) -> FeatureMaps:
         """Forward pass for attention pooling linear probe.