Float input support for segmentation metrics (#3198)

* Add logit support to segmentation metrics
* Improved num_classes logic
* Fix for mypy
* chlog

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Jirka Borovec <[email protected]>
Co-authored-by: Jirka B <[email protected]>
(cherry picked from commit 357c349)

Author: 3 people

CHANGELOG.md

@@ -21,6 +21,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 -
 
 
+- Float input support for segmentation metrics ([#3198](https://github.com/Lightning-AI/torchmetrics/pull/3198))
+
+
 ### Deprecated
 
 -

src/torchmetrics/functional/segmentation/dice.py

@@ -17,9 +17,8 @@
 from torch import Tensor
 from typing_extensions import Literal
 
-from torchmetrics.functional.segmentation.utils import _check_mixed_shape, _ignore_background
+from torchmetrics.functional.segmentation.utils import _segmentation_inputs_format
 from torchmetrics.utilities import rank_zero_warn
-from torchmetrics.utilities.checks import _check_same_shape
 from torchmetrics.utilities.compute import _safe_divide
 
 
@@ -56,25 +55,7 @@ def _dice_score_update(
     input_format: Literal["one-hot", "index", "mixed"] = "one-hot",
 ) -> tuple[Tensor, Tensor, Tensor]:
     """Update the state with the current prediction and 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}.")
-
-    if not include_background:
-        preds, target = _ignore_background(preds, target)
+    preds, target = _segmentation_inputs_format(preds, target, include_background, num_classes, input_format)
 
     reduce_axis = list(range(2, target.ndim))
     intersection = torch.sum(preds * target, dim=reduce_axis)

src/torchmetrics/functional/segmentation/generalized_dice.py

@@ -17,8 +17,7 @@
 from torch import Tensor
 from typing_extensions import Literal
 
-from torchmetrics.functional.segmentation.utils import _check_mixed_shape, _ignore_background
-from torchmetrics.utilities.checks import _check_same_shape
+from torchmetrics.functional.segmentation.utils import _segmentation_inputs_format
 from torchmetrics.utilities.compute import _safe_divide
 
 
@@ -55,25 +54,7 @@ def _generalized_dice_update(
     input_format: Literal["one-hot", "index", "mixed"] = "one-hot",
 ) -> Tuple[Tensor, Tensor]:
     """Update the state with the current prediction and 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}.")
-
-    if not include_background:
-        preds, target = _ignore_background(preds, target)
+    preds, target = _segmentation_inputs_format(preds, target, include_background, num_classes, input_format)
 
     reduce_axis = list(range(2, target.ndim))
     intersection = torch.sum(preds * target, dim=reduce_axis)

src/torchmetrics/functional/segmentation/hausdorff_distance.py

@@ -17,12 +17,7 @@
 import torch
 from torch import Tensor
 
-from torchmetrics.functional.segmentation.utils import (
-    _check_mixed_shape,
-    _ignore_background,
-    edge_surface_distance,
-)
-from torchmetrics.utilities.checks import _check_same_shape
+from torchmetrics.functional.segmentation.utils import _segmentation_inputs_format, edge_surface_distance
 
 
 def _hausdorff_distance_validate_args(
@@ -93,22 +88,8 @@ def hausdorff_distance(
 
     """
     _hausdorff_distance_validate_args(num_classes, include_background, distance_metric, spacing, directed, input_format)
-    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)
+    preds, target = _segmentation_inputs_format(preds, target, include_background, num_classes, input_format)
 
     distances = torch.zeros(preds.shape[0], preds.shape[1], device=preds.device)
 

src/torchmetrics/functional/segmentation/mean_iou.py

@@ -18,8 +18,7 @@
 from torch import Tensor
 from typing_extensions import Literal
 
-from torchmetrics.functional.segmentation.utils import _check_mixed_shape, _ignore_background
-from torchmetrics.utilities.checks import _check_same_shape
+from torchmetrics.functional.segmentation.utils import _segmentation_inputs_format
 from torchmetrics.utilities.compute import _safe_divide
 
 
@@ -51,8 +50,8 @@ def _mean_iou_validate_args(
     input_format: Literal["one-hot", "index", "mixed"] = "one-hot",
 ) -> None:
     """Validate the arguments of the metric."""
-    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 input_format in ["index"] and num_classes is None:
+        raise ValueError("Argument `num_classes` must be provided when `input_format` is 'index'.")
     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}."
@@ -76,33 +75,8 @@ def _mean_iou_update(
 ) -> tuple[Tensor, Tensor]:
     """Update the intersection and union counts for the mean IoU computation."""
     preds, target = _mean_iou_reshape_args(preds, target, input_format)
-    if input_format == "mixed":
-        _check_mixed_shape(preds, target)
-    else:
-        _check_same_shape(preds, target)
-
-    if input_format == "index":
-        if num_classes is None:
-            raise ValueError("Argument `num_classes` must be provided when `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 == "one-hot" and num_classes is None:
-        try:
-            num_classes = preds.shape[1]
-        except IndexError as err:
-            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)
+
+    preds, target = _segmentation_inputs_format(preds, target, include_background, num_classes, input_format)
 
     reduce_axis = list(range(2, preds.ndim))
     intersection = torch.sum(preds & target, dim=reduce_axis)
@@ -136,7 +110,8 @@ def mean_iou(
     Args:
         preds: Predictions from model
         target: Ground truth values
-        num_classes: Number of classes (required when input_format="index", optional when input_format="one-hot")
+        num_classes: Number of classes
+            (required when input_format="index", optional when input_format="one-hot" or "mixed")
         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.

src/torchmetrics/functional/segmentation/utils.py

@@ -49,6 +49,69 @@ def _check_mixed_shape(preds: Tensor, target: Tensor) -> None:
         )
 
 
+def _segmentation_inputs_format(
+    preds: Tensor,
+    target: Tensor,
+    include_background: bool,
+    num_classes: Optional[int] = None,
+    input_format: Literal["one-hot", "index", "mixed"] = "one-hot",
+) -> tuple[Tensor, Tensor]:
+    """Check and format inputs to the one-hot encodings."""
+    if input_format == "mixed":
+        _check_mixed_shape(preds, target)
+    else:
+        _check_same_shape(preds, target)
+
+    if input_format == "index":
+        if num_classes is None:
+            raise ValueError("Argument `num_classes` must be provided when `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 == "one-hot":
+        if num_classes is None:
+            num_classes = _get_num_classes(preds)
+        preds = _format_logits(preds, num_classes)
+        target = _format_logits(target, num_classes)
+    elif input_format == "mixed":
+        if preds.dim() == (target.dim() + 1):
+            if num_classes is None:
+                num_classes = _get_num_classes(preds)
+            preds = _format_logits(preds, num_classes)
+            target = torch.nn.functional.one_hot(target, num_classes=num_classes).movedim(-1, 1)
+        elif (preds.dim() + 1) == target.dim():
+            if num_classes is None:
+                num_classes = _get_num_classes(target)
+            target = _format_logits(target, num_classes)
+            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}.")
+
+    if not include_background:
+        preds, target = _ignore_background(preds, target)
+
+    return preds, target
+
+
+def _format_logits(tensor: Tensor, num_classes: int) -> Tensor:
+    """Transform logits or probabilities into integer one-hot encodings."""
+    if torch.is_floating_point(tensor):
+        tensor = tensor.argmax(dim=1)
+        tensor = torch.nn.functional.one_hot(tensor, num_classes=num_classes).movedim(-1, 1)
+    return tensor
+
+
+def _get_num_classes(tensor: Tensor) -> int:
+    """Get num classes from a tensor if it is not set."""
+    try:
+        num_classes = tensor.shape[1]
+    except IndexError as err:
+        raise IndexError(f"Cannot determine `num_classes` from tensor: {tensor}.") from err
+    if num_classes == 0:
+        raise ValueError(f"Expected argument `num_classes` to be a positive integer, but got {num_classes}.")
+    return num_classes
+
+
 def check_if_binarized(x: Tensor) -> None:
     """Check if tensor is binarized.
 

src/torchmetrics/segmentation/mean_iou.py

@@ -53,8 +53,8 @@ class MeanIoU(Metric):
           set to ``False``, the output will be a scalar tensor.
 
     Args:
-        num_classes: The number of classes in the segmentation problem. Required when input_format="index" or "mixed",
-            optional when input_format="one-hot".
+        num_classes: The number of classes in the segmentation problem. Required when input_format="index",
+            optional when input_format="one-hot" or "mixed".
         include_background: Whether to include the background class in the computation
         per_class: Whether to compute the IoU for each class separately. If set to ``False``, the metric will
             compute the mean IoU over all classes.
@@ -67,7 +67,7 @@ class MeanIoU(Metric):
         ValueError:
             If ``num_classes`` is not ``None`` or a positive integer
         ValueError:
-            If ``num_classes`` is not provided when ``input_format`` is ``"index"`` or ``"mixed"``
+            If ``num_classes`` is not provided when ``input_format`` is ``"index"``
         ValueError:
             If ``include_background`` is not a boolean
         ValueError:
@@ -132,7 +132,20 @@ def update(self, preds: Tensor, target: Tensor) -> None:
         """Update the state with the new data."""
         if not self._is_initialized:
             try:
-                self.num_classes = preds.shape[1]
+                if self.input_format == "one-hot":
+                    self.num_classes = preds.shape[1]
+                elif self.input_format == "mixed":
+                    if preds.dim() == (target.dim() + 1):
+                        self.num_classes = preds.shape[1]
+                    elif (preds.dim() + 1) == target.dim():
+                        self.num_classes = target.shape[1]
+                    else:
+                        raise ValueError(
+                            "Predictions and targets are expected to have the same shape,",
+                            f"got {preds.shape} and {target.shape}.",
+                        )
+                else:
+                    raise ValueError("Argument `num_classes` must be provided when `input_format` is 'index'.")
             except IndexError as err:
                 raise IndexError(f"Cannot determine `num_classes` from `preds` tensor: {preds}.") from err
 

tests/unittests/segmentation/inputs.py

@@ -48,3 +48,8 @@
     preds=torch.randint(0, NUM_CLASSES, (NUM_BATCHES, BATCH_SIZE, 32, 32)),
     target=to_one_hot(torch.randint(0, NUM_CLASSES, (NUM_BATCHES, BATCH_SIZE, 32, 32))),
 )
+
+_mixed_logits_input = _Input(
+    preds=(torch.rand((NUM_BATCHES, BATCH_SIZE, NUM_CLASSES, 32, 32)) * 12 - 6),
+    target=torch.randint(0, NUM_CLASSES, (NUM_BATCHES, BATCH_SIZE, 32, 32)),
+)

tests/unittests/segmentation/test_dice.py

@@ -28,6 +28,7 @@
     _index_input_2,
     _mixed_input_1,
     _mixed_input_2,
+    _mixed_logits_input,
     _one_hot_input_1,
     _one_hot_input_2,
 )
@@ -76,6 +77,7 @@ def _reference_dice_score(
         (_index_input_2.preds, _index_input_2.target, "index"),
         (_mixed_input_1.preds, _mixed_input_1.target, "mixed"),
         (_mixed_input_2.preds, _mixed_input_2.target, "mixed"),
+        (_mixed_logits_input.preds, _mixed_logits_input.target, "mixed"),
     ],
 )
 @pytest.mark.parametrize("include_background", [True, False])

tests/unittests/segmentation/test_generalized_dice_score.py

@@ -29,6 +29,7 @@
     _index_input_2,
     _mixed_input_1,
     _mixed_input_2,
+    _mixed_logits_input,
     _one_hot_input_1,
     _one_hot_input_2,
 )
