rslearn 0.0.20__py3-none-any.whl → 0.0.22__py3-none-any.whl

rslearn/data_sources/stac.py

@@ -0,0 +1,255 @@
+"""A partial data source implementation providing get_items using a STAC API."""
+
+import json
+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_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)
+            stac_items = self.client.search(
+                collections=[self.collection_name],
+                intersects=json.loads(shapely.to_geojson(wgs84_geometry.shp)),
+                date_time=wgs84_geometry.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/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.

rslearn/models/olmoearth_pretrain/model.py

@@ -159,20 +159,19 @@ class OlmoEarth(FeatureExtractor):
         that contains the distributed checkpoint. This is the format produced by
         pre-training runs in olmoearth_pretrain.
         """
-        # We avoid loading the train module here because it depends on running within
-        # olmo_core.
-        # Only pull in olmo_core when trying to load a distributed checkpoint to avoid dependency.
-        require_olmo_core("_load_model_from_checkpoint")
-        from olmo_core.distributed.checkpoint import load_model_and_optim_state
-
         with (checkpoint_upath / "config.json").open() as f:
             config_dict = json.load(f)
             model_config = Config.from_dict(config_dict["model"])
 
         model = model_config.build()
 
-        # Load the checkpoint.
+        # Load the checkpoint (requires olmo_core for distributed checkpoint loading).
         if not random_initialization:
+            require_olmo_core(
+                "_load_model_from_checkpoint with random_initialization=False"
+            )
+            from olmo_core.distributed.checkpoint import load_model_and_optim_state
+
             train_module_dir = checkpoint_upath / "model_and_optim"
             load_model_and_optim_state(str(train_module_dir), model)
             logger.info(f"loaded OlmoEarth encoder from {train_module_dir}")
@@ -242,7 +241,7 @@ class OlmoEarth(FeatureExtractor):
             present_modalities.append(modality)
             tensors = []
             for idx, inp in enumerate(context.inputs):
-                assert isinstance(inp, RasterImage)
+                assert isinstance(inp[modality], RasterImage)
                 tensors.append(inp[modality].image)
                 cur_timestamps = inp[modality].timestamps
                 if cur_timestamps is not None and len(cur_timestamps) > len(

rslearn/train/dataset.py

@@ -205,8 +205,7 @@ def read_raster_layer_for_data_input(
     group_idx: int,
     layer_config: LayerConfig,
     data_input: DataInput,
-    layer_data: WindowLayerData | None,
-) -> tuple[torch.Tensor, tuple[datetime, datetime] | None]:
+) -> torch.Tensor:
     """Read a raster layer for a DataInput.
 
     This scans the available rasters for the layer at the window to determine which
@@ -219,11 +218,9 @@ def read_raster_layer_for_data_input(
         group_idx: the item group.
         layer_config: the layer configuration.
         data_input: the DataInput that specifies the bands and dtype.
-        layer_data: the WindowLayerData associated with this layer and window.
 
     Returns:
-        RasterImage containing raster data and the timestamp associated
-            with that data.
+        Raster data as a tensor.
     """
     # See what different sets of bands we need to read to get all the
     # configured bands.
@@ -294,34 +291,46 @@ def read_raster_layer_for_data_input(
             src[src_indexes, :, :].astype(data_input.dtype.get_numpy_dtype())
         )
 
-    # add the timestamp. this is a tuple defining the start and end of the time range.
-    time_range = None
-    if layer_data is not None:
-        item = Item.deserialize(layer_data.serialized_item_groups[group_idx][0])
-        if item.geometry.time_range is not None:
-            # we assume if one layer data has a geometry & time range, all of them do
-            time_ranges = [
-                (
-                    datetime.fromisoformat(
-                        Item.deserialize(
-                            layer_data.serialized_item_groups[group_idx][idx]
-                        ).geometry.time_range[0]  # type: ignore
-                    ),
-                    datetime.fromisoformat(
-                        Item.deserialize(
-                            layer_data.serialized_item_groups[group_idx][idx]
-                        ).geometry.time_range[1]  # type: ignore
-                    ),
-                )
-                for idx in range(len(layer_data.serialized_item_groups[group_idx]))
-            ]
-            # take the min and max
-            time_range = (
-                min([t[0] for t in time_ranges]),
-                max([t[1] for t in time_ranges]),
+    return image
+
+
+def read_layer_time_range(
+    layer_data: WindowLayerData | None, group_idx: int
+) -> tuple[datetime, datetime] | None:
+    """Extract the combined time range from all items in a layer data group.
+
+    Returns the min start time and max end time across all items, or None if
+    no items have time ranges.
+
+    Raises:
+        ValueError: If some items have time_range and others don't.
+    """
+    if layer_data is None:
+        return None
+
+    serialized_items = layer_data.serialized_item_groups[group_idx]
+    if not serialized_items:
+        return None
+
+    first_item = Item.deserialize(serialized_items[0])
+    if first_item.geometry.time_range is None:
+        return None
+
+    # If the first item has a time_range, all items must have one
+    time_ranges: list[tuple[datetime, datetime]] = []
+    for serialized_item in serialized_items:
+        item = Item.deserialize(serialized_item)
+        if item.geometry.time_range is None:
+            raise ValueError(
+                f"Item '{item.name}' has no time_range, but first item does. "
+                "All items in a group must consistently have or lack time_range."
             )
+        time_ranges.append(item.geometry.time_range)
 
-    return image, time_range
+    return (
+        min(tr[0] for tr in time_ranges),
+        max(tr[1] for tr in time_ranges),
+    )
 
 
 def read_data_input(
@@ -378,17 +387,17 @@ def read_data_input(
         time_ranges: list[tuple[datetime, datetime] | None] = []
         for layer_name, group_idx in layers_to_read:
             layer_config = dataset.layers[layer_name]
-            image, time_range = read_raster_layer_for_data_input(
+            image = read_raster_layer_for_data_input(
                 window,
                 bounds,
                 layer_name,
                 group_idx,
                 layer_config,
                 data_input,
-                # some layers (e.g. "label_raster") won't have associated
-                # layer datas
-                layer_datas[layer_name] if layer_name in layer_datas else None,
             )
+            # some layers (e.g. "label_raster") won't have associated layer datas
+            layer_data = layer_datas.get(layer_name)
+            time_range = read_layer_time_range(layer_data, group_idx)
             if len(time_ranges) > 0:
                 if type(time_ranges[-1]) is not type(time_range):
                     raise ValueError(

rslearn/train/lightning_module.py

@@ -345,14 +345,14 @@ class RslearnLightningModule(L.LightningModule):
         )
 
         if self.visualize_dir:
-            for idx, (inp, target, output, metadata) in enumerate(
-                zip(inputs, targets, outputs, metadatas)
+            for inp, target, output, metadata in zip(
+                inputs, targets, outputs, metadatas
             ):
                 images = self.task.visualize(inp, target, output)
                 for image_suffix, image in images.items():
                     out_fname = os.path.join(
                         self.visualize_dir,
-                        f"{metadata['window_name']}_{metadata['bounds'][0]}_{metadata['bounds'][1]}_{image_suffix}.png",
+                        f"{metadata.window_name}_{metadata.patch_bounds[0]}_{metadata.patch_bounds[1]}_{image_suffix}.png",
                     )
                     Image.fromarray(image).save(out_fname)
 

rslearn/train/tasks/embedding.py

@@ -6,7 +6,7 @@ import numpy.typing as npt
 import torch
 from torchmetrics import MetricCollection
 
-from rslearn.models.component import FeatureMaps
+from rslearn.models.component import FeatureMaps, Predictor
 from rslearn.train.model_context import ModelContext, ModelOutput, SampleMetadata
 from rslearn.utils import Feature
 
@@ -83,7 +83,7 @@ class EmbeddingTask(Task):
         return MetricCollection({})
 
 
-class EmbeddingHead:
+class EmbeddingHead(Predictor):
     """Head for embedding task.
 
     It just adds a dummy loss to act as a Predictor.

rslearn/train/tasks/task.py

@@ -108,8 +108,10 @@ class BasicTask(Task):
         Returns:
             a dictionary mapping image name to visualization image
         """
-        image = input_dict["image"].cpu()
-        image = image[self.image_bands, :, :]
+        raster_image = input_dict["image"]
+        assert isinstance(raster_image, RasterImage)
+        # We don't really handle time series here, just use the first timestep.
+        image = raster_image.image.cpu()[self.image_bands, 0, :, :]
         if self.remap_values:
             factor = (self.remap_values[1][1] - self.remap_values[1][0]) / (
                 self.remap_values[0][1] - self.remap_values[0][0]

rslearn/utils/geometry.py

@@ -153,8 +153,8 @@ class ResolutionFactor:
         else:
             return Projection(
                 projection.crs,
-                projection.x_resolution // self.numerator,
-                projection.y_resolution // self.numerator,
+                projection.x_resolution / self.numerator,
+                projection.y_resolution / self.numerator,
             )
 
     def multiply_bounds(self, bounds: PixelBounds) -> PixelBounds:

rslearn/utils/stac.py

@@ -0,0 +1,173 @@
+"""STAC API client."""
+
+import logging
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any
+
+import requests
+
+logger = logging.getLogger(__name__)
+
+Bbox = tuple[float, float, float, float]
+
+
+@dataclass(frozen=True)
+class StacAsset:
+    """A STAC asset."""
+
+    href: str
+    title: str | None
+    type: str | None
+    roles: list[str] | None
+
+
+@dataclass(frozen=True)
+class StacItem:
+    """A STAC item."""
+
+    id: str
+    properties: dict[str, Any]
+    collection: str | None
+    bbox: Bbox | None
+    geometry: dict[str, Any] | None
+    assets: dict[str, StacAsset] | None
+    time_range: tuple[datetime, datetime] | None
+
+    @classmethod
+    def from_dict(cls, item: dict[str, Any]) -> "StacItem":
+        """Create a STAC item from the item dict returned from API."""
+        properties = item.get("properties", {})
+
+        # Parse bbox.
+        bbox: Bbox | None = None
+        if "bbox" in item:
+            if len(item["bbox"]) != 4:
+                raise NotImplementedError(
+                    f"got bbox with {len(item['bbox'])} coordinates but only 4 coordinates is implemented"
+                )
+            bbox = tuple(item["bbox"])
+
+        # Parse assets.
+        assets: dict[str, StacAsset] = {}
+        for name, asset in item.get("assets", {}).items():
+            assets[name] = StacAsset(
+                href=asset["href"],
+                title=asset.get("title"),
+                type=asset.get("type"),
+                roles=asset.get("roles"),
+            )
+
+        # Parse time range.
+        time_range: tuple[datetime, datetime] | None = None
+        if "start_datetime" in properties and "end_datetime" in properties:
+            time_range = (
+                datetime.fromisoformat(properties["start_datetime"]),
+                datetime.fromisoformat(properties["end_datetime"]),
+            )
+        elif "datetime" in properties:
+            ts = datetime.fromisoformat(properties["datetime"])
+            time_range = (ts, ts)
+
+        return cls(
+            id=item["id"],
+            properties=properties,
+            collection=item.get("collection"),
+            bbox=bbox,
+            geometry=item.get("geometry"),
+            assets=assets,
+            time_range=time_range,
+        )
+
+
+class StacClient:
+    """Limited functionality client for STAC APIs."""
+
+    def __init__(self, endpoint: str):
+        """Create a new StacClient.
+
+        Args:
+            endpoint: the STAC endpoint (base URL)
+        """
+        self.endpoint = endpoint
+        self.session = requests.Session()
+
+    def search(
+        self,
+        collections: list[str] | None = None,
+        bbox: Bbox | None = None,
+        intersects: dict[str, Any] | None = None,
+        date_time: datetime | tuple[datetime, datetime] | None = None,
+        ids: list[str] | None = None,
+        limit: int | None = None,
+        query: dict[str, Any] | None = None,
+    ) -> list[StacItem]:
+        """Execute a STAC item search.
+
+        We use the JSON POST API. Pagination is handled so the returned items are
+        concatenated across all available pages.
+
+        Args:
+            collections: only search within the provided collection(s).
+            bbox: only return features intersecting the provided bounding box.
+            intersects: only return features intersecting this GeoJSON geometry.
+            date_time: only return features that have a temporal property intersecting
+                the provided time range or timestamp.
+            ids: only return the provided item IDs.
+            limit: number of items per page. We will read all the pages.
+            query: query dict, if STAC query extension is supported by this API. See
+                https://github.com/stac-api-extensions/query.
+
+        Returns:
+            list of matching STAC items.
+        """
+        # Build JSON request data.
+        request_data: dict[str, Any] = {}
+        if collections is not None:
+            request_data["collections"] = collections
+        if bbox is not None:
+            request_data["bbox"] = bbox
+        if intersects is not None:
+            request_data["intersects"] = intersects
+        if date_time is not None:
+            if isinstance(date_time, tuple):
+                start_time = date_time[0].isoformat().replace("+00:00", "Z")
+                end_time = date_time[1].isoformat().replace("+00:00", "Z")
+                request_data["datetime"] = f"{start_time}/{end_time}"
+            else:
+                request_data["datetime"] = date_time.isoformat().replace("+00:00", "Z")
+        if ids is not None:
+            request_data["ids"] = ids
+        if limit is not None:
+            request_data["limit"] = limit
+        if query is not None:
+            request_data["query"] = query
+
+        # Handle pagination.
+        cur_url = self.endpoint + "/search"
+        items: list[StacItem] = []
+        while True:
+            logger.debug("Reading STAC items from %s", cur_url)
+            response = self.session.post(url=cur_url, json=request_data)
+            response.raise_for_status()
+            data = response.json()
+            for item_dict in data["features"]:
+                items.append(StacItem.from_dict(item_dict))
+
+            next_link = None
+            next_request_data: dict[str, Any] = {}
+            for link in data.get("links", []):
+                if "rel" not in link or link["rel"] != "next":
+                    continue
+                assert link["method"] == "POST"
+                next_link = link["href"]
+                next_request_data = link["body"]
+                break
+
+            if next_link is None:
+                break
+
+            cur_url = next_link
+            request_data = next_request_data
+
+        return items