diff --git a/monai/transforms/__init__.py b/monai/transforms/__init__.py index c7ac4b77e6f..62cf49937ec 100644 --- a/monai/transforms/__init__.py +++ b/monai/transforms/__init__.py @@ -298,6 +298,7 @@ KeepLargestConnectedComponent, LabelFilter, LabelToContour, + MarchingCubes, MeanEnsemble, ProbNMS, RemoveSmallObjects, @@ -335,6 +336,9 @@ LabelToContourD, LabelToContourd, LabelToContourDict, + MarchingCubesD, + MarchingCubesd, + MarchingCubesDict, MeanEnsembleD, MeanEnsembled, MeanEnsembleDict, diff --git a/monai/transforms/post/array.py b/monai/transforms/post/array.py index 3b5d38cf52d..d6c9735fb85 100644 --- a/monai/transforms/post/array.py +++ b/monai/transforms/post/array.py @@ -34,6 +34,7 @@ distance_transform_edt, fill_holes, get_largest_connected_component_mask, + get_marching_cubes_surface, get_unique_labels, remove_small_objects, ) @@ -63,6 +64,7 @@ "Invert", "GenerateHeatmap", "DistanceTransformEDT", + "MarchingCubes", ] @@ -644,7 +646,6 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: class Ensemble: - @staticmethod def get_stacked_torch(img: Sequence[NdarrayOrTensor] | NdarrayOrTensor) -> torch.Tensor: """Get either a sequence or single instance of np.ndarray/torch.Tensor. Return single torch.Tensor.""" @@ -1187,3 +1188,89 @@ 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) # type: ignore + + +class MarchingCubes(Transform): + """ + Extract a surface mesh from a 3D segmentation using marching cubes. + + Thin wrapper around :func:`monai.transforms.utils.get_marching_cubes_surface` + (`skimage.measure.marching_cubes`). The input is a channel-first volume with + shape ``(C, M, N, P)``; marching cubes runs per channel on the CPU (scikit-image + has no GPU kernel, CUDA inputs are moved to CPU first). + + Note: + The output is a mesh, not an image, so it cannot be composed with further + image transforms or inverted. Place it at the end of the pipeline. + For STL export, smoothing, or physical-space + mapping see ``monai.deploy`` ``STLConversionOperator``. + + Args: + level: isosurface value, defaults to 0.5 for binary masks. + spacing: voxel spacing along each spatial dim. A single number is used for + all axes. If ``None`` and the input is a MetaTensor, its pixdim is used, + otherwise unity spacing is assumed. + step_size: step size in voxels for marching cubes. Larger steps are faster + but yield coarser meshes. + allow_degenerate: allow degenerate triangles in the mesh. + method: one of ("lewiner", "lorensen"), see scikit-image docs. + return_normals_values: if ``True``, also return vertex normals and values, + i.e. ``(verts, faces, normals, values)`` matching scikit-image output. + Defaults to ``False`` (``(verts, faces)`` only). + + Example: + >>> import numpy as np + >>> from monai.transforms import MarchingCubes + >>> vol = np.zeros((1, 10, 10, 10), np.float32); vol[0, 3:7, 3:7, 3:7] = 1.0 + >>> verts, faces = MarchingCubes(level=0.5)(vol) + >>> verts.shape[1], faces.shape[1] + (3, 3) + """ + + backend = [TransformBackends.NUMPY] + + def __init__( + self, + level: float | None = 0.5, + spacing: Sequence[float] | float | None = None, + step_size: int = 1, + allow_degenerate: bool = True, + method: str = "lewiner", + return_normals_values: bool = False, + ) -> None: + super().__init__() + self.level = level + self.spacing = spacing + self.step_size = step_size + self.allow_degenerate = allow_degenerate + self.method = method + self.return_normals_values = return_normals_values + + def __call__(self, img: NdarrayOrTensor): + """ + Args: + img: channel-first volume with shape (C, M, N, P). + + Returns: + ``(vertices, faces)`` tuple for single-channel input, or a list of one + such tuple per channel for multi-channel input. Numpy arrays with shapes + ``(V, 3)`` and ``(F, 3)``. With ``return_normals_values=True`` each item + is ``(verts, faces, normals, values)`` instead. + + Raises: + ValueError: when ``img`` is not a channel-first 3D volume or has no channels. + """ + if img.ndim != 4 or img.shape[0] == 0: + raise ValueError(f"MarchingCubes requires a channel-first 3D volume (C, M, N, P), got shape {img.shape}.") + results = [] + for c in range(img.shape[0]): + verts, faces, normals, values = get_marching_cubes_surface( + img[c], + level=self.level, + spacing=self.spacing, + step_size=self.step_size, + allow_degenerate=self.allow_degenerate, + method=self.method, + ) + results.append((verts, faces, normals, values) if self.return_normals_values else (verts, faces)) + return results[0] if len(results) == 1 else results diff --git a/monai/transforms/post/dictionary.py b/monai/transforms/post/dictionary.py index 65fdd22b221..dc9f201f376 100644 --- a/monai/transforms/post/dictionary.py +++ b/monai/transforms/post/dictionary.py @@ -39,6 +39,7 @@ KeepLargestConnectedComponent, LabelFilter, LabelToContour, + MarchingCubes, MeanEnsemble, ProbNMS, RemoveSmallObjects, @@ -79,6 +80,9 @@ "LabelToContourD", "LabelToContourDict", "LabelToContourd", + "MarchingCubesD", + "MarchingCubesDict", + "MarchingCubesd", "MeanEnsembleD", "MeanEnsembleDict", "MeanEnsembled", @@ -270,6 +274,59 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, N return d +class MarchingCubesd(MapTransform): + """ + Dictionary-based wrapper of :py:class:`monai.transforms.MarchingCubes`. + + Note: the output stored at each key is a surface mesh ``(vertices, faces)``, + not an image, so it cannot be composed with further image transforms or + inverted. Place it at the end of the pipeline. + """ + + backend = MarchingCubes.backend + + def __init__( + self, + keys: KeysCollection, + level: float | None = 0.5, + spacing: Sequence[float] | float | None = None, + step_size: int = 1, + allow_degenerate: bool = True, + method: str = "lewiner", + return_normals_values: bool = False, + allow_missing_keys: bool = False, + ) -> None: + """ + Args: + keys: keys of the corresponding items to be transformed. + See also: :py:class:`monai.transforms.compose.MapTransform` + level: isosurface value, defaults to 0.5 for binary masks. + spacing: voxel spacing along each spatial dim. A single number is used + for all axes. If ``None`` and the input is a MetaTensor, its pixdim + is used, otherwise unity spacing is assumed. + step_size: step size in voxels for marching cubes. + allow_degenerate: allow degenerate triangles in the mesh. + method: one of ("lewiner", "lorensen"), see scikit-image docs. + return_normals_values: if ``True``, store ``(verts, faces, normals, values)``. + allow_missing_keys: don't raise exception if key is missing. + """ + super().__init__(keys, allow_missing_keys) + self.converter = MarchingCubes( + level=level, + spacing=spacing, + step_size=step_size, + allow_degenerate=allow_degenerate, + method=method, + return_normals_values=return_normals_values, + ) + + def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, NdarrayOrTensor]: + d = dict(data) + for key in self.key_iterator(d): + d[key] = self.converter(d[key]) + return d + + class RemoveSmallObjectsd(MapTransform): """ Dictionary-based wrapper of :py:class:`monai.transforms.RemoveSmallObjectsd`. @@ -1128,6 +1185,7 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Mapping[Hashable RemoveSmallObjectsD = RemoveSmallObjectsDict = RemoveSmallObjectsd LabelFilterD = LabelFilterDict = LabelFilterd LabelToContourD = LabelToContourDict = LabelToContourd +MarchingCubesD = MarchingCubesDict = MarchingCubesd MeanEnsembleD = MeanEnsembleDict = MeanEnsembled ProbNMSD = ProbNMSDict = ProbNMSd SaveClassificationD = SaveClassificationDict = SaveClassificationd diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py index 0b8a65b0fb3..b3eb0da62b9 100644 --- a/monai/transforms/utils.py +++ b/monai/transforms/utils.py @@ -110,6 +110,7 @@ "generate_spatial_bounding_box", "get_extreme_points", "get_largest_connected_component_mask", + "get_marching_cubes_surface", "keep_merge_components_with_points", "keep_components_with_positive_points", "convert_points_to_disc", @@ -1237,6 +1238,82 @@ def get_largest_connected_component_mask( return convert_to_dst_type(out, dst=img, dtype=out.dtype)[0] +def get_marching_cubes_surface( + volume: NdarrayTensor, + level: float | None = 0.5, + spacing: Sequence[float] | float | None = None, + step_size: int = 1, + allow_degenerate: bool = True, + method: str = "lewiner", + mask: NdarrayTensor | None = None, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """ + Extract a surface mesh from a 3D volume using `skimage.measure.marching_cubes`. + + NOTE: computation always runs on CPU via scikit-image (there is no cucim GPU + kernel for marching cubes). GPU tensors are moved to CPU first; the mesh is + tiny compared to the volume so the transfer cost is negligible. + + Args: + volume: 3D array with shape (M, N, P). For a channel-first image, pass one + channel at a time. + level: isosurface value. Defaults to 0.5 for binary masks. + spacing: voxel spacing along each dimension. If a single number, it is used + for all axes. If ``None`` and ``volume`` is a MetaTensor, the pixdim is + used, otherwise unity spacing is assumed. + step_size: step size in voxels. Larger steps yield coarser, faster meshes. + allow_degenerate: allow degenerate triangles in the mesh. + method: one of ("lewiner", "lorensen"), see scikit-image docs. + mask: optional boolean array of the same shape as ``volume``. Marching cubes + is computed only on ``True`` elements. + + Returns: + Tuple of (verts, faces, normals, values) as numpy arrays, matching + `skimage.measure.marching_cubes` output. Coordinate order matches the + input ``volume`` (M, N, P), scaled by ``spacing``. + + Raises: + RuntimeError: when scikit-image is not installed. + ValueError: when ``volume`` is not 3D, ``method`` is unsupported, or + ``mask`` shape does not match ``volume`` shape. + """ + if not has_measure: + raise RuntimeError("Skimage.measure required.") + look_up_option(method, ["lewiner", "lorensen"]) + volume_np, *_ = convert_data_type(volume, np.ndarray) + if volume_np.ndim != 3: + raise ValueError(f"marching cubes requires a 3D volume, got shape {volume_np.shape}.") + if spacing is not None: + spacing_t = ensure_tuple_rep(spacing, 3) + elif isinstance(volume, monai.data.MetaTensor): + try: + spacing_t = tuple(float(s) for s in volume.pixdim[-3:]) + if len(spacing_t) != 3: + raise ValueError + except Exception: + warnings.warn("Could not determine spacing from MetaTensor, assuming unity spacing.") + spacing_t = (1.0, 1.0, 1.0) + else: + spacing_t = (1.0, 1.0, 1.0) + mask_np: np.ndarray | None = None + if mask is not None: + mask_np, *_ = convert_data_type(mask, np.ndarray) + mask_np = np.asarray(mask_np, dtype=bool) + if mask_np.shape != volume_np.shape: + raise ValueError(f"mask shape {mask_np.shape} must match volume shape {volume_np.shape}.") + + verts, faces, normals, values = measure.marching_cubes( + volume_np, + level=level, + spacing=spacing_t, + step_size=step_size, + allow_degenerate=allow_degenerate, + method=method, + mask=mask_np, + ) + return verts, faces, normals, values + + def keep_merge_components_with_points( img_pos: NdarrayTensor, img_neg: NdarrayTensor, diff --git a/tests/transforms/test_marching_cubes.py b/tests/transforms/test_marching_cubes.py new file mode 100644 index 00000000000..a91c54dc847 --- /dev/null +++ b/tests/transforms/test_marching_cubes.py @@ -0,0 +1,119 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import unittest + +import numpy as np +import torch +from parameterized import parameterized + +from monai.data.meta_tensor import MetaTensor +from monai.transforms import MarchingCubes, MarchingCubesd +from monai.transforms.utils import get_marching_cubes_surface +from monai.utils import optional_import +from tests.test_utils import TEST_NDARRAYS, SkipIfNoModule + +measure, has_measure = optional_import("skimage.measure") + + +def _cube(channel_first=True): + vol = np.zeros((10, 10, 10), np.float32) + vol[3:7, 3:7, 3:7] = 1.0 + return vol[np.newaxis] if channel_first else vol + + +@SkipIfNoModule("skimage.measure") +class TestMarchingCubes(unittest.TestCase): + def test_matches_skimage(self): + vol = _cube() + verts, faces = MarchingCubes()(vol) + e_verts, e_faces, _, _ = measure.marching_cubes(vol[0], level=0.5) + self.assertTupleEqual(verts.shape, e_verts.shape) + self.assertTupleEqual(faces.shape, e_faces.shape) + np.testing.assert_allclose(verts, e_verts, rtol=1e-5) + np.testing.assert_array_equal(faces, e_faces) + self.assertEqual(verts.shape[1], 3) + self.assertEqual(faces.shape[1], 3) + + @parameterized.expand(TEST_NDARRAYS) + def test_input_types(self, im_type): + vol = im_type(_cube()) + verts, faces = MarchingCubes()(vol) + self.assertEqual(verts.shape[1], 3) + self.assertEqual(faces.shape[1], 3) + self.assertTrue((faces.max(0) < len(verts)).all()) + + def test_spacing_scales_vertices(self): + vol = _cube() + verts, _ = MarchingCubes()(vol) + verts2, _ = MarchingCubes(spacing=2.0)(vol) + np.testing.assert_allclose(verts2, verts * 2.0, rtol=1e-5) + + def test_metatensor_pixdim_spacing(self): + vol = _cube() + affine = torch.diag(torch.as_tensor([2.0, 3.0, 4.0, 1.0])) + verts, _ = MarchingCubes()(MetaTensor(vol, affine=affine)) + verts_ref, _ = MarchingCubes(spacing=(2.0, 3.0, 4.0))(vol) + np.testing.assert_allclose(verts, verts_ref, rtol=1e-5) + + def test_multi_channel_returns_list(self): + vol = np.concatenate([_cube(), _cube()], axis=0) + out = MarchingCubes()(vol) + self.assertIsInstance(out, list) + self.assertEqual(len(out), 2) + for verts, _faces in out: + self.assertEqual(verts.shape[1], 3) + + def test_return_normals_values(self): + verts, faces, normals, values = MarchingCubes(return_normals_values=True)(_cube()) + self.assertEqual(normals.shape, verts.shape) + self.assertEqual(len(values), len(verts)) + + def test_step_size(self): + verts, faces = MarchingCubes(step_size=2)(_cube()) + self.assertEqual(verts.shape[1], 3) + self.assertEqual(faces.shape[1], 3) + + def test_dict_wrapper(self): + vol = _cube() + out = MarchingCubesd(keys=["seg"])({"seg": vol})["seg"] + verts, _ = MarchingCubes()(vol) + np.testing.assert_allclose(out[0], verts, rtol=1e-5) + # allow missing keys + out = MarchingCubesd(keys=["missing"], allow_missing_keys=True)({"seg": vol}) + self.assertIn("seg", out) + + def test_util_matches_skimage(self): + vol = _cube(channel_first=False) + verts, faces, _, _ = get_marching_cubes_surface(vol, level=0.5) + e_verts, e_faces, _, _ = measure.marching_cubes(vol, level=0.5) + np.testing.assert_allclose(verts, e_verts, rtol=1e-5) + np.testing.assert_array_equal(faces, e_faces) + + def test_errors(self): + with self.assertRaises(ValueError): + MarchingCubes()(np.zeros((1, 10, 10), np.float32)) + with self.assertRaises(ValueError): + # empty channel axis + MarchingCubes()(np.zeros((0, 10, 10, 10), np.float32)) + with self.assertRaises(ValueError): + get_marching_cubes_surface(np.zeros((10, 10), np.float32)) + with self.assertRaises(ValueError): + get_marching_cubes_surface(np.zeros((10, 10, 10), np.float32), mask=np.ones((5, 5, 5), bool)) + with self.assertRaises(ValueError): + # empty volume has no isosurface + MarchingCubes()(np.zeros((1, 10, 10, 10), np.float32)) + + +if __name__ == "__main__": + unittest.main()