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

rslearn/dataset/storage/file.py

@@ -15,7 +15,7 @@ from rslearn.dataset.window import (
     get_window_layer_dir,
 )
 from rslearn.log_utils import get_logger
-from rslearn.utils.fsspec import open_atomic
+from rslearn.utils.fsspec import iter_nonhidden_subdirs, open_atomic
 from rslearn.utils.mp import star_imap_unordered
 
 from .storage import WindowStorage, WindowStorageFactory
@@ -77,8 +77,8 @@ class FileWindowStorage(WindowStorage):
         window_dirs = []
         if not groups:
             groups = []
-            for p in (self.path / "windows").iterdir():
-                groups.append(p.name)
+            for group_dir in iter_nonhidden_subdirs(self.path / "windows"):
+                groups.append(group_dir.name)
         for group in groups:
             group_dir = self.path / "windows" / group
             if not group_dir.exists():
@@ -86,16 +86,20 @@ class FileWindowStorage(WindowStorage):
                     f"Skipping group directory {group_dir} since it does not exist"
                 )
                 continue
+            if not group_dir.is_dir():
+                logger.warning(
+                    f"Skipping group path {group_dir} since it is not a directory"
+                )
+                continue
             if names:
-                cur_names = names
+                for window_name in names:
+                    window_dir = group_dir / window_name
+                    if not window_dir.is_dir():
+                        continue
+                    window_dirs.append(window_dir)
             else:
-                cur_names = []
-                for p in group_dir.iterdir():
-                    cur_names.append(p.name)
-
-            for window_name in cur_names:
-                window_dir = group_dir / window_name
-                window_dirs.append(window_dir)
+                for window_dir in iter_nonhidden_subdirs(group_dir):
+                    window_dirs.append(window_dir)
 
         if workers == 0:
             windows = [load_window(self, window_dir) for window_dir in window_dirs]
@@ -162,7 +166,7 @@ class FileWindowStorage(WindowStorage):
             return []
 
         completed_layers = []
-        for layer_dir in layers_directory.iterdir():
+        for layer_dir in iter_nonhidden_subdirs(layers_directory):
             layer_name, group_idx = get_layer_and_group_from_dir_name(layer_dir.name)
             if not self.is_layer_completed(group, name, layer_name, group_idx):
                 continue

rslearn/tile_stores/default.py

