rslearn 0.0.22__tar.gz → 0.0.23__tar.gz

{rslearn-0.0.22/rslearn.egg-info → rslearn-0.0.23}/PKG-INFO

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

{rslearn-0.0.22 → rslearn-0.0.23}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rslearn"
-version = "0.0.22"
+version = "0.0.23"
 description = "A library for developing remote sensing datasets and models"
 authors = [
     { name = "OlmoEarth Team" },

{rslearn-0.0.22 → rslearn-0.0.23}/rslearn/data_sources/planetary_computer.py

@@ -567,3 +567,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-0.0.22 → rslearn-0.0.23}/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
@@ -132,6 +133,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.
 
@@ -191,10 +210,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,
             )

{rslearn-0.0.22 → rslearn-0.0.23}/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-0.0.22 → rslearn-0.0.23}/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-0.0.22 → rslearn-0.0.23}/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-0.0.22 → rslearn-0.0.23}/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-0.0.22 → rslearn-0.0.23}/rslearn/train/tasks/segmentation.py

@@ -53,6 +53,7 @@ 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] = {},
@@ -74,6 +75,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
@@ -107,6 +110,7 @@ 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
@@ -237,29 +241,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 +290,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 +338,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 +348,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 +366,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 +375,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 +417,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 +467,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 +477,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 +529,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 +569,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 +594,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 +606,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 +654,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 +671,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-0.0.22 → rslearn-0.0.23/rslearn.egg-info}/PKG-INFO

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