decruft.py - does some math morphology operations to clean up a segmentation so that only the big connected regions are kept.

decruft.py

"""
Run this with a label map loaded in the slicer editor.

Change this path to be wherever you save this file and then paste 
the line into the slicer python console:

execfile('/Users/pieper/Dropbox/hacks/slicer/decruft.py')

This script will apply the algorithm and display the result.
You can use the editor's undo operation to get back to the 
label map before applying this algorithm.
"""
import vtkITK

fullyConnected = True
minimumSize = 0


def decruft(image,foregroundLabel,backgroundLabel,iterations=1):

  #
  # make binary mask foregroundLabel=1, backgroundLabel=0
  #

  binThresh = vtk.vtkImageThreshold()
  binThresh.SetInput( image )
  binThresh.ThresholdBetween(foregroundLabel,foregroundLabel)
  binThresh.SetInValue( 1 )
  binThresh.SetOutValue( 0 )
  binThresh.Update()

  #
  # first, erode iterations number of times
  #
  eroder = slicer.vtkImageErode()
  eroderImage = vtk.vtkImageData()
  eroderImage.DeepCopy(binThresh.GetOutput())
  eroder.SetInput(eroderImage)
  for iteration in range(iterations):
    eroder.SetForeground( 1 )
    eroder.SetBackground( 0 )
    eroder.SetNeighborTo8()
    eroder.Update()
    eroderImage.DeepCopy(eroder.GetOutput())


  #
  # now save only islands bigger than a specified size
  #

  # note that island operation happens in unsigned long space
  # but the slicer editor works in Short
  castIn = vtk.vtkImageCast()
  castIn.SetInput( eroderImage )
  castIn.SetOutputScalarTypeToUnsignedLong()

  # now identify the islands in the inverted volume
  # and find the pixel that corresponds to the background
  islandMath = vtkITK.vtkITKIslandMath()
  islandMath.SetInput( castIn.GetOutput() )
  islandMath.SetFullyConnected( fullyConnected )
  islandMath.SetMinimumSize( minimumSize )

  # note that island operation happens in unsigned long space
  # but the slicer editor works in Short
  castOut = vtk.vtkImageCast()
  castOut.SetInput( islandMath.GetOutput() )
  castOut.SetOutputScalarTypeToShort()

  castOut.Update()
  islandCount = islandMath.GetNumberOfIslands()
  islandOrigCount = islandMath.GetOriginalNumberOfIslands()
  ignoredIslands = islandOrigCount - islandCount
  print( "%d islands created (%d ignored)" % (islandCount, ignoredIslands) )

  #
  # now map everything back to 0 and 1
  #

  thresh = vtk.vtkImageThreshold()
  thresh.SetInput( castOut.GetOutput() )
  thresh.ThresholdByUpper(1)
  thresh.SetInValue( 1 )
  thresh.SetOutValue( 0 )
  thresh.Update()

  #
  # first, erode iterations plus one number of times
  #
  dilater = slicer.vtkImageErode()
  dilaterImage = vtk.vtkImageData()
  dilaterImage.DeepCopy(thresh.GetOutput())
  dilater.SetInput(dilaterImage)
  for iteration in range(1+iterations):
    dilater.SetForeground( 0 )
    dilater.SetBackground( 1 )
    dilater.SetNeighborTo8()
    dilater.Update()
    dilaterImage.DeepCopy(dilater.GetOutput())


  #
  # only keep pixels in both original and dilated result
  #

  logic = vtk.vtkImageLogic()
  logic.SetInput1(dilaterImage)
  logic.SetInput2(binThresh.GetOutput())
  if foregroundLabel == 0:
    logic.SetOperationToNand()
  logic.SetOutputTrueValue(1)
  logic.Update()



  image.DeepCopy(logic.GetOutput())

util = EditorLib.EditUtil.EditUtil()

slicer.modules.EditorWidget.toolsBox.undoRedo.saveState()

labelImage = util.getLabelImage()
label = util.getLabel()

decruft(labelImage,0,label,iterations=1)
decruft(labelImage,label,0,iterations=3)
util.markVolumeNodeAsModified(util.getLabelVolume())