@@ -15,6 +15,8 @@ from upath import UPath
 from rslearn.const import WGS84_PROJECTION
 from rslearn.utils.feature import Feature
 from rslearn.utils.fsspec import (
+    iter_nonhidden_files,
+    iter_nonhidden_subdirs,
     join_upath,
     open_atomic,
     open_rasterio_upath_reader,
@@ -129,7 +131,7 @@ class DefaultTileStore(TileStore):
             ValueError: if no file is found.
         """
         raster_dir = self._get_raster_dir(layer_name, item_name, bands)
-        for fname in raster_dir.iterdir():
+        for fname in iter_nonhidden_files(raster_dir):
             # Ignore completed sentinel files, bands files, as well as temporary files created by
             # open_atomic (in case this tile store is on local filesystem).
             if fname.name == COMPLETED_FNAME:
@@ -175,7 +177,7 @@ class DefaultTileStore(TileStore):
             return []
 
         bands: list[list[str]] = []
-        for raster_dir in item_dir.iterdir():
+        for raster_dir in iter_nonhidden_subdirs(item_dir):
             if not (raster_dir / BANDS_FNAME).exists():
                 # This is likely a legacy directory where the bands are only encoded in
                 # the directory name, so we have to rely on that.

rslearn/train/data_module.py

@@ -108,10 +108,10 @@ class RslearnDataModule(L.LightningDataModule):
         self.use_in_memory_all_crops_dataset = use_in_memory_all_crops_dataset
         self.index_mode = index_mode
         self.split_configs = {
-            "train": default_config.update(train_config),
-            "val": default_config.update(val_config),
-            "test": default_config.update(test_config),
-            "predict": default_config.update(predict_config),
+            "train": SplitConfig.merge_and_validate([default_config, train_config]),
+            "val": SplitConfig.merge_and_validate([default_config, val_config]),
+            "test": SplitConfig.merge_and_validate([default_config, test_config]),
+            "predict": SplitConfig.merge_and_validate([default_config, predict_config]),
         }
 
     def setup(
@@ -141,7 +141,7 @@ class RslearnDataModule(L.LightningDataModule):
                 task=self.task,
                 workers=self.init_workers,
                 name=self.name,
-                fix_patch_pick=(split != "train"),
+                fix_crop_pick=(split != "train"),
                 index_mode=self.index_mode,
             )
             logger.info(f"got {len(dataset)} examples in split {split}")
@@ -203,13 +203,16 @@ class RslearnDataModule(L.LightningDataModule):
         # Enable persistent workers unless we are using main process.
         persistent_workers = self.num_workers > 0
 
-        # If using all patches, limit number of workers to the number of windows.
+        # If using all crops, limit number of workers to the number of windows.
         # Otherwise it has to distribute the same window to different workers which can
         # cause issues for RslearnWriter.
         # If the number of windows is 0, then we can set positive number of workers
         # since they won't yield anything anyway.
         num_workers = self.num_workers
-        if split_config.load_all_crops and len(dataset.get_dataset_examples()) > 0:
+        if (
+            split_config.get_load_all_crops()
+            and len(dataset.get_dataset_examples()) > 0
+        ):
             num_workers = min(num_workers, len(dataset.get_dataset_examples()))
 
         kwargs: dict[str, Any] = dict(

rslearn/train/dataset.py

@@ -496,53 +496,6 @@ class SplitConfig:
             overlap_ratio: deprecated, use overlap_pixels instead
             load_all_patches: deprecated, use load_all_crops instead
         """
-        # Handle deprecated load_all_patches parameter
-        if load_all_patches is not None:
-            warnings.warn(
-                "load_all_patches is deprecated, use load_all_crops instead",
-                FutureWarning,
-                stacklevel=2,
-            )
-            if load_all_crops is not None:
-                raise ValueError(
-                    "Cannot specify both load_all_patches and load_all_crops"
-                )
-            load_all_crops = load_all_patches
-        # Handle deprecated patch_size parameter
-        if patch_size is not None:
-            warnings.warn(
-                "patch_size is deprecated, use crop_size instead",
-                FutureWarning,
-                stacklevel=2,
-            )
-            if crop_size is not None:
-                raise ValueError("Cannot specify both patch_size and crop_size")
-            crop_size = patch_size
-
-        # Normalize crop_size to tuple[int, int] | None
-        self.crop_size: tuple[int, int] | None = None
-        if crop_size is not None:
-            if isinstance(crop_size, int):
-                self.crop_size = (crop_size, crop_size)
-            else:
-                self.crop_size = crop_size
-
-        # Handle deprecated overlap_ratio parameter
-        if overlap_ratio is not None:
-            warnings.warn(
-                "overlap_ratio is deprecated, use overlap_pixels instead",
-                FutureWarning,
-                stacklevel=2,
-            )
-            if overlap_pixels is not None:
-                raise ValueError("Cannot specify both overlap_ratio and overlap_pixels")
-            if self.crop_size is None:
-                raise ValueError("overlap_ratio requires crop_size to be set")
-            overlap_pixels = round(self.crop_size[0] * overlap_ratio)
-
-        if overlap_pixels is not None and overlap_pixels < 0:
-            raise ValueError("overlap_pixels must be non-negative")
-
         self.groups = groups
         self.names = names
         self.tags = tags
@@ -555,13 +508,22 @@ class SplitConfig:
             output_layer_name_skip_inference_if_exists
         )
 
-        # Note that load_all_crops is handled by the RslearnDataModule rather than the
-        # ModelDataset.
-        self.load_all_crops = load_all_crops
-        self.overlap_pixels = overlap_pixels
+        # These have deprecated equivalents -- we store both raw values since we don't
+        # have a complete picture until the final merged SplitConfig is computed. We
+        # raise deprecation warnings in merge_and_validate and we disambiguate them in
+        # get_ functions (so the variables should never be accessed directly).
+        self._crop_size = crop_size
+        self._patch_size = patch_size
+        self._overlap_pixels = overlap_pixels
+        self._overlap_ratio = overlap_ratio
+        self._load_all_crops = load_all_crops
+        self._load_all_patches = load_all_patches
 
-    def update(self, other: "SplitConfig") -> "SplitConfig":
-        """Override settings in this SplitConfig with those in another.
+    def _merge(self, other: "SplitConfig") -> "SplitConfig":
+        """Merge settings from another SplitConfig into this one.
+
+        Args:
+            other: the config to merge in (its non-None values override self's)
 
         Returns:
             the resulting SplitConfig combining the settings.
@@ -574,9 +536,12 @@ class SplitConfig:
             num_patches=self.num_patches,
             transforms=self.transforms,
             sampler=self.sampler,
-            crop_size=self.crop_size,
-            overlap_pixels=self.overlap_pixels,
-            load_all_crops=self.load_all_crops,
+            crop_size=self._crop_size,
+            patch_size=self._patch_size,
+            overlap_pixels=self._overlap_pixels,
+            overlap_ratio=self._overlap_ratio,
+            load_all_crops=self._load_all_crops,
+            load_all_patches=self._load_all_patches,
             skip_targets=self.skip_targets,
             output_layer_name_skip_inference_if_exists=self.output_layer_name_skip_inference_if_exists,
         )
