Bug Fix in Image Generator - Now Allows FCN's to be used

keras/preprocessing/image.py

@@ -191,14 +191,17 @@ def flip_axis(x, axis):
     return x
 
 
-def array_to_img(x, dim_ordering='default', scale=True):
+def array_to_img(x, dim_ordering='default', scale=True, color_mode='rgb'):
     """Converts a 3D Numpy array to a PIL Image instance.
 
     # Arguments
         x: Input Numpy array.
         dim_ordering: Image data format.
         scale: Whether to rescale image values
             to be within [0, 255].
+        color_mode: 'bgr' to reverse the
+            channel dimensions, otherwise it picks
+            depending on the shape of X.
 
     # Returns
         A PIL Image instance.
@@ -232,6 +235,8 @@ def array_to_img(x, dim_ordering='default', scale=True):
             x /= x_max
         x *= 255
     if x.shape[2] == 3:
+        if color_mode == 'bgr':
+            x = x[:, :, ::-1]
         # RGB
         return pil_image.fromarray(x.astype('uint8'), 'RGB')
     elif x.shape[2] == 1:
@@ -241,12 +246,15 @@ def array_to_img(x, dim_ordering='default', scale=True):
         raise ValueError('Unsupported channel number: ', x.shape[2])
 
 
-def img_to_array(img, dim_ordering='default'):
+def img_to_array(img, dim_ordering='default', color_mode='rgb'):
     """Converts a PIL Image instance to a Numpy array.
 
     # Arguments
         img: PIL Image instance.
         dim_ordering: Image data format.
+        color_mode: 'bgr' to reverse the
+            channel dimensions, otherwise it picks
+            depending on the shape of X.
 
     # Returns
         A 3D Numpy array (float32).
@@ -262,6 +270,8 @@ def img_to_array(img, dim_ordering='default'):
     # or (channel, height, width)
     # but original PIL image has format (width, height, channel)
     x = np.asarray(img, dtype='float32')
+    if color_mode == 'bgr':
+        x = x[:, :, ::-1]
     if len(x.shape) == 3:
         if dim_ordering == 'th':
             x = x.transpose(2, 0, 1)
@@ -275,7 +285,7 @@ def img_to_array(img, dim_ordering='default'):
     return x
 
 
-def load_img(path, grayscale=False, target_size=None):
+def load_img(path, grayscale=False, target_size=(None, None)):
     """Loads an image into PIL format.
 
     # Arguments
@@ -298,7 +308,7 @@ def load_img(path, grayscale=False, target_size=None):
         img = img.convert('L')
     else:  # Ensure 3 channel even when loaded image is grayscale
         img = img.convert('RGB')
-    if target_size:
+    if target_size[0] and target_size[1]:  # Can't pass None in from iterator, hence tuple
         img = img.resize((target_size[1], target_size[0]))
     return img
 
@@ -742,9 +752,9 @@ def __init__(self, directory, image_data_generator,
         self.directory = directory
         self.image_data_generator = image_data_generator
         self.target_size = tuple(target_size)
-        if color_mode not in {'rgb', 'grayscale'}:
+        if color_mode not in {'rgb', 'bgr', 'grayscale'}:
             raise ValueError('Invalid color mode:', color_mode,
-                             '; expected "rgb" or "grayscale".')
+                             '; expected "rgb", "bgr" or "grayscale".')
         self.color_mode = color_mode
         self.dim_ordering = dim_ordering
         if self.color_mode == 'rgb':
@@ -777,6 +787,9 @@ def __init__(self, directory, image_data_generator,
             for subdir in sorted(os.listdir(directory)):
                 if os.path.isdir(os.path.join(directory, subdir)):
                     classes.append(subdir)
+        # User passed in None and no sub-folders, hence FCN. N.B.: Pass in classes=['.'] for improvements.
+        if not classes:
+            classes = [directory]
         self.nb_class = len(classes)
         self.class_indices = dict(zip(classes, range(len(classes))))
 
@@ -822,22 +835,30 @@ def next(self):
             index_array, current_index, current_batch_size = next(self.index_generator)
         # The transformation of images is not under thread lock
         # so it can be done in parallel
-        batch_x = np.zeros((current_batch_size,) + self.image_shape)
+        # If image_shape has None in position 1 it is intended to be used as a FCN
+        if self.image_shape[1]:
+            batch_x = np.zeros((current_batch_size,) + self.image_shape)
         grayscale = self.color_mode == 'grayscale'
         # build batch of image data
         for i, j in enumerate(index_array):
             fname = self.filenames[j]
             img = load_img(os.path.join(self.directory, fname),
                            grayscale=grayscale,
                            target_size=self.target_size)
-            x = img_to_array(img, dim_ordering=self.dim_ordering)
+            x = img_to_array(img, dim_ordering=self.dim_ordering, color_mode=self.color_mode)
+            # Allows shape to be None, but batch size must equal one.
+            if not self.image_shape[1] and i == 0:
+                if current_batch_size != 1:
+                    raise ValueError('Current batch size is: %s , should be size one for current shape'
+                                     % current_batch_size)
+                batch_x = np.zeros((current_batch_size,) + x.shape)
             x = self.image_data_generator.random_transform(x)
             x = self.image_data_generator.standardize(x)
             batch_x[i] = x
         # optionally save augmented images to disk for debugging purposes
         if self.save_to_dir:
             for i in range(current_batch_size):
-                img = array_to_img(batch_x[i], self.dim_ordering, scale=True)
+                img = array_to_img(batch_x[i], self.dim_ordering, scale=True, color_mode=self.color_mode)
                 fname = '{prefix}_{index}_{hash}.{format}'.format(prefix=self.save_prefix,
                                                                   index=current_index + i,
                                                                   hash=np.random.randint(1e4),

tests/keras/preprocessing/test_image.py

@@ -160,6 +160,33 @@ def test_directory_iterator(self):
         assert(sorted(dir_iterator.filenames) == sorted(filenames))
         shutil.rmtree(tmp_folder)
 
+    ''' Test for the directory iterator when we wish to use it for a FCN.'''
+    def test_directory_iterator_fcn(self):
+        num_classes = 1  # 1 class encompassing all masks
+        tmp_folder = tempfile.mkdtemp(prefix='test_images')
+
+        # no subfolders
+
+        # save the images in the paths
+        count = 0
+        filenames = []
+        for test_images in self.all_test_images:
+            for im in test_images:
+                filename = 'image-{}.jpg'.format(count)
+                filenames.append(filename)
+                im.save(os.path.join(tmp_folder, filename))
+                count += 1
+
+        # create iterator
+        generator = image.ImageDataGenerator()
+        dir_iterator = generator.flow_from_directory(tmp_folder, classes=['.'])
+
+        # check number of classes and images
+        assert (len(dir_iterator.class_indices) == num_classes)
+        assert (len(dir_iterator.classes) == count)
+        assert (sorted(dir_iterator.filenames) == sorted(filenames))
+        shutil.rmtree(tmp_folder)
+
     def test_img_utils(self):
         height, width = 10, 8