feat(transforms): add MarchingCubes surface extraction

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,

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,90 @@ 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, TransformBackends.TORCH]
+
+    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. Items have shapes ``(V, 3)``
+            and ``(F, 3)`` and match the input type: numpy arrays for array input, plain
+            ``torch.Tensor`` on the input device for tensor input. 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

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`.
@@ -1231,6 +1288,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

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,89 @@ 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[NdarrayOrTensor, NdarrayOrTensor, NdarrayOrTensor, NdarrayOrTensor]:
+    """
+    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 and the mesh is
+    moved back to the input device; 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) matching `skimage.measure.marching_cubes`
+        output. These are numpy arrays for array input, and plain (non-meta) ``torch.Tensor``
+        on the same device as ``volume`` for tensor input; ``faces`` keeps an integer dtype.
+        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(f"Expected 3 spatial spacings from MetaTensor pixdim, got {len(spacing_t)}.")
+        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,
+    )
+    if isinstance(volume, torch.Tensor):
+        verts, faces, normals, values = (
+            torch.as_tensor(np.ascontiguousarray(a), device=volume.device) for a in (verts, faces, normals, values)
+        )
+    return verts, faces, normals, values
+
+
 def keep_merge_components_with_points(
     img_pos: NdarrayTensor,
     img_neg: NdarrayTensor,