@@ -594,12 +559,18 @@ class SplitConfig:
             result.transforms = other.transforms
         if other.sampler:
             result.sampler = other.sampler
-        if other.crop_size:
-            result.crop_size = other.crop_size
-        if other.overlap_pixels is not None:
-            result.overlap_pixels = other.overlap_pixels
-        if other.load_all_crops is not None:
-            result.load_all_crops = other.load_all_crops
+        if other._crop_size is not None:
+            result._crop_size = other._crop_size
+        if other._patch_size is not None:
+            result._patch_size = other._patch_size
+        if other._overlap_pixels is not None:
+            result._overlap_pixels = other._overlap_pixels
+        if other._overlap_ratio is not None:
+            result._overlap_ratio = other._overlap_ratio
+        if other._load_all_crops is not None:
+            result._load_all_crops = other._load_all_crops
+        if other._load_all_patches is not None:
+            result._load_all_patches = other._load_all_patches
         if other.skip_targets is not None:
             result.skip_targets = other.skip_targets
         if other.output_layer_name_skip_inference_if_exists is not None:
@@ -608,17 +579,90 @@ class SplitConfig:
             )
         return result
 
+    @staticmethod
+    def merge_and_validate(configs: list["SplitConfig"]) -> "SplitConfig":
+        """Merge a list of SplitConfigs and validate the result.
+
+        Args:
+            configs: list of SplitConfig to merge. Later configs override earlier ones.
+
+        Returns:
+            the merged and validated SplitConfig.
+        """
+        if not configs:
+            return SplitConfig()
+
+        result = configs[0]
+        for config in configs[1:]:
+            result = result._merge(config)
+
+        # Emit deprecation warnings
+        if result._patch_size is not None:
+            warnings.warn(
+                "patch_size is deprecated, use crop_size instead",
+                FutureWarning,
+                stacklevel=2,
+            )
+        if result._overlap_ratio is not None:
+            warnings.warn(
+                "overlap_ratio is deprecated, use overlap_pixels instead",
+                FutureWarning,
+                stacklevel=2,
+            )
+        if result._load_all_patches is not None:
+            warnings.warn(
+                "load_all_patches is deprecated, use load_all_crops instead",
+                FutureWarning,
+                stacklevel=2,
+            )
+
+        # Check for conflicting parameters
+        if result._crop_size is not None and result._patch_size is not None:
+            raise ValueError("Cannot specify both crop_size and patch_size")
+        if result._overlap_pixels is not None and result._overlap_ratio is not None:
+            raise ValueError("Cannot specify both overlap_pixels and overlap_ratio")
+        if result._load_all_crops is not None and result._load_all_patches is not None:
+            raise ValueError("Cannot specify both load_all_crops and load_all_patches")
+
+        # Validate overlap_pixels is non-negative
+        if result._overlap_pixels is not None and result._overlap_pixels < 0:
+            raise ValueError("overlap_pixels must be non-negative")
+
+        # overlap_pixels requires load_all_crops.
+        if result.get_overlap_pixels() > 0 and not result.get_load_all_crops():
+            raise ValueError(
+                "overlap_pixels requires load_all_crops to be True since (overlap is only used during sliding window inference"
+            )
+
+        return result
+
     def get_crop_size(self) -> tuple[int, int] | None:
-        """Get crop size as tuple."""
-        return self.crop_size
+        """Get crop size as tuple, handling deprecated patch_size."""
+        size = self._crop_size if self._crop_size is not None else self._patch_size
+        if size is None:
+            return None
+        if isinstance(size, int):
+            return (size, size)
+        return size
 
     def get_overlap_pixels(self) -> int:
-        """Get the overlap pixels (default 0)."""
-        return self.overlap_pixels if self.overlap_pixels is not None else 0
+        """Get the overlap pixels (default 0), handling deprecated overlap_ratio."""
+        if self._overlap_pixels is not None:
+            return self._overlap_pixels
+        if self._overlap_ratio is not None:
+            crop_size = self.get_crop_size()
+            if crop_size is None:
+                raise ValueError("overlap_ratio requires crop_size to be set")
+            return round(crop_size[0] * self._overlap_ratio)
+        return 0
 
     def get_load_all_crops(self) -> bool:
