This notebook has been tested with

  • Python 3.5
  • Keras 2.0.6
  • Tensorflow 1.2.1

Use image generator to have more training data


In [ ]:
from keras.preprocessing.image import ImageDataGenerator, array_to_img, img_to_array, load_img
import concurrent.futures
import os

In [ ]:
# The variable 'number' gives the numbers of images generated from one image
def datagen(filename, destination, number):
    datagen = ImageDataGenerator(
            rotation_range=0.2,
            width_shift_range=0.2,
            height_shift_range=0.2,
            rescale=1./255,
            shear_range=0.2,
            zoom_range=0.2,
            horizontal_flip=True,
            fill_mode='nearest')
    
    img = load_img(filename)
    x = img_to_array(img) 
    x = x.reshape((1,) + x.shape)
    filename = filename.split('/')[1]
    
    i = 0 
    for batch in datagen.flow(x, batch_size=1,
                          save_to_dir=destination, save_prefix=filename, save_format='jpg'):
        i += 1
        if i > number:
            break

In [ ]:
# The source directory of original images
photoDir = 'middle_hair_man_train/'
photoList = os.listdir(photoDir)
# The target directory to save images generated
direction = 'middle_hair_man_train/'

# Use multi-threading to speed up
with concurrent.futures.ThreadPoolExecutor(max_workers=20) as executor:
    for photo in photoList:
        try:
            filename = photoDir+photo
            executor.submit(datagen, filename, direction, 20)
        except Exception as exc:
            print(exc)