Add "mixed" input format to segmentation metrics (#3176)

* Add `"mixed"` input format for segmentation metrics
* update changelog
* better naming of inputs
* Fixed mean_iou test
* Fixed another test

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Nicki Skafte <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: 4 people

CHANGELOG.md

@@ -32,6 +32,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added `Lip Vertex Error (LVE)` in multimodal domain ([3090](https://github.com/Lightning-AI/torchmetrics/pull/3090))
 
 
+- Added `mixed` input format to segmentation metrics ([3176](https://github.com/Lightning-AI/torchmetrics/pull/3176))
+
 ### Changed
 
 -

src/torchmetrics/functional/segmentation/dice.py

@@ -17,7 +17,7 @@
 from torch import Tensor
 from typing_extensions import Literal
 
-from torchmetrics.functional.segmentation.utils import _ignore_background
+from torchmetrics.functional.segmentation.utils import _check_mixed_shape, _ignore_background
 from torchmetrics.utilities import rank_zero_warn
 from torchmetrics.utilities.checks import _check_same_shape
 from torchmetrics.utilities.compute import _safe_divide
@@ -27,7 +27,7 @@ def _dice_score_validate_args(
     num_classes: int,
     include_background: bool,
     average: Optional[Literal["micro", "macro", "weighted", "none"]] = "micro",
-    input_format: Literal["one-hot", "index"] = "one-hot",
+    input_format: Literal["one-hot", "index", "mixed"] = "one-hot",
     aggregation_level: Optional[Literal["samplewise", "global"]] = "samplewise",
 ) -> None:
     """Validate the arguments of the metric."""
@@ -38,8 +38,10 @@ def _dice_score_validate_args(
     allowed_average = ["micro", "macro", "weighted", "none"]
     if average is not None and average not in allowed_average:
         raise ValueError(f"Expected argument `average` to be one of {allowed_average} or None, but got {average}.")
-    if input_format not in ["one-hot", "index"]:
-        raise ValueError(f"Expected argument `input_format` to be one of 'one-hot', 'index', but got {input_format}.")
+    if input_format not in ["one-hot", "index", "mixed"]:
+        raise ValueError(
+            f"Expected argument `input_format` to be one of 'one-hot', 'index', 'mixed', but got {input_format}."
+        )
     if aggregation_level not in ("samplewise", "global"):
         raise ValueError(
             f"Expected argument `aggregation_level` to be one of `samplewise`, `global`, but got {aggregation_level}"
@@ -51,14 +53,22 @@ def _dice_score_update(
     target: Tensor,
     num_classes: int,
     include_background: bool,
-    input_format: Literal["one-hot", "index"] = "one-hot",
+    input_format: Literal["one-hot", "index", "mixed"] = "one-hot",
 ) -> tuple[Tensor, Tensor, Tensor]:
     """Update the state with the current prediction and target."""
-    _check_same_shape(preds, target)
+    if input_format == "mixed":
+        _check_mixed_shape(preds, target)
+    else:
+        _check_same_shape(preds, target)
 
     if input_format == "index":
         preds = torch.nn.functional.one_hot(preds, num_classes=num_classes).movedim(-1, 1)
         target = torch.nn.functional.one_hot(target, num_classes=num_classes).movedim(-1, 1)
+    elif input_format == "mixed":
+        if preds.dim() == (target.dim() + 1):
+            target = torch.nn.functional.one_hot(target, num_classes=num_classes).movedim(-1, 1)
+        elif (preds.dim() + 1) == target.dim():
+            preds = torch.nn.functional.one_hot(preds, num_classes=num_classes).movedim(-1, 1)
 
     if preds.ndim < 3:
         raise ValueError(f"Expected both `preds` and `target` to have at least 3 dimensions, but got {preds.ndim}.")
@@ -117,7 +127,7 @@ def dice_score(
     num_classes: int,
     include_background: bool = True,
     average: Optional[Literal["micro", "macro", "weighted", "none"]] = "micro",
-    input_format: Literal["one-hot", "index"] = "one-hot",
+    input_format: Literal["one-hot", "index", "mixed"] = "one-hot",
     aggregation_level: Optional[Literal["samplewise", "global"]] = "samplewise",
 ) -> Tensor:
     """Compute the Dice score for semantic segmentation.
@@ -128,9 +138,10 @@ def dice_score(
         num_classes: Number of classes
         include_background: Whether to include the background class in the computation
         average: The method to average the dice score. Options are ``"micro"``, ``"macro"``, ``"weighted"``, ``"none"``
-          or ``None``. This determines how to average the dice score across different classes.
-        input_format: What kind of input the function receives. Choose between ``"one-hot"`` for one-hot encoded tensors
-          or ``"index"`` for index tensors
+            or ``None``. This determines how to average the dice score across different classes.
+        input_format: What kind of input the function receives.
+            Choose between ``"one-hot"`` for one-hot encoded tensors, ``"index"`` for index tensors
+            or ``"mixed"`` for one one-hot encoded and one index tensor
         aggregation_level: The level at which to aggregate the dice score. Options are ``"samplewise"`` or ``"global"``.
             For ``"samplewise"`` the dice score is computed for each sample and then averaged. For ``"global"`` the dice
             score is computed globally over all samples.

src/torchmetrics/functional/segmentation/generalized_dice.py

@@ -17,7 +17,7 @@
 from torch import Tensor
 from typing_extensions import Literal
 
-from torchmetrics.functional.segmentation.utils import _ignore_background
+from torchmetrics.functional.segmentation.utils import _check_mixed_shape, _ignore_background
 from torchmetrics.utilities.checks import _check_same_shape
 from torchmetrics.utilities.compute import _safe_divide
 
@@ -27,7 +27,7 @@ def _generalized_dice_validate_args(
     include_background: bool,
     per_class: bool,
     weight_type: Literal["square", "simple", "linear"],
-    input_format: Literal["one-hot", "index"],
+    input_format: Literal["one-hot", "index", "mixed"],
 ) -> None:
     """Validate the arguments of the metric."""
     if not isinstance(num_classes, int) or num_classes <= 0:
@@ -40,8 +40,10 @@ def _generalized_dice_validate_args(
         raise ValueError(
             f"Expected argument `weight_type` to be one of 'square', 'simple', 'linear', but got {weight_type}."
         )
-    if input_format not in ["one-hot", "index"]:
-        raise ValueError(f"Expected argument `input_format` to be one of 'one-hot', 'index', but got {input_format}.")
+    if input_format not in ["one-hot", "index", "mixed"]:
+        raise ValueError(
+            f"Expected argument `input_format` to be one of 'one-hot', 'index', 'mixed', but got {input_format}."
+        )
 
 
 def _generalized_dice_update(
@@ -50,14 +52,22 @@ def _generalized_dice_update(
     num_classes: int,
     include_background: bool,
     weight_type: Literal["square", "simple", "linear"] = "square",
-    input_format: Literal["one-hot", "index"] = "one-hot",
+    input_format: Literal["one-hot", "index", "mixed"] = "one-hot",
 ) -> Tuple[Tensor, Tensor]:
     """Update the state with the current prediction and target."""
-    _check_same_shape(preds, target)
+    if input_format == "mixed":
+        _check_mixed_shape(preds, target)
+    else:
+        _check_same_shape(preds, target)
 
     if input_format == "index":
         preds = torch.nn.functional.one_hot(preds, num_classes=num_classes).movedim(-1, 1)
         target = torch.nn.functional.one_hot(target, num_classes=num_classes).movedim(-1, 1)
+    elif input_format == "mixed":
+        if preds.dim() == (target.dim() + 1):
+            target = torch.nn.functional.one_hot(target, num_classes=num_classes).movedim(-1, 1)
+        elif (preds.dim() + 1) == target.dim():
+            preds = torch.nn.functional.one_hot(preds, num_classes=num_classes).movedim(-1, 1)
 
     if preds.ndim < 3:
         raise ValueError(f"Expected both `preds` and `target` to have at least 3 dimensions, but got {preds.ndim}.")
@@ -109,7 +119,7 @@ def generalized_dice_score(
     include_background: bool = True,
     per_class: bool = False,
     weight_type: Literal["square", "simple", "linear"] = "square",
-    input_format: Literal["one-hot", "index"] = "one-hot",
+    input_format: Literal["one-hot", "index", "mixed"] = "one-hot",
 ) -> Tensor:
     """Compute the Generalized Dice Score for semantic segmentation.
 
@@ -120,8 +130,9 @@ def generalized_dice_score(
         include_background: Whether to include the background class in the computation
         per_class: Whether to compute the score for each class separately, else average over all classes
         weight_type: Type of weight factor to apply to the classes. One of ``"square"``, ``"simple"``, or ``"linear"``
-        input_format: What kind of input the function receives. Choose between ``"one-hot"`` for one-hot encoded tensors
-            or ``"index"`` for index tensors
+        input_format: What kind of input the function receives.
+            Choose between ``"one-hot"`` for one-hot encoded tensors, ``"index"`` for index tensors
+            or ``"mixed"`` for one one-hot encoded and one index tensor
 
     Returns:
         The Generalized Dice Score

src/torchmetrics/functional/segmentation/hausdorff_distance.py

@@ -18,6 +18,7 @@
 from torch import Tensor
 
 from torchmetrics.functional.segmentation.utils import (
+    _check_mixed_shape,
     _ignore_background,
     edge_surface_distance,
 )
@@ -30,7 +31,7 @@ def _hausdorff_distance_validate_args(
     distance_metric: Literal["euclidean", "chessboard", "taxicab"] = "euclidean",
     spacing: Optional[Union[Tensor, list[float]]] = None,
     directed: bool = False,
-    input_format: Literal["one-hot", "index"] = "one-hot",
+    input_format: Literal["one-hot", "index", "mixed"] = "one-hot",
 ) -> None:
     """Validate the arguments of `hausdorff_distance` function."""
     if num_classes <= 0:
@@ -45,8 +46,10 @@ def _hausdorff_distance_validate_args(
         raise ValueError(f"Arg `spacing` must be a list or tensor, but got {type(spacing)}.")
     if not isinstance(directed, bool):
         raise ValueError(f"Expected argument `directed` must be a boolean, but got {directed}.")
-    if input_format not in ["one-hot", "index"]:
-        raise ValueError(f"Expected argument `input_format` to be one of 'one-hot', 'index', but got {input_format}.")
+    if input_format not in ["one-hot", "index", "mixed"]:
+        raise ValueError(
+            f"Expected argument `input_format` to be one of 'one-hot', 'index', 'mixed', but got {input_format}."
+        )
 
 
 def hausdorff_distance(
@@ -57,7 +60,7 @@ def hausdorff_distance(
     distance_metric: Literal["euclidean", "chessboard", "taxicab"] = "euclidean",
     spacing: Optional[Union[Tensor, list[float]]] = None,
     directed: bool = False,
-    input_format: Literal["one-hot", "index"] = "one-hot",
+    input_format: Literal["one-hot", "index", "mixed"] = "one-hot",
 ) -> Tensor:
     """Calculate `Hausdorff Distance`_ for semantic segmentation.
 
@@ -70,8 +73,9 @@ def hausdorff_distance(
           `"chessboard"` or `"taxicab"`
         spacing: spacing between pixels along each spatial dimension. If not provided the spacing is assumed to be 1
         directed: whether to calculate directed or undirected Hausdorff distance
-        input_format: What kind of input the function receives. Choose between ``"one-hot"`` for one-hot encoded tensors
-          or ``"index"`` for index tensors
+        input_format: What kind of input the function receives.
+            Choose between ``"one-hot"`` for one-hot encoded tensors, ``"index"`` for index tensors
+            or ``"mixed"`` for one one-hot encoded and one index tensor
 
     Returns:
         Hausdorff Distance for each class and batch element
@@ -89,11 +93,19 @@ def hausdorff_distance(
 
     """
     _hausdorff_distance_validate_args(num_classes, include_background, distance_metric, spacing, directed, input_format)
-    _check_same_shape(preds, target)
+    if input_format == "mixed":
+        _check_mixed_shape(preds, target)
+    else:
+        _check_same_shape(preds, target)
 
     if input_format == "index":
         preds = torch.nn.functional.one_hot(preds, num_classes=num_classes).movedim(-1, 1)
         target = torch.nn.functional.one_hot(target, num_classes=num_classes).movedim(-1, 1)
+    elif input_format == "mixed":
+        if preds.dim() == (target.dim() + 1):
+            target = torch.nn.functional.one_hot(target, num_classes=num_classes).movedim(-1, 1)
+        elif (preds.dim() + 1) == target.dim():
+            preds = torch.nn.functional.one_hot(preds, num_classes=num_classes).movedim(-1, 1)
 
     if not include_background:
         preds, target = _ignore_background(preds, target)

src/torchmetrics/functional/segmentation/mean_iou.py

@@ -18,15 +18,15 @@
 from torch import Tensor
 from typing_extensions import Literal
 
-from torchmetrics.functional.segmentation.utils import _ignore_background
+from torchmetrics.functional.segmentation.utils import _check_mixed_shape, _ignore_background
 from torchmetrics.utilities.checks import _check_same_shape
 from torchmetrics.utilities.compute import _safe_divide
 
 
 def _mean_iou_reshape_args(
     preds: Tensor,
     targets: Tensor,
-    input_format: Literal["one-hot", "index"] = "one-hot",
+    input_format: Literal["one-hot", "index", "mixed"] = "one-hot",
 ) -> Tuple[Tensor, Tensor]:
     """Reshape tensors to 3D if needed."""
     if input_format == "one-hot":
@@ -48,11 +48,11 @@ def _mean_iou_validate_args(
     num_classes: Optional[int],
     include_background: bool,
     per_class: bool,
-    input_format: Literal["one-hot", "index"] = "one-hot",
+    input_format: Literal["one-hot", "index", "mixed"] = "one-hot",
 ) -> None:
     """Validate the arguments of the metric."""
-    if input_format == "index" and num_classes is None:
-        raise ValueError("Argument `num_classes` must be provided when `input_format='index'`.")
+    if input_format in ["index", "mixed"] and num_classes is None:
+        raise ValueError("Argument `num_classes` must be provided when `input_format` is 'index' or 'mixed'.")
     if num_classes is not None and num_classes <= 0:
         raise ValueError(
             f"Expected argument `num_classes` must be `None` or a positive integer, but got {num_classes}."
@@ -61,20 +61,25 @@ def _mean_iou_validate_args(
         raise ValueError(f"Expected argument `include_background` must be a boolean, but got {include_background}.")
     if not isinstance(per_class, bool):
         raise ValueError(f"Expected argument `per_class` must be a boolean, but got {per_class}.")
-    if input_format not in ["one-hot", "index"]:
-        raise ValueError(f"Expected argument `input_format` to be one of 'one-hot', 'index', but got {input_format}.")
+    if input_format not in ["one-hot", "index", "mixed"]:
+        raise ValueError(
+            f"Expected argument `input_format` to be one of 'one-hot', 'index', 'mixed', but got {input_format}."
+        )
 
 
 def _mean_iou_update(
     preds: Tensor,
     target: Tensor,
     num_classes: Optional[int] = None,
     include_background: bool = False,
-    input_format: Literal["one-hot", "index"] = "one-hot",
+    input_format: Literal["one-hot", "index", "mixed"] = "one-hot",
 ) -> tuple[Tensor, Tensor]:
     """Update the intersection and union counts for the mean IoU computation."""
     preds, target = _mean_iou_reshape_args(preds, target, input_format)
-    _check_same_shape(preds, target)
+    if input_format == "mixed":
+        _check_mixed_shape(preds, target)
+    else:
+        _check_same_shape(preds, target)
 
     if input_format == "index":
         if num_classes is None:
@@ -88,6 +93,13 @@ def _mean_iou_update(
             raise IndexError(f"Cannot determine `num_classes` from `preds` tensor: {preds}.") from err
         if num_classes == 0:
             raise ValueError(f"Expected argument `num_classes` to be a positive integer, but got {num_classes}.")
+    elif input_format == "mixed":
+        if num_classes is None:
+            raise ValueError("Argument `num_classes` must be provided when `input_format='mixed'`.")
+        if preds.dim() == (target.dim() + 1):
+            target = torch.nn.functional.one_hot(target, num_classes=num_classes).movedim(-1, 1)
+        elif (preds.dim() + 1) == target.dim():
+            preds = torch.nn.functional.one_hot(preds, num_classes=num_classes).movedim(-1, 1)
 
     if not include_background:
         preds, target = _ignore_background(preds, target)
@@ -115,7 +127,7 @@ def mean_iou(
     num_classes: Optional[int] = None,
     include_background: bool = True,
     per_class: bool = False,
-    input_format: Literal["one-hot", "index"] = "one-hot",
+    input_format: Literal["one-hot", "index", "mixed"] = "one-hot",
 ) -> Tensor:
     """Calculates the mean Intersection over Union (mIoU) for semantic segmentation.
 
@@ -127,8 +139,9 @@ def mean_iou(
         num_classes: Number of classes (required when input_format="index", optional when input_format="one-hot")
         include_background: Whether to include the background class in the computation
         per_class: Whether to compute the IoU for each class separately, else average over all classes
-        input_format: What kind of input the function receives. Choose between ``"one-hot"`` for one-hot encoded tensors
-            or ``"index"`` for index tensors
+        input_format: What kind of input the function receives.
+            Choose between ``"one-hot"`` for one-hot encoded tensors, ``"index"`` for index tensors
+            or ``"mixed"`` for one one-hot encoded and one index tensor
 
     Returns:
         The mean IoU score

src/torchmetrics/functional/segmentation/utils.py

@@ -31,6 +31,24 @@ def _ignore_background(preds: Tensor, target: Tensor) -> tuple[Tensor, Tensor]:
     return preds, target
 
 
+def _check_mixed_shape(preds: Tensor, target: Tensor) -> None:
+    """Check that predictions and target have the same shape, else raise error."""
+    if preds.dim() == (target.dim() + 1):
+        if preds.shape[0] != target.shape[0] or preds.shape[2:] != target.shape[1:]:
+            raise RuntimeError(
+                f"Predictions and targets are expected to have the same shape, got {preds.shape} and {target.shape}."
+            )
+    elif (preds.dim() + 1) == target.dim():
+        if preds.shape[0] != target.shape[0] or preds.shape[1:] != target.shape[2:]:
+            raise RuntimeError(
+                f"Predictions and targets are expected to have the same shape, got {preds.shape} and {target.shape}."
+            )
+    else:
+        raise RuntimeError(
+            f"Predictions and targets are expected to have the same shape, got {preds.shape} and {target.shape}."
+        )
+
+
 def check_if_binarized(x: Tensor) -> None:
     """Check if tensor is binarized.
 

src/torchmetrics/segmentation/dice.py

@@ -70,8 +70,9 @@ class DiceScore(Metric):
         aggregation_level: The level at which to aggregate the dice score. Options are ``"samplewise"`` or ``"global"``.
             For ``"samplewise"`` the dice score is computed for each sample and then averaged. For ``"global"`` the dice
             score is computed globally over all samples.
-        input_format: What kind of input the function receives. Choose between ``"one-hot"`` for one-hot encoded tensors
-            or ``"index"`` for index tensors.
+        input_format: What kind of input the function receives.
+            Choose between ``"one-hot"`` for one-hot encoded tensors, ``"index"`` for index tensors
+            or ``"mixed"`` for one one-hot encoded and one index tensor
         zero_division: The value to return when there is a division by zero. Options are 1.0, 0.0, "warn" or "nan".
             Setting it to "warn" behaves like 0.0 but will also create a warning.
         kwargs: Additional keyword arguments, see :ref:`Metric kwargs` for more info.
@@ -84,7 +85,7 @@ class DiceScore(Metric):
         ValueError:
             If ``average`` is not one of ``"micro"``, ``"macro"``, ``"weighted"``, ``"none"`` or ``None``
         ValueError:
-            If ``input_format`` is not one of ``"one-hot"`` or ``"index"``
+            If ``input_format`` is not one of ``"one-hot"``, ``"index"`` or ``"mixed"``
 
     Example:
         >>> from torch import randint
@@ -116,7 +117,7 @@ def __init__(
         include_background: bool = True,
         average: Optional[Literal["micro", "macro", "weighted", "none"]] = "micro",
         aggregation_level: Optional[Literal["samplewise", "global"]] = "samplewise",
-        input_format: Literal["one-hot", "index"] = "one-hot",
+        input_format: Literal["one-hot", "index", "mixed"] = "one-hot",
         **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)