Tensorflow.keras如何加载自定义目录结构的图像分割数据集

不需要移动或拍平原有的目录结构，无侵入式的实现方案更优，具体适配两类接口的方法如下：

方案1：直接遍历子目录拼接数据集（推荐） #

适配 tf.keras.utils.image_dataset_from_directory 接口 #

核心逻辑是逐个加载a-z每个子目录的原图、掩码数据集，配对后合并为总训练集，完全无需修改源文件：

import tensorflow as tf
from pathlib import Path

# 自定义配置
TRAIN_ROOT = Path("../input/hand-seg-tr/train/")
TEST_PATH = "../input/hand-segmentation-test/test/"
IMAGE_SIZE = (256, 256)
BATCH_SIZE = 8

train_dataset = None
# 遍历所有字母子目录
for sub_folder in TRAIN_ROOT.iterdir():
    if not sub_folder.is_dir():
        continue
    # 加载当前子目录的原图
    img_ds = tf.keras.utils.image_dataset_from_directory(
        sub_folder / "images",
        label_mode=None,
        image_size=IMAGE_SIZE,
        batch_size=BATCH_SIZE,
        shuffle=False # 必须关闭，保证和掩码顺序对应
    )
    # 加载当前子目录的对应掩码
    mask_ds = tf.keras.utils.image_dataset_from_directory(
        sub_folder / "segmentation",
        label_mode=None,
        image_size=IMAGE_SIZE,
        batch_size=BATCH_SIZE,
        color_mode="grayscale",
        shuffle=False
    )
    # 配对当前子目录的原图和掩码
    paired_ds = tf.data.Dataset.zip((img_ds, mask_ds))
    # 合并到总训练集
    train_dataset = paired_ds if train_dataset is None else train_dataset.concatenate(paired_ds)

# 合并完成后全局打乱
train_dataset = train_dataset.shuffle(buffer_size=1000)

# 加载无标注测试集
test_dataset = tf.keras.utils.image_dataset_from_directory(
    TEST_PATH,
    label_mode=None,
    image_size=IMAGE_SIZE,
    batch_size=BATCH_SIZE,
    shuffle=False
)

适配 flow_from_directory 接口 #

逻辑和上述方案一致，逐个子目录生成原图、掩码生成器，配对后合并为统一的迭代器即可，不需要修改目录结构。

方案2：软链接构造统一目录（如需拍平结构时可选） #

如果确实需要使用单目录加载的逻辑，不需要物理复制移动文件，用软链接将所有a-z子目录下的图片链接到同一个统一目录即可，不会占用额外磁盘空间，也不会修改源文件：

• Linux/macOS 可直接用shell命令批量生成软链接
• Windows 系统也支持mklink命令创建软链接

拍平原目录方案的弊端 #

• 物理复制文件会占用额外磁盘空间，大体积数据集下耗时极长
• 不同子目录下可能存在同名文件，容易出现覆盖丢失的问题
• 修改源文件结构有出错风险，故障后恢复成本高

内容的提问来源于stack exchange，提问作者Dmitry Sokolov