rslearn 0.0.21__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/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

{rslearn-0.0.21.dist-info → rslearn-0.0.22.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rslearn
-Version: 0.0.21
+Version: 0.0.22
 Summary: A library for developing remote sensing datasets and models
 Author: OlmoEarth Team
 License:                                  Apache License
@@ -218,11 +218,13 @@ Requires-Dist: fsspec>=2025.10.0
 Requires-Dist: jsonargparse>=4.35.0
 Requires-Dist: lightning>=2.5.1.post0
 Requires-Dist: Pillow>=11.3
+Requires-Dist: pydantic>=2
 Requires-Dist: pyproj>=3.7
 Requires-Dist: python-dateutil>=2.9
 Requires-Dist: pytimeparse>=1.1
 Requires-Dist: rasterio>=1.4
 Requires-Dist: shapely>=2.1
+Requires-Dist: soilgrids>=0.1.4
 Requires-Dist: torch>=2.7.0
 Requires-Dist: torchvision>=0.22.0
 Requires-Dist: tqdm>=4.67
@@ -317,6 +319,7 @@ for how to setup these data sources.
 - Xyz (Slippy) Tiles (e.g., Mapbox tiles)
 - Planet Labs (PlanetScope, SkySat)
 - ESA WorldCover 2021
+- ISRIC SoilGrids (WCS)
 
 rslearn can also be used to easily mosaic, crop, and re-project any sets of local
 raster and vector files you may have.

{rslearn-0.0.21.dist-info → rslearn-0.0.22.dist-info}/RECORD

@@ -10,21 +10,24 @@ rslearn/config/__init__.py,sha256=n1qpZ0ImshTtLYl5mC73BORYyUcjPJyHiyZkqUY1hiY,47
 rslearn/config/dataset.py,sha256=AJlS-41B70E2Y4ST545J2P3Kz83jrFCRdkZecKZ7sQY,23255
 rslearn/data_sources/__init__.py,sha256=zzuZUxrlEIw84YpD2I0HJvCoLDB29LbmnKTXiJykzGU,660
 rslearn/data_sources/aws_landsat.py,sha256=0ZQtmd2NCnvLy4vFSB1AlmoguJbiQB_e_T4eS1tnW9Q,20443
-rslearn/data_sources/aws_open_data.py,sha256=lrHnMJTH3NAaRdNjxwCIxSq8rq90IvV4ho-qAG6Hdgc,29348
+rslearn/data_sources/aws_open_data.py,sha256=rDE9VZXXA669fD9jfOizxpN41lmvvJc-skZz3LettJU,29197
 rslearn/data_sources/aws_sentinel1.py,sha256=knBg2ZdwzGRIibUPAqnhYR-DbHqFL4tLsuRVrucTWU4,4745
+rslearn/data_sources/aws_sentinel2_element84.py,sha256=VkIAjCqzqi_m6StSQ-aY_AMtesFhkYYgwiUNkpBe7AA,13984
 rslearn/data_sources/climate_data_store.py,sha256=gwlbugwAT0YoOINOKNVlZuzpxUO5G0tfzwy4XVlURFg,18360
 rslearn/data_sources/copernicus.py,sha256=DLyXwnkFo2LzbxfLKHkOvHyZcNYFg5w-yXQPeBL67_w,35049
 rslearn/data_sources/data_source.py,sha256=lxdcxefETX5ui5uUdZn-ACSBhy6YVo4XDfk368EMD40,4862
 rslearn/data_sources/earthdaily.py,sha256=BNP7BK_vb9yQw6LaIaJ80vjCfBZ9TeALGkjHYIet_W0,18205
 rslearn/data_sources/earthdata_srtm.py,sha256=bwo8e_y9fFPliZ411tTWOiAUlEcb3AWBkeZxnkFY5SI,10633
 rslearn/data_sources/eurocrops.py,sha256=FQqdBcNl43fArq6rd_-iffVJyliIDbB0lIHvNdmtQBU,8663
-rslearn/data_sources/gcp_public_data.py,sha256=E3BhJqOgHMqwcfaxg7D47KNmfAwJDiDjH8-OdqNdjuI,36524
+rslearn/data_sources/gcp_public_data.py,sha256=qXWFQ4y0cCVBla68H4bpolG0FKW9_vLyTlA-0mJc1Y0,37349
 rslearn/data_sources/google_earth_engine.py,sha256=xdoSDjSVp6lPVPMv4UJZ6BRUozUA2hFbSTl1707TBoM,23523
 rslearn/data_sources/local_files.py,sha256=mo5W_BxBl89EPTIHNDEpXM6qBjrP225KK0PcmNgvJZQ,19090
 rslearn/data_sources/openstreetmap.py,sha256=TzZfouc2Z4_xjx2v_uv7aPn4tVW3flRVQN4qBfl507E,18161
 rslearn/data_sources/planet.py,sha256=6FWQ0bl1k3jwvwp4EVGi2qs3OD1QhnKOKP36mN4HELI,9446
 rslearn/data_sources/planet_basemap.py,sha256=e9R6FlagJjg8Z6Rc1dC6zK3xMkCohz8eohXqXmd29xg,9670
-rslearn/data_sources/planetary_computer.py,sha256=qgvGe_hQDxAIrRVB9szvJZL5IvE2jTo8UWuKCMzx7jM,29773
+rslearn/data_sources/planetary_computer.py,sha256=nTJ6Jh6CNBdCEIsn7G_xLQ0Nige5evdPdqLYmWTdDl4,20722
+rslearn/data_sources/soilgrids.py,sha256=rwO4goFPQ7lx420FvYBHYFXdihnZqn_-IjdqtxQ9j2g,12455
+rslearn/data_sources/stac.py,sha256=1qUbTD1fNSvBCX9QIXtyb9mGQ4K8ubRNIeEJs_I3QFU,9889
 rslearn/data_sources/usda_cdl.py,sha256=_WvxZkm0fbXfniRs6NT8iVCbTTmVPflDhsFT2ci6_Dk,6879
 rslearn/data_sources/usgs_landsat.py,sha256=kPOb3hsZe5-guUcFZZkwzcRpYZ3Zo7Bk4E829q_xiyU,18516
 rslearn/data_sources/utils.py,sha256=v_90ALOuts7RHNcx-j8o-aQ_aFjh8ZhXrmsaa9uEGDA,11651
@@ -46,7 +49,7 @@ rslearn/dataset/storage/file.py,sha256=g9HZ3CD4QcgyVNsBaXhjIKQgDOAeZ4R08sJ7ntx4w
 rslearn/dataset/storage/storage.py,sha256=DxZ7iwV938PiLwdQzb5EXSb4Mj8bRGmOTmA9fzq_Ge8,4840
 rslearn/models/__init__.py,sha256=_vWoF9d2Slah8-6XhYhdU4SRsy_CNxXjCGQTD2yvu3Q,22
 rslearn/models/anysat.py,sha256=nzk6hB83ltNFNXYRNA1rTvq2AQcAhwyvgBaZui1M37o,8107
-rslearn/models/attention_pooling.py,sha256=VeRoBLk0326Kpj780Pw8EiSFMU5_K9qg6HQO-B2r2PU,7044
+rslearn/models/attention_pooling.py,sha256=Kss1W_HYNzfX78RVDi18m0-cG2BoPsC240u0oEqJ0d4,7187
 rslearn/models/clip.py,sha256=QG1oUFqcVuNEhx7BNfJ1FnxIOMNUwRNBwXCe3CR6wFI,2415
 rslearn/models/component.py,sha256=uikFDzPYaW_LSXsrSsES1aup4IDIuWHsitWLpKgF7zU,3432
 rslearn/models/concatenate_features.py,sha256=Attemr5KurxlOojpclD0Pd5Cu2KHpNdpXe8jCSjpJ9U,3818
@@ -113,7 +116,7 @@ rslearn/train/__init__.py,sha256=fnJyY4aHs5zQqbDKSfXsJZXY_M9fbTsf7dRYaPwZr2M,30
 rslearn/train/all_patches_dataset.py,sha256=EVoYCmS3g4OfWPt5CZzwHVx9isbnWh5HIGA0RBqPFeA,21145
 rslearn/train/data_module.py,sha256=pgut8rEWHIieZ7RR8dUvhtlNqk0egEdznYF3tCvqdHg,23552
 rslearn/train/dataset.py,sha256=Jy1jU3GigfHaFeX9rbveX9bqy2Pd5Wh_vquD6_aFnS8,36522
-rslearn/train/lightning_module.py,sha256=HA9e-74oUZR5s7piQP9Mwxwz0vw0-p4HdmijfgBSpU0,14776
+rslearn/train/lightning_module.py,sha256=7WBAgdJhcMHueLKE2DthSFmNYvlNUh1dB4sibkqCsRA,14761
 rslearn/train/model_context.py,sha256=6o66BY6okBK-D5e0JUwPd7fxD_XehVaqxdQkJJKmQ3E,2580
 rslearn/train/optimizer.py,sha256=EKSqkmERalDA0bF32Gey7n6z69KLyaUWKlRsGJfKBmE,927
 rslearn/train/prediction_writer.py,sha256=rW0BUaYT_F1QqmpnQlbrLiLya1iBfC5Pb78G_NlF-vA,15956
@@ -126,12 +129,12 @@ rslearn/train/callbacks/peft.py,sha256=wEOKsS3RhsRaZTXn_Kz2wdsZdIiIaZPdCJWtdJBur
 rslearn/train/tasks/__init__.py,sha256=dag1u72x1-me6y0YcOubUo5MYZ0Tjf6-dOir9UeFNMs,75
 rslearn/train/tasks/classification.py,sha256=72ZBcbunMsdPYQN53S-4GfiLIDrr1X3Hni07dBJ0pu0,14261
 rslearn/train/tasks/detection.py,sha256=B0tfB7UGIbRtjnye3PhzLmfeQ4X7ImO3A-_LeNhBA54,21988
-rslearn/train/tasks/embedding.py,sha256=98ykdmfaxQjsH0UrUdTGmz1f0hCMPcNYTzt1YFqNQwQ,3869
+rslearn/train/tasks/embedding.py,sha256=NdJEAaDWlWYzvOBVf7eIHfFOzqTgavfFH1J1gMbAMVo,3891
 rslearn/train/tasks/multi_task.py,sha256=1ML9mZ-kM3JfElisLOWBUn4k12gsKTFjoYYgamnyxt8,6124
 rslearn/train/tasks/per_pixel_regression.py,sha256=znCLFaZbGx8lvIkntDXjcX7yy7giyyBdWN-TwTGaPV4,10197
 rslearn/train/tasks/regression.py,sha256=bVS_ApZSpbL0NaaM8Mu5Bsu4SBUyLpVtrPslulvvZHs,12695
 rslearn/train/tasks/segmentation.py,sha256=ie9ZV-sklLjQs35caiEglC1xff6dxeug_N-f_A8VosA,23034
-rslearn/train/tasks/task.py,sha256=B-z5XheXPTDaKauYzwze5ZD_O43R9uF76awgEk05bGE,3954
+rslearn/train/tasks/task.py,sha256=nMPunl9OlnOimr48saeTnwKMQ7Du4syGrwNKVQq4FL4,4110
 rslearn/train/transforms/__init__.py,sha256=BkCAzm4f-8TEhPIuyvCj7eJGh36aMkZFYlq-H_jkSvY,778
 rslearn/train/transforms/concatenate.py,sha256=hVVBaxIdk1Cx8JHPirj54TGpbWAJx5y_xD7k1rmGmT0,3166
 rslearn/train/transforms/crop.py,sha256=d5mrC8k7g4zz9J0RJrjz_pS1q06HfhUdzH6dUlL3Wqc,4668
@@ -147,7 +150,7 @@ rslearn/utils/__init__.py,sha256=GNvdTUmXakiEMnLdje7k1fe5aC7SFVqP757kbpN6Fzw,558
 rslearn/utils/array.py,sha256=RC7ygtPnQwU6Lb9kwORvNxatJcaJ76JPsykQvndAfes,2444
 rslearn/utils/feature.py,sha256=lsg0WThZDJzo1mrbaL04dXYI5G3x-n5FG9aEjj7uUaI,1649
 rslearn/utils/fsspec.py,sha256=h3fER_bkewzR9liEAULXguTIvXLUXA17pC_yZoWN5Tk,5902
-rslearn/utils/geometry.py,sha256=Z4alE4u0xa3fKGwoa4WQN5ZCjeqkPCJWfgcn8aw9t8I,21487
+rslearn/utils/geometry.py,sha256=S60uEvCQ3SyR2-3TJ1yUucFmGRM_rcoIT_gH83PN3RE,21485
 rslearn/utils/get_utm_ups_crs.py,sha256=kUrcyjCK7KWvuP1XR-nURPeRqYeRO-3L8QUJ1QTF9Ps,3599
 rslearn/utils/grid_index.py,sha256=hRmrtgpqN1pLa-djnZtgSXqKJlbgGyttGnCEmPLD0zo,2347
 rslearn/utils/jsonargparse.py,sha256=TRyZA151KzhjJlZczIHYguEA-YxCDYaZ2IwCRgx76nM,4791
@@ -156,12 +159,13 @@ rslearn/utils/raster_format.py,sha256=qZpbODF4I7BsOxf43O6vTmH2TSNw6N8PP0wsFUVvdI
 rslearn/utils/rtree_index.py,sha256=j0Zwrq3pXuAJ-hKpiRFQ7VNtvO3fZYk-Em2uBPAqfx4,6460
 rslearn/utils/spatial_index.py,sha256=eomJAUgzmjir8j9HZnSgQoJHwN9H0wGTjmJkMkLLfsU,762
 rslearn/utils/sqlite_index.py,sha256=YGOJi66544e6JNtfSft6YIlHklFdSJO2duxQ4TJ2iu4,2920
+rslearn/utils/stac.py,sha256=z93N5ZeEe1oUikX5ILMA5sQEZX276sAeMjsg0TShnSk,5776
 rslearn/utils/time.py,sha256=2ilSLG94_sxLP3y5RSV5L5CG8CoND_dbdzYEHVtN-I8,387
 rslearn/utils/vector_format.py,sha256=4ZDYpfBLLxguJkiIaavTagiQK2Sv4Rz9NumbHlq-3Lw,15041
-rslearn-0.0.21.dist-info/licenses/LICENSE,sha256=_99ZWPoLdlUbqZoSC5DF4ihiNwl5rTEmBaq2fACecdg,11352
-rslearn-0.0.21.dist-info/licenses/NOTICE,sha256=wLPr6rwV_jCg-xEknNGwhnkfRfuoOE9MZ-lru2yZyLI,5070
-rslearn-0.0.21.dist-info/METADATA,sha256=SUDvVod3HgoDdLYjRz9dwOL0c6V0SF-dTYuzsA4JJYk,37853
-rslearn-0.0.21.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-rslearn-0.0.21.dist-info/entry_points.txt,sha256=doTBQ57NT7nq-dgYGgTTw6mafcGWb_4PWYtYR4rGm50,46
-rslearn-0.0.21.dist-info/top_level.txt,sha256=XDKo90WBH8P9RQumHxo0giLJsoufT4r9odv-WE6Ahk4,8
-rslearn-0.0.21.dist-info/RECORD,,
+rslearn-0.0.22.dist-info/licenses/LICENSE,sha256=_99ZWPoLdlUbqZoSC5DF4ihiNwl5rTEmBaq2fACecdg,11352
+rslearn-0.0.22.dist-info/licenses/NOTICE,sha256=wLPr6rwV_jCg-xEknNGwhnkfRfuoOE9MZ-lru2yZyLI,5070
+rslearn-0.0.22.dist-info/METADATA,sha256=UArAfc_JYTffP8-cOwQf5mxh6XUtsRv5cwzFiLWNzLU,37936
+rslearn-0.0.22.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+rslearn-0.0.22.dist-info/entry_points.txt,sha256=doTBQ57NT7nq-dgYGgTTw6mafcGWb_4PWYtYR4rGm50,46
+rslearn-0.0.22.dist-info/top_level.txt,sha256=XDKo90WBH8P9RQumHxo0giLJsoufT4r9odv-WE6Ahk4,8
+rslearn-0.0.22.dist-info/RECORD,,