trjhtr

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.

The function 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[0], 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[1] == train1.shape[2] == 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[0], flattened_size)
 train2 = train2.reshape(train2.shape[0], flattened_size)
 test1 = test1.reshape(test1.shape[0], flattened_size)
 test2 = test2.reshape(test2.shape[0], 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[0], 1)
 self.y_test_ = self.y_test_.reshape(self.y_test_.shape[0], 1)

 if num_samples_last:
 # Code conversion for class
 self.X_train_.shape = (
 self.X_train_.shape[1], self.X_train_.shape[0])
 self.X_test_.shape = (self.X_test_.shape[1], self.X_test_.shape[0])
 self.y_train_.shape = (
 self.y_train_.shape[1], self.y_train_.shape[0])
 self.y_test_.shape = (self.y_test_.shape[1], self.y_test_.shape[0])

 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:

Augmenting images:

augment_images uses the augmentation functionality of Keras, but gives me more flexibility. I know Keras has save_to_dir and 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)

Getting data:

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.