Add conversion of 3D label maps to labeled meshes with SurfaceNets. (p…

…yvista#5176)

* Add conversion of 3D label maps to labeled meshes.

ImageData.contour_labeled is added and SurfaceNets3D algorithm from
a pre-release version of VTK is used. This is particularly useful for
visualization of medical images with their labels and the outputs of
multi-label segmentation maps.

Implements pyvista#5170

Note, version of vtk 9.3+ must be used which is currently in the pre-release
stage. Otherwise, the example is not generated and the test is not run.

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/

Note, silencing Mypy errors (as in pyvista#3837):
pyvista/core/filters/image_data.py:853: error: Argument 1 to "set_default_active_scalars" has incompatible type "ImageDataFilters"; expected "DataSet"  [arg-type]
pyvista/core/filters/image_data.py:854: error: "ImageDataFilters" has no attribute "active_scalars_info"  [attr-defined]
pyvista/core/filters/image_data.py:858: error: "ImageDataFilters" has no attribute "get_array_association"  [attr-defined]

* Refer to the contouring example

* Add tests for SurfaceNets contour extraction

* Update comments

* Update pyvista/core/filters/image_data.py

Co-authored-by: MatthewFlamm <39341281+MatthewFlamm@users.noreply.github.com>

* Update pyvista/core/filters/image_data.py

Co-authored-by: MatthewFlamm <39341281+MatthewFlamm@users.noreply.github.com>

* Update examples/01-filter/contouring.py

Co-authored-by: MatthewFlamm <39341281+MatthewFlamm@users.noreply.github.com>

* Update pyvista/core/filters/image_data.py

Co-authored-by: MatthewFlamm <39341281+MatthewFlamm@users.noreply.github.com>

* Use FieldAssiciation enums

* Use FieldAssociation enums correctly

* Add tests for improved coverage

* Add tests for improved coverage

* Fix quotes in tests

---------

Co-authored-by: MatthewFlamm <39341281+MatthewFlamm@users.noreply.github.com>
Co-authored-by: Bane Sullivan <banesullivan@gmail.com>

examples/01-filter/contouring.py

@@ -74,3 +74,15 @@
 pl.add_mesh(contours, **dargs)
 pl.add_mesh(arrows, **dargs)
 pl.show()
+
+###############################################################################
+# Contours from a label map
+# +++++++++++++++++++++++++
+#
+# Create labeled surfaces from 3D label maps (e.f. multi-label image segmentation)
+# using :func:`contour_labeled() <pyvista.ImageDataFilters.contour_labeled>`.
+# Requires VTK version 9.3
+if pv.vtk_version_info >= (9, 3):
+    label_map = pv.examples.download_frog_tissue()
+    mesh = label_map.contour_labeled()
+    mesh.plot(cmap="glasbey_warm", cpos="yx", show_scalar_bar=False)

pyvista/core/_vtk_core.py

@@ -421,3 +421,9 @@ def __init__(self):  # pragma: no cover
     from vtkmodules.vtkRenderingCore import vtkViewport
 except ImportError:  # pragma: no cover
     pass
+
+# 9.3+ imports
+try:
+    from vtkmodules.vtkFiltersCore import vtkSurfaceNets3D
+except ImportError:  # pragma: no cover
+    pass

pyvista/core/filters/image_data.py

@@ -1,5 +1,6 @@
 """Filters with a class to manage filters/algorithms for uniform grid datasets."""
 import collections.abc
+from typing import Literal, Optional
 
 import numpy as np
 
@@ -8,7 +9,7 @@
 from pyvista.core.errors import AmbiguousDataError, MissingDataError
 from pyvista.core.filters import _get_output, _update_alg
 from pyvista.core.filters.data_set import DataSetFilters
-from pyvista.core.utilities.arrays import set_default_active_scalars
+from pyvista.core.utilities.arrays import FieldAssociation, set_default_active_scalars
 from pyvista.core.utilities.helpers import wrap
 from pyvista.core.utilities.misc import abstract_class
 
@@ -772,3 +773,128 @@ def _flip_uniform(self, axis) -> 'pyvista.ImageData':
         alg.SetFilteredAxes(axis)
         alg.Update()
         return wrap(alg.GetOutput())
+
+    def contour_labeled(
+        self,
+        n_labels: Optional[int] = None,
+        smoothing: bool = False,
+        smoothing_num_iterations: int = 50,
+        smoothing_relaxation_factor: float = 0.5,
+        smoothing_constraint_distance: float = 1,
+        output_mesh_type: Literal['quads', 'triangles'] = 'quads',
+        output_style: Literal['default', 'boundary'] = 'default',
+        scalars: Optional[str] = None,
+        progress_bar: bool = False,
+    ) -> 'pyvista.PolyData':
+        """Generate labeled contours from 3D label maps.
+
+        SurfaceNets algorithm is used to extract contours preserving sharp
+        boundaries for the selected labels from the label maps.
+        Optionally, the boundaries can be smoothened to reduce the staircase
+        appearance in case of low resolution input label maps.
+
+        This filter requires that the :class:`ImageData` has integer point
+        scalars, such as multi-label maps generated from image segmentation.
+
+        .. note::
+           Requires ``vtk>=9.3.0``.
+
+        Parameters
+        ----------
+        n_labels : int, optional
+            Number of labels to be extracted (all are extracted if None is given).
+
+        smoothing : bool, default: False
+            Apply smoothing to the meshes.
+
+        smoothing_num_iterations : int, default: 50
+            Number of smoothing iterations.
+
+        smoothing_relaxation_factor : float, default: 0.5
+            Relaxation factor of the smoothing.
+
+        smoothing_constraint_distance : float, default: 1
+            Constraint distance of the smoothing.
+
+        output_mesh_type : str, default: 'quads'
+            Type of the output mesh. Must be either ``'quads'``, or ``'triangles'``.
+
+        output_style : str, default: 'default'
+            Style of the output mesh. Must be either ``'default'`` or ``'boundary'``.
+            When ``'default'`` is specified, the filter produces a mesh with both
+            interior and exterior polygons. When ``'boundary'`` is selected, only
+            polygons on the border with the background are produced (without interior
+            polygons). Note that style ``'selected'`` is currently not implemented.
+
+        scalars : str, optional
+            Name of scalars to process. Defaults to currently active scalars.
+
+        progress_bar : bool, default: False
+            Display a progress bar to indicate progress.
+
+        Returns
+        -------
+        pyvista.PolyData
+            :class:`pyvista.PolyData` Labeled mesh with the segments labeled.
+
+        References
+        ----------
+        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/
+
+        Examples
+        --------
+        See :ref:`contouring_example` for a full example using this filter.
+
+        """
+        if not hasattr(_vtk, 'vtkSurfaceNets3D'):  # pragma: no cover
+            from pyvista.core.errors import VTKVersionError
+
+            raise VTKVersionError('Surface nets 3D require VTK 9.3.0 or newer.')
+
+        alg = _vtk.vtkSurfaceNets3D()
+        if scalars is None:
+            set_default_active_scalars(self)  # type: ignore
+            field, scalars = self.active_scalars_info  # type: ignore
+            if field != FieldAssociation.POINT:
+                raise ValueError('If `scalars` not given, active scalars must be point array.')
+        else:
+            field = self.get_array_association(scalars, preference='point')  # type: ignore
+            if field != FieldAssociation.POINT:
+                raise ValueError(
+                    f'Can only process point data, given `scalars` are {field.name.lower()} data.'
+                )
+        alg.SetInputArrayToProcess(
+            0, 0, 0, field.value, scalars
+        )  # args: (idx, port, connection, field, name)
+        alg.SetInputData(self)
+        if n_labels is not None:
+            alg.GenerateLabels(n_labels, 1, n_labels)
+        if output_mesh_type == 'quads':
+            alg.SetOutputMeshTypeToQuads()
+        elif output_mesh_type == 'triangles':
+            alg.SetOutputMeshTypeToTriangles()
+        else:
+            raise ValueError(
+                f'Invalid output mesh type "{output_mesh_type}", use "quads" or "triangles"'
+            )
+        if output_style == 'default':
+            alg.SetOutputStyleToDefault()
+        elif output_style == 'boundary':
+            alg.SetOutputStyleToBoundary()
+        elif output_style == 'selected':
+            raise NotImplementedError(f'Output style "{output_style}" is not implemented')
+        else:
+            raise ValueError(f'Invalid output style "{output_style}", use "default" or "boundary"')
+        if smoothing:
+            alg.SmoothingOn()
+            alg.GetSmoother().SetNumberOfIterations(smoothing_num_iterations)
+            alg.GetSmoother().SetRelaxationFactor(smoothing_relaxation_factor)
+            alg.GetSmoother().SetConstraintDistance(smoothing_constraint_distance)
+        else:
+            alg.SmoothingOff()
+        _update_alg(alg, progress_bar, 'Performing Labeled Surface Extraction')
+        return wrap(alg.GetOutput())

tests/core/test_imagedata_filters.py

@@ -0,0 +1,127 @@
+import numpy as np
+import pytest
+
+import pyvista as pv
+from pyvista import examples
+
+VTK93 = pv.vtk_version_info >= (9, 3)
+
+
+@pytest.mark.skipif(not VTK93, reason='At least VTK 9.3 is required')
+def test_contour_labeled():
+    # Load a 3D label map (segmentation of a frog's tissue)
+    label_map = examples.download_frog_tissue()
+
+    # Extract surface for each label
+    mesh = label_map.contour_labeled()
+
+    assert label_map.point_data.active_scalars.max() == 29
+    assert 'BoundaryLabels' in mesh.cell_data
+    assert np.max(mesh['BoundaryLabels'][:, 0]) == 29
+
+
+@pytest.mark.skipif(not VTK93, reason='At least VTK 9.3 is required')
+def test_contour_labeled_with_smoothing():
+    # Load a 3D label map (segmentation of a frog's tissue)
+    label_map = examples.download_frog_tissue()
+
+    # Extract smooth surface for each label
+    mesh = label_map.contour_labeled(smoothing=True)
+    # this somehow mutates the object... also the n_labels is likely not correct
+
+    assert 'BoundaryLabels' in mesh.cell_data
+    assert np.max(mesh['BoundaryLabels'][:, 0]) == 29
+
+
+@pytest.mark.skipif(not VTK93, reason='At least VTK 9.3 is required')
+def test_contour_labeled_with_reduced_labels_count():
+    # Load a 3D label map (segmentation of a frog's tissue)
+    label_map = examples.download_frog_tissue()
+
+    # Extract surface for each label
+    mesh = label_map.contour_labeled(n_labels=2)
+    # this somehow mutates the object... also the n_labels is likely not correct
+
+    assert 'BoundaryLabels' in mesh.cell_data
+    assert np.max(mesh['BoundaryLabels'][:, 0]) == 2
+
+
+@pytest.mark.skipif(not VTK93, reason='At least VTK 9.3 is required')
+def test_contour_labeled_with_triangle_output_mesh():
+    # Load a 3D label map (segmentation of a frog's tissue)
+    label_map = examples.download_frog_tissue()
+
+    # Extract surface for each label
+    mesh = label_map.contour_labeled(scalars='MetaImage', output_mesh_type='triangles')
+
+    assert 'BoundaryLabels' in mesh.cell_data
+    assert np.max(mesh['BoundaryLabels'][:, 0]) == 29
+
+
+@pytest.mark.skipif(not VTK93, reason='At least VTK 9.3 is required')
+def test_contour_labeled_with_boundary_output_style():
+    # Load a 3D label map (segmentation of a frog's tissue)
+    label_map = examples.download_frog_tissue()
+
+    # Extract surface for each label
+    mesh = label_map.contour_labeled(output_style='boundary')
+
+    assert 'BoundaryLabels' in mesh.cell_data
+    assert np.max(mesh['BoundaryLabels'][:, 0]) == 29
+
+
+@pytest.mark.skipif(not VTK93, reason='At least VTK 9.3 is required')
+def test_contour_labeled_with_invalid_output_mesh_type():
+    # Load a 3D label map (segmentation of a frog's tissue)
+    label_map = examples.download_frog_tissue()
+
+    # Extract surface for each label
+    with pytest.raises(ValueError):
+        label_map.contour_labeled(output_mesh_type='invalid')
+
+
+@pytest.mark.skipif(not VTK93, reason='At least VTK 9.3 is required')
+def test_contour_labeled_with_invalid_output_style():
+    # Load a 3D label map (segmentation of a frog's tissue)
+    label_map = examples.download_frog_tissue()
+
+    # Extract surface for each label
+    with pytest.raises(NotImplementedError):
+        label_map.contour_labeled(output_style='selected')
+
+    with pytest.raises(ValueError):
+        label_map.contour_labeled(output_style='invalid')
+
+
+@pytest.mark.skipif(not VTK93, reason='At least VTK 9.3 is required')
+def test_contour_labeled_with_scalars():
+    # Load a 3D label map (segmentation of a frog's tissue)
+    # and create a new array with reduced number of labels
+    label_map = examples.download_frog_tissue()
+    label_map['labels'] = label_map['MetaImage'] // 2
+
+    # Extract surface for each label
+    mesh = label_map.contour_labeled(scalars='labels')
+
+    assert 'BoundaryLabels' in mesh.cell_data
+    assert np.max(mesh['BoundaryLabels'][:, 0]) == 14
+
+
+@pytest.mark.skipif(not VTK93, reason='At least VTK 9.3 is required')
+def test_contour_labeled_with_invalid_scalars():
+    # Load a 3D label map (segmentation of a frog's tissue)
+    label_map = examples.download_frog_tissue()
+
+    # Nonexistent scalar key
+    with pytest.raises(KeyError):
+        label_map.contour_labeled(scalars='nonexistent_key')
+
+    # Using cell data
+    label_map.cell_data['cell_data'] = np.zeros(label_map.n_cells)
+    with pytest.raises(ValueError, match='Can only process point data'):
+        label_map.contour_labeled(scalars='cell_data')
+
+    # When no scalas are given and active scalars are not point data
+    label_map.set_active_scalars('cell_data', preference='cell')
+    with pytest.raises(ValueError, match='active scalars must be point array'):
+        label_map.contour_labeled()