feat(metrics): add ClDiceMetric and compute_cldice (Fixes #4421)

monai/metrics/__init__.py

@@ -15,6 +15,7 @@
 from .active_learning_metrics import LabelQualityScore, VarianceMetric, compute_variance, label_quality_score
 from .average_precision import AveragePrecisionMetric, compute_average_precision
 from .calibration import CalibrationErrorMetric, CalibrationReduction, calibration_binning
+from .cldice import ClDiceMetric, compute_cldice
 from .confusion_matrix import ConfusionMatrixMetric, compute_confusion_matrix_metric, get_confusion_matrix
 from .cumulative_average import CumulativeAverage
 from .embedding_collapse import EmbeddingCollapseMetric, compute_embedding_collapse

monai/metrics/cldice.py

@@ -0,0 +1,226 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import torch
+
+from monai.losses.cldice import soft_erode, soft_open, soft_skel
+from monai.metrics.utils import do_metric_reduction, ignore_background
+from monai.utils import MetricReduction, look_up_option
+
+from .metric import CumulativeIterationMetric
+
+__all__ = ["ClDiceMetric", "compute_cldice"]
+
+
+def _compute_converged_skel(img: torch.Tensor, iter_: int) -> torch.Tensor:
+    """
+    Compute the morphological soft skeleton with at least ``iter_`` iterations, and if any thick
+    component still retains un-eroded interior mass after ``iter_`` erosions, continue skeletonization
+    until all components are reduced to their medial axes.
+
+    Args:
+        img: Binary or soft segmentation tensor of shape ``[B, C, H, W]`` or ``[B, C, H, W, D]``.
+        iter_: Minimum number of morphological skeletonization iterations.
+
+    Returns:
+        Skeletonized tensor of the same shape as ``img``.
+    """
+    skel = soft_skel(img, iter_)
+    eroded = img
+    for _ in range(iter_):
+        eroded = soft_erode(eroded)
+
+    if not torch.any(eroded > 0):
+        return skel
+
+    max_extra = max(0, (max(img.shape[2:]) + 1) // 2 - iter_)
+    for _ in range(max_extra):
+        eroded = soft_erode(eroded)
+        opened = soft_open(eroded)
+        delta = torch.relu(eroded - opened)
+        skel = skel + torch.relu(delta - skel * delta)
+        if not torch.any(eroded > 0):
+            break
+
+    return skel
+
+
+class ClDiceMetric(CumulativeIterationMetric):
+    """
+    Compute the centerline Dice (clDSC), Topology Precision (Tprec), or Topology Sensitivity (Tsens)
+    score between predictions and reference segmentations for 2D and 3D inputs.
+
+    Based on:
+        Shit et al., "clDice - a Novel Topology-Preserving Loss Function for Tubular Structure Segmentation", CVPR 2021.
+
+    Args:
+        include_background: Whether to include the first channel (background) in computation. Defaults to True.
+        iter_: Number of iterations for morphological skeletonization. Defaults to 3.
+        mode: Metric to return: ``"cldice"``, ``"tprec"``, or ``"tsens"``. Defaults to ``"cldice"``.
+        reduction: Reduction mode across batches (``"mean"``, ``"sum"``, ``"none"``, etc.). Defaults to ``"mean"``.
+        get_not_nans: Whether to return the ``not_nans`` count alongside the reduced metric. Defaults to False.
+        smooth: Smoothing constant to avoid division by zero. Defaults to 1e-5.
+    """
+
+    def __init__(
+        self,
+        include_background: bool = True,
+        iter_: int = 3,
+        mode: str = "cldice",
+        reduction: MetricReduction | str = MetricReduction.MEAN,
+        get_not_nans: bool = False,
+        smooth: float = 1e-5,
+    ) -> None:
+        """Initialize ClDiceMetric with skeletonization iterations, mode, and reduction settings."""
+        super().__init__()
+        self.include_background = include_background
+        self.iter_ = iter_
+        self.mode = look_up_option(mode.lower(), {"cldice", "tprec", "tsens"})
+        self.reduction = reduction
+        self.get_not_nans = get_not_nans
+        self.smooth = smooth
+
+    def _compute_tensor(self, y_pred: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        """
+        Compute the metric tensor for a single batch of predictions and targets.
+
+        Args:
+            y_pred: Predicted segmentation tensor of shape ``[B, C, H, W]`` or ``[B, C, H, W, D]``.
+            y: Reference segmentation tensor with the same shape as ``y_pred``.
+
+        Returns:
+            A tensor of shape ``[B, C]`` (or ``[B, C - 1]`` if ``include_background=False``) containing the scores.
+
+        Raises:
+            ValueError: When ``y_pred`` is not a 4D (2D) or 5D (3D) tensor, or shapes do not match.
+        """
+        if y_pred.ndim not in (4, 5):
+            raise ValueError(f"y_pred should be a 4D (2D) or 5D (3D) tensor, got shape {y_pred.shape}.")
+        return compute_cldice(
+            y_pred=y_pred,
+            y=y,
+            include_background=self.include_background,
+            iter_=self.iter_,
+            mode=self.mode,
+            smooth=self.smooth,
+        )
+
+    def aggregate(
+        self, reduction: MetricReduction | str | None = None
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        """
+        Aggregate the buffered metric results across iterations according to the specified reduction mode.
+
+        Args:
+            reduction: Optional override for the reduction mode (``"mean"``, ``"sum"``, ``"none"``, etc.).
+
+        Returns:
+            The aggregated metric tensor, or a tuple ``(aggregated_metric, not_nans)`` if ``get_not_nans=True``.
+
+        Raises:
+            ValueError: When the buffered data to aggregate is not a PyTorch Tensor.
+        """
+        data = self.get_buffer()
+        if not isinstance(data, torch.Tensor):
+            raise ValueError("The data to aggregate must be a PyTorch Tensor.")
+        f, not_nans = do_metric_reduction(data, reduction or self.reduction)
+        return (f, not_nans) if self.get_not_nans else f
+
+
+def compute_cldice(
+    y_pred: torch.Tensor,
+    y: torch.Tensor,
+    include_background: bool = True,
+    iter_: int = 3,
+    mode: str = "cldice",
+    smooth: float = 1e-5,
+) -> torch.Tensor:
+    """
+    Functional interface to compute ``cldice``, ``tprec``, or ``tsens`` per batch and per class.
+
+    If a reference mask channel is empty while the corresponding prediction is non-empty, ``NaN`` is returned
+    for that channel; if both reference and prediction masks are empty, ``1.0`` is returned. For thick
+    components where ``iter_`` is insufficient to reach the medial axis, skeletonization continues until
+    convergence before falling back to mask overlap if skeleton mass remains zero.
+
+    Args:
+        y_pred: Predicted segmentation tensor of shape ``[B, C, H, W]`` or ``[B, C, H, W, D]``.
+        y: Reference segmentation tensor with the same shape as ``y_pred``.
+        include_background: Whether to include the first channel (background) in computation. Defaults to True.
+        iter_: Number of iterations for morphological skeletonization. Defaults to 3.
+        mode: Metric to return: ``"cldice"``, ``"tprec"``, or ``"tsens"``. Defaults to ``"cldice"``.
+        smooth: Smoothing constant for numerical stability. Defaults to 1e-5.
+
+    Returns:
+        A tensor of shape ``[batch_size, num_classes]`` containing the computed metric values.
+
+    Raises:
+        ValueError: When ``mode`` is unsupported or ``y_pred`` and ``y`` have different shapes.
+    """
+    mode = look_up_option(mode.lower(), {"cldice", "tprec", "tsens"})
+
+    if y.shape != y_pred.shape:
+        raise ValueError(f"y_pred and y should have same shapes, got {y_pred.shape} and {y.shape}.")
+
+    if not include_background:
+        y_pred, y = ignore_background(y_pred=y_pred, y=y)
+
+    y = y.float()
+    y_pred = y_pred.float()
+
+    reduce_axis = list(range(2, y_pred.ndim))
+    skel_pred = _compute_converged_skel(y_pred, iter_)
+    skel_true = _compute_converged_skel(y, iter_)
+
+    y_sum = torch.sum(y, dim=reduce_axis)
+    y_pred_sum = torch.sum(y_pred, dim=reduce_axis)
+    skel_pred_sum = torch.sum(skel_pred, dim=reduce_axis)
+    skel_true_sum = torch.sum(skel_true, dim=reduce_axis)
+    intersection = torch.sum(y_pred * y, dim=reduce_axis)
+
+    fallback_tprec = torch.where(y_pred_sum > 0, intersection / y_pred_sum, torch.zeros_like(y_pred_sum))
+    fallback_tsens = torch.where(y_sum > 0, intersection / y_sum, torch.zeros_like(y_sum))
+
+    skel_pred_overlap = torch.sum(skel_pred * y, dim=reduce_axis)
+    skel_true_overlap = torch.sum(skel_true * y_pred, dim=reduce_axis)
+
+    tprec = torch.where(
+        skel_pred_sum > 0,
+        torch.where(
+            skel_pred_overlap > 0,
+            (skel_pred_overlap + smooth) / (skel_pred_sum + smooth),
+            torch.zeros_like(skel_pred_sum),
+        ),
+        fallback_tprec,
+    )
+    tsens = torch.where(
+        skel_true_sum > 0,
+        torch.where(
+            skel_true_overlap > 0,
+            (skel_true_overlap + smooth) / (skel_true_sum + smooth),
+            torch.zeros_like(skel_true_sum),
+        ),
+        fallback_tsens,
+    )
+
+    if mode == "tprec":
+        score = tprec
+    elif mode == "tsens":
+        score = tsens
+    else:
+        denom = tprec + tsens
+        score = torch.where(denom > 0, (2.0 * tprec * tsens) / denom, torch.zeros_like(denom))
+
+    nan_tensor = torch.tensor(float("nan"), device=y.device, dtype=y.dtype)
+    ones_tensor = torch.tensor(1.0, device=y.device, dtype=y.dtype)
+    return torch.where(y_sum > 0, score, torch.where(y_pred_sum > 0, nan_tensor, ones_tensor))

tests/metrics/test_cldice_metric.py

@@ -0,0 +1,109 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+
+from monai.metrics import ClDiceMetric, compute_cldice
+
+
+class TestClDiceMetric(unittest.TestCase):
+    """Unit tests for ClDiceMetric and compute_cldice across 2D and 3D inputs."""
+
+    def test_2d_and_3d_modes(self) -> None:
+        """Test partial-overlap 2D and 3D masks with distinct scores per mode and background option."""
+        y_2d = torch.zeros((2, 2, 16, 16))
+        y_pred_2d = torch.zeros((2, 2, 16, 16))
+        y_2d[:, 0, 2, 2:6] = 1.0
+        y_pred_2d[:, 0, 2, 2:5] = 1.0
+        y_pred_2d[:, 0, 12, 2:3] = 1.0
+
+        y_2d[:, 1, 5, 2:7] = 1.0
+        y_2d[:, 1, 10, 2:7] = 1.0
+        y_pred_2d[:, 1, 5, 2:10] = 1.0
+
+        y_3d = torch.zeros((2, 2, 12, 12, 12))
+        y_pred_3d = torch.zeros((2, 2, 12, 12, 12))
+        y_3d[:, 0, 2, 2, 2:6] = 1.0
+        y_pred_3d[:, 0, 2, 2, 2:5] = 1.0
+        y_pred_3d[:, 0, 8, 2, 2:3] = 1.0
+
+        y_3d[:, 1, 5, 2, 2:7] = 1.0
+        y_3d[:, 1, 9, 2, 2:7] = 1.0
+        y_pred_3d[:, 1, 5, 2, 2:10] = 1.0
+
+        cases = [
+            (y_pred_2d, y_2d, "tprec", False, 0.625),
+            (y_pred_2d, y_2d, "tsens", False, 0.500),
+            (y_pred_2d, y_2d, "cldice", False, 5.0 / 9.0),
+            (y_pred_2d, y_2d, "tprec", True, (0.75 + 0.625) / 2.0),
+            (y_pred_2d, y_2d, "tsens", True, (0.75 + 0.500) / 2.0),
+            (y_pred_2d, y_2d, "cldice", True, (0.75 + 5.0 / 9.0) / 2.0),
+            (y_pred_3d, y_3d, "cldice", False, 5.0 / 9.0),
+            (y_pred_3d, y_3d, "cldice", True, (0.75 + 5.0 / 9.0) / 2.0),
+        ]
+        for y_pred, y, mode, include_bg, expected in cases:
+            with self.subTest(ndim=y_pred.ndim, mode=mode, include_bg=include_bg):
+                metric = ClDiceMetric(include_background=include_bg, mode=mode, reduction="mean")
+                metric(y_pred, y)
+                result = metric.aggregate()
+                self.assertAlmostEqual(result.item(), expected, places=4)
+
+        with self.assertRaises(ValueError):
+            compute_cldice(torch.zeros((1, 1, 16, 16)), torch.zeros((1, 2, 16, 16)), include_background=False)
+
+    def test_solid_and_nested_thick_squares(self) -> None:
+        """Verify disjoint, identical, nested, and mixed thin-line + thick-square cases at iter_=3."""
+        y_pred_disjoint = torch.zeros((1, 1, 40, 40))
+        y_ref = torch.zeros((1, 1, 40, 40))
+        y_pred_disjoint[..., 5:15, 5:15] = 1.0
+        y_ref[..., 25:35, 25:35] = 1.0
+
+        y_pred_nested = torch.zeros((1, 1, 40, 40))
+        y_ref_nested = torch.zeros((1, 1, 40, 40))
+        y_pred_nested[..., 15:25, 15:25] = 1.0
+        y_ref_nested[..., 14:26, 14:26] = 1.0
+
+        for mode in ("cldice", "tprec", "tsens"):
+            disjoint_score = compute_cldice(y_pred_disjoint, y_ref, include_background=True, iter_=3, mode=mode)
+            identical_score = compute_cldice(y_ref, y_ref, include_background=True, iter_=3, mode=mode)
+            nested_score = compute_cldice(y_pred_nested, y_ref_nested, include_background=True, iter_=3, mode=mode)
+            self.assertEqual(disjoint_score.item(), 0.0)
+            self.assertEqual(identical_score.item(), 1.0)
+            self.assertEqual(nested_score.item(), 1.0)
+
+        # Mixed case: reference-matching thin line (mass 10) + disjoint 10x10 false-positive square (skel mass 4)
+        y_ref_line = torch.zeros((1, 1, 40, 40))
+        y_ref_line[..., 2, 2:12] = 1.0
+        y_pred_mixed = y_ref_line.clone()
+        y_pred_mixed[..., 20:30, 20:30] = 1.0
+
+        tprec_mixed = compute_cldice(y_pred_mixed, y_ref_line, include_background=True, iter_=3, mode="tprec")
+        cldice_mixed = compute_cldice(y_pred_mixed, y_ref_line, include_background=True, iter_=3, mode="cldice")
+        self.assertAlmostEqual(tprec_mixed.item(), 10.0 / 14.0, places=4)
+        self.assertAlmostEqual(cldice_mixed.item(), 10.0 / 12.0, places=4)
+
+    def test_empty_masks(self) -> None:
+        """Verify that two empty masks return 1.0 and empty reference with non-empty prediction returns NaN."""
+        y_pred = torch.zeros((1, 2, 16, 16))
+        y = torch.zeros((1, 2, 16, 16))
+        y_pred[0, 1, 4:8, 4:8] = 1.0
+
+        res = compute_cldice(y_pred, y, include_background=True)
+        self.assertEqual(res[0, 0].item(), 1.0)
+        self.assertTrue(torch.isnan(res[0, 1]))
+
+
+if __name__ == "__main__":
+    unittest.main()