Improved Lung Segmentation using Watershed

Improved Lung Segmentation using Watershed

The previously suggested Lung Segmentation methods in the challenge Kernels mainly involve thresholding the lung tissue based on its hounsfield value and using morphological dilation to include nodules in border regions. These methods have the severe drawback of also including lots of tissue that is neither lung, nor a region of interest. I coded up an algorithm based on the one presented in R Shojaii et al (2005, DOI: 10.1109/ICIP.2005.1530294) with some modifications, that I will present here. Some preprocessing and presentation code for this is also taken from the Kernels presented by Guido Zuidhof and ArnavJain. Lots of thanks to them for sharing their code.

The resulting CT Images are in HU and have the same (not necessarily equidistant) scale as the original scans.

#Required Imports and loading up a scan for processing as presented by Guide Zuidhof

%matplotlib inline

import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)
import dicom
import os
import scipy.ndimage as ndimage
import matplotlib.pyplot as plt

from skimage import measure,morphology,segmentation
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

# Some constants
INPUT_FOLDER='../input/sample_images/'
patients=os.listdir(INPUT_FOLDER)
patients.sort()

# Load the scans in given folder path
def load_scan(path):
    slices=[dicom.read_file(path+'/'+s) for s in os.listdir(path)]
    slices.sort(key=lambda x: int(x.InstanceNumber))
    try:
        slice_thickness=np.abs(slices[0].ImagePositionPatient[2]-slices[1].ImagePositionPatient[2])
    except:
        slice_thickness=np.abs(slices[0].SliceLocation-slices[1].SliceLocation)
    for s in slices:
        s.SliceThickness=slice_thickness
    return slices

def get_pixels_hu(scans):
    image=np.stack([s.pixel_array for s in scans])
    # Convert to int16(from sometimes int16),
    # should be possible as values should always be low enough(<32k)
    image=image.astype(np.int16)
    
    # Set outside-of-scan pixels to 0
    # The intercept is usually -1024, so air is approximately 0
    image[image==-2000]=0
    
    # Convert to Hounsfield units (HU)
    intercept=scans[0].RescaleIntercept
    slope=scans[0].RescaleSlope
    
    if slope!=1:
        image=slope*image.astype(np.float64)
        image=image.astype(np.int16)
        
    image+=np.int16(intercept)
    
    return np.array(image,dtype=np.int16)

test_patient_scans=load_scan(INPUT_FOLDER+patient[8])
test_patient_images=get_pixels_hu(test_patient_scans)
print("Original Slice")
plt.imshow(test_patient_images[65],cmap='gray')
plt.show()

Output:

Original Slice

In order to use marker based watershed segmentation, we need to identify two markers. An internal marker, that is definitely lung tissue and an external marker, that is definitely outside of our ROI. We’re starting by creating the internal marker by thresholding the Image and removing all regions but the biggest one. The external marker is created by morphological dilation of the internal marker with 2 different iterations and subtracting the results. A watershed marker is created superimposing the 2 markers with different grayscale values.

# Some of the starting Code is taken from ArnavJain, 
# since it's more readable than my own
def generate_markers(image):
    #Creation of the internal Marker
    marker_internal=image<-400
    marker_internal=segmentaion.clear_border(mark_internal)
    marker_internal_labels=measure.label(marker_internal)
    areas=[r.area for r in measure.regionprops(marker_internal_labels)]
    areas.sort()
    if len(areas)>2:
        for region in measure.regionprops(marker_internal_labels):
            if region.area<areas[-2]:
                for coordinates in region.coords:
                    marker_internal_labels[coordinates[0],coordinates[1]]=0
    marker_internal=marker_internal_labels>0
    #Creation of the external Marker
    external_a=ndimage.binary_dilation(marker_internal,iterations=10)
    external_b=ndimage.binary_dilation(marker_internal,iterations=55)
    marker_external=external_b^external_a
    #Creation of the Watershed Marker matrix
    marker_watershed=np.zeros((512,512),dtype=np.int)
    marker_watershed+=marker_internal*255
    marker_watershed+=marker_external*128
    
    return marker_internal,marker_external,marker_watershed

