Create trimap_module.py

trimap_module.py

@@ -0,0 +1,164 @@
+#!/usr/bin/env python
+import cv2, os, sys
+import numpy as np
+
+def extractImage(path):
+    # error handller if the intended path is not found
+    image = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
+    return image
+
+def checkImage(image):
+    """
+    Args:
+        image: input image to be checked
+    Returns:
+        binary image
+    Raises:
+        RGB image, grayscale image, all-black, and all-white image
+
+    """
+    if len(image.shape) > 2:
+        print("ERROR: non-binary image (RGB)"); sys.exit();
+
+    smallest = image.min(axis=0).min(axis=0) # lowest pixel value: 0 (black)
+    largest  = image.max(axis=0).max(axis=0) # highest pixel value: 1 (white)
+
+    if (smallest == 0 and largest == 0):
+        print("ERROR: non-binary image (all black)"); sys.exit()
+    elif (smallest == 255 and largest == 255):
+        print("ERROR: non-binary image (all white)"); sys.exit()
+    elif (smallest > 0 or largest < 255 ):
+        print("ERROR: non-binary image (grayscale)"); sys.exit()
+    else:
+        return True
+
+class Toolbox:
+    def __init__(self, image):
+        self.image = image
+
+    @property
+    def printImage(self):
+        """
+        Print image into a file for checking purpose
+        unitTest = Toolbox(image);
+        unitTest.printImage(image);
+        """
+        f = open("image_results.dat", "w+")
+        for i in range(0, self.image.shape[0]):
+            for j in range(0, self.image.shape[1]):
+                f.write("%d " %self.image[i,j])
+            f.write("\n")
+        f.close()
+        
+    @property
+    def displayImage(self):
+        """
+        Display the image on a window
+        Press any key to exit
+        """
+        cv2.imshow('Displayed Image', self.image)
+        cv2.waitKey(0)
+        cv2.destroyAllWindows()
+
+    def saveImage(self, title, extension):
+        """
+        Save as a specific image format (bmp, png, or jpeg)
+        """
+        cv2.imwrite("{}.{}".format(title,extension), self.image)
+
+    def morph_open(self, image, kernel):
+        """
+        Remove all white noises or speckles outside images
+        Need to tune the kernel size
+        Instruction:
+        unit01 = Toolbox(image);
+        kernel = np.ones( (9,9), np.uint8 );
+        morph  = unit01.morph_open(input_image, kernel);
+        """
+        bin_open = cv2.morphologyEx(self.image, cv2.MORPH_OPEN, kernel)
+        return bin_open
+
+    def morph_close(self, image, kernel):
+        """
+        Remove all black noises or speckles inside images
+        Need to tune the kernel size
+        Instruction:
+        unit01 = Toolbox(image);
+        kernel = np.ones( (11,11)_, np.uint8 );
+        morph  = unit01.morph_close(input_image, kernel);
+        """        
+        bin_close = cv2.morphologyEx(self.image, cv2.MORPH_CLOSE, kernel)
+        return bin_close
+
+
+def trimap(image, name, size, number, erosion=False):
+    """
+    This function creates a trimap based on simple dilation algorithm
+    Inputs [4]: a binary image (black & white only), name of the image, dilation pixels
+                the last argument is optional; i.e., how many iterations will the image get eroded
+    Output    : a trimap
+    """
+    checkImage(image)
+    row    = image.shape[0]
+    col    = image.shape[1]
+    pixels = 2*size + 1      ## Double and plus 1 to have an odd-sized kernel
+    kernel = np.ones((pixels,pixels),np.uint8)   ## Pixel of extension I get
+
+    if erosion is not False:
+        erosion = int(erosion)
+        erosion_kernel = np.ones((3,3), np.uint8)                     ## Design an odd-sized erosion kernel
+        image = cv2.erode(image, erosion_kernel, iterations=erosion)  ## How many erosion do you expect
+        image = np.where(image > 0, 255, image)                       ## Any gray-clored pixel becomes white (smoothing)
+        # Error-handler to prevent entire foreground annihilation
+        if cv2.countNonZero(image) == 0:
+            print("ERROR: foreground has been entirely eroded")
+            sys.exit()
+
+    dilation  = cv2.dilate(image, kernel, iterations = 1)
+
+    dilation  = np.where(dilation == 255, 127, dilation) 	## WHITE to GRAY
+    remake    = np.where(dilation != 127, 0, dilation)		## Smoothing
+    remake    = np.where(image > 127, 200, dilation)		## mark the tumor inside GRAY
+
+    remake    = np.where(remake < 127, 0, remake)		## Embelishment
+    remake    = np.where(remake > 200, 0, remake)		## Embelishment
+    remake    = np.where(remake == 200, 255, remake)		## GRAY to WHITE
+
+    #############################################
+    # Ensures only three pixel values available #
+    # TODO: Optimization with Cython            #
+    #############################################    
+    for i in range(0,row):
+        for j in range (0,col):
+            if (remake[i,j] != 0 and remake[i,j] != 255):
+                remake[i,j] = 127
+
+    path = "./images/results/"  ## Change the directory
+    new_name = '{}px_'.format(size) + name + '_{}.png'.format(number)
+    cv2.imwrite(os.path.join(path, new_name) , remake)
+
+
+#############################################
+###             TESTING SECTION           ###
+#############################################
+if __name__ == '__main__':
+    path  = "./images/test_images/test_image_11.png"
+    image = extractImage(path)
+
+    size = 10
+    number = path[-5]
+    title = "test_image"
+
+    unit01  = Toolbox(image);
+    kernel1 = np.ones( (11,11), np.uint8 )
+    unit01.displayImage
+    
+    opening = unit01.morph_close(image,kernel1)
+    trimap(opening, title, size, number, erosion=False)
+    unit02  = Toolbox(opening)
+    unit02.displayImage
+
+    ########################################################
+    ## Default instruction (no binary opening or closing  ##
+    ## trimap(image, title, size, number, erosion=False); ##
+    ########################################################