rslearn 0.0.26__py3-none-any.whl → 0.0.28__py3-none-any.whl

rslearn/train/prediction_writer.py

@@ -1,6 +1,7 @@
 """rslearn PredictionWriter implementation."""
 
 import json
+import warnings
 from collections.abc import Iterable, Sequence
 from dataclasses import dataclass
 from pathlib import Path
@@ -39,20 +40,20 @@ logger = get_logger(__name__)
 
 
 @dataclass
-class PendingPatchOutput:
-    """A patch output that hasn't been merged yet."""
+class PendingCropOutput:
+    """A crop output that hasn't been merged yet."""
 
     bounds: PixelBounds
     output: Any
 
 
-class PatchPredictionMerger:
-    """Base class for merging predictions from multiple patches."""
+class CropPredictionMerger:
+    """Base class for merging predictions from multiple crops."""
 
     def merge(
         self,
         window: Window,
-        outputs: Sequence[PendingPatchOutput],
+        outputs: Sequence[PendingCropOutput],
         layer_config: LayerConfig,
     ) -> Any:
         """Merge the outputs.
@@ -68,39 +69,60 @@ class PatchPredictionMerger:
         raise NotImplementedError
 
 
-class VectorMerger(PatchPredictionMerger):
+class VectorMerger(CropPredictionMerger):
     """Merger for vector data that simply concatenates the features."""
 
     def merge(
         self,
         window: Window,
-        outputs: Sequence[PendingPatchOutput],
+        outputs: Sequence[PendingCropOutput],
         layer_config: LayerConfig,
     ) -> list[Feature]:
         """Concatenate the vector features."""
         return [feat for output in outputs for feat in output.output]
 
 
-class RasterMerger(PatchPredictionMerger):
+class RasterMerger(CropPredictionMerger):
     """Merger for raster data that copies the rasters to the output."""
 
-    def __init__(self, padding: int | None = None, downsample_factor: int = 1):
+    def __init__(
+        self,
+        overlap_pixels: int | None = None,
+        downsample_factor: int = 1,
+        # Deprecated parameter (for backwards compatibility)
+        padding: int | None = None,
+    ):
         """Create a new RasterMerger.
 
         Args:
-            padding: the padding around the individual patch outputs to remove. This is
-                typically used when leveraging overlapping patches. Portions of outputs
-                at the border of the window will still be retained.
+            overlap_pixels: the number of pixels shared between adjacent crops during
+                sliding window inference. Half of this overlap is removed from each
+                crop during merging (except at window boundaries where the full crop
+                is retained).
             downsample_factor: the factor by which the rasters output by the task are
                 lower in resolution relative to the window resolution.
