Slicer test

interpolateMergedContours.py

import os
import sys
import argparse
from os.path import join, exists, isdir

"""
Usage: Slicer --no-main-window --python-script interpolateMergedContours.py -i labelmapsInDir -o labelmapsOutDir
"""

def interpolateMergedContours(inputImagePath, outputImagePath=None, outputMeshPath=None):
    # Import libraries
    import vtkITK
    import vtkSegmentationCorePython as vtkSegmentationCore
    segmentationsLogic = slicer.vtkSlicerSegmentationsModuleLogic()
    scene = slicer.mrmlScene
    # Load contours image
    inputLabelMapVolumeNode = slicer.util.loadLabelVolume(inputImagePath, returnNode=True)[1]
    # Create segmentation
    segmentationNode = slicer.vtkMRMLSegmentationNode()
    scene.AddNode(segmentationNode)
    segmentationNode.CreateDefaultDisplayNodes()
    displayNode = segmentationNode.GetDisplayNode()
    segmentation = segmentationNode.GetSegmentation()
    # Import label map to segmentation
    segment = segmentationsLogic.CreateSegmentFromLabelmapVolumeNode(inputLabelMapVolumeNode, segmentationNode)
    segmentation.AddSegment(segment)
    selectedSegmentIds = vtk.vtkStringArray()
    displayNode.GetVisibleSegmentIDs(selectedSegmentIds)
    mergedLabelmapGeometryImage = vtkSegmentationCore.vtkOrientedImageData()
    commonGeometryString = segmentation.DetermineCommonLabelmapGeometry(
        vtkSegmentationCore.vtkSegmentation.EXTENT_UNION_OF_EFFECTIVE_SEGMENTS,
        selectedSegmentIds)
    vtkSegmentationCore.vtkSegmentationConverter.DeserializeImageGeometry(
        commonGeometryString, mergedLabelmapGeometryImage)
    mergedImage = vtkSegmentationCore.vtkOrientedImageData()
    segmentationNode.GenerateMergedLabelmapForAllSegments(
        mergedImage,
        vtkSegmentationCore.vtkSegmentation.EXTENT_UNION_OF_EFFECTIVE_SEGMENTS,
        mergedLabelmapGeometryImage,
        selectedSegmentIds)
    outputLabelmap = vtkSegmentationCore.vtkOrientedImageData()
    interpolator = vtkITK.vtkITKMorphologicalContourInterpolator()
    interpolator.SetInputData(mergedImage)
    interpolator.Update()
    outputLabelmap.DeepCopy(interpolator.GetOutput())
    newSegment = vtkSegmentationCore.vtkSegment()
    newSegment.SetName(segment.GetName()+' interpolated')
    newSegment.SetColor(segment.GetColor())
    newSegmentID = segmentation.GenerateUniqueSegmentID("interpolated")
    segmentationNode.GetSegmentation().AddSegment(newSegment, newSegmentID)
    segmentationsLogic.SetBinaryLabelmapToSegment(outputLabelmap, segmentationNode, newSegmentID)
    CLOSED_SURFACE = vtkSegmentationCore.vtkSegmentationConverter.GetSegmentationClosedSurfaceRepresentationName()
    polyData = segmentationsLogic.CreateRepresentationForOneSegment(
            segmentation,
            newSegmentID,
            CLOSED_SURFACE)
    if outputMeshPath:
        writer = vtk.vtkPolyDataWriter()
        writer.SetInputDataObject(polyData)
        writer.SetFileName(outputMeshPath)
        writer.Write()
    if outputImagePath:
        displayNode.SetSegmentVisibility(selectedSegmentIds.GetValue(0), False)
        displayNode.GetVisibleSegmentIDs(selectedSegmentIds)
        outputLabelMapVolumeNode = slicer.vtkMRMLLabelMapVolumeNode()
        scene.AddNode(outputLabelMapVolumeNode)
        segmentationsLogic.ExportSegmentsToLabelmapNode(
            segmentationNode,
            selectedSegmentIds,
            outputLabelMapVolumeNode,
            inputLabelMapVolumeNode)
        slicer.util.saveNode(outputLabelMapVolumeNode, outputImagePath)


def main(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--input-folder", dest="input_folder", metavar="PATH",
                        required=True, help="Input folder with contours labelmaps")
    parser.add_argument("-o", "--output-folder", dest="output_folder", metavar="PATH",
                        required=True, help="Output folder for interpolated labelmaps")
    args = parser.parse_args(argv)
    if not isdir(args.output_folder):
        os.makedirs(args.output_folder)
    for filename in os.listdir(args.input_folder):
        inputImagePath = join(args.input_folder, filename)
        outputImagePath = join(args.output_folder, filename)
        outputMeshName = filename.split(os.extsep)[0] + '.vtk'
        outputMeshPath = join(args.output_folder, outputMeshName)
        interpolateMergedContours(
            inputImagePath,
            outputImagePath=outputImagePath,
            outputMeshPath=outputMeshPath)
    exit()


if __name__ == '__main__':
    main(sys.argv[1:])