如何将谷歌云盘自定义图像加载为符合要求的TensorFlow Pipeline格式

报错原因 #

tf.keras.utils.image_dataset_from_directory 返回的是 tf.data.Dataset 类型的批次数据集，并非 tfds.load 返回的带分类键的字典结构，因此不支持['A']这种下标访问方式。

调整步骤 #

1. 调整数据集目录结构 #

首先将谷歌云端硬盘中的自定义数据集按分类拆分存放，路径结构参考如下：

/content/drive/MyDrive/ColorGan/
├── trainA/    # A类训练图像，对应原逻辑中的train_horses
├── trainB/    # B类训练图像，对应原逻辑中的train_zebras
├── testA/     # A类测试图像
└── testB/     # B类测试图像

这种拆分方式加载效率更高，也更符合CycleGAN类任务的常规数据组织规范。

2. 修改数据加载代码 #

直接分别加载不同目录的数据集即可对齐原tfds.load的返回格式，代码示例如下：

import tensorflow as tf

# 加载A类训练集
train_horses = tf.keras.utils.image_dataset_from_directory(
    '/content/drive/MyDrive/ColorGan/trainA',
    labels=None,
    color_mode='rgb',
    batch_size=32,
    image_size=(256, 256),
    shuffle=True,
    seed=2000,
    interpolation='bilinear'
)

# 加载B类训练集
train_zebras = tf.keras.utils.image_dataset_from_directory(
    '/content/drive/MyDrive/ColorGan/trainB',
    labels=None,
    color_mode='rgb',
    batch_size=32,
    image_size=(256, 256),
    shuffle=True,
    seed=2000,
    interpolation='bilinear'
)

# 如需和原tfds返回的字典格式完全一致，可自行构造对应结构
dataset = {
    'trainA': train_horses,
    'trainB': train_zebras
}

可选替代方案 #

如果你不想调整现有目录结构，也可以对全量加载的数据集按标签过滤拆分：

# 你的原有加载逻辑
train_dataset=tf.keras.utils.image_dataset_from_directory(
    '/content/drive/MyDrive/ColorGan', labels='inferred', label_mode='int',
    class_names=['A', 'B'], # 显式指定类名保证标签顺序正确，A对应标签0，B对应标签1
    color_mode='rgb', batch_size=32, image_size=(256,256), 
    shuffle=True, seed=2000, validation_split=0.2, subset='training',
    interpolation='bilinear'
)

# 按标签拆分两类数据
train_horses = train_dataset.filter(lambda x, y: y == 0).map(lambda x, y: x)
train_zebras = train_dataset.filter(lambda x, y: y == 1).map(lambda x, y: x)

该方案需要额外的过滤计算，性能不如分开目录加载的方案。

内容的提问来源于stack exchange，提问作者Muhammad Afaq Saeed