#Show some example markers from the middle
test_patient_internal,test_patient_external,test_patient_watershed=generate_markers(test_patient_images[65])
print('Internal Marker')
plt.imshow(test_patient_internal,cmap='gray')
plt.show()
print('External Marker')
plt.imshow(test_patient_external,cmap='gray')
plt.show()
print('Watershed Marker')
plt.imshow(test_patient_watershed,cmap='gray')
plt.show()

Output:

Internal Marker

Output:

External Marker

Output:

Watershed Marker

Now we apply the marker based Watershed algorithm to find the precise border of the Lung located in the Black strip of the Watershed marker shown above. In order to do the algorithm we also need the Sobel-Gradient-Image of our original scan, which is calculated first.

In order to not miss nodules located next to the border regions a Black Top Hat Operation is performed to re-include those areas and areas surrounding the lung-hill. This is the main advantage of this method here over the methods from other kernels: Only areas that need re-inclusion get dilated, everywhere else the lung border stays precise.

def seperate_lungs(image):
    #Creation of the markers as shown above:
    marker_internal,marker_external,marker_watershed=generate_markers(image)
    
    #Creation of the Sobel-Gradient
    sobel_filtered_dx=ndimage.sobel(image,1)
    sobel_filtered_dy=ndimage.sobel(image,0)
    sobel_gradient=np.hypot(sobel_filtered_dx,sobel_filtered_dy)
    sobel_gradient*=255.0/np.max(sobel_gradient)
    
    #Watershed algorithm
    watershed=morphology.watershed(sobel_gradient,marker_watershed)
    
    #Reducing the image created by the Watershed algorithm to its outline
    outline=ndimage.morphological_gradient(watershed,size=(3,3))
    outline=outline.astype(bool)
    
    #Performing Black-Tophat Morphology for reinclusion
    #Creation of the disk-kernel and increasing its size a bit
    blackhat_struct=[[0,0,1,1,1,0,0],
                     [0,1,1,1,1,1,0],
                     [1,1,1,1,1,1,1],
                     [1,1,1,1,1,1,1],
                     [1,1,1,1,1,1,1],
                     [0,1,1,1,1,1,0],
                     [0,0,1,1,1,0,0]]
    blackhat_struct=ndimage.iterate_structure(blackhat_struct,8)
    #Perform the Black-Hat
    outline+=ndimage.black_tophat(outline,structure=blackhat_struct)
    
    #Use the internal marker and the Outline that 
    #was just created to generate the lungfilter
    lungfilter=np.bitwise_or(marker_internal,outline)
    #Close holes in the lungfilter
    #fill_holes is not used here, since in some slices 
    #the heart would be reincluded by accident
    lungfilter=ndimage.morphology.binary_closing(lungfilter,structure=np.ones((5,5)),iterations=3)
    
    #Apply the lungfilter (note the filtered areas being assigned -2000 HU)
    segmented=np.where(lungfilter==1,image,-2000*np.ones((512,512)))
    
    return segmented,lungfilter,outline,watershed,sobel_gradient,marker_internal,marker_external,marker_watershed

#Some Testcode
test_segmented,test_lungfilter,test_outline,test_watershed,test_sobel_gradient,test_marker_internal,test_marker_external,test_marker_watershed=seperate_lungs(test_patient_images[65])

print("Sobel Gradient")
plt.imshow(test_sobel_gradient,cmap='gray')
plt.show()
print("Watershed Image")
plt.imshow(test_watershed,cmap='gray')
plt.show()
print("Outline after reinclusion")
plt.imshow(test_outline,cmap='gray')
plt.show()
print("Lungfilter after closing")
plt.imshow(test_lungfilter,cmap='gray')
plt.show()
print("Segmented Lung")
plt.imshow(test_segmented,cmap='gray')
plt.show()

Output:

Sobel Gradient

Watershed Image

Outline after reinclusion

Lungfilter after closing

Segmented Lung

The resulting images of this code are still in the original dimensions of the CT Scan and in Hounsfield Units with the filtered areas being assigned -2000.

This method of lung segmentation preserves the original lung border very precisely while re-including possible nodule candidates in border regions. The main downside is the much longer processing time per patient.