如何使用flow_from_directory加载多目录下的多数据集以训练语义分割深度学习模型

没问题，我帮你搞定多数据集合并的生成器问题！你已经有了单个数据集的生成器代码，现在要把dataset1和dataset2合并起来训练，其实有两种常用的方法，我给你详细讲讲：

方法1：基于Keras ImageDataGenerator的自定义合并生成器 #

这种方法改动最小，在你现有代码的基础上，为每个数据集单独创建生成器，再写一个简单的生成器来迭代所有子生成器。

步骤1：为每个数据集创建单独的生成器 #

首先，分别为dataset1和dataset2创建图像、掩码生成器，注意要保证seed一致，这样图像和掩码的增强操作是同步的：

from tensorflow.keras.preprocessing.image import ImageDataGenerator

# 定义数据增强参数
img_data_gen_args = dict(horizontal_flip=True, vertical_flip=True, fill_mode='reflect')
image_datagen = ImageDataGenerator(**img_data_gen_args)
mask_datagen = ImageDataGenerator(**img_data_gen_args)

# ------------------- Dataset1 生成器 -------------------
train_img_path1 = "dataset1/train_images/"
train_mask_path1 = "dataset1/train_masks/"

image_gen1 = image_datagen.flow_from_directory(
    train_img_path1, 
    class_mode=None, 
    batch_size=16, 
    seed=123,
    target_size=(256, 256),  # 替换成你的图像实际尺寸
    color_mode="rgb"
)

mask_gen1 = mask_datagen.flow_from_directory(
    train_mask_path1, 
    class_mode=None, 
    batch_size=16, 
    seed=123,  # 和图像生成器seed一致，保证增强同步
    target_size=(256, 256),
    color_mode="grayscale"  # 掩码一般是单通道，记得设置这个
)

train_gen1 = zip(image_gen1, mask_gen1)

# ------------------- Dataset2 生成器 -------------------
train_img_path2 = "dataset2/train_images/"
train_mask_path2 = "dataset2/train_masks/"

image_gen2 = image_datagen.flow_from_directory(
    train_img_path2, 
    class_mode=None, 
    batch_size=16, 
    seed=123,
    target_size=(256, 256),
    color_mode="rgb"
)

mask_gen2 = mask_datagen.flow_from_directory(
    train_mask_path2, 
    class_mode=None, 
    batch_size=16, 
    seed=123,
    target_size=(256, 256),
    color_mode="grayscale"
)

train_gen2 = zip(image_gen2, mask_gen2)

步骤2：写一个合并生成器 #

这个生成器会循环遍历所有子生成器，每次返回一批数据：

def combined_generator(generator_list):
    while True:
        # 交替从每个数据集取数据，你也可以改成随机选择生成器
        for gen in generator_list:
            x, y = next(gen)
            yield x, y

# 合并两个数据集的生成器
combined_train_gen = combined_generator([train_gen1, train_gen2])

现在你就可以像使用单个生成器一样使用combined_train_gen了，比如：

x, y = next(combined_train_gen)

方法2：使用tf.data.Dataset（推荐） #

tf.data是TensorFlow官方推荐的数据管道，效率更高，支持多线程加载和更灵活的操作，适合大规模数据集。

步骤1：定义加载和增强函数 #

先写好加载图像、掩码，以及同步增强的函数：

import tensorflow as tf
import os

# 加载图像和掩码
def load_img_mask(img_path, mask_path):
    # 加载RGB图像
    img = tf.io.read_file(img_path)
    img = tf.image.decode_jpeg(img, channels=3)
    img = tf.image.convert_image_dtype(img, tf.float32)
    
    # 加载单通道掩码（假设是PNG格式）
    mask = tf.io.read_file(mask_path)
    mask = tf.image.decode_png(mask, channels=1)
    mask = tf.image.convert_image_dtype(mask, tf.float32)
    
    return img, mask

# 同步增强图像和掩码
def augment_img_mask(img, mask):
    # 随机水平翻转
    if tf.random.uniform(()) > 0.5:
        img = tf.image.flip_left_right(img)
        mask = tf.image.flip_left_right(mask)
    # 随机垂直翻转
    if tf.random.uniform(()) > 0.5:
        img = tf.image.flip_up_down(img)
        mask = tf.image.flip_up_down(mask)
    return img, mask

步骤2：创建单个数据集的管道 #

写一个函数来创建单个数据集的tf.data管道：

def create_tf_dataset(img_dir, mask_dir, img_size=(256,256), batch_size=16):
    # 获取所有训练图像路径
    img_paths = tf.data.Dataset.list_files(os.path.join(img_dir, "train", "*"), shuffle=False)
    
    # 根据图像路径生成对应的掩码路径（假设图像和掩码文件名前缀一致，后缀不同）
    def get_mask_path(img_path):
        img_filename = tf.strings.split(img_path, os.sep)[-1]
        img_prefix = tf.strings.split(img_filename, ".")[0]
        mask_path = tf.strings.join([mask_dir, "train", f"{img_prefix}.png"], separator=os.sep)
        return mask_path
    
    mask_paths = img_paths.map(get_mask_path)
    
    # 加载、预处理、增强数据
    dataset = tf.data.Dataset.zip((img_paths, mask_paths))
    dataset = dataset.map(load_img_mask, num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.map(lambda img, mask: (tf.image.resize(img, img_size), tf.image.resize(mask, img_size)), num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.map(augment_img_mask, num_parallel_calls=tf.data.AUTOTUNE)
    
    # 打乱、批量、预取
    dataset = dataset.shuffle(1000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
    return dataset

步骤3：合并两个数据集 #

最后创建两个数据集并合并：

# 创建dataset1和dataset2的管道
dataset1 = create_tf_dataset("dataset1/train_images/", "dataset1/train_masks/")
dataset2 = create_tf_dataset("dataset2/train_images/", "dataset2/train_masks/")

# 合并两个数据集
combined_dataset = dataset1.concatenate(dataset2)
# 可选：再次打乱合并后的数据集
combined_dataset = combined_dataset.shuffle(2000)

现在你可以直接把combined_dataset传入模型的fit方法中使用。

关键注意事项 #

• 图像尺寸统一：两个数据集的图像和掩码尺寸必须一致，否则生成器会报错，记得在target_size或resize中指定统一尺寸。
• 掩码通道数：掩码一般是单通道（二分类或多分类），所以在flow_from_directory中要设置color_mode="grayscale"，或者在tf.data加载时指定channels=1。
• 增强同步：无论是用ImageDataGenerator还是tf.data，都要保证图像和掩码的增强操作完全同步，否则会出现图像和掩码不匹配的问题。

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