A quick conversion of label maps to colored meshes with latest vtk.

surface_nets_3d_demo.py

"""A quick conversion of label maps to colored meshes.

This requires vtk > 9.3.
pip install --pre vtk colorcet pyvista colorcet
"""
from typing import Optional, Union

import pooch
import pyvista as pv
import vtk


def surface_nets(
    label_map: pv.ImageData,
    output_mesh_type: str = "quads",
    num_labels: Optional[int] = None,
    output_style: str = "boundary",
    smoothing: bool = False,
    num_iterations: int = 50,
    relaxation_factor: float = 0.5,
    constraint_distance: float = 1,
) -> pv.PolyData:
    """Extract surface meshes from a label map.

    Parameters
    ----------
    label_map : pv.ImageData
        The input label map.

    output_mesh_type : str, optional
        Mesh type - triangles or quads

    num_labels : int, optional
        Number of labels to be extracted. If not specified, all are taken.
    
    output_style : str, default: "boundary"
        Output style (boundary, selected, default)
    
    smoothing : bool, default: False
        Apply smoothing to the meshes
    
    num_iterations : int, default: 50
        Number of smoothing iterations

    relaxation_factor : float, default: 0.5
        Relaxation factor of the smoothing

    constraint_distance : float, default: 1
        Constraint distance of the smoothing

    See also:
    Sarah F. Frisken, SurfaceNets for Multi-Label Segmentations with Preservation of Sharp
    Boundaries, Journal of Computer Graphics Techniques (JCGT), vol. 11, no. 1, 34-54, 2022
    Available online http://jcgt.org/published/0011/01/03/

    https://www.kitware.com/really-fast-isocontouring/

    """
    if num_labels is None:
        num_labels = int(label_map.point_data.get_array(0).max())

    surface_nets = vtk.vtkSurfaceNets3D()
    surface_nets.SetInputData(label_map)
    surface_nets.GenerateLabels(num_labels, 1, num_labels)
    if output_mesh_type == "quads":
        surface_nets.SetOutputMeshTypeToQuads()
    elif output_mesh_type == "triangles":
        surface_nets.SetOutputMeshTypeToTriangles()
    if output_style == "boundary":
        surface_nets.SetOutputStyleToBoundary()
    elif output_style == "selected":
        surface_nets.SetOutputStyleToSelected()
    elif output_style == "default":
        surface_nets.SetOutputStyleToDefault()
    if smoothing:
        surface_nets.SmoothingOn()
        surface_nets.GetSmoother().SetNumberOfIterations(num_iterations)
        surface_nets.GetSmoother().SetRelaxationFactor(relaxation_factor)
        surface_nets.GetSmoother().SetConstraintDistance(constraint_distance)
    else:
        surface_nets.SmoothingOff()

    surface_nets.Update()

    return pv.wrap(surface_nets.GetOutput())


def run_example():
    # Using 1000 parcelation Schaeffer label map from the excellent neuroparc project:
    # https://github.com/neurodata/neuroparc
    file_name = pooch.retrieve(
        "https://github.com/neurodata/neuroparc/raw/master/atlases/label/Human/Schaefer1000_space-MNI152NLin6_res-1x1x1.nii.gz",
        known_hash="c3efe797aab3b3d9e705645bf29fac4e932c88dbe54ccaeb03982f11e66b3249",
    )
    label_map = pv.read(file_name)

    coarse = surface_nets(label_map, smoothing=False)
    smooth = surface_nets(label_map, smoothing=True, relaxation_factor=0.3)

    pl = pv.Plotter(shape=(1, 2))
    pl.subplot(0, 0)
    pl.add_mesh(coarse, cmap="glasbey_bw", show_scalar_bar=False)
    pl.subplot(0, 1)
    pl.add_mesh(smooth, cmap="glasbey_bw", show_scalar_bar=False)
    pl.link_views()
    pl.show()

if __name__ == "__main__":
    run_example()

The above outputs this:

brain-surfacenets-30fps.webm

See https://twitter.com/jmargeta/status/1719489527541100931 for more.