+            padding: deprecated, use overlap_pixels instead. The old padding value
+                equals overlap_pixels // 2.
         """
-        self.padding = padding
+        # Handle deprecated padding parameter
+        if padding is not None:
+            warnings.warn(
+                "padding is deprecated, use overlap_pixels instead. "
+                "Note: overlap_pixels = padding * 2",
+                FutureWarning,
+                stacklevel=2,
+            )
+            if overlap_pixels is not None:
+                raise ValueError("Cannot specify both padding and overlap_pixels")
+            overlap_pixels = padding * 2
+
+        self.overlap_pixels = overlap_pixels
         self.downsample_factor = downsample_factor
 
     def merge(
         self,
         window: Window,
-        outputs: Sequence[PendingPatchOutput],
+        outputs: Sequence[PendingCropOutput],
         layer_config: LayerConfig,
     ) -> npt.NDArray:
         """Merge the raster outputs."""
@@ -114,6 +136,12 @@ class RasterMerger(PatchPredictionMerger):
             dtype=layer_config.band_sets[0].dtype.get_numpy_dtype(),
         )
 
+        # Compute how many pixels to trim from each side.
+        # We remove half of the overlap from each side (not at window boundaries).
+        trim_pixels = (
+            self.overlap_pixels // 2 if self.overlap_pixels is not None else None
+        )
+
         # Ensure the outputs are sorted by height then width.
         # This way when we merge we can be sure that outputs that are lower or further
         # to the right will overwrite earlier outputs.
@@ -123,18 +151,18 @@ class RasterMerger(PatchPredictionMerger):
         for output in sorted_outputs:
             # So now we just need to compute the src_offset to copy.
             # If the output is not on the left or top boundary, then we should apply
-            # the padding (if set).
+            # the trim (if set).
             src = output.output
             src_offset = (
                 output.bounds[0] // self.downsample_factor,
                 output.bounds[1] // self.downsample_factor,
             )
-            if self.padding is not None and output.bounds[0] != window.bounds[0]:
-                src = src[:, :, self.padding :]
-                src_offset = (src_offset[0] + self.padding, src_offset[1])
-            if self.padding is not None and output.bounds[1] != window.bounds[1]:
-                src = src[:, self.padding :, :]
-                src_offset = (src_offset[0], src_offset[1] + self.padding)
+            if trim_pixels is not None and output.bounds[0] != window.bounds[0]:
+                src = src[:, :, trim_pixels:]
+                src_offset = (src_offset[0] + trim_pixels, src_offset[1])
+            if trim_pixels is not None and output.bounds[1] != window.bounds[1]:
+                src = src[:, trim_pixels:, :]
+                src_offset = (src_offset[0], src_offset[1] + trim_pixels)
 
             copy_spatial_array(
                 src=src,
@@ -162,7 +190,7 @@ class RslearnWriter(BasePredictionWriter):
         output_layer: str,
         path_options: dict[str, Any] | None = None,
         selector: list[str] | None = None,
-        merger: PatchPredictionMerger | None = None,
+        merger: CropPredictionMerger | None = None,
         output_path: str | Path | None = None,
         layer_config: LayerConfig | None = None,
         storage_config: StorageConfig | None = None,
@@ -175,7 +203,7 @@ class RslearnWriter(BasePredictionWriter):
             path_options: additional options for path to pass to fsspec
             selector: keys to access the desired output in the output dict if needed.
                 e.g ["key1", "key2"] gets output["key1"]["key2"]
-            merger: merger to use to merge outputs from overlapped patches.
+            merger: merger to use to merge outputs from overlapped crops.
             output_path: optional custom path for writing predictions. If provided,
                 predictions will be written to this path instead of deriving from dataset path.
             layer_config: optional layer configuration. If provided, this config will be
@@ -217,9 +245,9 @@ class RslearnWriter(BasePredictionWriter):
             self.merger = VectorMerger()
 
         # Map from window name to pending data to write.
-        # This is used when windows are split up into patches, so the data from all the
-        # patches of each window need to be reconstituted.
-        self.pending_outputs: dict[str, list[PendingPatchOutput]] = {}
+        # This is used when windows are split up into crops, so the data from all the
+        # crops of each window need to be reconstituted.
+        self.pending_outputs: dict[str, list[PendingCropOutput]] = {}
 
     def _get_layer_config_and_dataset_storage(
         self,
@@ -327,7 +355,7 @@ class RslearnWriter(BasePredictionWriter):
                 will be processed by the task to obtain a vector (list[Feature]) or
                 raster (npt.NDArray) output.
             metadatas: corresponding list of metadatas from the batch describing the
-                patches that were processed.
+                crops that were processed.
         """
         # Process the predictions into outputs that can be written.
         outputs: list = [
@@ -349,17 +377,17 @@ class RslearnWriter(BasePredictionWriter):
             )
             self.process_output(
                 window,
-                metadata.patch_idx,
-                metadata.num_patches_in_window,
-                metadata.patch_bounds,
+                metadata.crop_idx,
+                metadata.num_crops_in_window,
+                metadata.crop_bounds,
                 output,
             )
 
     def process_output(
         self,
         window: Window,
-        patch_idx: int,
-        num_patches: int,
+        crop_idx: int,
+        num_crops: int,
         cur_bounds: PixelBounds,
         output: npt.NDArray | list[Feature],
     ) -> None:
@@ -367,28 +395,28 @@ class RslearnWriter(BasePredictionWriter):
 
         Args:
             window: the window that the output pertains to.
