From 48a28731f95a1d9e56068ad23c65d88d3c0e1d48 Mon Sep 17 00:00:00 2001 From: Matthias Hadlich Date: Wed, 13 Sep 2023 17:22:08 +0200 Subject: [PATCH 01/23] Implement distance_transform_edt and the DistanceTransformEDT transform Signed-off-by: Matthias Hadlich --- monai/transforms/__init__.py | 1 + monai/transforms/post/array.py | 34 ++++++++ monai/transforms/utils.py | 37 ++++++++- tests/test_distance_transform_edt.py | 113 +++++++++++++++++++++++++++ 4 files changed, 182 insertions(+), 3 deletions(-) create mode 100644 tests/test_distance_transform_edt.py diff --git a/monai/transforms/__init__.py b/monai/transforms/__init__.py index eb8c5af19e..23239c0b45 100644 --- a/monai/transforms/__init__.py +++ b/monai/transforms/__init__.py @@ -287,6 +287,7 @@ RemoveSmallObjects, SobelGradients, VoteEnsemble, + DistanceTransformEDT, ) from .post.dictionary import ( ActivationsD, diff --git a/monai/transforms/post/array.py b/monai/transforms/post/array.py index df2e807a4b..d03329a3d0 100644 --- a/monai/transforms/post/array.py +++ b/monai/transforms/post/array.py @@ -35,6 +35,7 @@ get_largest_connected_component_mask, get_unique_labels, remove_small_objects, + distance_transform_edt, ) from monai.transforms.utils_pytorch_numpy_unification import unravel_index from monai.utils import TransformBackends, convert_data_type, convert_to_tensor, ensure_tuple, look_up_option @@ -53,6 +54,7 @@ "SobelGradients", "VoteEnsemble", "Invert", + "DistanceTransformEDT" ] @@ -936,3 +938,35 @@ 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, + force_scipy=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) + diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py index b769b8da2f..a25114cc9e 100644 --- a/monai/transforms/utils.py +++ b/monai/transforms/utils.py @@ -66,11 +66,11 @@ 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_tensor, convert_to_numpy 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 +124,7 @@ "reset_ops_id", "resolves_modes", "has_status_keys", + "distance_transform_edt", ] @@ -2012,7 +2013,7 @@ def has_status_keys(data: torch.Tensor, status_key: Any, default_message: str = Status keys are defined in :class:`TraceStatusKeys`. - This function also accepts: + This fun ction also accepts: * dictionaries of tensors * lists or tuples of tensors @@ -2051,5 +2052,35 @@ def has_status_keys(data: torch.Tensor, status_key: Any, default_message: str = return True, None +def distance_transform_edt(img: NdarrayOrTensor, sampling: 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. + + + """ + 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() diff --git a/tests/test_distance_transform_edt.py b/tests/test_distance_transform_edt.py new file mode 100644 index 0000000000..8e302f650f --- /dev/null +++ b/tests/test_distance_transform_edt.py @@ -0,0 +1,113 @@ +# 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. + +import unittest + +import numpy as np + +from parameterized import parameterized + +from monai.transforms import DistanceTransformEDT +from tests.utils import SkipIfNoModule, assert_allclose, optional_import +from tests.utils import HAS_CUPY, 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. , 1. , 1.4142, 2.2361, 3. ], + [ 0. , 0. , 1. , 2. , 2. ], + [ 0. , 1. , 1.4142, 1.4142, 1. ], + [ 0. , 1. , 1.4142, 1. , 0. ], + [ 0. , 1. , 1. , 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. , 2. , 2.828427, 4.472136, 6. ], + [0. , 0. , 2. , 4. , 4. ], + [0. , 2. , 2.828427, 2.828427, 2. ], + [0. , 2. , 2.828427, 2. , 0. ], + [0. , 2. , 2. , 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) + + @skip_if_no_cuda + @unittest.skipUnless(HAS_CUPY, "CuPy is required.") + @unittest.skipUnless(momorphology, "cuCIM transforms are required.") + @parameterized.expand(TEST_CASES) + 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(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(SAMPLING_TEST_CASES) + 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() From cbac352240b68b48fa913ad22139884f358015b6 Mon Sep 17 00:00:00 2001 From: Matthias Hadlich Date: Wed, 13 Sep 2023 17:27:17 +0200 Subject: [PATCH 02/23] Fix code style Signed-off-by: Matthias Hadlich --- monai/transforms/__init__.py | 2 +- monai/transforms/post/array.py | 15 ++---- monai/transforms/utils.py | 26 ++++++--- tests/test_distance_transform_edt.py | 80 ++++++++++++++-------------- 4 files changed, 64 insertions(+), 59 deletions(-) diff --git a/monai/transforms/__init__.py b/monai/transforms/__init__.py index 23239c0b45..6baaa3b1e6 100644 --- a/monai/transforms/__init__.py +++ b/monai/transforms/__init__.py @@ -277,6 +277,7 @@ from .post.array import ( Activations, AsDiscrete, + DistanceTransformEDT, FillHoles, Invert, KeepLargestConnectedComponent, @@ -287,7 +288,6 @@ RemoveSmallObjects, SobelGradients, VoteEnsemble, - DistanceTransformEDT, ) from .post.dictionary import ( ActivationsD, diff --git a/monai/transforms/post/array.py b/monai/transforms/post/array.py index d03329a3d0..6cf1419d3d 100644 --- a/monai/transforms/post/array.py +++ b/monai/transforms/post/array.py @@ -31,11 +31,11 @@ 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, remove_small_objects, - distance_transform_edt, ) from monai.transforms.utils_pytorch_numpy_unification import unravel_index from monai.utils import TransformBackends, convert_data_type, convert_to_tensor, ensure_tuple, look_up_option @@ -54,7 +54,7 @@ "SobelGradients", "VoteEnsemble", "Invert", - "DistanceTransformEDT" + "DistanceTransformEDT", ] @@ -939,27 +939,23 @@ def __call__(self, image: NdarrayOrTensor) -> torch.Tensor: 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(). + 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, - force_scipy=False, - ) -> None: + def __init__(self, sampling=None, force_scipy=False) -> None: super().__init__() self.force_scipy = force_scipy self.sampling = sampling - def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: """ Args: @@ -969,4 +965,3 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: 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) - diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py index a25114cc9e..1b425cd3a5 100644 --- a/monai/transforms/utils.py +++ b/monai/transforms/utils.py @@ -66,7 +66,13 @@ 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, convert_to_numpy +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") @@ -2052,20 +2058,25 @@ def has_status_keys(data: torch.Tensor, status_key: Any, default_message: str = return True, None -def distance_transform_edt(img: NdarrayOrTensor, sampling: float | list[float]=None, force_scipy: bool=False) -> NdarrayOrTensor: +def distance_transform_edt( + img: NdarrayOrTensor, sampling: float | list[float] = None, force_scipy: bool = False +) -> NdarrayOrTensor: """ - Euclidean distance transform, either GPU based with CuPy / cuCIM + Euclidean distance transform, either GPU based with CuPy / cuCIM or CPU based with scipy.ndimage. - Choice only depends on cuCIM being available. + 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. + 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. """ - distance_transform_edt, has_cucim = optional_import("cucim.core.operations.morphology", name="distance_transform_edt") - + 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 @@ -2081,6 +2092,5 @@ def distance_transform_edt(img: NdarrayOrTensor, sampling: float | list[float]=N return out - if __name__ == "__main__": print_transform_backends() diff --git a/tests/test_distance_transform_edt.py b/tests/test_distance_transform_edt.py index 8e302f650f..9c969591cd 100644 --- a/tests/test_distance_transform_edt.py +++ b/tests/test_distance_transform_edt.py @@ -9,15 +9,15 @@ # 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 -from tests.utils import SkipIfNoModule, assert_allclose, optional_import -from tests.utils import HAS_CUPY, skip_if_no_cuda +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") @@ -25,55 +25,55 @@ 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. , 1. , 1.4142, 2.2361, 3. ], - [ 0. , 0. , 1. , 2. , 2. ], - [ 0. , 1. , 1.4142, 1.4142, 1. ], - [ 0. , 1. , 1.4142, 1. , 0. ], - [ 0. , 1. , 1. , 0. , 0. ]]) - ], + 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. , 2. , 2.828427, 4.472136, 6. ], - [0. , 0. , 2. , 4. , 4. ], - [0. , 2. , 2.828427, 2.828427, 2. ], - [0. , 2. , 2.828427, 2. , 0. ], - [0. , 2. , 2. , 0. , 0. ]]), - ], + [ + 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), -], +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): +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) - + @skip_if_no_cuda @unittest.skipUnless(HAS_CUPY, "CuPy is required.") @unittest.skipUnless(momorphology, "cuCIM transforms are required.") From 150d59b396020d9771bd279b1630d172e4537892 Mon Sep 17 00:00:00 2001 From: Matthias Hadlich Date: Wed, 13 Sep 2023 18:52:43 +0200 Subject: [PATCH 03/23] Add test_distance_transform_edt to min_tests.py Signed-off-by: Matthias Hadlich --- tests/min_tests.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/min_tests.py b/tests/min_tests.py index 4c4e374311..a497f0d356 100644 --- a/tests/min_tests.py +++ b/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", From 2349f89209a9e5b52e7121e87a31988455ebe4af Mon Sep 17 00:00:00 2001 From: Matthias Hadlich Date: Wed, 13 Sep 2023 19:12:34 +0200 Subject: [PATCH 04/23] Add DistanceTransformEDTd Signed-off-by: Matthias Hadlich --- monai/transforms/__init__.py | 3 +++ monai/transforms/post/dictionary.py | 40 ++++++++++++++++++++++++++++ monai/transforms/utils.py | 5 ++-- tests/test_distance_transform_edt.py | 10 ++++++- 4 files changed, 54 insertions(+), 4 deletions(-) diff --git a/monai/transforms/__init__.py b/monai/transforms/__init__.py index 6baaa3b1e6..0eef4774cd 100644 --- a/monai/transforms/__init__.py +++ b/monai/transforms/__init__.py @@ -296,6 +296,9 @@ AsDiscreteD, AsDiscreted, AsDiscreteDict, + DistanceTransformEDTd, + DistanceTransformEDTD, + DistanceTransformEDTDict, Ensembled, EnsembleD, EnsembleDict, diff --git a/monai/transforms/post/dictionary.py b/monai/transforms/post/dictionary.py index 3fbfe46118..06e6a3371d 100644 --- a/monai/transforms/post/dictionary.py +++ b/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,41 @@ 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, force_scipy=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]) -> 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 +908,4 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, N VoteEnsembleD = VoteEnsembleDict = VoteEnsembled EnsembleD = EnsembleDict = Ensembled SobelGradientsD = SobelGradientsDict = SobelGradientsd +DistanceTransformEDTD = DistanceTransformEDTDict = DistanceTransformEDTd diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py index 1b425cd3a5..7f9fab4dd6 100644 --- a/monai/transforms/utils.py +++ b/monai/transforms/utils.py @@ -2069,9 +2069,8 @@ def distance_transform_edt( 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. - - + 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" diff --git a/tests/test_distance_transform_edt.py b/tests/test_distance_transform_edt.py index 9c969591cd..e2917db12d 100644 --- a/tests/test_distance_transform_edt.py +++ b/tests/test_distance_transform_edt.py @@ -16,7 +16,7 @@ import numpy as np from parameterized import parameterized -from monai.transforms import DistanceTransformEDT +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") @@ -74,6 +74,14 @@ def test_scipy_transform(self, input, expected_output): 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) + @skip_if_no_cuda @unittest.skipUnless(HAS_CUPY, "CuPy is required.") @unittest.skipUnless(momorphology, "cuCIM transforms are required.") From e3b3846771a29c1663fcd967b0f833eae743aec6 Mon Sep 17 00:00:00 2001 From: Matthias Hadlich Date: Wed, 13 Sep 2023 19:44:47 +0200 Subject: [PATCH 05/23] Update docs Signed-off-by: Matthias Hadlich --- docs/source/transforms.rst | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/docs/source/transforms.rst b/docs/source/transforms.rst index 051eab9e0e..c0b2bcace9 100644 --- a/docs/source/transforms.rst +++ b/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 From 62cbdcef0a3fb474d193d3957fd8d074e5c73925 Mon Sep 17 00:00:00 2001 From: Matthias Hadlich Date: Thu, 14 Sep 2023 11:04:24 +0200 Subject: [PATCH 06/23] Fix test Signed-off-by: Matthias Hadlich --- tests/test_distance_transform_edt.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_distance_transform_edt.py b/tests/test_distance_transform_edt.py index e2917db12d..ea0bb29504 100644 --- a/tests/test_distance_transform_edt.py +++ b/tests/test_distance_transform_edt.py @@ -82,10 +82,10 @@ def test_scipy_transformd(self, input, expected_output): 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.") - @parameterized.expand(TEST_CASES) def test_cucim_transform(self, input, expected_output): transform = DistanceTransformEDT() output = transform(input) @@ -95,12 +95,12 @@ def test_cucim_transform(self, input, expected_output): def test_sampling(self, sampling, input, expected_output): transform = DistanceTransformEDT(force_scipy=True, sampling=sampling) output = transform(input) - assert_allclose(cp.asnumpy(output), cp.asnumpy(expected_output), atol=1e-4, rtol=1e-4, type_test=False) + 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.") - @parameterized.expand(SAMPLING_TEST_CASES) def test_cucim_sampling(self, sampling, input, expected_output): transform = DistanceTransformEDT(sampling=sampling) output = transform(input) From f32841d77f79eb3b7d2b3bcd5f742e92c79d71cb Mon Sep 17 00:00:00 2001 From: Matthias Hadlich Date: Thu, 14 Sep 2023 11:04:50 +0200 Subject: [PATCH 07/23] Add typing Signed-off-by: Matthias Hadlich --- monai/transforms/post/array.py | 2 +- monai/transforms/post/dictionary.py | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/monai/transforms/post/array.py b/monai/transforms/post/array.py index 6cf1419d3d..10bfcc30b6 100644 --- a/monai/transforms/post/array.py +++ b/monai/transforms/post/array.py @@ -951,7 +951,7 @@ class DistanceTransformEDT(Transform): backend = [TransformBackends.NUMPY, TransformBackends.CUPY] - def __init__(self, sampling=None, force_scipy=False) -> None: + def __init__(self, sampling: float | list[float] = None, force_scipy: bool = False) -> None: super().__init__() self.force_scipy = force_scipy self.sampling = sampling diff --git a/monai/transforms/post/dictionary.py b/monai/transforms/post/dictionary.py index 06e6a3371d..a8bc4138c8 100644 --- a/monai/transforms/post/dictionary.py +++ b/monai/transforms/post/dictionary.py @@ -879,7 +879,11 @@ class DistanceTransformEDTd(MapTransform): backend = DistanceTransformEDT.backend def __init__( - self, keys: KeysCollection, allow_missing_keys: bool = False, sampling=None, force_scipy=False + self, + keys: KeysCollection, + allow_missing_keys: bool = False, + sampling: float | list[float] = None, + force_scipy: bool = False, ) -> None: super().__init__(keys, allow_missing_keys) self.force_scipy = force_scipy From 9a88bae835715622d583ebdcf772cc0c73b21005 Mon Sep 17 00:00:00 2001 From: Matthias Hadlich Date: Thu, 14 Sep 2023 11:42:41 +0200 Subject: [PATCH 08/23] Fix typing for sampling argument Signed-off-by: Matthias Hadlich --- monai/transforms/post/array.py | 2 +- monai/transforms/post/dictionary.py | 2 +- monai/transforms/utils.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/monai/transforms/post/array.py b/monai/transforms/post/array.py index 10bfcc30b6..5d578369e3 100644 --- a/monai/transforms/post/array.py +++ b/monai/transforms/post/array.py @@ -951,7 +951,7 @@ class DistanceTransformEDT(Transform): backend = [TransformBackends.NUMPY, TransformBackends.CUPY] - def __init__(self, sampling: float | list[float] = None, force_scipy: bool = False) -> None: + def __init__(self, sampling: None | float | list[float] = None, force_scipy: bool = False) -> None: super().__init__() self.force_scipy = force_scipy self.sampling = sampling diff --git a/monai/transforms/post/dictionary.py b/monai/transforms/post/dictionary.py index a8bc4138c8..d2518c4e50 100644 --- a/monai/transforms/post/dictionary.py +++ b/monai/transforms/post/dictionary.py @@ -882,7 +882,7 @@ def __init__( self, keys: KeysCollection, allow_missing_keys: bool = False, - sampling: float | list[float] = None, + sampling: None | float | list[float] = None, force_scipy: bool = False, ) -> None: super().__init__(keys, allow_missing_keys) diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py index 7f9fab4dd6..ecd79fe44d 100644 --- a/monai/transforms/utils.py +++ b/monai/transforms/utils.py @@ -2059,7 +2059,7 @@ def has_status_keys(data: torch.Tensor, status_key: Any, default_message: str = def distance_transform_edt( - img: NdarrayOrTensor, sampling: float | list[float] = None, force_scipy: bool = False + img: NdarrayOrTensor, sampling: None | float | list[float] = None, force_scipy: bool = False ) -> NdarrayOrTensor: """ Euclidean distance transform, either GPU based with CuPy / cuCIM From c4ed6c54fc282615f037309353d408b3ac1f8981 Mon Sep 17 00:00:00 2001 From: Matthias Hadlich Date: Thu, 14 Sep 2023 15:16:40 +0200 Subject: [PATCH 09/23] Fix typing return value Signed-off-by: Matthias Hadlich --- monai/transforms/post/dictionary.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monai/transforms/post/dictionary.py b/monai/transforms/post/dictionary.py index d2518c4e50..efbbb1b5f7 100644 --- a/monai/transforms/post/dictionary.py +++ b/monai/transforms/post/dictionary.py @@ -890,7 +890,7 @@ def __init__( self.sampling = sampling self.distance_transform = DistanceTransformEDT(sampling=self.sampling, force_scipy=self.force_scipy) - def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> NdarrayOrTensor: + 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]) From 6eccc4a2a1f3b14ff8597d38182a40aa87fd1b8b Mon Sep 17 00:00:00 2001 From: Matthias Hadlich Date: Fri, 15 Sep 2023 13:24:23 +0200 Subject: [PATCH 10/23] Update docs to match the code Signed-off-by: Matthias Hadlich --- monai/transforms/post/array.py | 4 +++- monai/transforms/utils.py | 10 ++++++---- tests/test_distance_transform_edt.py | 23 ++++++++++++----------- 3 files changed, 21 insertions(+), 16 deletions(-) diff --git a/monai/transforms/post/array.py b/monai/transforms/post/array.py index 5d578369e3..03cb38baf0 100644 --- a/monai/transforms/post/array.py +++ b/monai/transforms/post/array.py @@ -959,7 +959,9 @@ def __init__(self, sampling: None | float | list[float] = None, force_scipy: boo def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: """ Args: - img: shape must 2D or 3D for cupy, otherwise no restrictions + img: Input image on which the distance transform shall be run. + For CuPy shape must be (spatial_dim1, spatial_dim2 [, spatial_dim3]), so 2D or 3D input only. + If you need to run the transform on other shapes, use the ``force_scipy`` flag. Returns: An array with the same shape and data type as img diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py index ecd79fe44d..17414402ab 100644 --- a/monai/transforms/utils.py +++ b/monai/transforms/utils.py @@ -2064,10 +2064,14 @@ def distance_transform_edt( """ Euclidean distance transform, either GPU based with CuPy / cuCIM or CPU based with scipy.ndimage. - Choice only depends on cuCIM being available. + Choice depends on cuCIM being available or scipy can be forced with the ``force_scipy`` flag. + + Note that the runtime running on the CPU may be really depending on the inputs size. Args: - img: Input image on which the distance transform shall be run + img: Input image on which the distance transform shall be run. + For CuPy shape must be (spatial_dim1, spatial_dim2 [, spatial_dim3]), so 2D or 3D input only. + If you need to run the transform on other shapes, use the ``force_scipy`` flag. 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 @@ -2078,8 +2082,6 @@ def 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: diff --git a/tests/test_distance_transform_edt.py b/tests/test_distance_transform_edt.py index ea0bb29504..4758f34b1f 100644 --- a/tests/test_distance_transform_edt.py +++ b/tests/test_distance_transform_edt.py @@ -58,8 +58,9 @@ ] ] + RAISES_TEST_CASES = ( - [ + [ # Example 4D input. Should raise under CuPy 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 ) @@ -92,7 +93,7 @@ def test_cucim_transform(self, input, expected_output): 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): + def test_scipy_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) @@ -106,15 +107,15 @@ def test_cucim_sampling(self, sampling, input, expected_output): 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) + @parameterized.expand(RAISES_TEST_CASES) + @skip_if_no_cuda + @unittest.skipUnless(HAS_CUPY, "CuPy is required.") + @unittest.skipUnless(momorphology, "cuCIM transforms are required.") + 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): + transform(raises) if __name__ == "__main__": From 4e26689f7a9e8d710cbe18edf922eb4f71c8ab63 Mon Sep 17 00:00:00 2001 From: Matthias Hadlich Date: Tue, 19 Sep 2023 15:35:17 +0200 Subject: [PATCH 11/23] CuPy now allows 4D channel-wise input Signed-off-by: Matthias Hadlich --- monai/transforms/post/array.py | 3 ++- monai/transforms/post/dictionary.py | 4 +++- monai/transforms/utils.py | 11 +++++++++-- tests/test_distance_transform_edt.py | 25 ++++++++++++++++++++++--- 4 files changed, 36 insertions(+), 7 deletions(-) diff --git a/monai/transforms/post/array.py b/monai/transforms/post/array.py index 03cb38baf0..42d38c4b24 100644 --- a/monai/transforms/post/array.py +++ b/monai/transforms/post/array.py @@ -960,8 +960,9 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: """ Args: img: Input image on which the distance transform shall be run. - For CuPy shape must be (spatial_dim1, spatial_dim2 [, spatial_dim3]), so 2D or 3D input only. + For CuPy shape must be ([channel, ]spatial_dim1, spatial_dim2 [, spatial_dim3]), so 2D, 3D and 4D input only. If you need to run the transform on other shapes, use the ``force_scipy`` flag. + 4D input gets passed channel-wise to cupy. Returns: An array with the same shape and data type as img diff --git a/monai/transforms/post/dictionary.py b/monai/transforms/post/dictionary.py index efbbb1b5f7..fabe636baf 100644 --- a/monai/transforms/post/dictionary.py +++ b/monai/transforms/post/dictionary.py @@ -865,7 +865,9 @@ class DistanceTransformEDTd(MapTransform): 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(). + Note that the results of the two libraries can deviate, for details look into the cuCIM Docs about ``distance_transform_edt()``. + + Input has to 2D, 3D or 4D for CuPy. For details look into :py:class:`monai.transforms.DistanceTransformEDT`. Args: keys: keys of the corresponding items to model output. diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py index 17414402ab..dfb22d9786 100644 --- a/monai/transforms/utils.py +++ b/monai/transforms/utils.py @@ -2070,8 +2070,9 @@ def distance_transform_edt( Args: img: Input image on which the distance transform shall be run. - For CuPy shape must be (spatial_dim1, spatial_dim2 [, spatial_dim3]), so 2D or 3D input only. + For CuPy shape must be ([channel, ]spatial_dim1, spatial_dim2 [, spatial_dim3]), so 2D, 3D and 4D input only. If you need to run the transform on other shapes, use the ``force_scipy`` flag. + 4D input gets passed channel-wise to cupy. 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 @@ -2082,7 +2083,13 @@ def distance_transform_edt( if has_cp and has_cucim and not force_scipy: img_ = convert_to_cupy(img) - distance = distance_transform_edt(img_, sampling=sampling) + if img_.ndim == 4: + out = [] + for channel in img_: + out.append(distance_transform_edt(channel, sampling=sampling)) + distance = cp.stack(out) + else: + distance = distance_transform_edt(img_, sampling=sampling) else: if not has_ndimage: raise RuntimeError("scipy.ndimage required if cupy is not available") diff --git a/tests/test_distance_transform_edt.py b/tests/test_distance_transform_edt.py index 4758f34b1f..f89473d84d 100644 --- a/tests/test_distance_transform_edt.py +++ b/tests/test_distance_transform_edt.py @@ -37,7 +37,25 @@ [0.0, 1.0, 1.0, 0.0, 0.0], ] ), - ] + ], + [ # Example 4D input to test channel-wise CuPy + 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 = [ @@ -62,7 +80,8 @@ RAISES_TEST_CASES = ( [ # Example 4D input. Should raise under CuPy 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 + [[[[[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, ) ], ) @@ -112,7 +131,7 @@ def test_cucim_sampling(self, sampling, input, expected_output): @unittest.skipUnless(HAS_CUPY, "CuPy is required.") @unittest.skipUnless(momorphology, "cuCIM transforms are required.") def test_cucim_raises(self, raises): - """Currently only 2D and 3D images are supported by CuPy. This test checks for the according error message""" + """Currently only 2D, 3D and 4D images are supported by CuPy. This test checks for the according error message""" transform = DistanceTransformEDT() with self.assertRaises(NotImplementedError): transform(raises) From 7943ffdeab2e052d705c3c312f2fcdb9abd076e7 Mon Sep 17 00:00:00 2001 From: Wenqi Li Date: Tue, 19 Sep 2023 16:42:43 +0100 Subject: [PATCH 12/23] fixes format Signed-off-by: Wenqi Li --- monai/transforms/utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py index dfb22d9786..8e93a88f76 100644 --- a/monai/transforms/utils.py +++ b/monai/transforms/utils.py @@ -2096,8 +2096,7 @@ def distance_transform_edt( 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 + return convert_to_dst_type(distance, dst=img, dtype=distance.dtype)[0] if __name__ == "__main__": From 101cc6227b57b6ffd4d736c3bc1be39494994813 Mon Sep 17 00:00:00 2001 From: Matthias Hadlich Date: Sat, 23 Sep 2023 15:11:18 +0200 Subject: [PATCH 13/23] Update monai/transforms/post/array.py Co-authored-by: YunLiu <55491388+KumoLiu@users.noreply.github.com> Signed-off-by: Matthias Hadlich --- monai/transforms/post/array.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monai/transforms/post/array.py b/monai/transforms/post/array.py index 42d38c4b24..7e687e6203 100644 --- a/monai/transforms/post/array.py +++ b/monai/transforms/post/array.py @@ -960,7 +960,7 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: """ Args: img: Input image on which the distance transform shall be run. - For CuPy shape must be ([channel, ]spatial_dim1, spatial_dim2 [, spatial_dim3]), so 2D, 3D and 4D input only. + channel first array, must have shape: (num_channels, H[, W, ..., ]) If you need to run the transform on other shapes, use the ``force_scipy`` flag. 4D input gets passed channel-wise to cupy. From 24814c4e5fb405f8514722cad5acf375cc84a6dd Mon Sep 17 00:00:00 2001 From: Matthias Hadlich Date: Sat, 23 Sep 2023 15:12:03 +0200 Subject: [PATCH 14/23] Apply suggestions from code review Co-authored-by: YunLiu <55491388+KumoLiu@users.noreply.github.com> Signed-off-by: Matthias Hadlich --- monai/transforms/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py index 8e93a88f76..f5e81d0d33 100644 --- a/monai/transforms/utils.py +++ b/monai/transforms/utils.py @@ -2019,7 +2019,7 @@ def has_status_keys(data: torch.Tensor, status_key: Any, default_message: str = Status keys are defined in :class:`TraceStatusKeys`. - This fun ction also accepts: + This function also accepts: * dictionaries of tensors * lists or tuples of tensors @@ -2070,7 +2070,7 @@ def distance_transform_edt( Args: img: Input image on which the distance transform shall be run. - For CuPy shape must be ([channel, ]spatial_dim1, spatial_dim2 [, spatial_dim3]), so 2D, 3D and 4D input only. + channel first array, must have shape: (num_channels, H[, W, ..., ]) If you need to run the transform on other shapes, use the ``force_scipy`` flag. 4D input gets passed channel-wise to cupy. sampling: Spacing of elements along each dimension. If a sequence, must be of length equal to the input rank; From e98b10ac429fd801a3eebe059c8d02daf49b6199 Mon Sep 17 00:00:00 2001 From: Matthias Hadlich Date: Wed, 20 Sep 2023 15:11:40 +0200 Subject: [PATCH 15/23] Add test for 4D input Signed-off-by: Matthias Hadlich --- tests/test_distance_transform_edt.py | 54 ++++++++++++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/tests/test_distance_transform_edt.py b/tests/test_distance_transform_edt.py index f89473d84d..78349b5a47 100644 --- a/tests/test_distance_transform_edt.py +++ b/tests/test_distance_transform_edt.py @@ -56,6 +56,60 @@ ] ), ], + [ + np.array( + [ + [ + [0.0, 1.0, 1.0, 1.0, 1.0], + [0.0, 0.0, 1.0, 1.0, 1.0], + [0.0, 1.0, 1.0, 1.0, 1.0], + [0.0, 1.0, 1.0, 1.0, 0.0], + [0.0, 1.0, 1.0, 0.0, 0.0], + ], + [ + [0.0, 1.0, 1.0, 1.0, 1.0], + [0.0, 0.0, 1.0, 1.0, 1.0], + [0.0, 1.0, 1.0, 1.0, 1.0], + [0.0, 1.0, 1.0, 1.0, 0.0], + [0.0, 1.0, 1.0, 0.0, 0.0], + ], + [ + [0.0, 1.0, 1.0, 1.0, 1.0], + [0.0, 0.0, 1.0, 1.0, 1.0], + [0.0, 1.0, 1.0, 1.0, 1.0], + [0.0, 1.0, 1.0, 1.0, 0.0], + [0.0, 1.0, 1.0, 0.0, 0.0], + ], + ], + dtype=np.float32, + ), + np.array( + [ + [ + [0.0, 1.0, 1.4142135, 2.236068, 3.0], + [0.0, 0.0, 1.0, 2.0, 2.0], + [0.0, 1.0, 1.4142135, 1.4142135, 1.0], + [0.0, 1.0, 1.4142135, 1.0, 0.0], + [0.0, 1.0, 1.0, 0.0, 0.0], + ], + [ + [0.0, 1.0, 1.4142135, 2.236068, 3.0], + [0.0, 0.0, 1.0, 2.0, 2.0], + [0.0, 1.0, 1.4142135, 1.4142135, 1.0], + [0.0, 1.0, 1.4142135, 1.0, 0.0], + [0.0, 1.0, 1.0, 0.0, 0.0], + ], + [ + [0.0, 1.0, 1.4142135, 2.236068, 3.0], + [0.0, 0.0, 1.0, 2.0, 2.0], + [0.0, 1.0, 1.4142135, 1.4142135, 1.0], + [0.0, 1.0, 1.4142135, 1.0, 0.0], + [0.0, 1.0, 1.0, 0.0, 0.0], + ], + ], + dtype=np.float32, + ), + ], ] SAMPLING_TEST_CASES = [ From 6d19ae4123df7f697930e1c13795ca250d9abaa8 Mon Sep 17 00:00:00 2001 From: Matthias Hadlich Date: Sun, 24 Sep 2023 10:14:25 +0200 Subject: [PATCH 16/23] Remove force_scipy flag Signed-off-by: Matthias Hadlich --- monai/transforms/post/array.py | 10 +++++----- monai/transforms/post/dictionary.py | 5 +---- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/monai/transforms/post/array.py b/monai/transforms/post/array.py index 7e687e6203..f7e3a0e5f3 100644 --- a/monai/transforms/post/array.py +++ b/monai/transforms/post/array.py @@ -951,9 +951,8 @@ class DistanceTransformEDT(Transform): backend = [TransformBackends.NUMPY, TransformBackends.CUPY] - def __init__(self, sampling: None | float | list[float] = None, force_scipy: bool = False) -> None: + def __init__(self, sampling: None | float | list[float] = None) -> None: super().__init__() - self.force_scipy = force_scipy self.sampling = sampling def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: @@ -961,10 +960,11 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: Args: img: Input image on which the distance transform shall be run. channel first array, must have shape: (num_channels, H[, W, ..., ]) - If you need to run the transform on other shapes, use the ``force_scipy`` flag. - 4D input gets passed channel-wise to cupy. + Input gets passed channel-wise to the distance-transform, thus results from this function may differ + from directly calling ``distance_transform_edt()`` in CuPy or scipy. + 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) + return distance_transform_edt(img=img, sampling=self.sampling) diff --git a/monai/transforms/post/dictionary.py b/monai/transforms/post/dictionary.py index fabe636baf..e6eda18d6d 100644 --- a/monai/transforms/post/dictionary.py +++ b/monai/transforms/post/dictionary.py @@ -874,7 +874,6 @@ class DistanceTransformEDTd(MapTransform): 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 """ @@ -885,12 +884,10 @@ def __init__( 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) + self.distance_transform = DistanceTransformEDT(sampling=self.sampling) def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Mapping[Hashable, NdarrayOrTensor]: d = dict(data) From 1f60c62b279123b68c906b76817a185f37fe18fa Mon Sep 17 00:00:00 2001 From: Matthias Hadlich Date: Sun, 24 Sep 2023 10:14:51 +0200 Subject: [PATCH 17/23] Remove force_scipy flag Signed-off-by: Matthias Hadlich --- tests/test_distance_transform_edt.py | 39 ++++++++++++++++------------ 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/tests/test_distance_transform_edt.py b/tests/test_distance_transform_edt.py index 78349b5a47..63b6eff536 100644 --- a/tests/test_distance_transform_edt.py +++ b/tests/test_distance_transform_edt.py @@ -16,6 +16,8 @@ import numpy as np from parameterized import parameterized +import torch + from monai.transforms import DistanceTransformEDT, DistanceTransformEDTd from tests.utils import HAS_CUPY, assert_allclose, optional_import, skip_if_no_cuda @@ -26,16 +28,16 @@ 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 + ([[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], - ] + ]] ), ], [ # Example 4D input to test channel-wise CuPy @@ -116,16 +118,16 @@ [ 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 + ([[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], - ] + ]] ), ] ] @@ -144,40 +146,44 @@ class TestDistanceTransformEDT(unittest.TestCase): @parameterized.expand(TEST_CASES) def test_scipy_transform(self, input, expected_output): - transform = DistanceTransformEDT(force_scipy=True) + transform = DistanceTransformEDT() 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) + transform = DistanceTransformEDTd(keys=("to_transform",)) 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(SAMPLING_TEST_CASES) + def test_scipy_sampling(self, sampling, input, expected_output): + transform = DistanceTransformEDT(sampling=sampling) + output = transform(input) + 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): + input_ = torch.tensor(input, device='cuda') transform = DistanceTransformEDT() - output = transform(input) + 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_scipy_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): + input_ = torch.tensor(input, device='cuda') transform = DistanceTransformEDT(sampling=sampling) - output = transform(input) + output = transform(input_) assert_allclose(cp.asnumpy(output), cp.asnumpy(expected_output), atol=1e-4, rtol=1e-4, type_test=False) @parameterized.expand(RAISES_TEST_CASES) @@ -186,9 +192,10 @@ def test_cucim_sampling(self, sampling, input, expected_output): @unittest.skipUnless(momorphology, "cuCIM transforms are required.") def test_cucim_raises(self, raises): """Currently only 2D, 3D and 4D images are supported by CuPy. This test checks for the according error message""" + input_ = torch.tensor(raises, device='cuda') transform = DistanceTransformEDT() with self.assertRaises(NotImplementedError): - transform(raises) + transform(input_) if __name__ == "__main__": From 3cccb546b1fb80e3f4832b8e9de470bd9188fe50 Mon Sep 17 00:00:00 2001 From: Matthias Hadlich Date: Tue, 26 Sep 2023 20:35:26 +0200 Subject: [PATCH 18/23] Rework distance_transform_edt to include more parameters Signed-off-by: Matthias Hadlich --- monai/transforms/post/array.py | 2 + monai/transforms/utils.py | 128 ++++++++++++++++++++++++++++----- 2 files changed, 111 insertions(+), 19 deletions(-) diff --git a/monai/transforms/post/array.py b/monai/transforms/post/array.py index f7e3a0e5f3..3ebf0e3c24 100644 --- a/monai/transforms/post/array.py +++ b/monai/transforms/post/array.py @@ -962,6 +962,8 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: channel first array, must have shape: (num_channels, H[, W, ..., ]) Input gets passed channel-wise to the distance-transform, thus results from this function may differ from directly calling ``distance_transform_edt()`` in CuPy or scipy. + 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. Returns: diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py index f5e81d0d33..38d8f3f59e 100644 --- a/monai/transforms/utils.py +++ b/monai/transforms/utils.py @@ -18,7 +18,7 @@ from contextlib import contextmanager from functools import lru_cache, wraps from inspect import getmembers, isclass -from typing import Any +from typing import Any, Tuple import numpy as np import torch @@ -2059,44 +2059,134 @@ def has_status_keys(data: torch.Tensor, status_key: Any, default_message: str = def distance_transform_edt( - img: NdarrayOrTensor, sampling: None | float | list[float] = None, force_scipy: bool = False -) -> NdarrayOrTensor: + img: NdarrayOrTensor, + sampling: None | float | list[float] = None, + return_distances: bool = True, + return_indices: bool = False, + distances: NdarrayOrTensor | None=None, + indices: NdarrayOrTensor | None=None, + block_params:Tuple[int,int,int] | None=None, + float64_distances: bool=False, +) -> Tuple[NdarrayOrTensor, NdarrayOrTensor]: """ Euclidean distance transform, either GPU based with CuPy / cuCIM - or CPU based with scipy.ndimage. - Choice depends on cuCIM being available or scipy can be forced with the ``force_scipy`` flag. + or CPU based with scipy. + To use the GPU implementation, make sure cuCIM is available and that the data is a `torch.tensor` on a GPU device. - Note that the runtime running on the CPU may be really depending on the inputs size. + For details about the implementation, check out the + `SciPy`_ + and `cuCIM`_ docs. Args: img: Input image on which the distance transform shall be run. channel first array, must have shape: (num_channels, H[, W, ..., ]) - If you need to run the transform on other shapes, use the ``force_scipy`` flag. - 4D input gets passed channel-wise to cupy. + Input gets passed channel-wise to the distance-transform, thus results from this function may differ + from directly calling ``distance_transform_edt()`` in CuPy or scipy. 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 + return_distances: Whether to calculate the distance transform. + return_indices: Whether to calculate the feature transform. + distances: An output array to store the calculated distance transform, instead of returning it. `return_distances` must be True. + indices: An output array to store the calculated feature transform, instead of returning it. `return_indicies` must be True. + block_params: This parameter is specific to cuCIM and does not exist in SciPy. For details, look here + `distance_transform_edt`_ + float64_distances: This parameter is specific to cuCIM and does not exist in SciPy. + If True, use double precision in the distance computation (to match SciPy behavior). + Otherwise, single precision will be used for efficiency. + + Returns: + distances: The calculated distance transform. Returned only when `return_distances` is True and `distances` is not supplied. + It will have the same shape as image. For cuCIM: Will have dtype torch.float64 if float64_distances is True, otherwise it will have dtype torch.float32. + For scipy: Will have dtype np.float64. + indices: The calculated feature transform. It has an image-shaped array for each dimension of the image. + Returned only when `return_indices` is True and `indices` is not supplied. dtype np.float64. + """ distance_transform_edt, has_cucim = optional_import( "cucim.core.operations.morphology", name="distance_transform_edt" ) + use_cp = has_cp and has_cucim and isinstance(img, torch.Tensor) and img.device.type == torch.device("cuda").type - if has_cp and has_cucim and not force_scipy: + if not return_distances and not return_indices: + raise RuntimeError("Neither return_distances nor return_indices True") + distances_original, indices_original = distances, indices + distances, indices = None, None + if use_cp: + distances_, indices_ = None, None + if return_distances: + dtype = torch.float64 if float64_distances else torch.float32 + if distances is None: + distances = torch.zeros_like(img, dtype=dtype) + else: + if not isinstance(distances, torch.Tensor) and distances.device != img.device: + raise TypeError("distances must be a torch.Tensor on the same device as img") + if not distances.dtype == dtype: + raise TypeError("distances must be a torch.Tensor of dtype float32 or float64") + distances_ = convert_to_cupy(distances) + if return_indices: + dtype = torch.int32 + if indices is None: + indices = torch.zeros((img.dim(),) + img.shape, dtype=dtype) + else: + if not isinstance(indices, torch.Tensor) and indices.device != img.device: + raise TypeError("indices must be a torch.Tensor on the same device as img") + if not indices.dtype == dtype: + raise TypeError("indices must be a torch.Tensor of dtype int32") + indices_ = convert_to_cupy(indices) img_ = convert_to_cupy(img) - if img_.ndim == 4: - out = [] - for channel in img_: - out.append(distance_transform_edt(channel, sampling=sampling)) - distance = cp.stack(out) - else: - distance = distance_transform_edt(img_, sampling=sampling) + for channel_idx in range(img_.shape[0]): + distance_transform_edt( + img_[channel_idx], + sampling=sampling, + return_distances=return_distances, + return_indices=return_indices, + distances=distances_[channel_idx] if distances_ is not None else None, + indices=indices_[channel_idx] if indices_ is not None else None, + block_params=block_params, + float64_distances=float64_distances, + ) 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) + if return_distances: + if distances is None: + distances = np.zeros_like(img_, dtype=np.float64) + else: + if not isinstance(distances, np.ndarray): + raise TypeError("distances must be a numpy.ndarray") + if not distances.dtype == np.float64: + raise TypeError("distances must be a numpy.ndarray of dtype float64") + if return_indices: + if indices is None: + indices = np.zeros((img_.ndim,) + img_.shape, dtype=np.int32) + else: + if not isinstance(indices, np.ndarray): + raise TypeError("indices must be a numpy.ndarray") + if not indices.dtype == np.int32: + raise TypeError("indices must be a numpy.ndarray of dtype int32") + + for channel_idx in range(img_.shape[0]): + ndimage.distance_transform_edt( + img_[channel_idx], + sampling=sampling, + return_distances=return_distances, + return_indices=return_indices, + distances=distances[channel_idx] if distances is not None else None, + indices=indices[channel_idx] if indices is not None else None, + ) + + r_vals = [] + if return_distances and distances_original is None: + r_vals.append(distances) + if return_indices and indices_original is None: + r_vals.append(indices) + if not r_vals: + return None + if len(r_vals) == 1: + return r_vals[0] + return tuple(r_vals) - return convert_to_dst_type(distance, dst=img, dtype=distance.dtype)[0] if __name__ == "__main__": From 582efb4bc156f75b086adf05b2bafb0bb52be9f2 Mon Sep 17 00:00:00 2001 From: Matthias Hadlich Date: Tue, 26 Sep 2023 20:37:17 +0200 Subject: [PATCH 19/23] Code styling Signed-off-by: Matthias Hadlich --- monai/transforms/post/dictionary.py | 5 +--- monai/transforms/utils.py | 32 ++++++++++----------- tests/test_distance_transform_edt.py | 43 ++++++++++++++-------------- 3 files changed, 39 insertions(+), 41 deletions(-) diff --git a/monai/transforms/post/dictionary.py b/monai/transforms/post/dictionary.py index e6eda18d6d..66f60dec09 100644 --- a/monai/transforms/post/dictionary.py +++ b/monai/transforms/post/dictionary.py @@ -880,10 +880,7 @@ class DistanceTransformEDTd(MapTransform): backend = DistanceTransformEDT.backend def __init__( - self, - keys: KeysCollection, - allow_missing_keys: bool = False, - sampling: None | float | list[float] = None, + self, keys: KeysCollection, allow_missing_keys: bool = False, sampling: None | float | list[float] = None ) -> None: super().__init__(keys, allow_missing_keys) self.sampling = sampling diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py index 38d8f3f59e..78f43cb3af 100644 --- a/monai/transforms/utils.py +++ b/monai/transforms/utils.py @@ -2063,18 +2063,18 @@ def distance_transform_edt( sampling: None | float | list[float] = None, return_distances: bool = True, return_indices: bool = False, - distances: NdarrayOrTensor | None=None, - indices: NdarrayOrTensor | None=None, - block_params:Tuple[int,int,int] | None=None, - float64_distances: bool=False, + distances: NdarrayOrTensor | None = None, + indices: NdarrayOrTensor | None = None, + block_params: Tuple[int, int, int] | None = None, + float64_distances: bool = False, ) -> Tuple[NdarrayOrTensor, NdarrayOrTensor]: """ Euclidean distance transform, either GPU based with CuPy / cuCIM or CPU based with scipy. To use the GPU implementation, make sure cuCIM is available and that the data is a `torch.tensor` on a GPU device. - For details about the implementation, check out the - `SciPy`_ + For details about the implementation, check out the + `SciPy`_ and `cuCIM`_ docs. Args: @@ -2086,19 +2086,20 @@ def distance_transform_edt( if a single number, this is used for all axes. If not specified, a grid spacing of unity is implied. return_distances: Whether to calculate the distance transform. return_indices: Whether to calculate the feature transform. - distances: An output array to store the calculated distance transform, instead of returning it. `return_distances` must be True. + distances: An output array to store the calculated distance transform, instead of returning it. + `return_distances` must be True. indices: An output array to store the calculated feature transform, instead of returning it. `return_indicies` must be True. - block_params: This parameter is specific to cuCIM and does not exist in SciPy. For details, look here + block_params: This parameter is specific to cuCIM and does not exist in SciPy. For details, look here `distance_transform_edt`_ - float64_distances: This parameter is specific to cuCIM and does not exist in SciPy. - If True, use double precision in the distance computation (to match SciPy behavior). - Otherwise, single precision will be used for efficiency. + float64_distances: This parameter is specific to cuCIM and does not exist in SciPy. + If True, use double precision in the distance computation (to match SciPy behavior). + Otherwise, single precision will be used for efficiency. Returns: - distances: The calculated distance transform. Returned only when `return_distances` is True and `distances` is not supplied. - It will have the same shape as image. For cuCIM: Will have dtype torch.float64 if float64_distances is True, otherwise it will have dtype torch.float32. - For scipy: Will have dtype np.float64. - indices: The calculated feature transform. It has an image-shaped array for each dimension of the image. + distances: The calculated distance transform. Returned only when `return_distances` is True and `distances` is not supplied. + It will have the same shape as image. For cuCIM: Will have dtype torch.float64 if float64_distances is True, + otherwise it will have dtype torch.float32. For scipy: Will have dtype np.float64. + indices: The calculated feature transform. It has an image-shaped array for each dimension of the image. Returned only when `return_indices` is True and `indices` is not supplied. dtype np.float64. """ @@ -2188,6 +2189,5 @@ def distance_transform_edt( return tuple(r_vals) - if __name__ == "__main__": print_transform_backends() diff --git a/tests/test_distance_transform_edt.py b/tests/test_distance_transform_edt.py index 63b6eff536..73df324a20 100644 --- a/tests/test_distance_transform_edt.py +++ b/tests/test_distance_transform_edt.py @@ -14,9 +14,8 @@ import unittest import numpy as np -from parameterized import parameterized - import torch +from parameterized import parameterized from monai.transforms import DistanceTransformEDT, DistanceTransformEDTd from tests.utils import HAS_CUPY, assert_allclose, optional_import, skip_if_no_cuda @@ -31,13 +30,15 @@ ([[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], - ]] + [ + [ + [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], + ] + ] ), ], [ # Example 4D input to test channel-wise CuPy @@ -121,13 +122,15 @@ ([[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], - ]] + [ + [ + [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], + ] + ] ), ] ] @@ -164,24 +167,22 @@ def test_scipy_sampling(self, sampling, input, expected_output): output = transform(input) 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): - input_ = torch.tensor(input, device='cuda') + input_ = torch.tensor(input, device="cuda") 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) @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): - input_ = torch.tensor(input, device='cuda') + input_ = torch.tensor(input, device="cuda") 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) @@ -192,7 +193,7 @@ def test_cucim_sampling(self, sampling, input, expected_output): @unittest.skipUnless(momorphology, "cuCIM transforms are required.") def test_cucim_raises(self, raises): """Currently only 2D, 3D and 4D images are supported by CuPy. This test checks for the according error message""" - input_ = torch.tensor(raises, device='cuda') + input_ = torch.tensor(raises, device="cuda") transform = DistanceTransformEDT() with self.assertRaises(NotImplementedError): transform(input_) From 2bb17e144274f7de672602de6f306e84b1e0424e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 26 Sep 2023 18:37:48 +0000 Subject: [PATCH 20/23] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- monai/transforms/utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py index 78f43cb3af..d5e2f4e040 100644 --- a/monai/transforms/utils.py +++ b/monai/transforms/utils.py @@ -18,7 +18,7 @@ from contextlib import contextmanager from functools import lru_cache, wraps from inspect import getmembers, isclass -from typing import Any, Tuple +from typing import Any import numpy as np import torch @@ -2065,9 +2065,9 @@ def distance_transform_edt( return_indices: bool = False, distances: NdarrayOrTensor | None = None, indices: NdarrayOrTensor | None = None, - block_params: Tuple[int, int, int] | None = None, + block_params: tuple[int, int, int] | None = None, float64_distances: bool = False, -) -> Tuple[NdarrayOrTensor, NdarrayOrTensor]: +) -> tuple[NdarrayOrTensor, NdarrayOrTensor]: """ Euclidean distance transform, either GPU based with CuPy / cuCIM or CPU based with scipy. From aa31e319058d8824d3c83fef19d10caf5f1a0a55 Mon Sep 17 00:00:00 2001 From: Matthias Hadlich Date: Tue, 26 Sep 2023 20:42:39 +0200 Subject: [PATCH 21/23] DCO Remediation Commit for Matthias Hadlich I, Matthias Hadlich , hereby add my Signed-off-by to this commit: 101cc6227b57b6ffd4d736c3bc1be39494994813 I, Matthias Hadlich , hereby add my Signed-off-by to this commit: 24814c4e5fb405f8514722cad5acf375cc84a6dd Signed-off-by: Matthias Hadlich --- tests/test_distance_transform_edt.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_distance_transform_edt.py b/tests/test_distance_transform_edt.py index 73df324a20..f8845b916c 100644 --- a/tests/test_distance_transform_edt.py +++ b/tests/test_distance_transform_edt.py @@ -135,7 +135,6 @@ ] ] - RAISES_TEST_CASES = ( [ # Example 4D input. Should raise under CuPy np.array( From 94db8eaca880e0cb63fca0e5f8c272ac65c3f330 Mon Sep 17 00:00:00 2001 From: Matthias Hadlich Date: Wed, 27 Sep 2023 15:33:28 +0200 Subject: [PATCH 22/23] Final fixes Signed-off-by: Matthias Hadlich --- monai/transforms/post/array.py | 20 +++++++++------ monai/transforms/post/dictionary.py | 22 +++++++++++----- monai/transforms/utils.py | 38 ++++++++++++++++------------ tests/test_distance_transform_edt.py | 4 +-- 4 files changed, 52 insertions(+), 32 deletions(-) diff --git a/monai/transforms/post/array.py b/monai/transforms/post/array.py index 3ebf0e3c24..f7dc10120b 100644 --- a/monai/transforms/post/array.py +++ b/monai/transforms/post/array.py @@ -943,10 +943,14 @@ def __call__(self, image: NdarrayOrTensor) -> torch.Tensor: class DistanceTransformEDT(Transform): """ Applies the Euclidean distance transform on the input. + Either GPU based with CuPy / cuCIM or CPU based with scipy. + To use the GPU implementation, make sure cuCIM is available and that the data is a `torch.tensor` on a GPU device. - 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(). + Note that the results of the libraries can differ, so stick to one if possible. + For details, check out the `SciPy`_ and `cuCIM`_ documentation and / or :func:`monai.transforms.utils.distance_transform_edt`. + + .. _SciPy: https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.distance_transform_edt.html + .. _cuCIM: https://docs.rapids.ai/api/cucim/nightly/api/#cucim.core.operations.morphology.distance_transform_edt """ backend = [TransformBackends.NUMPY, TransformBackends.CUPY] @@ -959,13 +963,13 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: """ Args: img: Input image on which the distance transform shall be run. - channel first array, must have shape: (num_channels, H[, W, ..., ]) - Input gets passed channel-wise to the distance-transform, thus results from this function may differ - from directly calling ``distance_transform_edt()`` in CuPy or scipy. - sampling: Spacing of elements along each dimension. If a sequence, must be of length equal to the input rank; + Has to be a channel first array, must have shape: (num_channels, H, W [,D]). + Can be of any type but will be converted into binary: 1 wherever image equates to True, 0 elsewhere. + Input gets passed channel-wise to the distance-transform, thus results from this function will differ + from directly calling ``distance_transform_edt()`` in CuPy or SciPy. + sampling: Spacing of elements along each dimension. If a sequence, must be of length equal to the input rank -1; if a single number, this is used for all axes. If not specified, a grid spacing of unity is implied. - Returns: An array with the same shape and data type as img """ diff --git a/monai/transforms/post/dictionary.py b/monai/transforms/post/dictionary.py index 66f60dec09..393f161917 100644 --- a/monai/transforms/post/dictionary.py +++ b/monai/transforms/post/dictionary.py @@ -862,19 +862,29 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, N class DistanceTransformEDTd(MapTransform): """ Applies the Euclidean distance transform on the input. + Either GPU based with CuPy / cuCIM or CPU based with scipy. + To use the GPU implementation, make sure cuCIM is available and that the data is a `torch.tensor` on a GPU device. - Either GPU based with CuPy / cuCIM or CPU based with scipy.ndimage. - Choice only depends on cuCIM being available. - Note that the results of the two libraries can deviate, for details look into the cuCIM Docs about ``distance_transform_edt()``. + Note that the results of the libraries can differ, so stick to one if possible. + For details, check out the `SciPy`_ and `cuCIM`_ documentation and / or :func:`monai.transforms.utils.distance_transform_edt`. - Input has to 2D, 3D or 4D for CuPy. For details look into :py:class:`monai.transforms.DistanceTransformEDT`. + + Note on the input shape: + Has to be a channel first array, must have shape: (num_channels, H, W [,D]). + Can be of any type but will be converted into binary: 1 wherever image equates to True, 0 elsewhere. + Input gets passed channel-wise to the distance-transform, thus results from this function will differ + from directly calling ``distance_transform_edt()`` in CuPy or SciPy. Args: - keys: keys of the corresponding items to model output. + keys: keys of the corresponding items to be transformed. 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; + sampling: Spacing of elements along each dimension. If a sequence, must be of length equal to the input rank -1; if a single number, this is used for all axes. If not specified, a grid spacing of unity is implied. + .. _SciPy: https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.distance_transform_edt.html + .. _cuCIM: https://docs.rapids.ai/api/cucim/nightly/api/#cucim.core.operations.morphology.distance_transform_edt + + """ backend = DistanceTransformEDT.backend diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py index d5e2f4e040..7e34ece521 100644 --- a/monai/transforms/utils.py +++ b/monai/transforms/utils.py @@ -2065,32 +2065,34 @@ def distance_transform_edt( return_indices: bool = False, distances: NdarrayOrTensor | None = None, indices: NdarrayOrTensor | None = None, + *, block_params: tuple[int, int, int] | None = None, float64_distances: bool = False, -) -> tuple[NdarrayOrTensor, NdarrayOrTensor]: +) -> None | NdarrayOrTensor | tuple[NdarrayOrTensor, NdarrayOrTensor]: """ - Euclidean distance transform, either GPU based with CuPy / cuCIM - or CPU based with scipy. + Euclidean distance transform, either GPU based with CuPy / cuCIM or CPU based with scipy. To use the GPU implementation, make sure cuCIM is available and that the data is a `torch.tensor` on a GPU device. - For details about the implementation, check out the - `SciPy`_ - and `cuCIM`_ docs. + Note that the results of the libraries can differ, so stick to one if possible. + For details, check out the `SciPy`_ and `cuCIM`_ documentation. + + .. _SciPy: https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.distance_transform_edt.html + .. _cuCIM: https://docs.rapids.ai/api/cucim/nightly/api/#cucim.core.operations.morphology.distance_transform_edt Args: img: Input image on which the distance transform shall be run. - channel first array, must have shape: (num_channels, H[, W, ..., ]) - Input gets passed channel-wise to the distance-transform, thus results from this function may differ - from directly calling ``distance_transform_edt()`` in CuPy or scipy. - sampling: Spacing of elements along each dimension. If a sequence, must be of length equal to the input rank; + Has to be a channel first array, must have shape: (num_channels, H, W [,D]). + Can be of any type but will be converted into binary: 1 wherever image equates to True, 0 elsewhere. + Input gets passed channel-wise to the distance-transform, thus results from this function will differ + from directly calling ``distance_transform_edt()`` in CuPy or SciPy. + sampling: Spacing of elements along each dimension. If a sequence, must be of length equal to the input rank -1; if a single number, this is used for all axes. If not specified, a grid spacing of unity is implied. return_distances: Whether to calculate the distance transform. return_indices: Whether to calculate the feature transform. distances: An output array to store the calculated distance transform, instead of returning it. `return_distances` must be True. indices: An output array to store the calculated feature transform, instead of returning it. `return_indicies` must be True. - block_params: This parameter is specific to cuCIM and does not exist in SciPy. For details, look here - `distance_transform_edt`_ + block_params: This parameter is specific to cuCIM and does not exist in SciPy. For details, look into `cuCIM`_. float64_distances: This parameter is specific to cuCIM and does not exist in SciPy. If True, use double precision in the distance computation (to match SciPy behavior). Otherwise, single precision will be used for efficiency. @@ -2098,7 +2100,7 @@ def distance_transform_edt( Returns: distances: The calculated distance transform. Returned only when `return_distances` is True and `distances` is not supplied. It will have the same shape as image. For cuCIM: Will have dtype torch.float64 if float64_distances is True, - otherwise it will have dtype torch.float32. For scipy: Will have dtype np.float64. + otherwise it will have dtype torch.float32. For SciPy: Will have dtype np.float64. indices: The calculated feature transform. It has an image-shaped array for each dimension of the image. Returned only when `return_indices` is True and `indices` is not supplied. dtype np.float64. @@ -2106,10 +2108,14 @@ def distance_transform_edt( distance_transform_edt, has_cucim = optional_import( "cucim.core.operations.morphology", name="distance_transform_edt" ) - use_cp = has_cp and has_cucim and isinstance(img, torch.Tensor) and img.device.type == torch.device("cuda").type + use_cp = has_cp and has_cucim and isinstance(img, torch.Tensor) and img.device.type == "cuda" if not return_distances and not return_indices: raise RuntimeError("Neither return_distances nor return_indices True") + + if not (img.ndim >= 3 and img.ndim <= 4): + raise RuntimeError("Wrong input dimensionality. Use (num_channels, H, W [,D])") + distances_original, indices_original = distances, indices distances, indices = None, None if use_cp: @@ -2117,7 +2123,7 @@ def distance_transform_edt( if return_distances: dtype = torch.float64 if float64_distances else torch.float32 if distances is None: - distances = torch.zeros_like(img, dtype=dtype) + distances = torch.zeros_like(img, dtype=dtype) # type: ignore else: if not isinstance(distances, torch.Tensor) and distances.device != img.device: raise TypeError("distances must be a torch.Tensor on the same device as img") @@ -2127,7 +2133,7 @@ def distance_transform_edt( if return_indices: dtype = torch.int32 if indices is None: - indices = torch.zeros((img.dim(),) + img.shape, dtype=dtype) + indices = torch.zeros((img.dim(),) + img.shape, dtype=dtype) # type: ignore else: if not isinstance(indices, torch.Tensor) and indices.device != img.device: raise TypeError("indices must be a torch.Tensor on the same device as img") diff --git a/tests/test_distance_transform_edt.py b/tests/test_distance_transform_edt.py index f8845b916c..83b9348348 100644 --- a/tests/test_distance_transform_edt.py +++ b/tests/test_distance_transform_edt.py @@ -191,10 +191,10 @@ def test_cucim_sampling(self, sampling, input, expected_output): @unittest.skipUnless(HAS_CUPY, "CuPy is required.") @unittest.skipUnless(momorphology, "cuCIM transforms are required.") def test_cucim_raises(self, raises): - """Currently only 2D, 3D and 4D images are supported by CuPy. This test checks for the according error message""" + """Currently images of shape a certain shape are supported. This test checks for the according error message""" input_ = torch.tensor(raises, device="cuda") transform = DistanceTransformEDT() - with self.assertRaises(NotImplementedError): + with self.assertRaises(RuntimeError): transform(input_) From d0075c5517aabe94532c517854f85d333d75a0c3 Mon Sep 17 00:00:00 2001 From: Matthias Hadlich Date: Wed, 27 Sep 2023 15:47:19 +0200 Subject: [PATCH 23/23] Fix typing Signed-off-by: Matthias Hadlich --- monai/transforms/post/array.py | 2 +- monai/transforms/utils.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/monai/transforms/post/array.py b/monai/transforms/post/array.py index f7dc10120b..f10dd21642 100644 --- a/monai/transforms/post/array.py +++ b/monai/transforms/post/array.py @@ -973,4 +973,4 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: Returns: An array with the same shape and data type as img """ - return distance_transform_edt(img=img, sampling=self.sampling) + return distance_transform_edt(img=img, sampling=self.sampling) # type: ignore diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py index 7e34ece521..cdf58514e6 100644 --- a/monai/transforms/utils.py +++ b/monai/transforms/utils.py @@ -2192,7 +2192,7 @@ def distance_transform_edt( return None if len(r_vals) == 1: return r_vals[0] - return tuple(r_vals) + return tuple(r_vals) # type: ignore if __name__ == "__main__":