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

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,19 +356,16 @@ 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 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/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,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/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.23.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rslearn
-Version: 0.0.21
+Version: 0.0.23
 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.