如何从tf.keras.preprocessing.image_dataset_from_directory()？浏览和修改创建的数据集

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.preprocessing.image import ImageDataGenerator

# 创建数据集
train_dataset = tf.keras.preprocessing.image_dataset_from_directory(
    directory='path/to/train/directory',
    labels='inferred',
    image_size=(128, 128),
    batch_size=32,
    validation_split=0.2,
    subset='training'
)

# 浏览数据集
for images, labels in train_dataset:
    # 对数据集进行操作，如打印图像和标签
    print(images.shape)
    print(labels.shape)
    break

# 修改数据集
data_augmentation = keras.Sequential([
    tf.keras.layers.experimental.preprocessing.Rescaling(1./255),
    tf.keras.layers.experimental.preprocessing.RandomRotation(0.2),
    tf.keras.layers.experimental.preprocessing.RandomZoom(0.2)
])

train_dataset = train_dataset.map(lambda x, y: (data_augmentation(x, training=True), y))

# 再次浏览修改后的数据集
for images, labels in train_dataset:
    print(images.shape)
    print(labels.shape)
    break