--- a
+++ b/2 - KMeans & Morphology Methods - Rim et al/SegmentationFunctions.py
@@ -0,0 +1,362 @@
+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
+from skimage.measure import label, regionprops
+	
+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)
+	
+#	print(len(labels), 'Labels: \n', labels)
+#	print(len(centers), 'Centers: \n', centers)
+#	print('Center labels and their points:', center_labels)
+	
+	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=showOutput)
+	
+	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
+
+def chullForegroundMask(img, showOutput=0):
+	'''
+	Generate a convex hull of foreground mask
+	Input: Whole chest CT image in grayscale format (512x512)
+	Output: Convex hull of foreground mask in binary format (512x512)
+	'''
+	
+	centroid_clusters, segmented_img = imgKMeans(img, 2, showOutput=showOutput)
+	
+	fg_threshold = sum(centroid_clusters.values())[0] / 2
+	
+	retval, fg_mask = cv2.threshold(img, fg_threshold, 255, cv2.THRESH_BINARY) 
+	
+	fg_mask_opened = opening(fg_mask, square(3))
+	fg_mask_opened2 = opening(fg_mask_opened, disk(4))
+	
+	fg_mask = fg_mask_opened2
+	
+	ch_fg_mask = convex_hull_image(fg_mask)
+	
+	if showOutput:
+		fig, axes = plt.subplots(1, 3, sharex=True, sharey=True, figsize=(16, 10))
+		axes = axes.ravel()
+
+		axes[0].set_title('On Opening with Square SE (3)')
+		axes[0].imshow(fg_mask_opened, cmap='gray')
+
+		axes[1].set_title('On Opening with Disk SE (4)')
+		axes[1].imshow(fg_mask_opened2, cmap='gray')
+
+		axes[2].set_title('Convex Hull of Foreground Mask')
+		axes[2].imshow(ch_fg_mask, cmap='gray')
+	
+	return fg_mask, ch_fg_mask, fg_threshold
+
+def chullLungMask(img, ch_fg_mask, fg_threshold, showOutput=0):
+	'''
+	Generate a convex hull of lung mask
+	Input: Whole chest CT image in grayscale format (512x512),
+		   Convex hull of foreground mask in binary format (512x512)
+	Output: Convex hull of lung mask in binary format (512x512),
+			Intermediate heart mask in binary format (512x512)
+	'''
+	
+	enhanced = img.copy()
+	
+	for i in range(img.shape[0]):
+		for j in range(img.shape[1]):
+			
+			if ch_fg_mask[i][j] == 0:
+				enhanced[i][j] = 255
+		 
+	retval, initial_lung_mask = cv2.threshold(enhanced, fg_threshold, 255, cv2.THRESH_BINARY_INV) 
+	
+	lung_mask_op1 = opening(initial_lung_mask, square(2))
+	lung_mask_op1cl1 = closing(lung_mask_op1, disk(8))	
+	lung_mask_op2cl1 = opening(lung_mask_op1cl1, square(8))
+	lung_mask_op3cl1 = opening(lung_mask_op2cl1, disk(8))
+	lung_mask = closing(lung_mask_op3cl1, disk(16))
+
+	ch_lung_mask = convex_hull_image(lung_mask)
+	
+	initial_int_heart_mask = ch_lung_mask * np.invert(lung_mask)
+	int_heart_mask_op1 = opening(initial_int_heart_mask, square(3))
+	int_heart_mask = opening(int_heart_mask_op1, disk(3))
+	
+	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_op1, cmap='gray')
+
+		axes[4].set_title('On Closing with Disk SE (8)')
+		axes[4].imshow(lung_mask_op1cl1, cmap='gray')
+
+		axes[5].set_title('On Opening with Square SE (8)')
+		axes[5].imshow(lung_mask_op2cl1, cmap='gray')
+
+		axes[6].set_title('On Opening with Disk SE (8)')
+		axes[6].imshow(lung_mask_op3cl1, cmap='gray')
+
+		axes[7].set_title('On Closing with Disk SE (16)')
+		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 lung_mask, ch_lung_mask, int_heart_mask
+
+def chullSpineMask(img, int_heart_mask, showOutput=0):
+	
+	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) 
+	
+	spine_mask_cl1 = closing(initial_spine_mask, disk(20))
+	spine_mask = opening(spine_mask_cl1, square(4))	
+	
+	label_img = label(spine_mask)
+	spine_regions = regionprops(label_img)
+
+	# Center point of the spine
+	y0, x0 = spine_regions[0].centroid
+	
+	orientation = spine_regions[0].orientation
+
+	# Top-most point of the spine
+	x2 = x0 - math.sin(orientation) * 0.6 * spine_regions[0].axis_major_length
+	y2 = y0 - math.cos(orientation) * 0.6 * spine_regions[0].axis_major_length
+
+	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_cl1, 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')
+		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')
+		
+		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 segmentLungs(img, lung_mask, showOutput=0): 
+	
+	seg_lungs = cv2.bitwise_and(img, img, mask=lung_mask)
+	
+	if showOutput:
+		plt.figure(figsize=(10, 5))
+		plt.title('Segmented Lungs')
+		plt.imshow(seg_lungs, cmap='gray')
+	return seg_lungs
+
+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
+	masked = masked / (maxValue - minValue)
+	
+	return masked
+
+def getColoredMasks(img, heart_mask, lung_mask, showOutput=0):
+	heart_mask_color = np.array([256, 0, 0])
+	lung_mask_color = np.array([0, 256, 0])
+
+	heart_colored = applyMaskColor(heart_mask, heart_mask_color)
+	lung_colored = applyMaskColor(lung_mask, lung_mask_color)
+	colored_masks = heart_colored + lung_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, colored_masks