--- a
+++ b/1 - Methods with Improved Results/SegmentationFunctions.py
@@ -0,0 +1,458 @@
+import os
+import cv2 
+import math
+
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+from matplotlib.patches import Rectangle
+
+from sklearn.cluster import KMeans
+from skimage.morphology import erosion, opening, closing, square, \
+                               disk, convex_hull_image, remove_small_holes
+from skimage.measure import label, regionprops
+
+from skimage.filters import sobel
+from skimage.feature import canny
+from scipy import ndimage as ndi
+from skimage.segmentation import watershed
+    
+SMALL_FONT = 13
+MEDIUM_FONT = 15
+LARGE_FONT = 18
+
+plt.rc('font', size=SMALL_FONT)          # controls default text sizes
+plt.rc('axes', titlesize=SMALL_FONT)     # fontsize of the axes title
+plt.rc('axes', labelsize=MEDIUM_FONT)    # fontsize of the x and y labels
+plt.rc('xtick', labelsize=SMALL_FONT)    # fontsize of the tick labels
+plt.rc('ytick', labelsize=SMALL_FONT)    # fontsize of the tick labels
+plt.rc('legend', fontsize=MEDIUM_FONT)   # legend fontsize
+plt.rc('figure', titlesize=LARGE_FONT)   # fontsize of the figure title
+
+plt.rcParams["figure.figsize"] = (10, 5)
+
+def readSortedSlices(path):
+    
+    slices = []
+    for s in os.listdir(path):
+        slices.append(path + '/' + s)       
+    slices.sort(key = lambda s: int(s[s.find('_') + 1 : s.find('.')]))
+    ID = slices[0][slices[0].find('/') + 1 : slices[0].find('_')]
+    print('CT scan of Patient %s consists of %d slices.' % (ID, len(slices)))  
+    return (slices, ID)
+
+def getSliceImages(slices):
+    
+    return list(map(readImg, slices))
+
+def readImg(path, showOutput=0):
+    
+    img = cv2.imread(path)
+    img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+    
+    if showOutput:
+        plt.title('A CT Scan Image Slice')
+        plt.imshow(img, cmap='gray')
+    return img
+    
+def imgKMeans(img, K, showOutput=0, showHistogram=0):
+    '''
+    Apply KMeans on an image with the number of clusters K
+    Input: Image, Number of clusters K
+    Output: Dictionary of cluster center labels and points, Output segmented image
+    '''
+
+    imgflat = np.reshape(img, img.shape[0] * img.shape[1]).reshape(-1, 1)
+        
+    kmeans = KMeans(n_clusters=K, verbose=0)
+    
+    kmmodel = kmeans.fit(imgflat)
+    
+    labels = kmmodel.labels_
+    centers = kmmodel.cluster_centers_
+    center_labels = dict(zip(np.arange(K), centers))
+    
+    output = np.array([center_labels[label] for label in labels])
+    output = output.reshape(img.shape[0], img.shape[1]).astype(int)
+    
+    if showOutput:
+
+        fig, axes = plt.subplots(1, 2, sharex=True, sharey=True, figsize=(10, 5))
+        axes = axes.ravel()
+
+        axes[0].imshow(img, cmap='gray')
+        axes[0].set_title('Original Image')
+
+        axes[1].imshow(output)
+        axes[1].set_title('Image after KMeans (K = ' + str(K) + ')')
+    
+    return center_labels, output
+
+def preprocessImage(img, showOutput=0):
+    '''
+    Preprocess the image by applying truncated thresholding using KMeans
+    Input: Image
+    Output: Preprocessed image
+    '''
+    centroids, segmented_img = imgKMeans(img, 3, showOutput=0)
+    
+    sorted_center_values = sorted([i[0] for i in centroids.values()])
+    threshold = (sorted_center_values[-1] + sorted_center_values[-2]) / 2
+    
+    retval, procImg = cv2.threshold(img, threshold, 255, cv2.THRESH_TOZERO) 
+    
+    if showOutput:
+        
+        fig, axes = plt.subplots(1, 2, sharex=True, sharey=True, figsize=(10, 5))
+        axes = axes.ravel()
+
+        axes[0].imshow(img, cmap='gray')
+        axes[0].set_title('Original Image')
+
+        axes[1].imshow(procImg, cmap='gray')
+        axes[1].set_title('Processed Image - After Thresholding')
+    
+    return procImg, threshold
+
+
+def getForegroundMask(img, fg_threshold, showOutput=0):
+    
+    retval, init_fg_mask = cv2.threshold(img, fg_threshold, 255, cv2.THRESH_BINARY)
+
+    # Morphological operations to clean the mask
+    fg_mask_opened = opening(init_fg_mask, square(3))
+    fg_mask_opened2 = opening(fg_mask_opened, disk(4))
+    
+    # Perform edge-based segmentation of the foreground...
+    
+    # Detect contours that delineate the foreground with the Canny edge detector
+    edges = canny(fg_mask_opened2) # Background is uniform - edges are on the boundary/inside ROI
+    
+    # Fill the inner part of the boundary using morphology ops
+    fg_mask = ndi.binary_fill_holes(fg_mask_opened2)
+    
+    # Plot all steps
+    if showOutput:
+        fig, axes = plt.subplots(2, 3, sharex=True, sharey=True, figsize=(16, 10))
+        axes = axes.ravel()
+
+        axes[0].set_title('Original Image')
+        axes[0].imshow(img, cmap='gray')
+        axes[1].set_title('On Performing Thresholding\'s')
+        axes[1].imshow(init_fg_mask, cmap='gray')
+        axes[2].set_title('On Opening with Square SE (3)')
+        axes[2].imshow(fg_mask_opened, cmap='gray')
+        axes[3].set_title('On Opening with Disk SE (4)')
+        axes[3].imshow(fg_mask_opened2, cmap='gray')
+        axes[4].set_title('Outer Boundary Delineation with Canny\'s')
+        axes[4].imshow(edges, cmap='gray')
+        axes[5].set_title('Foreground Mask')
+        axes[5].imshow(fg_mask, cmap='gray')
+
+    return fg_mask
+
+
+def getLungTracheaMasks(img, fg_mask, fg_threshold, showOutput=0):
+    
+    # Distinguish black pixels of the background from those of the lungs
+    enhanced = img.copy()
+    
+    for i in range(img.shape[0]):
+        for j in range(img.shape[1]):
+            
+            if fg_mask[i][j] == 0:
+                enhanced[i][j] = 255
+         
+    # Extract lungs from the foreground mask
+    retval, initial_lung_mask = cv2.threshold(enhanced, fg_threshold, 255, cv2.THRESH_BINARY_INV) 
+    
+    # Clean up the lung mask with morphological operations
+    lung_mask_op = opening(initial_lung_mask, square(2))
+    lung_mask_opcl = closing(lung_mask_op, disk(6))
+    lung_mask_opclrm = ndi.binary_fill_holes(lung_mask_opcl)
+
+    # Get connected components of the segmented image and label them
+    label_img = label(lung_mask_opclrm)
+    lung_regions = regionprops(label_img)
+    
+    # Upon experimentation: areas of regions < 1500 are wind pipe structures
+    trachea_labels = []
+    for i in lung_regions:
+        if i.area < 1500:
+            trachea_labels.append(i.label)
+            
+    # Create trachea mask as a summation of all those regions
+    trachea_mask = np.zeros(img.shape, dtype=np.uint8)
+    for row in range(label_img.shape[0]):
+        for col in range(label_img.shape[1]):
+            if label_img[row][col] in trachea_labels:
+                trachea_mask[row][col] = 255
+            
+    # Lung mask is made of all the other regions
+    lung_mask = lung_mask_opclrm * np.invert(trachea_mask) 
+    
+    # Lung mask is all black? Convex hull set to 0, since convex hull op on empty img errors out
+    if sum(sum(lung_mask)) > 0:
+        ch_lung_mask = convex_hull_image(lung_mask)
+    else:
+        ch_lung_mask = lung_mask.copy()
+    
+    
+    initial_int_heart_mask = ch_lung_mask * np.invert(lung_mask) * np.invert(trachea_mask)
+    
+    int_heart_mask_op1 = opening(initial_int_heart_mask, square(5))
+    int_heart_mask_op2 = opening(int_heart_mask_op1, disk(4))
+    
+    heart_label_img = label(int_heart_mask_op2)
+    heart_regions = regionprops(heart_label_img)
+    
+    areas = {}
+    for i in heart_regions:
+        areas[i.label] = i.area
+    
+    if areas:
+        heart_label = max(areas, key=areas.get)
+        int_heart_mask = np.where(heart_label_img==heart_label, np.uint8(255), np.uint8(0))
+    else:
+        int_heart_mask = np.zeros(img.shape, dtype=np.uint8)
+        
+    if showOutput:
+        
+        fig, axes = plt.subplots(4, 3, sharex=True, sharey=True, figsize=(20, 20))
+        axes = axes.ravel()
+
+        axes[0].set_title('Original Image')
+        axes[0].imshow(img, cmap='gray')
+
+        axes[1].set_title('Enhanced Image')
+        axes[1].imshow(enhanced, cmap='gray')
+
+        axes[2].set_title('Initial Lung Mask')
+        axes[2].imshow(initial_lung_mask, cmap='gray')
+
+        axes[3].set_title('On Opening with Square SE (2)')
+        axes[3].imshow(lung_mask_op, cmap='gray')
+
+        axes[4].set_title('On Closing with Disk SE (6)')
+        axes[4].imshow(lung_mask_opcl, cmap='gray')
+
+        axes[5].set_title('On Filling Regions')
+        axes[5].imshow(lung_mask_opclrm, cmap='gray')
+
+        axes[6].set_title('Trachea Mask/Primary Bronchi')
+        axes[6].imshow(trachea_mask, cmap='gray')
+
+        axes[7].set_title('Lung Mask')
+        axes[7].imshow(lung_mask, cmap='gray')
+
+        axes[8].set_title('Convex Hull of Lung Mask')
+        axes[8].imshow(ch_lung_mask, cmap='gray')
+
+        axes[9].set_title('Initial Intermediate Heart Mask')
+        axes[9].imshow(initial_int_heart_mask, cmap='gray')
+
+        axes[10].set_title('On Opening with Square SE (3)')
+        axes[10].imshow(int_heart_mask_op1, cmap='gray')
+
+        axes[11].set_title('Intermediate heart mask')
+        axes[11].imshow(int_heart_mask, cmap='gray')
+    
+    return trachea_mask, lung_mask, ch_lung_mask, int_heart_mask
+
+    
+def chullSpineMask(img, int_heart_mask, showOutput=0):
+    
+    # If no heart
+    if not int_heart_mask.any():
+        return int_heart_mask, int_heart_mask
+    
+    int_heart_pixel = img.copy()
+    
+    for i in range(img.shape[0]):
+        for j in range(img.shape[1]):
+            if int_heart_mask[i][j] == 0:
+                int_heart_pixel[i][j] = 0
+                
+    centroids, segmented_heart_img = imgKMeans(int_heart_pixel, 3, showOutput=0)
+    
+    spine_threshold = (max(centroids.values()))[0]
+
+    retval, initial_spine_mask = cv2.threshold(int_heart_pixel, spine_threshold, 255, cv2.THRESH_BINARY) 
+    
+    bone_mask = closing(initial_spine_mask, disk(20))  
+    
+    label_spine = label(bone_mask)
+    spine_regions = regionprops(label_spine)
+
+    # Assumption: The spine's area is greater than that of any calcium deposits
+    labels = []
+    areas = {}
+    geometric_measures = {}
+    for i in spine_regions:
+        labels.append(i.label)
+        areas[i.label] = i.area
+        geometric_measures[i.label] = [i.centroid, i.orientation, i.axis_major_length]
+    
+    spine_label = max(areas, key=areas.get)
+    labels.remove(spine_label)
+    spine_mask = np.where(label_spine==spine_label, np.uint8(255), np.uint8(0))
+    
+#     if labels:
+#         calcium_deposit_mask = np.where(label_spine==labels[0], np.uint8(255), np.uint8(0))
+#     else:
+#         calcium_deposit_mask = np.zeros(img.shape, dtype=np.uint8)
+    
+    label_heart = label(int_heart_mask)
+    heart_regions = regionprops(label_heart)
+    heart_region_area = heart_regions[0].area
+    spine_region_area = areas[spine_label] 
+    
+    frac_heart = (heart_region_area - spine_region_area)/heart_region_area
+    
+    if frac_heart < 0.5:
+        heart_mask = np.zeros(img.shape, dtype=np.uint8)
+        make_spine_mask = 0
+    else:
+        make_spine_mask = 1
+        
+        # Center point of the spine - get the centroid 
+        y0, x0 = geometric_measures[spine_label][0]
+
+        orientation = geometric_measures[spine_label][1]
+
+        # Top-most point of the spine
+        # top_most_coordinate = centroid - (slightly_more_than_half * axis_major_length * sin(angle))
+        x2 = x0 - math.sin(orientation) * 0.6 * geometric_measures[spine_label][2]
+        y2 = y0 - math.cos(orientation) * 0.6 * geometric_measures[spine_label][2]
+
+        chull_spine_mask = spine_mask.copy()
+
+        # Vertical axis
+        for i in range(math.ceil(y2), img.shape[1]):
+
+            if i > math.ceil(y0):
+                # Horizontal axis
+                for j in range(img.shape[0]):
+                    chull_spine_mask[i][j] = 255
+            else:
+                # Horizontal axis
+                for j in range(math.ceil(x0)):
+                    chull_spine_mask[i][j] = 255
+
+
+        heart_mask = int_heart_mask * np.invert(chull_spine_mask)
+    
+    if showOutput:
+        
+        fig, axes = plt.subplots(3, 3, sharex=True, sharey=True, figsize=(15, 15))
+        axes = axes.ravel()
+
+        axes[0].set_title('Intermediate Heart Mask')
+        axes[0].imshow(int_heart_mask, cmap='gray')
+
+        axes[1].set_title('Intermediate Heart Segment')
+        axes[1].imshow(int_heart_pixel, cmap='gray')
+        
+        axes[2].set_title('Intermediate Heart Segment on K-Means (K = 3)')
+        axes[2].imshow(segmented_heart_img)
+
+        axes[3].set_title('Spine Mask')
+        axes[3].imshow(initial_spine_mask, cmap='gray')
+        
+        axes[4].set_title('On Closing with Disk SE (20)')
+        axes[4].imshow(spine_mask, cmap='gray')
+        
+        axes[5].set_title('On Opening with Square SE (4)')
+        axes[5].imshow(spine_mask, cmap='gray')
+        
+        axes[6].set_title('Centroid and uppermost point')
+        axes[6].imshow(spine_mask, cmap='gray')
+        
+        if make_spine_mask:
+            axes[6].plot((x0, x2), (y0, y2), '-r', linewidth=1.5)
+            axes[6].plot(x0, y0, '.g', markersize=5)
+            axes[6].plot(x2, y2, '.b', markersize=5)
+        
+            axes[7].set_title('Convex Hull of Spine Mask')
+            axes[7].imshow(chull_spine_mask, cmap='gray')
+        else:
+            axes[7].set_title('Convex Hull of Spine Mask')
+            axes[7].imshow(chull_spine_mask, cmap='gray')
+        
+        axes[8].set_title('Heart Mask')
+        axes[8].imshow(heart_mask, cmap='gray')
+        
+    return spine_mask, heart_mask
+
+def segmentHeart(img, heart_mask, showOutput=0):
+
+	seg_heart = cv2.bitwise_and(img, img, mask=heart_mask)
+	
+	if showOutput:
+		plt.figure(figsize=(10, 5))
+		plt.title('Segmented Heart')
+		plt.imshow(seg_heart, cmap='gray')
+	return seg_heart
+	
+def segmentHeartLungsTrachea(img, heart_mask, lung_mask, trachea_mask, showOutput=0):
+
+    seg_heart = cv2.bitwise_and(img, img, mask=heart_mask)
+    seg_lungs = cv2.bitwise_and(img, img, mask=lung_mask)
+    seg_trachea = cv2.bitwise_and(img, img, mask=trachea_mask)
+    
+    if showOutput:
+        fig, [ax1, ax2, ax3] = plt.subplots(1, 3, figsize=(12, 6), sharex=True, sharey=False)
+        
+        ax1.set_title('Segmented Heart')
+        ax1.imshow(seg_heart, cmap='gray')
+        
+        ax2.set_title('Segmented Lungs')
+        ax2.imshow(seg_lungs, cmap='gray')
+        
+        ax3.set_title('Segmented Trachea')
+        ax3.imshow(seg_trachea, cmap='gray')
+        
+    return seg_heart, seg_lungs, seg_trachea
+    
+
+def applyMaskColor(mask, mask_color):
+    
+    masked = np.concatenate(([mask[ ... , np.newaxis] * color for color in mask_color]), axis=2)
+    
+    # Matplotlib expects color intensities to range from 0 to 1 if a float
+    maxValue = np.amax(masked)
+    minValue = np.amin(masked)
+
+    # Therefore, scale the color image accordingly
+    if maxValue - minValue == 0:
+        return masked
+    else:
+        masked = masked / (maxValue - minValue)
+    
+    return masked
+
+def getColoredMasks(img, heart_mask, lung_mask, trachea_mask, showOutput=0):
+    heart_mask_color = np.array([256, 0, 0])
+    lung_mask_color = np.array([0, 256, 0])
+    trachea_mask_color = np.array([0, 0, 256])
+
+    heart_colored = applyMaskColor(heart_mask, heart_mask_color)
+    lung_colored = applyMaskColor(lung_mask, lung_mask_color)
+    trachea_colored = applyMaskColor(trachea_mask, trachea_mask_color)
+    
+    colored_masks = heart_colored + lung_colored + trachea_colored
+
+    if showOutput:
+        fig, axes = plt.subplots(2, 2, figsize=(10, 10))
+        ax = axes.ravel()
+        
+        ax[0].set_title("Original Image")
+        ax[0].imshow(img, cmap='gray')
+        ax[1].set_title("Heart Mask")
+        ax[1].imshow(heart_colored)
+        ax[2].set_title("Lung Mask")
+        ax[2].imshow(lung_colored)
+        ax[3].set_title("Masks")
+        ax[3].imshow(colored_masks)
+    
+    return heart_colored, lung_colored, trachea_colored, colored_masks