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

rslearn/data_sources/planetary_computer.py

@@ -3,7 +3,7 @@
 import os
 import tempfile
 import xml.etree.ElementTree as ET
-from datetime import timedelta
+from datetime import datetime, timedelta
 from typing import Any
 
 import affine
@@ -12,6 +12,7 @@ import planetary_computer
 import rasterio
 import requests
 from rasterio.enums import Resampling
+from typing_extensions import override
 from upath import UPath
 
 from rslearn.config import LayerConfig
@@ -24,11 +25,104 @@ from rslearn.tile_stores import TileStore, TileStoreWithLayer
 from rslearn.utils.fsspec import join_upath
 from rslearn.utils.geometry import PixelBounds, Projection, STGeometry
 from rslearn.utils.raster_format import get_raster_projection_and_bounds
+from rslearn.utils.stac import StacClient, StacItem
 
 from .copernicus import get_harmonize_callback
 
 logger = get_logger(__name__)
 
+# Max limit accepted by Planetary Computer API.
+PLANETARY_COMPUTER_LIMIT = 1000
+
+
+class PlanetaryComputerStacClient(StacClient):
+    """A StacClient subclass that handles Planetary Computer's pagination limits.
+
+    Planetary Computer STAC API does not support standard pagination and has a max
+    limit of 1000. If the initial query returns 1000 items, this client paginates
+    by sorting by ID and using gt (greater than) queries to fetch subsequent pages.
+    """
+
+    @override
+    def search(
+        self,
+        collections: list[str] | None = None,
+        bbox: tuple[float, float, float, float] | 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,
+        sortby: list[dict[str, str]] | None = None,
+    ) -> list[StacItem]:
+        # We will use sortby for pagination, so the caller must not set it.
+        if sortby is not None:
+            raise ValueError("sortby must not be set for PlanetaryComputerStacClient")
+
+        # First, try a simple query with the PC limit to detect if pagination is needed.
+        # We always use PLANETARY_COMPUTER_LIMIT for the request because PC doesn't
+        # support standard pagination, and we need to detect when we hit the limit
+        # to switch to ID-based pagination.
+        # We could just start sorting by ID here and do pagination, but we treate it as
+        # a special case to avoid sorting since that seems to speed up the query.
+        stac_items = super().search(
+            collections=collections,
+            bbox=bbox,
+            intersects=intersects,
+            date_time=date_time,
+            ids=ids,
+            limit=PLANETARY_COMPUTER_LIMIT,
+            query=query,
+        )
+
+        # If we got fewer than the PC limit, we have all the results.
+        if len(stac_items) < PLANETARY_COMPUTER_LIMIT:
+            return stac_items
+
+        # We hit the limit, so we need to paginate by ID.
+        # Re-fetch with sorting by ID to ensure consistent ordering for pagination.
+        logger.debug(
+            "Initial request returned %d items (at limit), switching to ID pagination",
+            len(stac_items),
+        )
+
+        all_items: list[StacItem] = []
+        last_id: str | None = None
+
+        while True:
+            # Build query with id > last_id if we're paginating.
+            combined_query: dict[str, Any] = dict(query) if query else {}
+            if last_id is not None:
+                combined_query["id"] = {"gt": last_id}
+
+            stac_items = super().search(
+                collections=collections,
+                bbox=bbox,
+                intersects=intersects,
+                date_time=date_time,
+                ids=ids,
+                limit=PLANETARY_COMPUTER_LIMIT,
+                query=combined_query if combined_query else None,
+                sortby=[{"field": "id", "direction": "asc"}],
+            )
+
+            all_items.extend(stac_items)
+
+            # If we got fewer than the limit, we've fetched everything.
+            if len(stac_items) < PLANETARY_COMPUTER_LIMIT:
+                break
+
+            # Otherwise, paginate using the last item's ID.
+            last_id = stac_items[-1].id
+            logger.debug(
+                "Got %d items, paginating with id > %s",
+                len(stac_items),
+                last_id,
+            )
+
+        logger.debug("Total items fetched: %d", len(all_items))
+        return all_items
+
 
 class PlanetaryComputer(StacDataSource, TileStore):
     """Modality-agnostic data source for data on Microsoft Planetary Computer.
@@ -100,6 +194,10 @@ class PlanetaryComputer(StacDataSource, TileStore):
             required_assets=required_assets,
             cache_dir=cache_upath,
         )
+
+        # Replace the client with PlanetaryComputerStacClient to handle PC's pagination limits.
+        self.client = PlanetaryComputerStacClient(self.STAC_ENDPOINT)
+
         self.asset_bands = asset_bands
         self.timeout = timeout
         self.skip_items_missing_assets = skip_items_missing_assets
@@ -567,3 +665,53 @@ class Naip(PlanetaryComputer):
             context=context,
             **kwargs,
         )
+
+
+class CopDemGlo30(PlanetaryComputer):
+    """A data source for Copernicus DEM GLO-30 (30m) on Microsoft Planetary Computer.
+
+    See https://planetarycomputer.microsoft.com/dataset/cop-dem-glo-30.
+    """
+
+    COLLECTION_NAME = "cop-dem-glo-30"
+    DATA_ASSET = "data"
+
+    def __init__(
+        self,
+        band_name: str = "DEM",
+        context: DataSourceContext = DataSourceContext(),
+        **kwargs: Any,
+    ):
+        """Initialize a new CopDemGlo30 instance.
+
+        Args:
+            band_name: band name to use if the layer config is missing from the
+                context.
+            context: the data source context.
+            kwargs: additional arguments to pass to PlanetaryComputer.
+        """
+        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")
+            band_name = context.layer_config.band_sets[0].bands[0]
+
+        super().__init__(
+            collection_name=self.COLLECTION_NAME,
+            asset_bands={self.DATA_ASSET: [band_name]},
+            # Skip since all items should have the same asset(s).
+            skip_items_missing_assets=True,
+            context=context,
+            **kwargs,
+        )
+
+    def _stac_item_to_item(self, stac_item: Any) -> SourceItem:
+        # Copernicus DEM is static; ignore item timestamps so it matches any window.
+        item = super()._stac_item_to_item(stac_item)
+        item.geometry = STGeometry(item.geometry.projection, item.geometry.shp, None)
+        return item
+
+    def _get_search_time_range(self, geometry: STGeometry) -> None:
+        # Copernicus DEM is static; do not filter STAC searches by time.
+        return None

rslearn/data_sources/stac.py

@@ -1,6 +1,7 @@
 """A partial data source implementation providing get_items using a STAC API."""
 
 import json
+from datetime import datetime
 from typing import Any
 
 import shapely
@@ -11,6 +12,7 @@ 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.fsspec import open_atomic
 from rslearn.utils.geometry import STGeometry
 from rslearn.utils.stac import StacClient, StacItem
 
@@ -132,6 +134,24 @@ class StacDataSource(ItemLookupDataSource[SourceItem]):
 
         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.
 
@@ -168,7 +188,7 @@ class StacDataSource(ItemLookupDataSource[SourceItem]):
 
         # Finally we cache it if cache_dir is set.
         if cache_fname is not None:
-            with cache_fname.open("w") as f:
+            with open_atomic(cache_fname, "w") as f:
                 json.dump(item.serialize(), f)
 
         return item
@@ -191,10 +211,11 @@ class StacDataSource(ItemLookupDataSource[SourceItem]):
             # 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=wgs84_geometry.time_range,
+                date_time=search_time_range,
                 query=self.query,
                 limit=self.limit,
             )
@@ -239,7 +260,7 @@ class StacDataSource(ItemLookupDataSource[SourceItem]):
                     cache_fname = self.cache_dir / f"{item.name}.json"
                     if cache_fname.exists():
                         continue
-                    with cache_fname.open("w") as f:
+                    with open_atomic(cache_fname, "w") as f:
                         json.dump(item.serialize(), f)
 
             cur_groups = match_candidate_items_to_window(

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/simple_time_series.py

@@ -180,7 +180,7 @@ class SimpleTimeSeries(FeatureExtractor):
         # want to pass 2 timesteps to the model.
         # TODO is probably to make this behaviour clearer but lets leave it like
         # this for now to not break things.
-        num_timesteps = images.shape[1] // image_channels
+        num_timesteps = image_channels // images.shape[1]
         batched_timesteps = images.shape[2] // num_timesteps
         images = rearrange(
             images,

rslearn/train/lightning_module.py

@@ -210,11 +210,30 @@ class RslearnLightningModule(L.LightningModule):
             # Fail silently for single-dataset case, which is okay
             pass
 
+    def on_validation_epoch_end(self) -> None:
+        """Compute and log validation metrics at epoch end.
+
+        We manually compute and log metrics here (instead of passing the MetricCollection
+        to log_dict) because MetricCollection.compute() properly flattens dict-returning
+        metrics, while log_dict expects each metric to return a scalar tensor.
+        """
+        metrics = self.val_metrics.compute()
+        self.log_dict(metrics)
+        self.val_metrics.reset()
+
     def on_test_epoch_end(self) -> None:
-        """Optionally save the test metrics to a file."""
+        """Compute and log test metrics at epoch end, optionally save to file.
+
+        We manually compute and log metrics here (instead of passing the MetricCollection
+        to log_dict) because MetricCollection.compute() properly flattens dict-returning
+        metrics, while log_dict expects each metric to return a scalar tensor.
+        """
+        metrics = self.test_metrics.compute()
+        self.log_dict(metrics)
+        self.test_metrics.reset()
+
         if self.metrics_file:
             with open(self.metrics_file, "w") as f:
-                metrics = self.test_metrics.compute()
                 metrics_dict = {k: v.item() for k, v in metrics.items()}
                 json.dump(metrics_dict, f, indent=4)
                 logger.info(f"Saved metrics to {self.metrics_file}")
@@ -300,9 +319,6 @@ class RslearnLightningModule(L.LightningModule):
             sync_dist=True,
         )
         self.val_metrics.update(outputs, targets)
-        self.log_dict(
-            self.val_metrics, batch_size=batch_size, on_epoch=True, sync_dist=True
-        )
 
     def test_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None:
         """Compute the test loss and additional metrics.
