CODEWITHSUNDEEP
pythontensorflowkerastensorflow2.0
user2874583
user2874583

Reputation: 535

How can I combine ImageDataGenerator with TensorFlow datasets in TF2?

I have a TF dataset to classify cats and dogs:

import tensorflow_datasets as tfds
SPLIT_WEIGHTS = (8, 1, 1)
splits = tfds.Split.TRAIN.subsplit(weighted=SPLIT_WEIGHTS)

(raw_train, raw_validation, raw_test), metadata = tfds.load(
    'cats_vs_dogs', split=list(splits),
    with_info=True, as_supervised=True)

In the example they use some image augmentation with a map function. I was wondering if that could also be done with the nice ImageDataGenerator class such as described here:

from tensorflow.keras.preprocessing.image import ImageDataGenerator
train_image_generator = ImageDataGenerator(rescale=1./255) # Generator for our training data
train_data_gen = train_image_generator.flow_from_directory(batch_size=batch_size,
                                                           directory=train_dir,
                                                           shuffle=True,
                                                           target_size=(IMG_HEIGHT, IMG_WIDTH),
                                                           class_mode='binary')

The problem I'm facing is that I can only see 3 ways to use the ImageDataGenerator: pandas dataframe, numpy array and directory of images. Is there any way to also use a Tensorflow dataset and combine these methods?

Upvotes: 15

Views: 6754

Answers (3)

moejoe
moejoe

Reputation: 137

One way could be to add the data augmentation via keras preprocessing layers, making the augmentation part of your model (see here: https://www.tensorflow.org/tutorials/images/data_augmentation#data_augmentation_2).

Example from the docs:

data_augmentation = tf.keras.Sequential([
  layers.RandomFlip("horizontal_and_vertical"),
  layers.RandomRotation(0.2),
])

Upvotes: 1

ssp
ssp

Reputation: 1710

One idea is you could create a generator wrapper function that uses your tfds Dataset to load multiples of your batch size. Then pass those images, labels to ImageDataGenerator's flow method which would yield augmented data at a rate of your desired batch size.

For example:

def tfds_imgen(ds, imgen, batch_size, batches_per):
    for images, labels in ds:
        flow_ = imgen.flow(images, labels, batch_size=batch_size)
        for _ in range(batches_per):
            yield next(flow_)

raw_train_ds = tfds.load(
    'cats_vs_dogs', split='train',
    batch_size=SOME_MULTIPLE_OF_32,
    as_supervised=True)

imgen = ImageDataGenerator(...)

train_ds = tfds_imgen(
    raw_train_ds.as_numpy_iterator(), imgen,
    batch_size=32, batches_per=SOME_MULTIPLE_OF_32 // 32)

Upvotes: 3

sebastian-sz
sebastian-sz

Reputation: 1498

Yes, it is but it's a bit tricky.
Keras ImageDataGenerator works on numpy.arrays and not on tf.Tensor's so we have to use Tensorflow's numpy_function. This will allow us to perform operations on tf.data.Dataset content just like it was numpy arrays.

First, let's declare the function that we will .map over our dataset (assuming your dataset consists of image, label pairs):

# We will take 1 original image and create 5 augmented images:
HOW_MANY_TO_AUGMENT = 5

def augment(image, label):

  # Create generator and fit it to an image
  img_gen = tf.keras.preprocessing.image.ImageDataGenerator(rescale=1./255)
  img_gen.fit(image)

  # We want to keep original image and label
  img_results = [(image/255.).astype(np.float32)] 
  label_results = [label]

  # Perform augmentation and keep the labels
  augmented_images = [next(img_gen.flow(image)) for _ in range(HOW_MANY_TO_AUGMENT)]
  labels = [label for _ in range(HOW_MANY_TO_AUGMENT)]

  # Append augmented data and labels to original data
  img_results.extend(augmented_images)
  label_results.extend(labels)

  return img_results, label_results

Now, in order to use this function inside tf.data.Dataset we must declare a numpy_function:

def py_augment(image, label):
  func = tf.numpy_function(augment, [image, label], [tf.float32, tf.int32])
  return func

py_augment can be safely used like:

augmented_dataset_ds = image_label_dataset.map(py_augment)

The image part in dataset is now in shape (HOW_MANY_TO_AUGMENT, image_height, image_width, channels). To convert it to simple (1, image_height, image_width, channels) you can simply use unbatch:

unbatched_augmented_dataset_ds = augmented_dataset_ds.unbatch()

So the entire section looks like this:

HOW_MANY_TO_AUGMENT = 5

def augment(image, label):

  # Create generator and fit it to an image
  img_gen = tf.keras.preprocessing.image.ImageDataGenerator(rescale=1./255)
  img_gen.fit(image)

  # We want to keep original image and label
  img_results = [(image/255.).astype(np.float32)] 
  label_results = [label]

  # Perform augmentation and keep the labels
  augmented_images = [next(img_gen.flow(image)) for _ in range(HOW_MANY_TO_AUGMENT)]
  labels = [label for _ in range(HOW_MANY_TO_AUGMENT)]

  # Append augmented data and labels to original data
  img_results.extend(augmented_images)
  label_results.extend(labels)

  return img_results, label_results

def py_augment(image, label):
  func = tf.numpy_function(augment, [image, label], [tf.float32, tf.int32])
  return func

unbatched_augmented_dataset_ds = augmented_dataset_ds.map(py_augment).unbatch()

# Iterate over the dataset for preview:
for image, label in unbatched_augmented_dataset_ds:
    ...

Upvotes: 7

Related Questions