-            patch_idx: the index of this patch for the window.
-            num_patches: the total number of patches to be processed for the window.
-            cur_bounds: the bounds of the current patch.
+            crop_idx: the index of this crop for the window.
+            num_crops: the total number of crops to be processed for the window.
+            cur_bounds: the bounds of the current crop.
             output: the output data.
         """
-        # Incorporate the output into our list of pending patch outputs.
+        # Incorporate the output into our list of pending crop outputs.
         if window.name not in self.pending_outputs:
             self.pending_outputs[window.name] = []
-        self.pending_outputs[window.name].append(PendingPatchOutput(cur_bounds, output))
+        self.pending_outputs[window.name].append(PendingCropOutput(cur_bounds, output))
         logger.debug(
-            f"Stored PendingPatchOutput for patch #{patch_idx}/{num_patches} at window {window.name}"
+            f"Stored PendingCropOutput for crop #{crop_idx}/{num_crops} at window {window.name}"
         )
 
-        if patch_idx < num_patches - 1:
+        if crop_idx < num_crops - 1:
             return
 
-        # This is the last patch so it's time to write it.
+        # This is the last crop so it's time to write it.
         # First get the pending output and clear it.
         pending_output = self.pending_outputs[window.name]
         del self.pending_outputs[window.name]
 
-        # Merge outputs from overlapped patches if merger is set.
+        # Merge outputs from overlapped crops if merger is set.
         logger.debug(f"Merging and writing for window {window.name}")
         merged_output = self.merger.merge(window, pending_output, self.layer_config)
 

rslearn/train/tasks/classification.py

@@ -16,6 +16,7 @@ from torchmetrics.classification import (
 )
 
 from rslearn.models.component import FeatureVector, Predictor
+from rslearn.train.metrics import ConfusionMatrixMetric
 from rslearn.train.model_context import (
     ModelContext,
     ModelOutput,
@@ -44,6 +45,7 @@ class ClassificationTask(BasicTask):
         f1_metric_kwargs: dict[str, Any] = {},
         positive_class: str | None = None,
         positive_class_threshold: float = 0.5,
+        enable_confusion_matrix: bool = False,
         **kwargs: Any,
     ):
         """Initialize a new ClassificationTask.
@@ -69,6 +71,8 @@ class ClassificationTask(BasicTask):
             positive_class: positive class name.
             positive_class_threshold: threshold for classifying the positive class in
                 binary classification (default 0.5).
+            enable_confusion_matrix: whether to compute confusion matrix (default false).
+                If true, it requires wandb to be initialized for logging.
             kwargs: other arguments to pass to BasicTask
         """
         super().__init__(**kwargs)
@@ -84,6 +88,7 @@ class ClassificationTask(BasicTask):
         self.f1_metric_kwargs = f1_metric_kwargs
         self.positive_class = positive_class
         self.positive_class_threshold = positive_class_threshold
+        self.enable_confusion_matrix = enable_confusion_matrix
 
         if self.positive_class_threshold != 0.5:
             # Must be binary classification
@@ -201,7 +206,7 @@ class ClassificationTask(BasicTask):
         feature = Feature(
             STGeometry(
                 metadata.projection,
-                shapely.Point(metadata.patch_bounds[0], metadata.patch_bounds[1]),
+                shapely.Point(metadata.crop_bounds[0], metadata.crop_bounds[1]),
                 None,
             ),
             {
@@ -278,6 +283,14 @@ class ClassificationTask(BasicTask):
                 )
                 metrics["f1"] = ClassificationMetric(MulticlassF1Score(**kwargs))
 
+        if self.enable_confusion_matrix:
+            metrics["confusion_matrix"] = ClassificationMetric(
+                ConfusionMatrixMetric(
+                    num_classes=len(self.classes),
+                    class_names=self.classes,
+                ),
+            )
+
         return MetricCollection(metrics)
 
 

rslearn/train/tasks/detection.py

@@ -128,7 +128,7 @@ class DetectionTask(BasicTask):
         if not load_targets:
             return {}, {}
 
-        bounds = metadata.patch_bounds
+        bounds = metadata.crop_bounds
 
         boxes = []
         class_labels = []
@@ -244,10 +244,10 @@ class DetectionTask(BasicTask):
         features = []
         for box, class_id, score in zip(boxes, class_ids, scores):
             shp = shapely.box(
-                metadata.patch_bounds[0] + float(box[0]),
-                metadata.patch_bounds[1] + float(box[1]),
-                metadata.patch_bounds[0] + float(box[2]),
-                metadata.patch_bounds[1] + float(box[3]),
+                metadata.crop_bounds[0] + float(box[0]),
+                metadata.crop_bounds[1] + float(box[1]),
+                metadata.crop_bounds[0] + float(box[2]),
+                metadata.crop_bounds[1] + float(box[3]),
             )
             geom = STGeometry(metadata.projection, shp, None)
             properties: dict[str, Any] = {

rslearn/train/tasks/per_pixel_regression.py

@@ -149,22 +149,28 @@ class PerPixelRegressionHead(Predictor):
     """Head for per-pixel regression task."""
 
     def __init__(
-        self, loss_mode: Literal["mse", "l1"] = "mse", use_sigmoid: bool = False
+        self,
+        loss_mode: Literal["mse", "l1", "huber"] = "mse",
+        use_sigmoid: bool = False,
+        huber_delta: float = 1.0,
     ):
-        """Initialize a new RegressionHead.
+        """Initialize a new PerPixelRegressionHead.
 
         Args:
-            loss_mode: the loss function to use, either "mse" (default) or "l1".
+            loss_mode: the loss function to use: "mse" (default), "l1", or "huber".
             use_sigmoid: whether to apply a sigmoid activation on the output. This
                 requires targets to be between 0-1.
+            huber_delta: delta parameter for Huber loss (only used when
+                loss_mode="huber").
         """
         super().__init__()
 
