Merge branch 'feature/voc-masks' into develop

AngusG · Jan 24, 2017 · d5d00a9 · d5d00a9
2 parents 834fbc4 4d6986f
commit d5d00a9
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 9 deletions.
diff --git a/VOClabelcolormap.py b/VOClabelcolormap.py
@@ -0,0  1,45 @@
 """ 
 Python implementation of the color map function for the PASCAL VOC data set. 
 Official Matlab version can be found in the PASCAL VOC devkit 
 http://host.robots.ox.ac.uk/pascal/VOC/voc2012/index.html#devkit
 """
 import numpy as np
 from skimage.io import imshow
 import matplotlib.pyplot as plt
 
 def color_map(N=256, normalized=False):
     def bitget(byteval, idx):
         return ((byteval & (1 << idx)) != 0)
 
     dtype = 'float32' if normalized else 'uint8'
     cmap = np.zeros((N, 3), dtype=dtype)
     for i in range(N):
         r = g = b = 0
         c = i
         for j in range(8):
             r = r | (bitget(c, 0) << 7-j)
             g = g | (bitget(c, 1) << 7-j)
             b = b | (bitget(c, 2) << 7-j)
             c = c >> 3
 
         cmap[i] = np.array([r, g, b])
 
     cmap = cmap/255 if normalized else cmap
     return cmap
 
 
 def color_map_viz():
     labels = ['background', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor', 'void']
     nclasses = 21
     row_size = 50
     col_size = 500
     cmap = color_map()
     array = np.empty((row_size*(nclasses 1), col_size, cmap.shape[1]), dtype=cmap.dtype)
     for i in range(nclasses):
         array[i*row_size:i*row_size row_size, :] = cmap[i]
     array[nclasses*row_size:nclasses*row_size row_size, :] = cmap[-1]
 
     imshow(array)
     plt.yticks([row_size*i row_size/2 for i in range(nclasses 1)], labels)
     plt.xticks([])
     plt.show()
diff --git a/truth_and_crop.py b/truth_and_crop.py
@@ -11,20  11,23 @@
 import sys
 from PyQt4 import QtCore, QtGui, uic
 from PyQt4.QtGui import *
 from VOClabelcolormap import color_map
 
 qtCreatorFile = "truth_and_crop_qt4.ui"
 
 # Constants
 APP_NAME = 'Truth and Crop'
 IMAGES_OUT_DIR = 'images/'
-MASKS_OUT_DIR = 'masks/'
 INT_MASKS_OUT_DIR = 'masks/'
 RGB_MASKS_OUT_DIR = 'PASCALVOC-style-masks/'
 VALID_EXT = '.JPG'  # File extension to consider valid when searching for prv/next image
 IMAGE_EXT = '.jpg'  # Output file extension
 MASK_EXT = '_mask.jpg'
 PX_INTENSITY = 0.4
 N_CHANNELS = 2
 
 # Constants - class labels
 NCLASSES = 4
 CLASS_OTHER = 0
 CLASS_MUSSEL = 1
 CLASS_CIONA = 2
@@ -139,13  142,15 @@ def __handle_crop_btn(self, event):
     def __handle_done_btn(self, event):
 
         image_path = os.path.join(self.outputFolder, IMAGES_OUT_DIR)
-        mask_path = os.path.join(self.outputFolder, MASKS_OUT_DIR)
         int_mask_path = os.path.join(self.outputFolder, INT_MASKS_OUT_DIR)
         rgb_mask_path = os.path.join(self.outputFolder, RGB_MASKS_OUT_DIR)
 
         if not os.path.exists(image_path):
             os.makedirs(image_path)
-
-        if not os.path.exists(mask_path):
-            os.makedirs(mask_path)
         if not os.path.exists(int_mask_path):
             os.makedirs(int_mask_path)
         if not os.path.exists(rgb_mask_path):
             os.makedirs(rgb_mask_path)
 
         self.original = cv2.cvtColor(self.original, cv2.COLOR_RGB2BGR)
 
@@ -160,13  165,24 @@ def __handle_done_btn(self, event):
             # Set all pixels in super_px to p_class.
             self.segmentation_mask[self.segments == super_px] = p_class
 
         # Make PASCAL fmt segmentation_mask as well
         cmap = color_map(NCLASSES)
         height, width, __ = self.original.shape
 
         # Initialize empty RGB array
         array = np.empty((height, width, cmap.shape[1]), dtype=cmap.dtype)
         #array = np.zeros((height, width, cmap.shape[1]), dtype=cmap.dtype)
         array = cv2.cvtColor(array, cv2.COLOR_RGB2BGR)
 
         # Convert integers in segmentation_mask to rgb vals
         for i in range(NCLASSES):
             array[self.segmentation_mask == i] = cmap[i] 
 
         for i, (x, y) in enumerate(crop_list):
 
             # Detailed cropped image suffix.
             details = self.__generate_image_details(img_name, i, x, y)
 
-            height, width, __ = self.original.shape
-
             # if y - self.w > 0 and y   self.w < height and x - self.w > 0
             # and x   self.w < width:
             y_lwr = y - self.w > 0
@@ -177,13  193,17 @@ def __handle_done_btn(self, event):
 
                 cropped_image = self.original[
                     y - self.w:y   self.w, x - self.w:x   self.w, :]
-                cropped_mask = self.segmentation_mask[
                 cropped_int_mask = self.segmentation_mask[
                     y - self.w:y   self.w, x - self.w:x   self.w]
                 cropped_rgb_mask = array[
                     y - self.w:y   self.w, x - self.w:x   self.w]
 
                 cv2.imwrite(os.path.join(
                     image_path, details   IMAGE_EXT), cropped_image)
                 cv2.imwrite(os.path.join(
-                    mask_path, details   IMAGE_EXT), cropped_mask)
                     int_mask_path, details   IMAGE_EXT), cropped_int_mask)
                 cv2.imwrite(os.path.join(
                     rgb_mask_path, details   IMAGE_EXT), cropped_rgb_mask)
 
                 print('Success: cropped image at x=%d,y=%d with wnd=%d' %
                       (x, y, self.w))