diff --git a/monai/metrics/__init__.py b/monai/metrics/__init__.py index f55f92db1b..f575cac854 100644 --- a/monai/metrics/__init__.py +++ b/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 diff --git a/monai/metrics/cldice.py b/monai/metrics/cldice.py new file mode 100644 index 0000000000..7ef0e51574 --- /dev/null +++ b/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)) diff --git a/tests/metrics/test_cldice_metric.py b/tests/metrics/test_cldice_metric.py new file mode 100644 index 0000000000..9e9dd9f52a --- /dev/null +++ b/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()