Implement distance_transform_edt and the DistanceTransformEDT transform

docs/source/transforms.rst

@@ -602,6 +602,12 @@ Post-processing
     :members:
     :special-members: __call__
 
+`DistanceTransformEDT`
+"""""""""""""""""""""""""""""""
+.. autoclass:: DistanceTransformEDT
+    :members:
+    :special-members: __call__
+
 `RemoveSmallObjects`
 """"""""""""""""""""
 .. image:: https://raw.githubusercontent.com/Project-MONAI/DocImages/main/transforms/RemoveSmallObjects.png
@@ -1640,6 +1646,12 @@ Post-processing (Dict)
     :members:
     :special-members: __call__
 
+`DistanceTransformEDTd`
+""""""""""""""""""""""""""""""""
+.. autoclass:: DistanceTransformEDTd
+    :members:
+    :special-members: __call__
+
 `RemoveSmallObjectsd`
 """""""""""""""""""""
 .. image:: https://raw.githubusercontent.com/Project-MONAI/DocImages/main/transforms/RemoveSmallObjectsd.png

monai/transforms/__init__.py

@@ -277,6 +277,7 @@
 from .post.array import (
     Activations,
     AsDiscrete,
+    DistanceTransformEDT,
     FillHoles,
     Invert,
     KeepLargestConnectedComponent,
@@ -295,6 +296,9 @@
     AsDiscreteD,
     AsDiscreted,
     AsDiscreteDict,
+    DistanceTransformEDTd,
+    DistanceTransformEDTD,
+    DistanceTransformEDTDict,
     Ensembled,
     EnsembleD,
     EnsembleDict,

monai/transforms/post/array.py

@@ -31,6 +31,7 @@
 from monai.transforms.utility.array import ToTensor
 from monai.transforms.utils import (
     convert_applied_interp_mode,
+    distance_transform_edt,
     fill_holes,
     get_largest_connected_component_mask,
     get_unique_labels,
@@ -53,6 +54,7 @@
     "SobelGradients",
     "VoteEnsemble",
     "Invert",
+    "DistanceTransformEDT",
 ]
 
 
@@ -936,3 +938,30 @@ def __call__(self, image: NdarrayOrTensor) -> torch.Tensor:
         grads = convert_to_dst_type(grads.squeeze(0), image_tensor)[0]
 
         return grads
+
+
+class DistanceTransformEDT(Transform):
+    """
+    Applies the Euclidean distance transform on the input.
+
+    Either GPU based with CuPy / cuCIM or CPU based with scipy.ndimage.
+    Choice only depends on cuCIM being available.
+    Note that the calculations can deviate, for details look into the cuCIM about distance_transform_edt().
+    """
+
+    backend = [TransformBackends.NUMPY, TransformBackends.CUPY]
+
+    def __init__(self, sampling: None | float | list[float] = None, force_scipy: bool = False) -> None:
+        super().__init__()
+        self.force_scipy = force_scipy
+        self.sampling = sampling
+
+    def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
+        """
+        Args:
+            img: shape must 2D or 3D for cupy, otherwise no restrictions
+
+        Returns:
+            An array with the same shape and data type as img
+        """
+        return distance_transform_edt(img=img, sampling=self.sampling, force_scipy=self.force_scipy)

monai/transforms/post/dictionary.py

@@ -33,6 +33,7 @@
 from monai.transforms.post.array import (
     Activations,
     AsDiscrete,
+    DistanceTransformEDT,
     FillHoles,
     KeepLargestConnectedComponent,
     LabelFilter,
@@ -91,6 +92,9 @@
     "VoteEnsembleD",
     "VoteEnsembleDict",
     "VoteEnsembled",
+    "DistanceTransformEDTd",
+    "DistanceTransformEDTD",
+    "DistanceTransformEDTDict",
 ]
 
 DEFAULT_POST_FIX = PostFix.meta()
@@ -855,6 +859,45 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, N
         return d
 
 
+class DistanceTransformEDTd(MapTransform):
+    """
+    Applies the Euclidean distance transform on the input.
+
+    Either GPU based with CuPy / cuCIM or CPU based with scipy.ndimage.
+    Choice only depends on cuCIM being available.
+    Note that the calculations can deviate, for details look into the cuCIM about distance_transform_edt().
+
+    Args:
+        keys: keys of the corresponding items to model output.
+        allow_missing_keys: don't raise exception if key is missing.
+        sampling: Spacing of elements along each dimension. If a sequence, must be of length equal to the input rank;
+            if a single number, this is used for all axes. If not specified, a grid spacing of unity is implied.
+        force_scipy: Force the CPU based scipy implementation of the euclidean distance transform
+
+    """
+
+    backend = DistanceTransformEDT.backend
+
+    def __init__(
+        self,
+        keys: KeysCollection,
+        allow_missing_keys: bool = False,
+        sampling: None | float | list[float] = None,
+        force_scipy: bool = False,
+    ) -> None:
+        super().__init__(keys, allow_missing_keys)
+        self.force_scipy = force_scipy
+        self.sampling = sampling
+        self.distance_transform = DistanceTransformEDT(sampling=self.sampling, force_scipy=self.force_scipy)
+
+    def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Mapping[Hashable, NdarrayOrTensor]:
+        d = dict(data)
+        for key in self.key_iterator(d):
+            d[key] = self.distance_transform(img=d[key])
+
+        return d
+
+
 ActivationsD = ActivationsDict = Activationsd
 AsDiscreteD = AsDiscreteDict = AsDiscreted
 FillHolesD = FillHolesDict = FillHolesd
@@ -869,3 +912,4 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, N
 VoteEnsembleD = VoteEnsembleDict = VoteEnsembled
 EnsembleD = EnsembleDict = Ensembled
 SobelGradientsD = SobelGradientsDict = SobelGradientsd
+DistanceTransformEDTD = DistanceTransformEDTDict = DistanceTransformEDTd

monai/transforms/utils.py

@@ -66,11 +66,17 @@
     pytorch_after,
 )
 from monai.utils.enums import TransformBackends
-from monai.utils.type_conversion import convert_data_type, convert_to_cupy, convert_to_dst_type, convert_to_tensor
+from monai.utils.type_conversion import (
+    convert_data_type,
+    convert_to_cupy,
+    convert_to_dst_type,
+    convert_to_numpy,
+    convert_to_tensor,
+)
 
 measure, has_measure = optional_import("skimage.measure", "0.14.2", min_version)
 morphology, has_morphology = optional_import("skimage.morphology")
-ndimage, _ = optional_import("scipy.ndimage")
+ndimage, has_ndimage = optional_import("scipy.ndimage")
 cp, has_cp = optional_import("cupy")
 cp_ndarray, _ = optional_import("cupy", name="ndarray")
 exposure, has_skimage = optional_import("skimage.exposure")
@@ -124,6 +130,7 @@
     "reset_ops_id",
     "resolves_modes",
     "has_status_keys",
+    "distance_transform_edt",
 ]
 
 
@@ -2012,7 +2019,7 @@ def has_status_keys(data: torch.Tensor, status_key: Any, default_message: str =
 
     Status keys are defined in :class:`TraceStatusKeys<monai.utils.enums.TraceStatusKeys>`.
 
-    This function also accepts:
+    This fun    ction also accepts:
 
     * dictionaries of tensors
     * lists or tuples of tensors
@@ -2051,5 +2058,38 @@ def has_status_keys(data: torch.Tensor, status_key: Any, default_message: str =
     return True, None
 
 
+def distance_transform_edt(
+    img: NdarrayOrTensor, sampling: None | float | list[float] = None, force_scipy: bool = False
+) -> NdarrayOrTensor:
+    """
+    Euclidean distance transform, either GPU based with CuPy / cuCIM
+    or CPU based with scipy.ndimage.
+    Choice only depends on cuCIM being available.
+
+    Args:
+        img: Input image on which the distance transform shall be run
+        sampling: Spacing of elements along each dimension. If a sequence, must be of length equal to the input rank;
+            if a single number, this is used for all axes. If not specified, a grid spacing of unity is implied.
+        force_scipy: Force the CPU based scipy implementation of the euclidean distance transform
+    """
+    distance_transform_edt, has_cucim = optional_import(
+        "cucim.core.operations.morphology", name="distance_transform_edt"
+    )
+
+    if has_cp and has_cucim and not force_scipy:
+        img_ = convert_to_cupy(img)
+        # Only accepts 2D and 3D input as of 09-2023
+        # TODO Add check and switch to scipy then?
+        distance = distance_transform_edt(img_, sampling=sampling)
+    else:
+        if not has_ndimage:
+            raise RuntimeError("scipy.ndimage required if cupy is not available")
+        img_ = convert_to_numpy(img)
+        distance = ndimage.distance_transform_edt(img_, sampling=sampling)
+
+    out = convert_to_dst_type(distance, dst=img, dtype=distance.dtype)[0]
+    return out
+
+
 if __name__ == "__main__":
     print_transform_backends()

tests/min_tests.py

@@ -61,6 +61,7 @@ def run_testsuit():
         "test_deepgrow_transforms",
         "test_detect_envelope",
         "test_dints_network",
+        "test_distance_transform_edt",
         "test_efficientnet",
         "test_ensemble_evaluator",
         "test_ensure_channel_first",

tests/test_distance_transform_edt.py

@@ -0,0 +1,121 @@
+# 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
+from parameterized import parameterized
+
+from monai.transforms import DistanceTransformEDT, DistanceTransformEDTd
+from tests.utils import HAS_CUPY, assert_allclose, optional_import, skip_if_no_cuda
+
+momorphology, has_cucim = optional_import("cucim.core.operations.morphology")
+ndimage, has_ndimage = optional_import("scipy.ndimage")
+cp, _ = optional_import("cupy")
+
+TEST_CASES = [
+    [
+        np.array(
+            ([0, 1, 1, 1, 1], [0, 0, 1, 1, 1], [0, 1, 1, 1, 1], [0, 1, 1, 1, 0], [0, 1, 1, 0, 0]), dtype=np.float32
+        ),
+        np.array(
+            [
+                [0.0, 1.0, 1.4142, 2.2361, 3.0],
+                [0.0, 0.0, 1.0, 2.0, 2.0],
+                [0.0, 1.0, 1.4142, 1.4142, 1.0],
+                [0.0, 1.0, 1.4142, 1.0, 0.0],
+                [0.0, 1.0, 1.0, 0.0, 0.0],
+            ]
+        ),
+    ]
+]
+
+SAMPLING_TEST_CASES = [
+    [
+        2,
+        np.array(
+            ([0, 1, 1, 1, 1], [0, 0, 1, 1, 1], [0, 1, 1, 1, 1], [0, 1, 1, 1, 0], [0, 1, 1, 0, 0]), dtype=np.float32
+        ),
+        np.array(
+            [
+                [0.0, 2.0, 2.828427, 4.472136, 6.0],
+                [0.0, 0.0, 2.0, 4.0, 4.0],
+                [0.0, 2.0, 2.828427, 2.828427, 2.0],
+                [0.0, 2.0, 2.828427, 2.0, 0.0],
+                [0.0, 2.0, 2.0, 0.0, 0.0],
+            ]
+        ),
+    ]
+]
+
+RAISES_TEST_CASES = (
+    [
+        np.array(
+            [[[[0, 1, 1, 1, 1], [0, 0, 1, 1, 1], [0, 1, 1, 1, 1], [0, 1, 1, 1, 0], [0, 1, 1, 0, 0]]]], dtype=np.float32
+        )
+    ],
+)
+
+
+class TestDistanceTransformEDT(unittest.TestCase):
+    @parameterized.expand(TEST_CASES)
+    def test_scipy_transform(self, input, expected_output):
+        transform = DistanceTransformEDT(force_scipy=True)
+        output = transform(input)
+        assert_allclose(output, expected_output, atol=1e-4, rtol=1e-4, type_test=False)
+
+    @parameterized.expand(TEST_CASES)
+    def test_scipy_transformd(self, input, expected_output):
+        transform = DistanceTransformEDTd(keys=("to_transform",), force_scipy=True)
+        data = {"to_transform": input}
+        data_ = transform(data)
+        output = data_["to_transform"]
+        assert_allclose(output, expected_output, atol=1e-4, rtol=1e-4, type_test=False)
+
+    @parameterized.expand(TEST_CASES)
+    @skip_if_no_cuda
+    @unittest.skipUnless(HAS_CUPY, "CuPy is required.")
+    @unittest.skipUnless(momorphology, "cuCIM transforms are required.")
+    def test_cucim_transform(self, input, expected_output):
+        transform = DistanceTransformEDT()
+        output = transform(input)
+        assert_allclose(cp.asnumpy(output), cp.asnumpy(expected_output), atol=1e-4, rtol=1e-4, type_test=False)
+
+    @parameterized.expand(SAMPLING_TEST_CASES)
+    def test_sampling(self, sampling, input, expected_output):
+        transform = DistanceTransformEDT(force_scipy=True, sampling=sampling)
+        output = transform(input)
+        assert_allclose(output, expected_output, atol=1e-4, rtol=1e-4, type_test=False)
+
+    @parameterized.expand(SAMPLING_TEST_CASES)
+    @skip_if_no_cuda
+    @unittest.skipUnless(HAS_CUPY, "CuPy is required.")
+    @unittest.skipUnless(momorphology, "cuCIM transforms are required.")
+    def test_cucim_sampling(self, sampling, input, expected_output):
+        transform = DistanceTransformEDT(sampling=sampling)
+        output = transform(input)
+        assert_allclose(cp.asnumpy(output), cp.asnumpy(expected_output), atol=1e-4, rtol=1e-4, type_test=False)
+
+    # @skip_if_no_cuda
+    # @unittest.skipUnless(HAS_CUPY, "CuPy is required.")
+    # @unittest.skipUnless(momorphology, "cuCIM transforms are required.")
+    # @parameterized.expand(RAISES_TEST_CASES)
+    # def test_cucim_raises(self, raises):
+    #     """Currently only 2D and 3D images are supported by CuPy. This test checks for the according error message"""
+    #     transform = DistanceTransformEDT()
+    #     with self.assertRaises(NotImplementedError):
+    #         output = transform(raises)
+
+
+if __name__ == "__main__":
+    unittest.main()