Skip to content

GenerateModelsFromLabels

vtk-examples/Python/Medical/GenerateModelsFromLabels

Description

This example uses vtkDiscreteFlyingEdges3D or vtkDiscreteMarchingCubes to create vtkPolyData models from a 3D volume that contains discrete labels. These volumes are normally the output of a segmentation algorithm. The polydata for each label will be output into a separate file.

You can load these files into ParaView, where they will appear as a series of time steps. You can then single step through displaying the polydate from each file making up the series.

If you want to see the segmentation results as cube models, see the example GenerateCubesFromLabels

The input volume must be in MetaIO format.

Once you generate the models, you can view them with Paraview

Note

This original source code for this example is here.

Other languages

See (Cxx)

Question

If you have a question about this example, please use the VTK Discourse Forum

Code

GenerateModelsFromLabels.py

#!/usr/bin/env python

import os
import sys

# noinspection PyUnresolvedReferences
import vtkmodules.vtkInteractionStyle
# noinspection PyUnresolvedReferences
import vtkmodules.vtkRenderingOpenGL2
from vtkmodules.vtkCommonCore import (
    VTK_VERSION_NUMBER,
    vtkVersion
)
from vtkmodules.vtkCommonDataModel import (
    vtkDataObject,
    vtkDataSetAttributes
)
from vtkmodules.vtkFiltersCore import (
    vtkMaskFields,
    vtkThreshold,
    vtkWindowedSincPolyDataFilter
)
from vtkmodules.vtkFiltersGeneral import (
    vtkDiscreteFlyingEdges3D,
    vtkDiscreteMarchingCubes
)
from vtkmodules.vtkFiltersGeometry import vtkGeometryFilter
from vtkmodules.vtkIOImage import vtkMetaImageReader
from vtkmodules.vtkIOXML import vtkXMLPolyDataWriter
from vtkmodules.vtkImagingStatistics import vtkImageAccumulate


def main():
    # vtkDiscreteFlyingEdges3D was introduced in VTK >= 8.2
    use_flying_edges = vtk_version_ok(8, 2, 0)

    file_name, start_label, end_label = get_program_parameters()
    if start_label > end_label:
        end_label, start_label = start_label, end_label

    # Create all of the classes we will need
    reader = vtkMetaImageReader()
    histogram = vtkImageAccumulate()
    if use_flying_edges:
        try:
            using_marching_cubes = False
            discrete_cubes = vtkDiscreteFlyingEdges3D()
        except AttributeError:
            using_marching_cubes = True
            discrete_cubes = vtkDiscreteMarchingCubes()
    else:
        using_marching_cubes = True
        discrete_cubes = vtkDiscreteMarchingCubes()
    smoother = vtkWindowedSincPolyDataFilter()
    selector = vtkThreshold()
    scalars_off = vtkMaskFields()
    geometry = vtkGeometryFilter()
    writer = vtkXMLPolyDataWriter()

    # Define all of the variables
    file_prefix = 'Label'
    smoothing_iterations = 15
    pass_band = 0.001
    feature_angle = 120.0

    # Generate models from labels
    # 1) Read the meta file
    # 2) Generate a histogram of the labels
    # 3) Generate models from the labeled volume
    # 4) Smooth the models
    # 5) Output each model into a separate file

    reader.SetFileName(file_name)

    histogram.SetInputConnection(reader.GetOutputPort())
    histogram.SetComponentExtent(0, end_label, 0, 0, 0, 0)
    histogram.SetComponentOrigin(0, 0, 0)
    histogram.SetComponentSpacing(1, 1, 1)
    histogram.Update()

    discrete_cubes.SetInputConnection(reader.GetOutputPort())
    discrete_cubes.GenerateValues(end_label - start_label + 1, start_label, end_label)

    smoother.SetInputConnection(discrete_cubes.GetOutputPort())
    smoother.SetNumberOfIterations(smoothing_iterations)
    smoother.BoundarySmoothingOff()
    smoother.FeatureEdgeSmoothingOff()
    smoother.SetFeatureAngle(feature_angle)
    smoother.SetPassBand(pass_band)
    smoother.NonManifoldSmoothingOn()
    smoother.NormalizeCoordinatesOn()
    smoother.Update()

    selector.SetInputConnection(smoother.GetOutputPort())
    if use_flying_edges:
        if using_marching_cubes:
            selector.SetInputArrayToProcess(0, 0, 0, vtkDataObject().FIELD_ASSOCIATION_CELLS,
                                            vtkDataSetAttributes().SCALARS)
        else:
            selector.SetInputArrayToProcess(0, 0, 0, vtkDataObject().FIELD_ASSOCIATION_POINTS,
                                            vtkDataSetAttributes().SCALARS)
    else:
        selector.SetInputArrayToProcess(0, 0, 0, vtkDataObject().FIELD_ASSOCIATION_CELLS,
                                        vtkDataSetAttributes().SCALARS)

    # Strip the scalars from the output
    scalars_off.SetInputConnection(selector.GetOutputPort())
    scalars_off.CopyAttributeOff(vtkMaskFields().POINT_DATA,
                                 vtkDataSetAttributes().SCALARS)
    scalars_off.CopyAttributeOff(vtkMaskFields().CELL_DATA,
                                 vtkDataSetAttributes().SCALARS)

    geometry.SetInputConnection(scalars_off.GetOutputPort())

    writer.SetInputConnection(geometry.GetOutputPort())

    for i in range(start_label, end_label + 1):
        # see if the label exists, if not skip it
        frequency = histogram.GetOutput().GetPointData().GetScalars().GetTuple1(i)
        if frequency == 0.0:
            continue

        # select the cells for a given label
        selector.SetLowerThreshold(i)
        selector.SetUpperThreshold(i)

        # output the polydata
        output_fn = '{:s}{:d}.vtp'.format(file_prefix, i)
        print('{:s} writing {:s}'.format(os.path.basename(sys.argv[0]), output_fn))

        writer.SetFileName(output_fn)
        writer.Write()


def get_program_parameters():
    import argparse
    description = 'Creates vtkPolyData models from a 3D volume that contains discrete labels.'
    epilogue = '''
These volumes are normally the output of a segmentation algorithm.
The polydata for each label will be output into a separate file.
    '''
    parser = argparse.ArgumentParser(description=description, epilog=epilogue,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('filename', help='Input volume e.g. Frog/frogtissue.mhd.')
    parser.add_argument('startlabel', type=int, help='The starting label in the input volume, e,g, 1.')
    parser.add_argument('endlabel', type=int, help='The ending label in the input volume e.g. 29')
    args = parser.parse_args()
    return args.filename, args.startlabel, args.endlabel


def vtk_version_ok(major, minor, build):
    """
    Check the VTK version.

    :param major: Major version.
    :param minor: Minor version.
    :param build: Build version.
    :return: True if the requested VTK version is greater or equal to the actual VTK version.
    """
    needed_version = 10000000000 * int(major) + 100000000 * int(minor) + int(build)
    try:
        vtk_version_number = VTK_VERSION_NUMBER
    except AttributeError:  # as error:
        ver = vtkVersion()
        vtk_version_number = 10000000000 * ver.GetVTKMajorVersion() + 100000000 * ver.GetVTKMinorVersion() \
                             + ver.GetVTKBuildVersion()
    if vtk_version_number >= needed_version:
        return True
    else:
        return False


if __name__ == '__main__':
    main()