watershed.py

import SimpleITK as sitk
import sys
from os.path import expanduser, join

'''
 Watershed.py
 Jan Witowski
 Garage of Complexity
'''

# 1) Loading
dicomPath = join(expanduser("~"), "Desktop", "dir")
reader = sitk.ImageSeriesReader()
seriesIDread = reader.GetGDCMSeriesIDs(dicomPath)[0]
print("len:", len(seriesIDread))
dicomFilenames = reader.GetGDCMSeriesFileNames(dicomPath, seriesIDread)
reader.SetFileNames(dicomFilenames)
imgSeries = reader.Execute()
sitk.Show(imgSeries, "Before pre-processing")

# 2) Pre-processing
## Edge features
sigma = imgSeries.GetSpacing()[0]
level = 17 # 15, 16, 17 work best
imgFeature = sitk.GradientMagnitude(imgSeries)
#sitk.Show(imgFeature, "EdgeFeatures")

# 3) Watershed segmentation
'''Oversegmentation
imgWS = sitk.MorphologicalWatershed(imgFeature, level=0, markWatershedLine=True, fullyConnected=False)
sitk.Show(sitk.LabelToRGB(imgWS), "WS oversegmentation")
'''

imgWS = sitk.MorphologicalWatershed(imgFeature, level=level, markWatershedLine=True, fullyConnected=False)
# Notice the behavior without sitk.LabelToRGB() - the image is all blurry and segmentation doesn't look like it should!!!
sitk.Show(sitk.LabelToRGB(imgWS), "WS segmentation")



# 4) Region growing
seed = [130, 140, 133] #slice 133
LowerTh, UpperTh = 120,160

segmentationFilter = sitk.ConnectedThresholdImageFilter() # New segmentation Filter instance
segmentationFilter.SetLower( float(LowerTh) ) # Arg -> lower threshold
segmentationFilter.SetUpper( float(UpperTh) ) # Arg -> upper threshold
segmentationFilter.SetReplaceValue( 255 )
segmentationFilter.AddSeed( seed ) # Arg -> Add seed to Filter instance

# Run the segmentation filter
image = segmentationFilter.Execute( imgWS ) # Apply filter (Segmentate)
image[seed]=255
sitk.Show(sitk.LabelToRGB(image), "After region growing")