@@ -49,8 +50,14 @@ def _reference_generalized_dice(
         target = torch.nn.functional.one_hot(target, num_classes=NUM_CLASSES).movedim(-1, 1)
     elif input_format == "mixed":
         if preds.dim() == (target.dim() + 1):
+            if torch.is_floating_point(preds):
+                preds = preds.argmax(dim=1)
+                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 (preds.dim() + 1) == target.dim():
+            if torch.is_floating_point(target):
+                target = target.argmax(dim=1)
+                target = torch.nn.functional.one_hot(target, num_classes=NUM_CLASSES).movedim(-1, 1)
             preds = torch.nn.functional.one_hot(preds, num_classes=NUM_CLASSES).movedim(-1, 1)
     monai_extra_arg = {"sum_over_classes": True} if RequirementCache("monai>=1.4.0") else {}
     val = compute_generalized_dice(preds, target, include_background=include_background, **monai_extra_arg)
@@ -68,6 +75,7 @@ def _reference_generalized_dice(
         (_index_input_2.preds, _index_input_2.target, "index"),
         (_mixed_input_1.preds, _mixed_input_1.target, "mixed"),
         (_mixed_input_2.preds, _mixed_input_2.target, "mixed"),
+        (_mixed_logits_input.preds, _mixed_logits_input.target, "mixed"),
     ],
 )
 @pytest.mark.parametrize("include_background", [True, False])