@@ -340,9 +356,6 @@ class RslearnLightningModule(L.LightningModule):
             sync_dist=True,
         )
         self.test_metrics.update(outputs, targets)
-        self.log_dict(
-            self.test_metrics, batch_size=batch_size, on_epoch=True, sync_dist=True
-        )
 
         if self.visualize_dir:
             for inp, target, output, metadata in zip(

rslearn/train/tasks/multi_task.py

@@ -118,13 +118,16 @@ class MultiTask(Task):
 
     def get_metrics(self) -> MetricCollection:
         """Get metrics for this task."""
-        metrics = []
+        # Flatten metrics into a single dict with task_name/ prefix to avoid nested
+        # MetricCollections. Nested collections cause issues because MetricCollection
+        # has postfix=None which breaks MetricCollection.compute().
+        all_metrics = {}
         for task_name, task in self.tasks.items():
-            cur_metrics = {}
             for metric_name, metric in task.get_metrics().items():
-                cur_metrics[metric_name] = MetricWrapper(task_name, metric)
-            metrics.append(MetricCollection(cur_metrics, prefix=f"{task_name}/"))
-        return MetricCollection(metrics)
+                all_metrics[f"{task_name}/{metric_name}"] = MetricWrapper(
+                    task_name, metric
+                )
+        return MetricCollection(all_metrics)
 
 
 class MetricWrapper(Metric):

rslearn/train/tasks/per_pixel_regression.py

@@ -100,7 +100,7 @@ class PerPixelRegressionTask(BasicTask):
             raise ValueError(
                 f"PerPixelRegressionTask output must be an HW tensor, but got shape {raw_output.shape}"
             )
-        return (raw_output / self.scale_factor).cpu().numpy()
+        return (raw_output[None, :, :] / self.scale_factor).cpu().numpy()
 
     def visualize(
         self,

rslearn/train/tasks/segmentation.py

@@ -53,11 +53,14 @@ class SegmentationTask(BasicTask):
         enable_accuracy_metric: bool = True,
         enable_miou_metric: bool = False,
         enable_f1_metric: bool = False,
+        report_metric_per_class: bool = False,
         f1_metric_thresholds: list[list[float]] = [[0.5]],
         metric_kwargs: dict[str, Any] = {},
         miou_metric_kwargs: dict[str, Any] = {},
         prob_scales: list[float] | None = None,
         other_metrics: dict[str, Metric] = {},
+        output_probs: bool = False,
+        output_class_idx: int | None = None,
         **kwargs: Any,
     ) -> None:
         """Initialize a new SegmentationTask.
@@ -74,6 +77,8 @@ class SegmentationTask(BasicTask):
             enable_accuracy_metric: whether to enable the accuracy metric (default
                 true).
             enable_f1_metric: whether to enable the F1 metric (default false).
+            report_metric_per_class: whether to report chosen metrics for each class, in
+                addition to the average score across classes.
             enable_miou_metric: whether to enable the mean IoU metric (default false).
             f1_metric_thresholds: list of list of thresholds to apply for F1 metric.
                 Each inner list is used to initialize a separate F1 metric where the
@@ -89,6 +94,10 @@ class SegmentationTask(BasicTask):
                 this is only applied during prediction, not when computing val or test
                 metrics.
             other_metrics: additional metrics to configure on this task.
+            output_probs: if True, output raw softmax probabilities instead of class IDs
+                during prediction.
+            output_class_idx: if set along with output_probs, only output the probability
+                for this specific class index (single-channel output).
             kwargs: additional arguments to pass to BasicTask
         """
         super().__init__(**kwargs)
@@ -107,11 +116,14 @@ class SegmentationTask(BasicTask):
         self.enable_accuracy_metric = enable_accuracy_metric
         self.enable_f1_metric = enable_f1_metric
         self.enable_miou_metric = enable_miou_metric
+        self.report_metric_per_class = report_metric_per_class
         self.f1_metric_thresholds = f1_metric_thresholds
         self.metric_kwargs = metric_kwargs
         self.miou_metric_kwargs = miou_metric_kwargs
         self.prob_scales = prob_scales
         self.other_metrics = other_metrics
+        self.output_probs = output_probs
+        self.output_class_idx = output_class_idx
 
     def process_inputs(
         self,
@@ -167,7 +179,9 @@ class SegmentationTask(BasicTask):
             metadata: metadata about the patch being read
 
         Returns:
-            CHW numpy array with one channel, containing the predicted class IDs.
+            CHW numpy array. If output_probs is False, returns one channel with
+            predicted class IDs. If output_probs is True, returns softmax probabilities
+            (num_classes channels, or 1 channel if output_class_idx is set).
         """
         if not isinstance(raw_output, torch.Tensor) or len(raw_output.shape) != 3:
             raise ValueError("the output for SegmentationTask must be a CHW tensor")
@@ -179,6 +193,15 @@ class SegmentationTask(BasicTask):
                     self.prob_scales, device=raw_output.device, dtype=raw_output.dtype
                 )[:, None, None]
             )
+
+        if self.output_probs:
+            # Return raw softmax probabilities
+            probs = raw_output.cpu().numpy()
+            if self.output_class_idx is not None:
+                # Return only the specified class probability
+                return probs[self.output_class_idx : self.output_class_idx + 1, :, :]
+            return probs
+
         classes = raw_output.argmax(dim=0).cpu().numpy()
         return classes[None, :, :]
 
@@ -237,29 +260,41 @@ class SegmentationTask(BasicTask):
                     # Metric name can't contain "." so change to ",".
                     suffix = "_" + str(thresholds[0]).replace(".", ",")
 
+                # Create one metric per type - it returns a dict with "avg" and optionally per-class keys
                 metrics["F1" + suffix] = SegmentationMetric(
-                    F1Metric(num_classes=self.num_classes, score_thresholds=thresholds)
+                    F1Metric(
+                        num_classes=self.num_classes,
+                        score_thresholds=thresholds,
+                        report_per_class=self.report_metric_per_class,
+                    ),
                 )
                 metrics["precision" + suffix] = SegmentationMetric(
                     F1Metric(
                         num_classes=self.num_classes,
                         score_thresholds=thresholds,
                         metric_mode="precision",
-                    )
+                        report_per_class=self.report_metric_per_class,
+                    ),
                 )
                 metrics["recall" + suffix] = SegmentationMetric(
                     F1Metric(
                         num_classes=self.num_classes,
                         score_thresholds=thresholds,
                         metric_mode="recall",
-                    )
+                        report_per_class=self.report_metric_per_class,
+                    ),
                 )
 
         if self.enable_miou_metric:
-            miou_metric_kwargs: dict[str, Any] = dict(num_classes=self.num_classes)
+            miou_metric_kwargs: dict[str, Any] = dict(
+                num_classes=self.num_classes,
+                report_per_class=self.report_metric_per_class,
+            )
             if self.nodata_value is not None:
                 miou_metric_kwargs["nodata_value"] = self.nodata_value
             miou_metric_kwargs.update(self.miou_metric_kwargs)
+
+            # Create one metric - it returns a dict with "avg" and optionally per-class keys
             metrics["mean_iou"] = SegmentationMetric(
                 MeanIoUMetric(**miou_metric_kwargs),
                 pass_probabilities=False,
@@ -274,6 +309,20 @@ class SegmentationTask(BasicTask):
 class SegmentationHead(Predictor):
     """Head for segmentation task."""
 
+    def __init__(self, weights: list[float] | None = None, dice_loss: bool = False):
+        """Initialize a new SegmentationTask.
+
+        Args:
+            weights: weights for cross entropy loss (Tensor of size C)
+            dice_loss: weather to add dice loss to cross entropy
+        """
+        super().__init__()
+        if weights is not None:
+            self.register_buffer("weights", torch.Tensor(weights))
+        else:
+            self.weights = None
+        self.dice_loss = dice_loss
+
     def forward(
         self,
         intermediates: Any,
@@ -308,7 +357,7 @@ class SegmentationHead(Predictor):
             labels = torch.stack([target["classes"] for target in targets], dim=0)
             mask = torch.stack([target["valid"] for target in targets], dim=0)
             per_pixel_loss = torch.nn.functional.cross_entropy(
-                logits, labels, reduction="none"
+                logits, labels, weight=self.weights, reduction="none"
             )
             mask_sum = torch.sum(mask)
             if mask_sum > 0:
@@ -318,6 +367,9 @@ class SegmentationHead(Predictor):
                 # If there are no valid pixels, we avoid dividing by zero and just let
                 # the summed mask loss be zero.
                 losses["cls"] = torch.sum(per_pixel_loss * mask)
+            if self.dice_loss:
+                dice_loss = DiceLoss()(outputs, labels, mask)
+                losses["dice"] = dice_loss
 
         return ModelOutput(
             outputs=outputs,
@@ -333,6 +385,7 @@ class SegmentationMetric(Metric):
         metric: Metric,
         pass_probabilities: bool = True,
         class_idx: int | None = None,
+        output_key: str | None = None,
     ):
         """Initialize a new SegmentationMetric.
 
@@ -341,12 +394,19 @@ class SegmentationMetric(Metric):
                 classes from the targets and masking out invalid pixels.
             pass_probabilities: whether to pass predicted probabilities to the metric.
                 If False, argmax is applied to pass the predicted classes instead.
-            class_idx: if metric returns value for multiple classes, select this class.
+            class_idx: if set, return only this class index's value. For backward
+                compatibility with configs using standard torchmetrics. Internally
+                converted to output_key="cls_{class_idx}".
+            output_key: if the wrapped metric returns a dict (or a tensor that gets
+                converted to a dict), return only this key's value. For standard
+                torchmetrics with average=None, tensors are converted to dicts with
+                keys "cls_0", "cls_1", etc. If None, the full dict is returned.
         """
         super().__init__()
         self.metric = metric
         self.pass_probablities = pass_probabilities
         self.class_idx = class_idx
+        self.output_key = output_key
 
     def update(
         self, preds: list[Any] | torch.Tensor, targets: list[dict[str, Any]]
@@ -376,10 +436,32 @@ class SegmentationMetric(Metric):
         self.metric.update(preds, labels)
 
     def compute(self) -> Any:
-        """Returns the computed metric."""
+        """Returns the computed metric.
+
+        If the wrapped metric returns a multi-element tensor (e.g., standard torchmetrics
+        with average=None), it is converted to a dict with keys like "cls_0", "cls_1", etc.
+        This allows uniform handling via output_key for both standard torchmetrics and
+        custom dict-returning metrics.
+        """
         result = self.metric.compute()
+
+        # Convert multi-element tensors to dict for uniform handling.
+        # This supports standard torchmetrics with average=None which return per-class tensors.
+        if isinstance(result, torch.Tensor) and result.ndim >= 1:
+            result = {f"cls_{i}": result[i] for i in range(len(result))}
+
+        if self.output_key is not None:
+            if not isinstance(result, dict):
+                raise TypeError(
+                    f"output_key is set to '{self.output_key}' but metric returned "
+                    f"{type(result).__name__} instead of dict"
+                )
+            return result[self.output_key]
         if self.class_idx is not None:
-            result = result[self.class_idx]
+            # For backward compatibility: class_idx can index into the converted dict
+            if isinstance(result, dict):
+                return result[f"cls_{self.class_idx}"]
+            return result[self.class_idx]
         return result
 
     def reset(self) -> None:
@@ -404,6 +486,7 @@ class F1Metric(Metric):
         num_classes: int,
         score_thresholds: list[float],
         metric_mode: str = "f1",
+        report_per_class: bool = False,
     ):
         """Create a new F1Metric.
 
@@ -413,11 +496,14 @@ class F1Metric(Metric):
                 metric is the best F1 across score thresholds.
             metric_mode: set to "precision" or "recall" to return that instead of F1
                 (default "f1")
+            report_per_class: whether to include per-class scores in the output dict.
+                If False, only returns the "avg" key.
         """
         super().__init__()
         self.num_classes = num_classes
         self.score_thresholds = score_thresholds
         self.metric_mode = metric_mode
+        self.report_per_class = report_per_class
 
         assert self.metric_mode in ["f1", "precision", "recall"]
 
@@ -462,9 +548,10 @@ class F1Metric(Metric):
         """Compute metric.
 
         Returns:
-            the best F1 score across score thresholds and classes.
+            dict with "avg" key containing mean score across classes.
+            If report_per_class is True, also includes "cls_N" keys for each class N.
         """
-        best_scores = []
+        cls_best_scores = {}
 
         for cls_idx in range(self.num_classes):
             best_score = None
@@ -501,9 +588,12 @@ class F1Metric(Metric):
                 if best_score is None or score > best_score:
                     best_score = score
 
-            best_scores.append(best_score)
+            cls_best_scores[f"cls_{cls_idx}"] = best_score
 
-        return torch.mean(torch.stack(best_scores))
+        report_scores = {"avg": torch.mean(torch.stack(list(cls_best_scores.values())))}
+        if self.report_per_class:
+            report_scores.update(cls_best_scores)
+        return report_scores
 
 
 class MeanIoUMetric(Metric):
@@ -523,7 +613,7 @@ class MeanIoUMetric(Metric):
         num_classes: int,
         nodata_value: int | None = None,
         ignore_missing_classes: bool = False,
-        class_idx: int | None = None,
+        report_per_class: bool = False,
     ):
         """Create a new MeanIoUMetric.
 
@@ -535,15 +625,14 @@ class MeanIoUMetric(Metric):
             ignore_missing_classes: whether to ignore classes that don't appear in
                 either the predictions or the ground truth. If false, the IoU for a
                 missing class will be 0.
-            class_idx: only compute and return the IoU for this class. This option is
-                provided so the user can get per-class IoU results, since Lightning
-                only supports scalar return values from metrics.
+            report_per_class: whether to include per-class IoU scores in the output dict.
+                If False, only returns the "avg" key.
         """
         super().__init__()
         self.num_classes = num_classes
         self.nodata_value = nodata_value
         self.ignore_missing_classes = ignore_missing_classes
-        self.class_idx = class_idx
+        self.report_per_class = report_per_class
 
         self.add_state(
             "intersections", default=torch.zeros(self.num_classes), dist_reduce_fx="sum"
@@ -584,9 +673,11 @@ class MeanIoUMetric(Metric):
         """Compute metric.
 
         Returns:
-            the mean IoU across classes.
+            dict with "avg" containing the mean IoU across classes.
+            If report_per_class is True, also includes "cls_N" keys for each valid class N.
         """
-        per_class_scores = []
+        cls_scores = {}
+        valid_scores = []
 
         for cls_idx in range(self.num_classes):
             # Check if nodata_value is set and is one of the classes
@@ -599,6 +690,56 @@ class MeanIoUMetric(Metric):
             if union == 0 and self.ignore_missing_classes:
                 continue
 
-            per_class_scores.append(intersection / union)
+            score = intersection / union
+            cls_scores[f"cls_{cls_idx}"] = score
+            valid_scores.append(score)
+
+        report_scores = {"avg": torch.mean(torch.stack(valid_scores))}
+        if self.report_per_class:
+            report_scores.update(cls_scores)
+        return report_scores
+
+
+class DiceLoss(torch.nn.Module):
+    """Mean Dice Loss for segmentation.
+
+    This is the mean of the per-class dice loss (1 - 2*intersection / union scores).
+    The per-class intersection is the number of pixels across all examples where
+    the predicted label and ground truth label are both that class, and the per-class
+    union is defined similarly.
+    """
+
+    def __init__(self, smooth: float = 1e-7):
+        """Initialize a new DiceLoss."""
+        super().__init__()
+        self.smooth = smooth
+
+    def forward(
+        self, inputs: torch.Tensor, targets: torch.Tensor, mask: torch.Tensor
+    ) -> torch.Tensor:
+        """Compute Dice Loss.
+
+        Returns:
+            the mean Dicen Loss across classes
+        """
+        num_classes = inputs.shape[1]
+        targets_one_hot = (
+            torch.nn.functional.one_hot(targets, num_classes)
+            .permute(0, 3, 1, 2)
+            .float()
+        )
+
+        # Expand mask to [B, C, H, W]
+        mask = mask.unsqueeze(1).expand_as(inputs)
+
+        dice_per_class = []
+        for c in range(num_classes):
+            pred_c = inputs[:, c] * mask[:, c]
+            target_c = targets_one_hot[:, c] * mask[:, c]
+
+            intersection = (pred_c * target_c).sum()
+            union = pred_c.sum() + target_c.sum()
+            dice_c = (2.0 * intersection + self.smooth) / (union + self.smooth)
+            dice_per_class.append(dice_c)
 
-        return torch.mean(torch.stack(per_class_scores))
+        return 1 - torch.stack(dice_per_class).mean()

rslearn/utils/raster_format.py

@@ -476,6 +476,7 @@ class GeotiffRasterFormat(RasterFormat):
         bounds: PixelBounds,
         array: npt.NDArray[Any],
         fname: str | None = None,
+        nodata_val: int | float | None = None,
     ) -> None:
         """Encodes raster data.
 
@@ -485,6 +486,7 @@ class GeotiffRasterFormat(RasterFormat):
             bounds: the bounds of the raster data in the projection
             array: the raster data
             fname: override the filename to save as
+            nodata_val: set the nodata value when writing the raster.
         """
         if fname is None:
             fname = self.fname
@@ -520,6 +522,9 @@ class GeotiffRasterFormat(RasterFormat):
             profile["tiled"] = True
             profile["blockxsize"] = self.block_size
             profile["blockysize"] = self.block_size
+        # Set nodata_val if provided.
+        if nodata_val is not None:
+            profile["nodata"] = nodata_val
 
         profile.update(self.geotiff_options)
 
@@ -535,6 +540,7 @@ class GeotiffRasterFormat(RasterFormat):
         bounds: PixelBounds,
         resampling: Resampling = Resampling.bilinear,
         fname: str | None = None,
+        nodata_val: int | float | None = None,
     ) -> npt.NDArray[Any]:
         """Decodes raster data.
 
@@ -544,6 +550,16 @@ class GeotiffRasterFormat(RasterFormat):
             bounds: the bounds to read in the given projection.
             resampling: resampling method to use in case resampling is needed.
             fname: override the filename to read from
+            nodata_val: override the nodata value in the raster when reading. Pixels in
+                bounds that are not present in the source raster will be initialized to
+                this value. Note that, if the raster specifies a nodata value, and
+                some source pixels have that value, they will still be read under their
+                original value; overriding the nodata value is primarily useful if the
+                user wants out of bounds pixels to have a different value from the
+                source pixels, e.g. if the source data has background and foreground
+                classes (with background being nodata) but we want to read it in a
+                different projection and have out of bounds pixels be a third "invalid"
+                value.
 
         Returns:
             the raster data
@@ -561,6 +577,7 @@ class GeotiffRasterFormat(RasterFormat):
                 width=bounds[2] - bounds[0],
                 height=bounds[3] - bounds[1],
                 resampling=resampling,
+                src_nodata=nodata_val,
             ) as vrt:
                 return vrt.read()
 

rslearn/utils/stac.py

@@ -101,6 +101,7 @@ class StacClient:
         ids: list[str] | None = None,
         limit: int | None = None,
         query: dict[str, Any] | None = None,
+        sortby: list[dict[str, str]] | None = None,
     ) -> list[StacItem]:
         """Execute a STAC item search.
 
@@ -117,6 +118,7 @@ class StacClient:
             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.
+            sortby: list of sort specifications, e.g. [{"field": "id", "direction": "asc"}].
 
         Returns:
             list of matching STAC items.
@@ -142,6 +144,8 @@ class StacClient:
             request_data["limit"] = limit
         if query is not None:
             request_data["query"] = query
+        if sortby is not None:
+            request_data["sortby"] = sortby
 
         # Handle pagination.
         cur_url = self.endpoint + "/search"

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rslearn
-Version: 0.0.22
+Version: 0.0.24
 Summary: A library for developing remote sensing datasets and models
 Author: OlmoEarth Team
 License:                                  Apache License

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

@@ -3,7 +3,7 @@ rslearn/arg_parser.py,sha256=Go1MyEflcau_cziirmNd7Yhxa0WtXTAljIVE4f5H1GE,1194
 rslearn/const.py,sha256=FUCfsvFAs-QarEDJ0grdy0C1HjUjLpNFYGo5I2Vpc5Y,449
 rslearn/lightning_cli.py,sha256=1eTeffUlFqBe2KnyuYyXJdNKYQClCA-PV1xr0vJyJao,17972
 rslearn/log_utils.py,sha256=unD9gShiuO7cx5Nnq8qqVQ4qrbOOwFVgcHxN5bXuiAo,941
-rslearn/main.py,sha256=jRMYeU3-QvYSkTAJB69S1mHEft7-5_-RomzX1B-b8GM,28581
+rslearn/main.py,sha256=rrDEoa0xCkDflH-HN2SaHt0hb-rLfXWP-kJKISZAe9s,28714
 rslearn/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rslearn/template_params.py,sha256=Vop0Ha-S44ctCa9lvSZRjrMETznJZlR5y_gJrVIwrPg,791
 rslearn/config/__init__.py,sha256=n1qpZ0ImshTtLYl5mC73BORYyUcjPJyHiyZkqUY1hiY,474
@@ -25,9 +25,9 @@ rslearn/data_sources/local_files.py,sha256=mo5W_BxBl89EPTIHNDEpXM6qBjrP225KK0Pcm
 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=nTJ6Jh6CNBdCEIsn7G_xLQ0Nige5evdPdqLYmWTdDl4,20722
+rslearn/data_sources/planetary_computer.py,sha256=k6CO5Yim2I-frlD8r2_uBo0CQFw89mN3_5mrv0Xk2WU,26449
 rslearn/data_sources/soilgrids.py,sha256=rwO4goFPQ7lx420FvYBHYFXdihnZqn_-IjdqtxQ9j2g,12455
-rslearn/data_sources/stac.py,sha256=1qUbTD1fNSvBCX9QIXtyb9mGQ4K8ubRNIeEJs_I3QFU,9889
+rslearn/data_sources/stac.py,sha256=Xty1JDueAAonNVLRo8vfNBhlHrVLjhmZ-uRBYbrGvtA,10753
 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
@@ -69,7 +69,7 @@ rslearn/models/prithvi.py,sha256=J45eC1pd4l5AGlr19Qjrjrw5PPwvYE9bNM5qCFoznmg,403
 rslearn/models/resize_features.py,sha256=U7ZIVwwToJJnwchFG59wLWWP9eikHDB_1c4OtpubxHU,1693
 rslearn/models/sam2_enc.py,sha256=WZOtlp0FjaVztW4gEVIcsFQdKArS9iblRODP0b6Oc8M,3641
 rslearn/models/satlaspretrain.py,sha256=2R48ulbtd44Qy2FYJCkllE2Wk35eZxkc79ruSgkmgcQ,3384
-rslearn/models/simple_time_series.py,sha256=Nfk5E3d9W-4AyLQiy-P8p-JvxmFYE3FBrvOgttjXSMw,14678
+rslearn/models/simple_time_series.py,sha256=farQwt_nJVyAbgaM2UzdyqpDuIO0SLmHr9e9EVPSWCE,14678
 rslearn/models/singletask.py,sha256=9DM9a9-Mv3vVQqRhPOIXG2HHuVqVa_zuvgafeeYh4r0,1903
 rslearn/models/ssl4eo_s12.py,sha256=DOlpIj6NfjIlWyJ27m9Xo8TMlovBDstFq0ARnmAJ6qY,3919
 rslearn/models/swin.py,sha256=Xqr3SswbHP6IhwT2atZMAPF2TUzQqfxvihksb8WSeRo,6065
@@ -116,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=7WBAgdJhcMHueLKE2DthSFmNYvlNUh1dB4sibkqCsRA,14761
+rslearn/train/lightning_module.py,sha256=V4YoEg9PrwrgG4q9Dmv_9OBrSIK-SRPzjWtZRIfmPFg,15366
 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
@@ -130,10 +130,10 @@ rslearn/train/tasks/__init__.py,sha256=dag1u72x1-me6y0YcOubUo5MYZ0Tjf6-dOir9UeFN
 rslearn/train/tasks/classification.py,sha256=72ZBcbunMsdPYQN53S-4GfiLIDrr1X3Hni07dBJ0pu0,14261
 rslearn/train/tasks/detection.py,sha256=B0tfB7UGIbRtjnye3PhzLmfeQ4X7ImO3A-_LeNhBA54,21988
 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/multi_task.py,sha256=32hvwyVsHqt7N_M3zXsTErK1K7-0-BPHzt7iGNehyaI,6314
+rslearn/train/tasks/per_pixel_regression.py,sha256=Clrod6LQGjgNC0IAR4HLY7eCGWMHj2mk4d4moZCl4Qc,10209
 rslearn/train/tasks/regression.py,sha256=bVS_ApZSpbL0NaaM8Mu5Bsu4SBUyLpVtrPslulvvZHs,12695
-rslearn/train/tasks/segmentation.py,sha256=ie9ZV-sklLjQs35caiEglC1xff6dxeug_N-f_A8VosA,23034
+rslearn/train/tasks/segmentation.py,sha256=LZeuveHhMQsjNOQfMcwqSI4Ux3k9zfa58A2eZHSif8Y,29391
 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
@@ -155,17 +155,17 @@ rslearn/utils/get_utm_ups_crs.py,sha256=kUrcyjCK7KWvuP1XR-nURPeRqYeRO-3L8QUJ1QTF
 rslearn/utils/grid_index.py,sha256=hRmrtgpqN1pLa-djnZtgSXqKJlbgGyttGnCEmPLD0zo,2347
 rslearn/utils/jsonargparse.py,sha256=TRyZA151KzhjJlZczIHYguEA-YxCDYaZ2IwCRgx76nM,4791
 rslearn/utils/mp.py,sha256=XYmVckI5TOQuCKc49NJyirDJyFgvb4AI-gGypG2j680,1399
-rslearn/utils/raster_format.py,sha256=qZpbODF4I7BsOxf43O6vTmH2TSNw6N8PP0wsFUVvdIw,26267
+rslearn/utils/raster_format.py,sha256=fwotJBadwqYSdK8UokiKOk1pOF8JMim3kP_VwLWivPI,27382
 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/stac.py,sha256=c8NwOCKWvUwA-FSKlxZn-t7RZYweuye53OufT0bAK4A,5996
 rslearn/utils/time.py,sha256=2ilSLG94_sxLP3y5RSV5L5CG8CoND_dbdzYEHVtN-I8,387
 rslearn/utils/vector_format.py,sha256=4ZDYpfBLLxguJkiIaavTagiQK2Sv4Rz9NumbHlq-3Lw,15041
-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,,
+rslearn-0.0.24.dist-info/licenses/LICENSE,sha256=_99ZWPoLdlUbqZoSC5DF4ihiNwl5rTEmBaq2fACecdg,11352
+rslearn-0.0.24.dist-info/licenses/NOTICE,sha256=wLPr6rwV_jCg-xEknNGwhnkfRfuoOE9MZ-lru2yZyLI,5070
+rslearn-0.0.24.dist-info/METADATA,sha256=gV5mgeYPYiKWrEu7D8acOubWvg76Nn_4ICvlD7iTpcs,37936
+rslearn-0.0.24.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
+rslearn-0.0.24.dist-info/entry_points.txt,sha256=doTBQ57NT7nq-dgYGgTTw6mafcGWb_4PWYtYR4rGm50,46
+rslearn-0.0.24.dist-info/top_level.txt,sha256=XDKo90WBH8P9RQumHxo0giLJsoufT4r9odv-WE6Ahk4,8
+rslearn-0.0.24.dist-info/RECORD,,

{rslearn-0.0.22.dist-info → rslearn-0.0.24.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.1)
 Root-Is-Purelib: true
 Tag: py3-none-any