将图像数据集分为特征和标签

Question

我正在使用ImageDataGenerator加载我自己的rgb图像数据集，我想将训练数据和测试数据分离为标签和特征。有没有办法做到这一点？。我正在使用tensorflow 2.1 in python。

data_gen = tf.keras.preprocessing.image.ImageDataGenerator(
    rescale=1. / 255)

train_data = data_gen.flow_from_directory(
    directory='data/train',
    shuffle=True,
    target_size=(224, 224),
    class_mode='categorical']