diff --git a/docs/source/metrics.rst b/docs/source/metrics.rst index ebd745e795f..b070f7169c1 100644 --- a/docs/source/metrics.rst +++ b/docs/source/metrics.rst @@ -123,6 +123,13 @@ Metrics .. autoclass:: AbsoluteVolumeDifferenceMetric :members: +`Centerline Dice` +----------------- +.. autofunction:: compute_centerline_dice + +.. autoclass:: CenterlineDiceMetric + :members: + `PanopticQualityMetric` ----------------------- .. autofunction:: compute_panoptic_quality diff --git a/monai/metrics/__init__.py b/monai/metrics/__init__.py index f55f92db1b8..d2cb8d620ed 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 .centerline_dice import CenterlineDiceMetric, compute_centerline_dice 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/centerline_dice.py b/monai/metrics/centerline_dice.py new file mode 100644 index 00000000000..6ca02763d00 --- /dev/null +++ b/monai/metrics/centerline_dice.py @@ -0,0 +1,255 @@ +# 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 + +from typing import Any + +import numpy as np +import torch + +from monai.metrics.utils import do_metric_reduction, ignore_background, is_binary_tensor +from monai.utils import MetricReduction, optional_import + +from .metric import CumulativeIterationMetric + +__all__ = ["CenterlineDiceMetric", "compute_centerline_dice"] + +skeletonize, has_skimage = optional_import("skimage.morphology", name="skeletonize") + +_METRIC_NAMES = { + "cldsc": "cldsc", + "centerline_dice": "cldsc", + "topology_precision": "topology_precision", + "tprec": "topology_precision", + "topology_sensitivity": "topology_sensitivity", + "tsens": "topology_sensitivity", +} + + +class CenterlineDiceMetric(CumulativeIterationMetric): + """ + Compute the centreline Dice score (clDSC), topology precision, or topology sensitivity between + predicted and reference segmentation masks, based on skeletonization of the binary segmentations. + These metrics evaluate the connectivity of segmentations and are particularly useful for tubular + structures such as vessels. + + This implements the metrics defined in: + + Shit et al. "clDice -- a Novel Topology-Preserving Loss Function for Tubular Structure + Segmentation." CVPR 2021. https://arxiv.org/abs/2003.07311 + + With :math:`S_{Pred}` and :math:`S_{Ref}` the skeletons of the prediction and the reference + segmentation, the topology precision and topology sensitivity are defined as: + + .. math:: + T_{prec}(S_{Pred}, Ref) = \\frac{|S_{Pred} \\cap Ref|}{|S_{Pred}|}, \\qquad + T_{sens}(S_{Ref}, Pred) = \\frac{|S_{Ref} \\cap Pred|}{|S_{Ref}|} + + and the centreline Dice score is their harmonic mean: + + .. math:: + clDSC = \\frac{2 \\cdot T_{prec} \\cdot T_{sens}}{T_{prec} + T_{sens}} + + The class- and batch sample-wise values can be aggregated with the function `aggregate`. + + The skeletons are computed with ``skimage.morphology.skeletonize`` (2D inputs) and its + ``method="lee"`` variant (3D inputs), matching the reference implementation of the paper + (https://github.com/jocpae/clDice). ``scikit-image`` is an optional dependency imported with + ``optional_import`` and is required to use this metric. + + The inputs ``y_pred`` and ``y`` are expected to be binarized one-hot tensors with shape + BCHW[D]. If they contain continuous values (e.g. sigmoid outputs), binarize them first with a + suitable threshold transform. + + Example: + .. code-block:: python + + import torch + from monai.metrics import CenterlineDiceMetric + + y_pred = torch.randint(0, 2, (4, 3, 64, 64)).float() + y = torch.randint(0, 2, (4, 3, 64, 64)).float() + + metric = CenterlineDiceMetric(metric_name="cldsc", include_background=False) + metric(y_pred, y) # accumulate + result = metric.aggregate() # shape: (n_classes - 1,) after mean reduction + metric.reset() + + Args: + metric_name: which metric to compute. One of ``"cldsc"`` (alias ``"centerline_dice"``), + ``"topology_precision"`` (alias ``"tprec"``), ``"topology_sensitivity"`` (alias + ``"tsens"``). Defaults to ``"cldsc"``. + include_background: whether to include the first channel (index 0), which is by convention + assumed to be background, in the computation. Defaults to ``True``. + reduction: define mode of reduction to the metrics, will only apply reduction on `not-nan` + values, available reduction modes: {``"none"``, ``"mean"``, ``"sum"``, ``"mean_batch"``, + ``"sum_batch"``, ``"mean_channel"``, ``"sum_channel"``}, default to ``"mean"``. + If "none", will not do reduction. + get_not_nans: whether to return the `not_nans` count. Defaults to ``False``. + `not_nans` is the number of batch samples for which not all class-specific values were + nan values. If set to ``True``, the function `aggregate` will return both the aggregated + value and the `not_nans` count. + + Raises: + ValueError: if ``metric_name`` is not one of the supported metrics. + """ + + def __init__( + self, + metric_name: str = "cldsc", + include_background: bool = True, + reduction: MetricReduction | str = MetricReduction.MEAN, + get_not_nans: bool = False, + ) -> None: + super().__init__() + if metric_name not in _METRIC_NAMES: + raise ValueError(f"metric_name must be one of {sorted(_METRIC_NAMES)}, got '{metric_name}'.") + self.metric_name = _METRIC_NAMES[metric_name] + self.include_background = include_background + self.reduction = reduction + self.get_not_nans = get_not_nans + + def _compute_tensor(self, y_pred: torch.Tensor, y: torch.Tensor, **kwargs: Any) -> torch.Tensor: # type: ignore[override] + """ + Args: + y_pred: binarized prediction tensor, shape BCHW[D]. + y: binarized reference tensor, shape BCHW[D]. + + Raises: + ValueError: when ``y_pred`` or ``y`` are not one-hot encoded BCHW[D] tensors. + """ + return compute_centerline_dice( + y_pred=y_pred, y=y, metric_name=self.metric_name, include_background=self.include_background + ) + + def aggregate( + self, reduction: MetricReduction | str | None = None + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """ + Aggregates the output of `_compute_tensor`. + + Args: + reduction: define mode of reduction to the metrics, will only apply reduction on + `not-nan` values, available reduction modes: {``"none"``, ``"mean"``, ``"sum"``, + ``"mean_batch"``, ``"sum_batch"``, ``"mean_channel"``, ``"sum_channel"``}, + default to `self.reduction`. if "none", will not do reduction. + + Returns: + If `get_not_nans` is set to ``True``, this function returns the aggregated value and the + `not_nans` count, otherwise only the aggregated value. + """ + 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 _skeletonize(mask: np.ndarray) -> np.ndarray: + """ + Compute the binary skeleton of a non-empty 2D or 3D mask (Zhang-Suen for 2D, Lee et al. for 3D). + + ``skimage``'s Lee method can return an empty skeleton for a non-empty 3D mask (e.g. small solid + blobs); in that case the mask itself is returned as its own skeleton, so a non-empty mask always + keeps at least one skeleton voxel. + + Args: + mask: binary 2D or 3D mask. + + Returns: + np.ndarray: the skeleton of ``mask``, or ``mask`` itself when it is empty or its + skeletonization is empty. + """ + if not mask.any(): + return mask + skel = skeletonize(mask) if mask.ndim == 2 else skeletonize(mask, method="lee") + return skel if skel.any() else mask + + +def compute_centerline_dice( + y_pred: torch.Tensor, y: torch.Tensor, metric_name: str = "cldsc", include_background: bool = True +) -> torch.Tensor: + """ + Compute the centreline Dice score (clDSC), topology precision, or topology sensitivity for a + batch of binarized segmentations. See :py:class:`CenterlineDiceMetric` for the definitions. + + In the case of a class being present neither in the predicted segmentation nor in the reference + segmentation, a nan value is returned for this class. In the case of a class being present in + only one of predicted segmentation or reference segmentation, 0 is returned for this class + (a missing prediction cannot cover the reference skeleton and vice versa). When the skeleton of + a non-empty mask is empty (``skimage``'s Lee method can produce this for small solid 3D blobs), + the mask itself is used as its own skeleton, so that identical masks always score 1. + + Args: + y_pred: binarized prediction tensor, one-hot encoded, shape BCHW[D]. + y: binarized reference tensor, one-hot encoded, shape BCHW[D]. + metric_name: which metric to compute. One of ``"cldsc"`` (alias ``"centerline_dice"``), + ``"topology_precision"`` (alias ``"tprec"``), ``"topology_sensitivity"`` (alias + ``"tsens"``). Defaults to ``"cldsc"``. + include_background: whether to include the first channel (background) in the computation. + Defaults to ``True``. + + Returns: + PyTorch tensor of shape [B, C] with the metric value for each batch sample and class. + + Raises: + ValueError: if ``y_pred`` and/or ``y`` are not PyTorch tensors, do not have four or five + dimensions, have different shapes, or if ``metric_name`` is unknown. + """ + if metric_name not in _METRIC_NAMES: + raise ValueError(f"metric_name must be one of {sorted(_METRIC_NAMES)}, got '{metric_name}'.") + metric = _METRIC_NAMES[metric_name] + + if not isinstance(y_pred, torch.Tensor) or not isinstance(y, torch.Tensor): + raise ValueError("y_pred and y must be PyTorch Tensor.") + + if y_pred.ndimension() not in (4, 5) or y.ndimension() not in (4, 5): + raise ValueError("y_pred and y should be one-hot encoded: [B,C,H,W] or [B,C,H,W,D].") + + if y_pred.shape != y.shape: + raise ValueError( + f"y_pred and y should have same shape, but instead, shapes are {y_pred.shape} (y_pred) and {y.shape} (y)." + ) + + is_binary_tensor(y_pred, "y_pred") + is_binary_tensor(y, "y") + + if not include_background: + y_pred, y = ignore_background(y_pred=y_pred, y=y) + + batch_size, n_class = y_pred.shape[:2] + result = torch.empty((batch_size, n_class), device=y_pred.device, dtype=torch.float) + + y_pred_np = y_pred.detach().cpu().numpy() > 0 + y_np = y.detach().cpu().numpy() > 0 + + for b, c in np.ndindex(batch_size, n_class): + pred, ref = y_pred_np[b, c], y_np[b, c] + if not ref.any() and not pred.any(): + # the class is neither present in the prediction nor in the reference segmentation + result[b, c] = torch.tensor(np.nan) + continue + skel_pred = _skeletonize(pred) + skel_ref = _skeletonize(ref) + n_skel_pred, n_skel_ref = int(skel_pred.sum()), int(skel_ref.sum()) + # an empty prediction cannot cover any reference skeleton (and vice versa): score 0 + t_prec = float(np.logical_and(skel_pred, ref).sum()) / n_skel_pred if n_skel_pred > 0 else 0.0 + t_sens = float(np.logical_and(skel_ref, pred).sum()) / n_skel_ref if n_skel_ref > 0 else 0.0 + if metric == "topology_precision": + result[b, c] = t_prec + elif metric == "topology_sensitivity": + result[b, c] = t_sens + else: + result[b, c] = 0.0 if (t_prec + t_sens) == 0 else 2.0 * t_prec * t_sens / (t_prec + t_sens) + + return result diff --git a/tests/metrics/test_centerline_dice.py b/tests/metrics/test_centerline_dice.py new file mode 100644 index 00000000000..aae81350358 --- /dev/null +++ b/tests/metrics/test_centerline_dice.py @@ -0,0 +1,168 @@ +# 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 numpy as np +import torch +from parameterized import parameterized + +from monai.metrics import CenterlineDiceMetric, compute_centerline_dice +from tests.test_utils import SkipIfNoModule + + +def _line_2d(row: int, cols: tuple[int, int], size: int = 7) -> np.ndarray: + """One-pixel-wide horizontal line; its skeleton is the line itself.""" + img = np.zeros((size, size), dtype=np.float32) + img[row, cols[0] : cols[1]] = 1.0 + return img + + +def _to_bc(img: np.ndarray) -> torch.Tensor: + return torch.tensor(img[None, None]) + + +# hand-computed cases with one-pixel-wide lines (skeleton == line) +REF_LINE = _line_2d(3, (1, 6)) # 5 skeleton pixels +IDENTICAL = [_to_bc(REF_LINE), _to_bc(REF_LINE), {"metric_name": "cldsc"}, [[1.0]]] +IDENTICAL_TPRECS = [_to_bc(REF_LINE), _to_bc(REF_LINE), {"metric_name": "tprec"}, [[1.0]]] +IDENTICAL_TSENS = [_to_bc(REF_LINE), _to_bc(REF_LINE), {"metric_name": "tsens"}, [[1.0]]] +# prediction is a 3-pixel subset of the 5-pixel reference line: +# T_prec = 3/3 = 1, T_sens = 3/5 = 0.6, clDSC = 2*1*0.6/1.6 = 0.75 +SUBSET = [_to_bc(_line_2d(3, (2, 5))), _to_bc(REF_LINE), {"metric_name": "cldsc"}, [[0.75]]] +SUBSET_TPRECS = [_to_bc(_line_2d(3, (2, 5))), _to_bc(REF_LINE), {"metric_name": "topology_precision"}, [[1.0]]] +SUBSET_TSENS = [_to_bc(_line_2d(3, (2, 5))), _to_bc(REF_LINE), {"metric_name": "topology_sensitivity"}, [[0.6]]] +# prediction is a superset of the reference: T_prec = 3/5 = 0.6, T_sens = 1 +SUPERSET = [_to_bc(REF_LINE), _to_bc(_line_2d(3, (2, 5))), {"metric_name": "cldsc"}, [[0.75]]] +# disjoint lines: no overlap at all +DISJOINT = [_to_bc(_line_2d(1, (1, 6))), _to_bc(REF_LINE), {"metric_name": "cldsc"}, [[0.0]]] +# class absent from both: nan; present in only one: 0 +EMPTY_EMPTY = [_to_bc(np.zeros((7, 7), dtype=np.float32)), _to_bc(np.zeros((7, 7), dtype=np.float32)), {}, [[np.nan]]] +EMPTY_PRED = [_to_bc(np.zeros((7, 7), dtype=np.float32)), _to_bc(REF_LINE), {}, [[0.0]]] +EMPTY_REF = [_to_bc(REF_LINE), _to_bc(np.zeros((7, 7), dtype=np.float32)), {}, [[0.0]]] +# alias +ALIAS = [_to_bc(REF_LINE), _to_bc(REF_LINE), {"metric_name": "centerline_dice"}, [[1.0]]] + + +@SkipIfNoModule("skimage.morphology") +class TestComputeCenterlineDice(unittest.TestCase): + @parameterized.expand( + [ + IDENTICAL, + IDENTICAL_TPRECS, + IDENTICAL_TSENS, + SUBSET, + SUBSET_TPRECS, + SUBSET_TSENS, + SUPERSET, + DISJOINT, + EMPTY_EMPTY, + EMPTY_PRED, + EMPTY_REF, + ALIAS, + ] + ) + def test_value_2d(self, y_pred, y, params, expected): + result = compute_centerline_dice(y_pred=y_pred, y=y, **params) + np.testing.assert_allclose(result.cpu().numpy(), expected, rtol=1e-5) + + def test_value_3d(self): + # one-voxel-wide line along the last axis; Lee skeletonization keeps it intact + line = np.zeros((7, 7, 7), dtype=np.float32) + line[3, 3, 1:6] = 1.0 + ref = _to_bc(line) + pred = _to_bc(line.copy()) + result = compute_centerline_dice(y_pred=pred, y=ref) + np.testing.assert_allclose(result.cpu().numpy(), [[1.0]], rtol=1e-5) + # 3-voxel subset: T_prec = 1, T_sens = 3/5 + sub = np.zeros((7, 7, 7), dtype=np.float32) + sub[3, 3, 2:5] = 1.0 + result = compute_centerline_dice(y_pred=_to_bc(sub), y=ref, metric_name="cldsc") + np.testing.assert_allclose(result.cpu().numpy(), [[0.75]], rtol=1e-5) + + def test_solid_blob_fallback(self): + # skimage's Lee method returns an empty skeleton for small solid 3D blobs such as a 2x2x2 + # cube; the mask fallback must still give identical masks a score of 1 (all metrics) and + # disjoint solid blobs a clDSC of 0 + cube = np.zeros((7, 7, 7), dtype=np.float32) + cube[2:4, 2:4, 2:4] = 1.0 + for params in ({"metric_name": "cldsc"}, {"metric_name": "tprec"}, {"metric_name": "tsens"}): + result = compute_centerline_dice(y_pred=_to_bc(cube), y=_to_bc(cube), **params) + np.testing.assert_allclose(result.cpu().numpy(), [[1.0]], rtol=1e-5) + shifted = np.zeros((7, 7, 7), dtype=np.float32) + shifted[4:6, 4:6, 4:6] = 1.0 + result = compute_centerline_dice(y_pred=_to_bc(cube), y=_to_bc(shifted)) + np.testing.assert_allclose(result.cpu().numpy(), [[0.0]], rtol=1e-5) + + def test_thick_structure(self): + # identical thick masks score 1 regardless of the skeleton shape + thick = np.zeros((9, 9), dtype=np.float32) + thick[3:6, 2:7] = 1.0 + result = compute_centerline_dice(y_pred=_to_bc(thick), y=_to_bc(thick), metric_name="cldsc") + np.testing.assert_allclose(result.cpu().numpy(), [[1.0]], rtol=1e-5) + # 3x5 bar vs an internal 5-pixel line. The Zhang-Suen skeleton of the bar is the 4-pixel + # shape {(3,5), (4,2), (4,3), (4,4)}, of which 3 pixels lie on the prediction line: + # T_prec = 5/5 = 1, T_sens = 3/4 = 0.75, clDSC = 2*1*0.75/1.75 = 0.8571... + thin = np.zeros((9, 9), dtype=np.float32) + thin[4, 2:7] = 1.0 + result = compute_centerline_dice(y_pred=_to_bc(thin), y=_to_bc(thick), metric_name="cldsc") + np.testing.assert_allclose(result.cpu().numpy(), [[6.0 / 7.0]], rtol=1e-5) + + def test_multi_class_include_background(self): + # channel 0 empty in both (nan), channel 1 identical lines (1.0) + y_pred = torch.zeros(1, 2, 7, 7) + y = torch.zeros(1, 2, 7, 7) + y_pred[0, 1] = torch.tensor(REF_LINE) + y[0, 1] = torch.tensor(REF_LINE) + result = compute_centerline_dice(y_pred=y_pred, y=y, include_background=True) + self.assertTrue(np.isnan(result[0, 0].item())) + np.testing.assert_allclose(result[0, 1].item(), 1.0, rtol=1e-5) + result = compute_centerline_dice(y_pred=y_pred, y=y, include_background=False) + self.assertEqual(result.shape, (1, 1)) + np.testing.assert_allclose(result[0, 0].item(), 1.0, rtol=1e-5) + + def test_invalid_inputs(self): + with self.assertRaises(ValueError): + compute_centerline_dice(y_pred=_to_bc(REF_LINE), y=_to_bc(REF_LINE), metric_name="unknown") + with self.assertRaises(ValueError): + compute_centerline_dice(y_pred=np.zeros((1, 1, 7, 7)), y=_to_bc(REF_LINE)) # type: ignore[arg-type] + with self.assertRaises(ValueError): + compute_centerline_dice(y_pred=torch.zeros(1, 1, 7), y=torch.zeros(1, 1, 7)) + with self.assertRaises(ValueError): + compute_centerline_dice(y_pred=_to_bc(REF_LINE), y=_to_bc(np.zeros((9, 9), dtype=np.float32))) + with self.assertRaises(ValueError): + CenterlineDiceMetric(metric_name="unknown") + + def test_metric_class(self): + metric = CenterlineDiceMetric(metric_name="cldsc", include_background=True) + metric(y_pred=_to_bc(REF_LINE), y=_to_bc(REF_LINE)) + metric(y_pred=_to_bc(_line_2d(3, (2, 5))), y=_to_bc(REF_LINE)) + result = metric.aggregate(reduction="mean") + np.testing.assert_allclose(result.item(), (1.0 + 0.75) / 2, rtol=1e-5) + # reduction "none" returns the accumulated [B, C] buffer + metric.reset() + metric(y_pred=_to_bc(REF_LINE), y=_to_bc(REF_LINE)) + result = metric.aggregate(reduction="none") + np.testing.assert_allclose(result.cpu().numpy(), [[1.0]], rtol=1e-5) + + def test_metric_class_not_nans(self): + metric = CenterlineDiceMetric(include_background=True, get_not_nans=True) + metric(y_pred=_to_bc(np.zeros((7, 7), dtype=np.float32)), y=_to_bc(np.zeros((7, 7), dtype=np.float32))) + metric(y_pred=_to_bc(REF_LINE), y=_to_bc(REF_LINE)) + result, not_nans = metric.aggregate(reduction="mean") + np.testing.assert_allclose(result.item(), 1.0, rtol=1e-5) + self.assertEqual(not_nans.item(), 1) + + +if __name__ == "__main__": + unittest.main()