带有 tf.data.Data.from_generator 的图像数据生成器- Keras

在Keras中，我们可以使用ImageDataGenerator来生成增强图像以进行训练。但是，如果我们需要完全控制如何读取和增强图像数据，我们可以使用tf.data.Dataset.from_generator函数。

该函数接受一个生成器函数作为输入，将其转换为tf.data.Dataset对象。生成器函数应不带参数，并在每次调用时返回一个元组，包含输入和目标图像。

下面是一个使用tf.data.Dataset.from_generator函数生成的数据集例子：

import tensorflow as tf
import numpy as np

def image_generator(image_file_names, label_file_names):
    for i in range(len(image_file_names)):
        # 读取图像和标签
        image = plt.imread(image_file_names[i])
        label = plt.imread(label_file_names[i])
        # 图像增强
        image = augmentation_function(image)
        # 将标签转换为one-hot编码
        label_one_hot = tf.one_hot(label, depth=2)
        # 返回输入和目标图像元组
        yield image, label_one_hot

# 图像文件和标签文件的路径
image_file_names = [...]
label_file_names = [...]

# 生成数据集，并进行shuffle和batch
dataset = tf.data.Dataset.from_generator(
    image_generator,
    args=[image_file_names, label_file_names],
    output_types=(tf.float32, tf.float32),
    output_shapes=([256, 256, 3], [256, 256, 2])
).shuffle(1000).batch(32)

# 使用数据集进行训练
model.fit(
    dataset,
    epochs=10,
    steps_per_epoch=len(image_file_names) // 32
)

在上面的例子中，我们定义了一个image_generator函数，该函数每次返回一个包含增强后的输入和目标图像的元组。然后，我们使用from_generator函数将其转换为tf.data.Dataset对象，并在参数中指定输出类型和形状。最后，我们使用生成的数据集进行模型训练。