TensorFlow图像分割自定义数据集的文件夹结构与加载方法咨询

核心问题说明 #

tf.keras.utils.ImageFolder（你代码里的tf.data_folder.ImageFolder应为笔误）是为图像分类任务设计的工具，它会将子文件夹映射为类别标签，无法直接处理图像-掩码配对的分割任务，这是你报错的根本原因。

推荐的文件夹结构 #

你当前的结构其实可以保留，只需调整加载逻辑；如果要更贴近官方教程风格，也可以简化为以下结构（核心是保证图像与掩码文件名一一对应）：

path/to/dataset/
  train/
    images/
      1.png
      2.png
      ...
    masks/
      1.jpg
      2.jpg
      ...
  test/
    images/
      6.png
      ...
    masks/
      6.jpg
      ...

正确的数据集加载方式 #

通过tf.data.Dataset手动配对图像与掩码，以下是完整实现代码：

import tensorflow as tf
import os

# 定义图像加载与预处理函数
def load_and_preprocess_image(img_path):
    # 读取图像
    img = tf.io.read_file(img_path)
    # 根据图像格式解码（png用decode_png，jpg用decode_jpeg）
    img = tf.image.decode_png(img, channels=3)
    # 转换为float32并归一化到[0,1]
    img = tf.image.convert_image_dtype(img, tf.float32)
    # 调整图像尺寸（根据你的模型需求修改）
    img = tf.image.resize(img, (256, 256))
    return img

# 定义掩码加载与预处理函数
def load_and_preprocess_mask(mask_path):
    mask = tf.io.read_file(mask_path)
    # 掩码一般为单通道，这里按你的格式用decode_jpeg
    mask = tf.image.decode_jpeg(mask, channels=1)
    mask = tf.image.convert_image_dtype(mask, tf.float32)
    # 掩码缩放用最近邻插值，避免类别混淆
    mask = tf.image.resize(mask, (256, 256), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
    return mask

# 加载训练集路径（确保文件名排序后一一对应）
train_img_dir = "path/to/image/train/image"
train_mask_dir = "path/to/image/train/segmentation_mask"
train_img_paths = sorted([os.path.join(train_img_dir, f) for f in os.listdir(train_img_dir) if f.endswith(('.png', '.jpg'))])
train_mask_paths = sorted([os.path.join(train_mask_dir, f) for f in os.listdir(train_mask_dir) if f.endswith(('.png', '.jpg'))])

# 构建训练数据集
train_dataset = tf.data.Dataset.from_tensor_slices((train_img_paths, train_mask_paths))
train_dataset = train_dataset.map(
    lambda img_p, mask_p: (load_and_preprocess_image(img_p), load_and_preprocess_mask(mask_p)),
    num_parallel_calls=tf.data.AUTOTUNE
)

# 加载测试集路径
test_img_dir = "path/to/image/test/image"
test_mask_dir = "path/to/image/test/segmentation_mask"
test_img_paths = sorted([os.path.join(test_img_dir, f) for f in os.listdir(test_img_dir) if f.endswith(('.png', '.jpg'))])
test_mask_paths = sorted([os.path.join(test_mask_dir, f) for f in os.listdir(test_mask_dir) if f.endswith(('.png', '.jpg'))])

# 构建测试数据集
test_dataset = tf.data.Dataset.from_tensor_slices((test_img_paths, test_mask_paths))
test_dataset = test_dataset.map(
    lambda img_p, mask_p: (load_and_preprocess_image(img_p), load_and_preprocess_mask(mask_p)),
    num_parallel_calls=tf.data.AUTOTUNE
)

# 后续添加批量、打乱、预取等优化操作
train_dataset = train_dataset.shuffle(buffer_size=100).batch(8).prefetch(tf.data.AUTOTUNE)
test_dataset = test_dataset.batch(8).prefetch(tf.data.AUTOTUNE)

关键注意事项 #

• 必须保证图像与掩码的文件名完全对应（如1.png对应1.jpg），排序后才能正确配对；
• 掩码预处理时使用NEAREST_NEIGHBOR插值，防止掩码的类别标签被模糊处理；
• 根据你的实际图像格式（png/jpg）调整decode_png或decode_jpeg，同时修改尺寸适配你的模型输入。

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