如何在 TF2 中将 ImageDataGenerator 与 TensorFlow 数据集结合起来？

Question

我有一个 TF 数据集来对猫和狗进行分类：

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)

在这个例子中，他们使用了一些带有地图功能的图像增强。我想知道是否也可以使用这里ImageDataGenerator描述的好类来完成：

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')

我面临的问题是我只能看到3 种使用方式ImageDataGenerator：pandas 数据框、numpy 数组和图像目录。有没有办法也使用 Tensorflow 数据集并结合这些方法？

Answer 1

是的，它是，但它有点棘手。
KerasImageDataGenerator适用于numpy.arrays 而不是tf.Tensor's 所以我们必须使用 Tensorflow 的numpy_function。这将允许我们对tf.data.Dataset像对 numpy 数组一样内容。

首先，让我们声明我们将.map在我们的数据集上使用的函数（假设您的数据集由图像、标签对组成）：

# 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

现在，为了在里面使用这个函数，tf.data.Dataset我们必须声明一个numpy_function：

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

py_augment 可以安全地使用，例如：

augmented_dataset_ds = image_label_dataset.map(py_augment)

image数据集中的部分现在是 shape (HOW_MANY_TO_AUGMENT, image_height, image_width, channels)。要将其转换为简单，(1, image_height, image_width, channels)您只需使用unbatch：

unbatched_augmented_dataset_ds = augmented_dataset_ds.unbatch()

所以整个部分看起来像这样：

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:
    ...