-        """Returns whether loading all patches is enabled (default False)."""
-        return True if self.load_all_crops is True else False
+        """Returns whether loading all crops is enabled (default False)."""
+        if self._load_all_crops is not None:
+            return self._load_all_crops
+        if self._load_all_patches is not None:
+            return self._load_all_patches
+        return False
 
     def get_skip_targets(self) -> bool:
         """Returns whether skip_targets is enabled (default False)."""
@@ -697,7 +741,7 @@ class ModelDataset(torch.utils.data.Dataset):
         task: Task,
         workers: int,
         name: str | None = None,
-        fix_patch_pick: bool = False,
+        fix_crop_pick: bool = False,
         index_mode: IndexMode = IndexMode.OFF,
     ) -> None:
         """Instantiate a new ModelDataset.
@@ -709,7 +753,7 @@ class ModelDataset(torch.utils.data.Dataset):
             task: the task to train on
             workers: number of workers to use for initializing the dataset
             name: name of the dataset
-            fix_patch_pick: if True, fix the patch pick to be the same every time
+            fix_crop_pick: if True, fix the crop pick to be the same every time
                 for a given window. Useful for testing (default: False)
             index_mode: controls dataset index caching behavior (default: IndexMode.OFF)
         """
@@ -718,14 +762,14 @@ class ModelDataset(torch.utils.data.Dataset):
         self.inputs = inputs
         self.task = task
         self.name = name
-        self.fix_patch_pick = fix_patch_pick
+        self.fix_crop_pick = fix_crop_pick
         if split_config.transforms:
             self.transforms = Sequential(*split_config.transforms)
         else:
             self.transforms = rslearn.train.transforms.transform.Identity()
 
         # Get normalized crop size from the SplitConfig.
-        # But if load all patches is enabled, this is handled by AllCropsDataset, so
+        # But if load_all_crops is enabled, this is handled by AllCropsDataset, so
         # here we instead load the entire windows.
         if split_config.get_load_all_crops():
             self.crop_size = None
