I built a class to help me handle image data to use in machine learning. I thought that there would be a pre-existing package that did what I wanted but I couldn't find it so I wrote this. I am not intentionally trying to re-invent the wheel so if there's something that already does this please let me know (although I would still be interested in how I could make this better).
The main goal is the have a class of functions that can read images files from directories and convert them into training and testing sets ready for machine learning. I want to have the flexibility to return the data in any of the following forms:
- grayscale or rgb
- flattened vectors or not
- any square image size
- rescaled, standardized, or not
- labels as either column vector or not ( (n,) or (n, 1) )
- number of samples either as first or last in ndarray
I use this class to accept directories in the following format:
data │ └───train │ └───image1 │ │ im1.jpg │ │ im2.jpg │ └───image2 │ im1.jpg │ im2.jpg └───test │ └───image1 │ │ im3.jpg │ │ im4.jpg │ └───image2 │ im3.jpg │ im4.jpg
then return data in any format I want, including (N, l, w, 3), (N, l*w*1), (l*w*3, N), etc.
train_test_sets is the meat of the class, but there are other helper functions as well.
from PIL import Image import os import numpy as np from keras.preprocessing.image import array_to_img, img_to_array, load_img import random class Gather_Data(object): def __init__(self): self.X_train_ = None self.X_test_ = None self.y_train_ = None self.y_test_ = None self.image_size_ = None self.num_train_images_ = None self.num_test_images_ = None def get_filenames(self, path): ''' Returns list of filenames in a path ''' # os.path.join will add the trailing slash if it's not already there files = [file for file in os.listdir( path) if os.path.isfile(os.path.join(path, file))] return files def get_images(self, path, result_format='list of PIL images', new_size=0, grayscale=True): ''' Accepts a path to a directory of images and returns an ndarray of shape N, H, W, c where N is the number of images H is the height of the images W is the width of the images\ c is 3 if RGB and 1 if grayscale result can be "ndarray" for a single large ndarray, "list of ndarrays", or list of PIL Images (PIL.Image.Image) If a new_size is added, it must be square This function also allows the images to be resized, but forces square ''' files = self.get_filenames(path) images =  for file in files: image = Image.open(os.path.join(path, file)) if grayscale: image = image.convert("L") if new_size != 0: image = image.resize((new_size, new_size), Image.ANTIALIAS) if result_format == 'ndarray' or result_format == 'list of ndarrays': image = np.array(image) images.append(image) if result_format == 'ndarray': return np.asarray(images) else: return images def make_dir_if_needed(self, folder): ''' Checks if a directory already exists and if not creates it ''' if not os.path.isdir(folder): os.makedirs(folder) def augment_images(self, original_file, output_path, output_prefix, image_number, datagen, count=10): ''' This function works on a single image at a time. It works best by enumerating a list of file names and passing the file and index. original_file must be the full path to the file, not just the filename The image_number should be the index from the enumeration e.g.: for index, file in enumerate(train_files): augment_images(os.path.join(train_path, file), output_path, str(index), datagen, count=10) ''' self.make_dir_if_needed(output_path) # load image to array image = img_to_array(load_img(original_file)) # set_trace() # reshape to array rank 4 image = image.reshape((1,) + image.shape) # let's create infinite flow of images images_flow = datagen.flow(image, batch_size=1) for index, new_images in enumerate(images_flow): if index >= count: break # we access only first image because of batch_size=1 new_image = array_to_img(new_images, scale=True) output_filename = output_path + output_prefix + image_number + \ '-' + str(index+1) + '.jpg' new_image.save(output_filename) def train_test_sets(self, input1_training_path, input2_training_path, input1_testing_path, input2_testing_path, new_size=256, grayscale=False, num_samples_last=False, standardization='normalize', seed=None, verbose=False, y_as_column_vector=False, flatten=True): ''' This assumes the data arrives in the form (N, H*W*c) where c is color color is 3 for RGB or 1 for grayscale To leave the images at their original size pass `new_size = 0` ''' # Get an ndarray of each group of images # Array should be N * H * W * c train1 = self.get_images( input1_training_path, result_format='ndarray', new_size=new_size, grayscale=grayscale) train2 = self.get_images( input2_training_path, result_format='ndarray', new_size=new_size, grayscale=grayscale) test1 = self.get_images( input1_testing_path, result_format='ndarray', new_size=new_size, grayscale=grayscale) test2 = self.get_images( input2_testing_path, result_format='ndarray', new_size=new_size, grayscale=grayscale) self.image_size_ = (new_size, new_size) # make sure the image is square assert train1.shape == train1.shape == new_size # Now we have an array of images N * W * H * 3 or N * W * H * 1 if flatten: if verbose: print("flattening") # Flatten the arrays if grayscale: flattened_size = new_size * new_size else: flattened_size = new_size * new_size * 3 train1 = train1.reshape(train1.shape, flattened_size) train2 = train2.reshape(train2.shape, flattened_size) test1 = test1.reshape(test1.shape, flattened_size) test2 = test2.reshape(test2.shape, flattened_size) # Combine the two different inputs into a single training set training_images = np.concatenate((train1, train2), axis=0) # Do same for testing set testing_images = np.concatenate((test1, test2), axis=0) # Get the number of training and testing examples self.num_train_images_ = len(training_images) self.num_test_images_ = len(testing_images) # Create labels training_labels = np.concatenate( (np.zeros(len(train1)), np.ones(len(train2)))) testing_labels = np.concatenate( (np.zeros(len(test1)), np.ones(len(test2)))) # Zip the images and labels together so they can be shuffled together if verbose: print("zipping") train_zipped = list(zip(training_images, training_labels)) test_zipped = list(zip(testing_images, testing_labels)) if verbose: print("shuffling") # Now shuffle both random.seed(seed) random.shuffle(train_zipped) random.shuffle(test_zipped) self.X_train_, self.y_train_ = zip(*train_zipped) self.X_test_, self.y_test_ = zip(*test_zipped) # Convert tuples back to ndarrays self.X_train_ = np.asarray(self.X_train_) self.X_test_ = np.asarray(self.X_test_) self.y_train_ = np.asarray(self.y_train_) self.y_test_ = np.asarray(self.y_test_) if standardization == 'normalize': if verbose: print("standardizing") # Standardize the values self.X_train_ = (self.X_train_ - self.X_train_.mean() ) / self.X_train_.std() # Use the train mean and standard deviation self.X_test_ = (self.X_test_ - self.X_train_.mean() ) / self.X_train_.std() elif standardization == 'rescale': if verbose: print("standardizing") # Standardize the values self.X_train_ = self.X_train_ / 255. # Use the train mean and standard deviation self.X_test_ = self.X_test_ / 255. if y_as_column_vector: # Reshape the y to matrix them n X 1 matricies self.y_train_ = self.y_train_.reshape(self.y_train_.shape, 1) self.y_test_ = self.y_test_.reshape(self.y_test_.shape, 1) if num_samples_last: # Code conversion for class self.X_train_.shape = ( self.X_train_.shape, self.X_train_.shape) self.X_test_.shape = (self.X_test_.shape, self.X_test_.shape) self.y_train_.shape = ( self.y_train_.shape, self.y_train_.shape) self.y_test_.shape = (self.y_test_.shape, self.y_test_.shape) def dataset_parameters(self): ''' Returns the parameters of the dataset ''' try: print("X_train shape: " + str(self.X_train_.shape)) print("y_train shape: " + str(self.y_train_.shape)) print("X_test shape: " + str(self.X_test_.shape)) print("y_test shape: " + str(self.y_test_.shape)) print("Number of training examples: " + str(self.num_train_images_)) print("Number of testing examples: " + str(self.num_test_images_)) print("Each image is of size: " + str(self.image_size_)) except AttributeError: print("Error: The data has not been input or is incorrectly configured.")
Here are some sample use cases:
augment_images uses the augmentation functionality of Keras, but gives me more flexibility. I know Keras has
save_prefix arguments, but I wanted to control exactly which images were augmented, how many times they were augmented, and what their files names are.
my_data = Gather_Data() image_path = 'img/' aug_path = 'aug/' filenames = my_data.get_filenames(image_path) # This import is down here because it would normally be in a separate file from keras.preprocessing.image import ImageDataGenerator datagen = ImageDataGenerator( rotation_range=45 ) for index, file in enumerate(filenames): my_data.augment_images(os.path.join(image_path, file), aug_path, 'augmented_image', str(index), datagen, count=2)
im1_train_path = 'data/train/im1/' im2_train_path = 'data/train/im2/' im1_test_path = 'data/test/im1/' im2_test_path = 'data/test/im2/' datagen = ImageDataGenerator( rescale=1. / 255, shear_range=0.2, zoom_range=0.2, horizontal_flip=True) my_data.train_test_sets(im1_train_path, im2_train_path, im1_test_path, im2_test_path) X_train = my_data.X_train_ y_train = my_data.y_train_ X_test = my_data.X_test_ y_test = my_data.y_test_ my_data.dataset_parameters()
I'm not particularly concerned with line length or blank lines, but I'd really appreciate any suggestions.