是的,但它有点棘手。
Keras ImageDataGenerator 适用于 numpy.arrays 而不是 tf.Tensor,因此我们必须使用 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 部分现已成型
(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:
...