rslearn 0.0.27__tar.gz → 0.0.28__tar.gz

{rslearn-0.0.27/rslearn.egg-info → rslearn-0.0.28}/PKG-INFO

@@ -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 → rslearn-0.0.28}/pyproject.toml

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

{rslearn-0.0.27 → rslearn-0.0.28}/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-0.0.27 → rslearn-0.0.28}/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-0.0.27 → rslearn-0.0.28}/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-0.0.27 → rslearn-0.0.28}/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-0.0.27 → rslearn-0.0.28}/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-0.0.28/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()