-        if loss_mode not in ["mse", "l1"]:
-            raise ValueError("invalid loss mode")
+        if loss_mode not in ["mse", "l1", "huber"]:
+            raise ValueError(f"invalid loss mode {loss_mode}")
 
         self.loss_mode = loss_mode
         self.use_sigmoid = use_sigmoid
+        self.huber_delta = huber_delta
 
     def forward(
         self,
@@ -217,8 +223,15 @@ class PerPixelRegressionHead(Predictor):
                 scores = torch.square(outputs - labels)
             elif self.loss_mode == "l1":
                 scores = torch.abs(outputs - labels)
+            elif self.loss_mode == "huber":
+                scores = torch.nn.functional.huber_loss(
+                    outputs,
+                    labels,
+                    reduction="none",
+                    delta=self.huber_delta,
+                )
             else:
-                assert False
+                raise ValueError(f"unknown loss mode {self.loss_mode}")
 
             # Compute average but only over valid pixels.
             mask_total = mask.sum()

rslearn/train/tasks/regression.py

@@ -130,7 +130,7 @@ class RegressionTask(BasicTask):
         feature = Feature(
             STGeometry(
                 metadata.projection,
-                shapely.Point(metadata.patch_bounds[0], metadata.patch_bounds[1]),
+                shapely.Point(metadata.crop_bounds[0], metadata.crop_bounds[1]),
                 None,
             ),
             {
@@ -196,18 +196,24 @@ class RegressionHead(Predictor):
     """Head for regression task."""
 
     def __init__(
-        self, loss_mode: Literal["mse", "l1"] = "mse", use_sigmoid: bool = False
+        self,
+        loss_mode: Literal["mse", "l1", "huber"] = "mse",
+        use_sigmoid: bool = False,
+        huber_delta: float = 1.0,
     ):
         """Initialize a new RegressionHead.
 
         Args:
-            loss_mode: the loss function to use, either "mse" (default) or "l1".
+            loss_mode: the loss function to use: "mse" (default), "l1", or "huber".
             use_sigmoid: whether to apply a sigmoid activation on the output. This
                 requires targets to be between 0-1.
+            huber_delta: delta parameter for Huber loss (only used when
+                loss_mode="huber").
         """
         super().__init__()
         self.loss_mode = loss_mode
         self.use_sigmoid = use_sigmoid
+        self.huber_delta = huber_delta
 
     def forward(
         self,
@@ -251,6 +257,16 @@ class RegressionHead(Predictor):
                 losses["regress"] = torch.mean(torch.square(outputs - labels) * mask)
             elif self.loss_mode == "l1":
                 losses["regress"] = torch.mean(torch.abs(outputs - labels) * mask)
+            elif self.loss_mode == "huber":
+                losses["regress"] = torch.mean(
+                    torch.nn.functional.huber_loss(
+                        outputs,
+                        labels,
+                        reduction="none",
+                        delta=self.huber_delta,
+                    )
+                    * mask
+                )
             else:
                 raise ValueError(f"unknown loss mode {self.loss_mode}")
 

rslearn/train/tasks/segmentation.py

@@ -10,6 +10,7 @@ import torchmetrics.classification
 from torchmetrics import Metric, MetricCollection
 
 from rslearn.models.component import FeatureMaps, Predictor
+from rslearn.train.metrics import ConfusionMatrixMetric
 from rslearn.train.model_context import (
     ModelContext,
     ModelOutput,
@@ -43,6 +44,8 @@ class SegmentationTask(BasicTask):
         other_metrics: dict[str, Metric] = {},
         output_probs: bool = False,
         output_class_idx: int | None = None,
+        enable_confusion_matrix: bool = False,
+        class_names: list[str] | None = None,
         **kwargs: Any,
     ) -> None:
         """Initialize a new SegmentationTask.
@@ -80,6 +83,10 @@ class SegmentationTask(BasicTask):
                 during prediction.
             output_class_idx: if set along with output_probs, only output the probability
                 for this specific class index (single-channel output).
+            enable_confusion_matrix: whether to compute confusion matrix (default false).
+                If true, it requires wandb to be initialized for logging.
+            class_names: optional list of class names for labeling confusion matrix axes.
+                If not provided, classes will be labeled as "class_0", "class_1", etc.
             kwargs: additional arguments to pass to BasicTask
         """
         super().__init__(**kwargs)
@@ -106,6 +113,8 @@ class SegmentationTask(BasicTask):
         self.other_metrics = other_metrics
         self.output_probs = output_probs
         self.output_class_idx = output_class_idx
+        self.enable_confusion_matrix = enable_confusion_matrix
+        self.class_names = class_names
 
     def process_inputs(
         self,
@@ -285,6 +294,14 @@ class SegmentationTask(BasicTask):
         if self.other_metrics:
             metrics.update(self.other_metrics)
 
+        if self.enable_confusion_matrix:
+            metrics["confusion_matrix"] = SegmentationMetric(
+                ConfusionMatrixMetric(
+                    num_classes=self.num_classes,
+                    class_names=self.class_names,
+                ),
+            )
+
         return MetricCollection(metrics)
 
 

rslearn/utils/__init__.py

@@ -7,6 +7,7 @@ from .geometry import (
     PixelBounds,
     Projection,
     STGeometry,
+    get_global_raster_bounds,
     is_same_resolution,
     shp_intersects,
 )
@@ -23,6 +24,7 @@ __all__ = (
     "Projection",
     "STGeometry",
     "daterange",
+    "get_global_raster_bounds",
     "get_utm_ups_crs",
     "is_same_resolution",
     "logger",

rslearn/utils/fsspec.py

@@ -2,7 +2,7 @@
 
 import os
 import tempfile
-from collections.abc import Generator
+from collections.abc import Generator, Iterator
 from contextlib import contextmanager
 from typing import Any
 
@@ -16,6 +16,56 @@ from rslearn.log_utils import get_logger
 logger = get_logger(__name__)
 
 
+def iter_nonhidden(path: UPath) -> Iterator[UPath]:
+    """Iterate over non-hidden entries in a directory.
+
+    Hidden entries are those whose basename starts with "." (e.g. ".DS_Store").
+
+    Args:
+        path: the directory to iterate.
+
+    Yields:
+        non-hidden UPath entries in the directory.
+    """
+    try:
+        it = path.iterdir()
+    except (FileNotFoundError, NotADirectoryError):
+        return
+
+    for p in it:
+        if p.name.startswith("."):
+            continue
+        yield p
+
+
+def iter_nonhidden_subdirs(path: UPath) -> Iterator[UPath]:
+    """Iterate over non-hidden subdirectories in a directory.
+
+    Args:
+        path: the directory to iterate.
+
+    Yields:
+        non-hidden subdirectories in the directory.
+    """
+    for p in iter_nonhidden(path):
+        if p.is_dir():
+            yield p
+
+
+def iter_nonhidden_files(path: UPath) -> Iterator[UPath]:
+    """Iterate over non-hidden files in a directory.
+
+    Args:
+        path: the directory to iterate.
+
+    Yields:
+        non-hidden files in the directory.
+    """
+    for p in iter_nonhidden(path):
+        if p.is_file():
+            yield p
+
+
 @contextmanager
 def get_upath_local(
     path: UPath, extra_paths: list[UPath] = []

rslearn/utils/geometry.py

@@ -116,6 +116,27 @@ class Projection:
 WGS84_PROJECTION = Projection(CRS.from_epsg(WGS84_EPSG), 1, 1)
 
 
+def get_global_raster_bounds(projection: Projection) -> PixelBounds:
+    """Get very large pixel bounds for a global raster in the given projection.
+
+    This is useful for data sources that cover the entire world and don't want to
+    compute exact bounds in arbitrary projections (which can fail for projections
+    like UTM that only cover part of the world).
+
+    Args:
+        projection: the projection to get bounds in.
+
+    Returns:
+        Pixel bounds that will intersect with any reasonable window. We assume that the
+        absolute value of CRS coordinates is at most 2^32, and adjust it based on the
+        resolution in the Projection in case very fine-grained resolutions are used.
+    """
+    crs_bound = 2**32
+    pixel_bound_x = int(crs_bound / abs(projection.x_resolution))
+    pixel_bound_y = int(crs_bound / abs(projection.y_resolution))
+    return (-pixel_bound_x, -pixel_bound_y, pixel_bound_x, pixel_bound_y)
+
+
 class ResolutionFactor:
     """Multiplier for the resolution in a Projection.