@@ -952,7 +996,7 @@ class ModelDataset(torch.utils.data.Dataset):
         """Get a list of examples in the dataset.
 
         If load_all_crops is False, this is a list of Windows. Otherwise, this is a
-        list of (window, crop_bounds, (crop_idx, # patches)) tuples.
+        list of (window, crop_bounds, (crop_idx, # crops)) tuples.
         """
         if self.dataset_examples is None:
             logger.debug(
@@ -985,7 +1029,7 @@ class ModelDataset(torch.utils.data.Dataset):
         """
         dataset_examples = self.get_dataset_examples()
         example = dataset_examples[idx]
-        rng = random.Random(idx if self.fix_patch_pick else None)
+        rng = random.Random(idx if self.fix_crop_pick else None)
 
         # Select bounds to read.
         if self.crop_size:

rslearn/train/lightning_module.py

@@ -6,12 +6,14 @@ from typing import Any
 
 import lightning as L
 import torch
+import wandb
 from lightning.pytorch.utilities.types import OptimizerLRSchedulerConfig
 from PIL import Image
 from upath import UPath
 
 from rslearn.log_utils import get_logger
 
+from .metrics import NonScalarMetricOutput
 from .model_context import ModelContext, ModelOutput
 from .optimizer import AdamW, OptimizerFactory
 from .scheduler import PlateauScheduler, SchedulerFactory
@@ -210,15 +212,53 @@ class RslearnLightningModule(L.LightningModule):
             # Fail silently for single-dataset case, which is okay
             pass
 
+    def _log_non_scalar_metric(self, name: str, value: NonScalarMetricOutput) -> None:
+        """Log a non-scalar metric to wandb.
+
+        Args:
+            name: the metric name (e.g., "val_confusion_matrix")
+            value: the non-scalar metric output
+        """
+        # The non-scalar metrics are logging directly without Lightning
+        # So we need to skip logging during sanity check.
+        if self.trainer.sanity_checking:
+            return
+
+        # Wandb is required for logging non-scalar metrics.
+        if not wandb.run:
+            logger.warning(
+                f"Weights & Biases is not initialized, skipping logging of {name}"
+            )
+            return
+
+        value.log_to_wandb(name)
+
     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.
+
+        Non-scalar metrics (like confusion matrices) are logged separately using
+        logger-specific APIs.
         """
         metrics = self.val_metrics.compute()
-        self.log_dict(metrics)
+
+        # Separate scalar and non-scalar metrics
+        scalar_metrics = {}
+        for k, v in metrics.items():
+            if isinstance(v, NonScalarMetricOutput):
+                self._log_non_scalar_metric(k, v)
+            elif isinstance(v, torch.Tensor) and v.dim() > 0 and v.numel() > 1:
+                raise ValueError(
+                    f"Metric '{k}' returned a non-scalar tensor with shape {v.shape}. "
+                    "Wrap it in a NonScalarMetricOutput subclass."
+                )
+            else:
+                scalar_metrics[k] = v
+
+        self.log_dict(scalar_metrics)
         self.val_metrics.reset()
 
     def on_test_epoch_end(self) -> None:
@@ -227,14 +267,30 @@ class RslearnLightningModule(L.LightningModule):
         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.
+
+        Non-scalar metrics (like confusion matrices) are logged separately.
         """
         metrics = self.test_metrics.compute()
-        self.log_dict(metrics)
+
+        # Separate scalar and non-scalar metrics
+        scalar_metrics = {}
+        for k, v in metrics.items():
+            if isinstance(v, NonScalarMetricOutput):
+                self._log_non_scalar_metric(k, v)
+            elif isinstance(v, torch.Tensor) and v.dim() > 0 and v.numel() > 1:
+                raise ValueError(
+                    f"Metric '{k}' returned a non-scalar tensor with shape {v.shape}. "
+                    "Wrap it in a NonScalarMetricOutput subclass."
+                )
+            else:
+                scalar_metrics[k] = v
+
+        self.log_dict(scalar_metrics)
         self.test_metrics.reset()
 
         if self.metrics_file:
             with open(self.metrics_file, "w") as f:
-                metrics_dict = {k: v.item() for k, v in metrics.items()}
+                metrics_dict = {k: v.item() for k, v in scalar_metrics.items()}
                 json.dump(metrics_dict, f, indent=4)
                 logger.info(f"Saved metrics to {self.metrics_file}")
 

rslearn/train/metrics.py

@@ -0,0 +1,162 @@
+"""Metric output classes for non-scalar metrics."""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+import torch
+import wandb
+from torchmetrics import Metric
+
+from rslearn.log_utils import get_logger
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class NonScalarMetricOutput(ABC):
+    """Base class for non-scalar metric outputs that need special logging.
+
+    Subclasses should implement the log_to_wandb method to define how the metric
+    should be logged (only supports logging to Weights & Biases).
+    """
+
+    @abstractmethod
+    def log_to_wandb(self, name: str) -> None:
+        """Log this metric to wandb.
+
+        Args:
+            name: the metric name
+        """
+        pass
+
+
+@dataclass
+class ConfusionMatrixOutput(NonScalarMetricOutput):
+    """Confusion matrix metric output.
+
+    Args:
+        confusion_matrix: confusion matrix of shape (num_classes, num_classes)
+            where cm[i, j] is the count of samples with true label i and predicted
+            label j.
+        class_names: optional list of class names for axis labels
+    """
+
+    confusion_matrix: torch.Tensor
+    class_names: list[str] | None = None
+
+    def _expand_confusion_matrix(self) -> tuple[list[int], list[int]]:
+        """Expand confusion matrix to (preds, labels) pairs for wandb.
+
+        Returns:
+            Tuple of (preds, labels) as lists of integers.
+        """
+        cm = self.confusion_matrix.detach().cpu()
+
+        # Handle extra dimensions from distributed reduction
+        if cm.dim() > 2:
+            cm = cm.sum(dim=0)
+
+        total = cm.sum().item()
+        if total == 0:
+            return [], []
+
+        preds = []
+        labels = []
+        for true_label in range(cm.shape[0]):
+            for pred_label in range(cm.shape[1]):
+                count = cm[true_label, pred_label].item()
+                if count > 0:
+                    preds.extend([pred_label] * int(count))
+                    labels.extend([true_label] * int(count))
+
+        return preds, labels
+
+    def log_to_wandb(self, name: str) -> None:
+        """Log confusion matrix to wandb.
+
+        Args:
+            name: the metric name (e.g., "val_confusion_matrix")
+        """
+        preds, labels = self._expand_confusion_matrix()
+
+        if len(preds) == 0:
+            logger.warning(f"No samples to log for {name}")
+            return
+
+        num_classes = self.confusion_matrix.shape[0]
+        if self.class_names is None:
+            class_names = [str(i) for i in range(num_classes)]
+        else:
+            class_names = self.class_names
+
+        wandb.log(
+            {
+                name: wandb.plot.confusion_matrix(
+                    preds=preds,
+                    y_true=labels,
+                    class_names=class_names,
+                    title=name,
+                ),
+            },
+        )
+
+
+class ConfusionMatrixMetric(Metric):
+    """Confusion matrix metric that works on flattened inputs.
+
+    Expects preds of shape (N, C) and labels of shape (N,).
+    Should be wrapped by ClassificationMetric or SegmentationMetric
+    which handle the task-specific preprocessing.
+
+    Args:
+        num_classes: number of classes
+        class_names: optional list of class names for labeling
+    """
+
+    def __init__(
+        self,
+        num_classes: int,
+        class_names: list[str] | None = None,
+    ):
+        """Initialize a new ConfusionMatrixMetric.
+
+        Args:
+            num_classes: number of classes
+            class_names: optional list of class names for labeling
+        """
+        super().__init__()
+        self.num_classes = num_classes
+        self.class_names = class_names
+        self.add_state(
+            "confusion_matrix",
+            default=torch.zeros(num_classes, num_classes, dtype=torch.long),
+            dist_reduce_fx="sum",
+        )
+
+    def update(self, preds: torch.Tensor, labels: torch.Tensor) -> None:
+        """Update metric.
+
+        Args:
+            preds: predictions of shape (N, C) - probabilities
+            labels: ground truth of shape (N,) - class indices
+        """
+        if len(preds) == 0:
+            return
+
+        pred_classes = preds.argmax(dim=1)  # (N,)
+
+        for true_label in range(self.num_classes):
+            for pred_label in range(self.num_classes):
+                count = ((labels == true_label) & (pred_classes == pred_label)).sum()
+                self.confusion_matrix[true_label, pred_label] += count
+
+    def compute(self) -> ConfusionMatrixOutput:
+        """Returns the confusion matrix wrapped in ConfusionMatrixOutput."""
+        return ConfusionMatrixOutput(
+            confusion_matrix=self.confusion_matrix,
+            class_names=self.class_names,
+        )
+
+    def reset(self) -> None:
+        """Reset metric."""
+        super().reset()

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
@@ -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/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

@@ -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/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-0.0.27.dist-info → rslearn-0.0.28.dist-info}/METADATA

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

{rslearn-0.0.27.dist-info → rslearn-0.0.28.dist-info}/RECORD

@@ -47,7 +47,7 @@ rslearn/dataset/materialize.py,sha256=VoL5Qf5pGcQV4QMlO5vrcu7w0Sl1NdIRLUVk0kSCMO
 rslearn/dataset/remap.py,sha256=6MaImsY02GNACpvRM81RvWmjZWRfAHxo_R3Ox6XLF6A,2723
 rslearn/dataset/window.py,sha256=X4q8YzcSOTtwKxCPf71QLMoyKUtYMSnZu0kPnmVSUx4,10644
 rslearn/dataset/storage/__init__.py,sha256=R50AVV5LH2g7ol0-jyvGcB390VsclXGbJXz4fmkn9as,52
-rslearn/dataset/storage/file.py,sha256=g9HZ3CD4QcgyVNsBaXhjIKQgDOAeZ4R08sJ7ntx4wo8,6815
+rslearn/dataset/storage/file.py,sha256=GJJgH_eHknLbMQwoC3mOoXJKl6Ha3oNXzz62FIEMWlg,7130
 rslearn/dataset/storage/storage.py,sha256=DxZ7iwV938PiLwdQzb5EXSb4Mj8bRGmOTmA9fzq_Ge8,4840
 rslearn/models/__init__.py,sha256=_vWoF9d2Slah8-6XhYhdU4SRsy_CNxXjCGQTD2yvu3Q,22
 rslearn/models/anysat.py,sha256=nzk6hB83ltNFNXYRNA1rTvq2AQcAhwyvgBaZui1M37o,8107
@@ -112,14 +112,15 @@ rslearn/models/presto/__init__.py,sha256=eZrB-XKi_vYqZhpyAOwppJi4dRuMtYVAdbq7KRy
 rslearn/models/presto/presto.py,sha256=fkyHB85Hfx5L-4yejSFAFv83gk9VFqAR1GTgggtq0EA,11049
 rslearn/models/presto/single_file_presto.py,sha256=-P00xjhj9dx3O6HqWpQmG9dPk_i6bT_t8vhX4uQm5tA,30242
 rslearn/tile_stores/__init__.py,sha256=-cW1J7So60SEP5ZLHCPdaFBV5CxvV3QlOhaFnUkhTJ0,1675
-rslearn/tile_stores/default.py,sha256=PYaDNvBxhJTDKJGw0EjDTSE1OKajR7_iJpMbOjj-mE8,15054
+rslearn/tile_stores/default.py,sha256=AG2j0FCNi_4cnXqLjRIef5wMqMJ5_YtSkTIhk7qJQVQ,15134
 rslearn/tile_stores/tile_store.py,sha256=9AeYduDYPp_Ia2NMlq6osptpz_AFGIOQcLJrqZ_m-z0,10469
 rslearn/train/__init__.py,sha256=fnJyY4aHs5zQqbDKSfXsJZXY_M9fbTsf7dRYaPwZr2M,30
 rslearn/train/all_crops_dataset.py,sha256=CWnqbSjRXJZQsudljvpA07oldiP4fZTmjwrT0sjVnq4,21399
-rslearn/train/data_module.py,sha256=yPShftkHJ2bhJ4carwYYb9c3PkkP7ArzXQyu37EuAxk,23718
-rslearn/train/dataset.py,sha256=AAkdHX3q8VD1Geq9yIatiPkM5blA2luVPPwDmXDB6Z8,43284
+rslearn/train/data_module.py,sha256=G1TRhXg8SPewYy0BTZN5KpeLPK72qIaH15ePfUwrxgM,23865
+rslearn/train/dataset.py,sha256=vCmm6yrW2bAc5A94aBwQe-SOIGdVcZYMM2oBYRq2_sw,45253
 rslearn/train/dataset_index.py,sha256=S5iXhQga5gnnkDqThXXlyjIwkJBPVWiUfDPx3iVs-pw,5306
-rslearn/train/lightning_module.py,sha256=m0aIGk5xO5y12DEiwSl6eAko6X-gQ78_Wsbvz4Hb_NE,15364
+rslearn/train/lightning_module.py,sha256=n4hasJBVlAmMhvf2yaFo0gy1vGz5haQkJpZdCSKlJ8A,17482
+rslearn/train/metrics.py,sha256=RknMf2n09D5XBCf0YM4Zmm0XI-pFbRtsmbY51ipVMPk,4799
 rslearn/train/model_context.py,sha256=8DMWGj5xCRmRDo_38lkhkUMHfK_yg3XZrUJQIz5a1vA,3200
 rslearn/train/optimizer.py,sha256=EKSqkmERalDA0bF32Gey7n6z69KLyaUWKlRsGJfKBmE,927
 rslearn/train/prediction_writer.py,sha256=cRFehEtr0iBuVqzE69a0B4Lvb8ywxLeyon34KWI86H0,16961
@@ -130,13 +131,13 @@ rslearn/train/callbacks/freeze_unfreeze.py,sha256=8fIzBMhCKKjpTffIeAdhdSjsBd8NjT
 rslearn/train/callbacks/gradients.py,sha256=4YqCf0tBb6E5FnyFYbveXfQFlgNPyxIXb2FCWX4-6qs,5075
 rslearn/train/callbacks/peft.py,sha256=wEOKsS3RhsRaZTXn_Kz2wdsZdIiIaZPdCJWtdJBurT8,4156
 rslearn/train/tasks/__init__.py,sha256=dag1u72x1-me6y0YcOubUo5MYZ0Tjf6-dOir9UeFNMs,75
-rslearn/train/tasks/classification.py,sha256=H-Ayqm59IxwrczC8lUV5J5vg-JILhQhTiVlyaTpBs2k,14259
+rslearn/train/tasks/classification.py,sha256=_3cRa8ojd9sG2ELRW_BvByZh2YFdCBaklR8Kv9LAgOY,14864
 rslearn/train/tasks/detection.py,sha256=uDMGtsCMSk9OGXn-vpFKBAyHyVN0ji2NCfqBgg1BQyw,21725
 rslearn/train/tasks/embedding.py,sha256=NdJEAaDWlWYzvOBVf7eIHfFOzqTgavfFH1J1gMbAMVo,3891
 rslearn/train/tasks/multi_task.py,sha256=32hvwyVsHqt7N_M3zXsTErK1K7-0-BPHzt7iGNehyaI,6314
-rslearn/train/tasks/per_pixel_regression.py,sha256=njShN-U9fx3SPcCxGgbDlZAp3DT_GlTt0BRZS416gnw,10387
-rslearn/train/tasks/regression.py,sha256=_TGlj3PA14Iye0duf25TcGZQpFXU9fGfNilzpNuPS78,12693
-rslearn/train/tasks/segmentation.py,sha256=dn1yo1dIArKvW9Giw8-LZyIZ87q76eslL0mk58GyApo,29663
+rslearn/train/tasks/per_pixel_regression.py,sha256=3m_BTP2akadYe3IuAlCG2bd_alfNyom55-pFrI2q4PE,10928
+rslearn/train/tasks/regression.py,sha256=TzHL42gm3aIdev0R7_uz_TSYbAwSvQPjCD42y1p9_7Y,13269
+rslearn/train/tasks/segmentation.py,sha256=b9XS09EQvum89eoW3vWqFMKuCRtznODteKIr1hFnIz4,30531
 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=S8f1svzwb5UmeAgzXe4Af_hFvt5o0tQctIE6t3QYuPI,2625
@@ -153,7 +154,7 @@ rslearn/utils/__init__.py,sha256=GZc1erpEfXTc32yjEDbt5rnMrnXEBY7WVm3v4NlwwWY,620
 rslearn/utils/array.py,sha256=RC7ygtPnQwU6Lb9kwORvNxatJcaJ76JPsykQvndAfes,2444
 rslearn/utils/colors.py,sha256=ELY9_buH06TOVPLrDAyf2S0G--ZiOxnnP8Ujim6_3ig,369
 rslearn/utils/feature.py,sha256=lsg0WThZDJzo1mrbaL04dXYI5G3x-n5FG9aEjj7uUaI,1649
-rslearn/utils/fsspec.py,sha256=h3fER_bkewzR9liEAULXguTIvXLUXA17pC_yZoWN5Tk,5902
+rslearn/utils/fsspec.py,sha256=TcEUgXKvsmtKHv5JVOI2Vp4WNfVNeTok0x4JgZaD1iw,7052
 rslearn/utils/geometry.py,sha256=VzLoxtwdV3uC3szowT-bGuCFF6ge8eK0m01lq8q-01Q,22423
 rslearn/utils/get_utm_ups_crs.py,sha256=kUrcyjCK7KWvuP1XR-nURPeRqYeRO-3L8QUJ1QTF9Ps,3599
 rslearn/utils/grid_index.py,sha256=hRmrtgpqN1pLa-djnZtgSXqKJlbgGyttGnCEmPLD0zo,2347
@@ -175,10 +176,10 @@ rslearn/vis/render_sensor_image.py,sha256=D0ynK6ABPV046970lIKwF98klpSCtrsUvZTwtZ
 rslearn/vis/render_vector_label.py,sha256=ncwgRKCYCJCK1-wTpjgksOiDDebku37LpAyq6wsg4jg,14939
 rslearn/vis/utils.py,sha256=Zop3dEmyaXUYhPiGdYzrTO8BRXWscP2dEZy2myQUnNk,2765
 rslearn/vis/vis_server.py,sha256=kIGnhTy-yfu5lBOVCoo8VVG259i974JPszudCePbzfI,20157
-rslearn-0.0.27.dist-info/licenses/LICENSE,sha256=_99ZWPoLdlUbqZoSC5DF4ihiNwl5rTEmBaq2fACecdg,11352
-rslearn-0.0.27.dist-info/licenses/NOTICE,sha256=wLPr6rwV_jCg-xEknNGwhnkfRfuoOE9MZ-lru2yZyLI,5070
-rslearn-0.0.27.dist-info/METADATA,sha256=_6aHsGpgH_T2MQSb9qKeYshP6zdzwu3CmI2sBhD3dqk,38714
-rslearn-0.0.27.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-rslearn-0.0.27.dist-info/entry_points.txt,sha256=doTBQ57NT7nq-dgYGgTTw6mafcGWb_4PWYtYR4rGm50,46
-rslearn-0.0.27.dist-info/top_level.txt,sha256=XDKo90WBH8P9RQumHxo0giLJsoufT4r9odv-WE6Ahk4,8
-rslearn-0.0.27.dist-info/RECORD,,
+rslearn-0.0.28.dist-info/licenses/LICENSE,sha256=_99ZWPoLdlUbqZoSC5DF4ihiNwl5rTEmBaq2fACecdg,11352
+rslearn-0.0.28.dist-info/licenses/NOTICE,sha256=wLPr6rwV_jCg-xEknNGwhnkfRfuoOE9MZ-lru2yZyLI,5070
+rslearn-0.0.28.dist-info/METADATA,sha256=OSGXg3yVyndUAZYL9EwvCl095O_Vfknpxlhx8O5dLmQ,38714
+rslearn-0.0.28.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+rslearn-0.0.28.dist-info/entry_points.txt,sha256=doTBQ57NT7nq-dgYGgTTw6mafcGWb_4PWYtYR4rGm50,46
+rslearn-0.0.28.dist-info/top_level.txt,sha256=XDKo90WBH8P9RQumHxo0giLJsoufT4r9odv-WE6Ahk4,8
+rslearn-0.0.28.dist-info/RECORD,,