--- a
+++ b/lung_segmentation.py
@@ -0,0 +1,218 @@
+import numpy as np
+import os
+
+from skimage import measure, morphology, segmentation
+from skimage.morphology import ball, disk, binary_erosion, binary_closing
+from skimage.measure import label, regionprops
+from skimage.filters import roberts
+from skimage.segmentation import clear_border
+
+from scipy import ndimage as ndi
+import matplotlib.pyplot as plt
+import scipy.misc
+
+import config
+
+
+# Optimal threshold, found in an article about segmentation algorithm using
+# morphological operations. This is in HU units.
+threshold = -420
+nodule_threshold = -180
+
+
+class SegmentationAlgorithm(object):
+    def __init__(self, threshold):
+        self._threshold = threshold
+
+    def get_segmented_lungs(self, plot=False):
+        pass
+
+    def apply_threshold(self, scan):
+        scan[scan < nodule_threshold] = 0
+        return scan
+
+    def get_lung_nodules_candidates(self, patient_imgs):
+        nodules = [self.apply_threshold(scan) for scan in patient_imgs]
+        return np.stack([nodule for nodule in nodules if nodule.any()])
+
+    def get_slices_with_nodules(self, patient_imgs):
+        return np.stack([slice_ for slice_ in patient_imgs if 
+            self._has_nodule(slice_)])
+
+    def _has_nodule(self, scan):
+        scan_copy = scan.copy()
+        scan_copy[scan_copy < nodule_threshold] = 0
+        return scan_copy.any()
+
+
+class MorphologicalSegmentation(SegmentationAlgorithm):
+    def __init__(self, threshold=-420):
+        super(MorphologicalSegmentation, self).__init__(threshold)
+
+    def get_segmented_lungs(self, im, plot=False):
+        '''
+        This funtion segments the lungs from the given 2D slice.
+        '''
+        '''
+        Step 1: Convert into a binary image. 
+        '''
+        binary = im < self._threshold
+        if plot == True:
+            plt.imshow(binary, cmap=plt.cm.bone)
+            plt.show()
+        '''
+        Step 2: Remove the blobs connected to the border of the image.
+        '''
+        cleared = clear_border(binary)
+        if plot == True:
+            plt.imshow(cleared, cmap=plt.cm.bone)
+            plt.show()
+        '''
+        Step 3: Closure operation with a disk of radius 2. This operation is 
+        to keep nodules attached to the lung wall.
+        '''
+        selem = disk(2)
+        binary = binary_closing(cleared, selem)
+        if plot == True:
+            plt.imshow(binary, cmap=plt.cm.bone)
+            plt.show()
+        '''    
+        Step 4: Label the image.
+        '''
+        label_image = label(binary)
+        if plot == True:
+            plt.imshow(label_image, cmap=plt.cm.bone)
+            plt.show() 
+        '''
+        Step 5: Keep the labels with 2 largest areas.
+        '''
+        areas = [r.area for r in regionprops(label_image)]
+
+        areas.sort()
+        if len(areas) > 2:
+            for region in regionprops(label_image):
+                if region.area < areas[-2]:
+                    for coordinates in region.coords:                
+                           label_image[coordinates[0], coordinates[1]] = 0
+        binary = label_image > 0
+        if plot == True:
+            plt.imshow(binary, cmap=plt.cm.bone)
+            plt.show() 
+        '''
+        Step 6: Closure operation with a disk of radius 2. This operation is 
+        to keep nodules attached to the lung wall.
+        '''
+        selem = disk(12)
+        binary = binary_closing(binary, selem)
+        if plot == True:
+            plt.imshow(binary, cmap=plt.cm.bone)
+            plt.show()
+        '''
+        Step 7: Fill in the small holes inside the binary mask of lungs.
+        '''
+        edges = roberts(binary)
+        binary = ndi.binary_fill_holes(edges)
+        if plot == True:
+            plt.imshow(binary, cmap=plt.cm.bone)
+            plt.show()
+        '''
+        Step 8: Erosion operation with a disk of radius 2. This operation is 
+        seperate the lung nodules attached to the blood vessels.
+        '''
+        selem = disk(2)
+        binary = binary_erosion(binary, selem)
+        if plot == True:
+            plt.imshow(binary, cmap=plt.cm.bone)
+            plt.show() 
+       
+        '''
+        Step 9: Superimpose the binary mask on the input image.
+        '''
+        get_high_vals = binary == 0
+        im[get_high_vals] = 0
+        if plot == True:
+            plt.imshow(im, cmap=plt.cm.bone)
+            plt.show() 
+            
+        return im
+
+
+class WatershedSegmentation(SegmentationAlgorithm):
+    def __init__(self, threshold=-400):
+        super(WatershedSegmentation, self).__init__(threshold)
+
+    def get_segmented_lungs(self, image, plot=False):
+        # TODO: might add the logic for plotting the filters applied in the process
+        #Creation of the markers as shown above:
+        marker_internal, marker_external, marker_watershed = self.generate_markers(image)
+        
+        #Creation of the Sobel-Gradient
+        sobel_filtered_dx = ndi.sobel(image, 1)
+        sobel_filtered_dy = ndi.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 = ndi.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 = ndi.iterate_structure(blackhat_struct, 8)
+        #Perform the Black-Hat
+        outline += ndi.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 = ndi.morphology.binary_closing(lungfilter, structure=np.ones((5,5)), iterations=3)
+        
+        #Apply the lungfilter (the filtered areas being assigned 0 HU
+        segmented = np.where(lungfilter == 1, image, np.zeros(image.shape))
+
+        return segmented
+
+
+    def generate_markers(self, image):
+        #Creation of the internal Marker
+        marker_internal = image < self._threshold
+        marker_internal = segmentation.clear_border(marker_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 = ndi.binary_dilation(marker_internal, iterations=10)
+        external_b = ndi.binary_dilation(marker_internal, iterations=55)
+        marker_external = external_b ^ external_a
+        #Creation of the Watershed Marker matrix
+        marker_watershed = np.zeros(image.shape, dtype=np.int)
+        marker_watershed += marker_internal * 255
+        marker_watershed += marker_external * 128
+        
+        return marker_internal, marker_external, marker_watershed
+
+
+def get_segmentation_algorithm():
+    if config.SEGMENTATION_ALGO == config.MORPHOLOGICAL_OPERATIONS:
+        return MorphologicalSegmentation()
+    if config.SEGMENTATION_ALGO == config.WATERSHED:
+        return WatershedSegmentation()
+
+    return WatershedSegmentation()
\ No newline at end of file