如何使用TensorFlow ImageDataGenerator实现多输入图像数据处理

原生ImageDataGenerator默认支持RGB（3通道）、RGBA（4通道）、灰度图（1通道）读取，针对自定义4通道输入需求，不需要改动框架源码，通过自定义生成器包装即可完美适配类目录结构的数据集，完整保留原有的自动标签生成、数据增强、乱序、批次加载能力。

实现步骤 #

1. 基础配置 #

先按常规方式初始化ImageDataGenerator，配置你需要的数据增强参数：

import numpy as np
from tensorflow.keras.preprocessing.image import ImageDataGenerator

# 按需配置数据增强规则，和普通图像分类场景用法一致
train_datagen = ImageDataGenerator(
    rescale=1./255,
    rotation_range=12,
    width_shift_range=0.08,
    height_shift_range=0.08,
    zoom_range=0.1,
    horizontal_flip=False # 词汇对应手势类图像不要开水平翻转，会改变语义
)
val_datagen = ImageDataGenerator(rescale=1./255)

2. 按场景选择生成逻辑 #

场景A：单文件自带4通道（如RGBA格式图像）

直接将flow_from_directory的color_mode参数设为"rgba"即可，不需要额外封装，生成器会直接输出形状为(batch_size, img_h, img_w, 4)的4通道张量，自动生成20类的one-hot标签：

train_gen = train_datagen.flow_from_directory(
    directory="./train", # 替换为训练集根目录，内部包含20个类名对应的子文件夹
    target_size=(224, 224), # 替换为模型要求的输入尺寸
    color_mode="rgba",
    batch_size=32,
    class_mode="categorical",
    shuffle=True,
    seed=42
)

场景B：需要自定义拼接第4通道（如3通道RGB加深度/特征通道）

先初始化基础生成器读取原始3通道数据，再通过自定义生成器拼接第4通道：

base_train_gen = train_datagen.flow_from_directory(
    directory="./train",
    target_size=(224, 224),
    color_mode="rgb",
    batch_size=32,
    class_mode="categorical",
    shuffle=True,
    seed=42
)

def four_channel_wrapper(gen, img_h=224, img_w=224):
    for x_rgb, y in gen:
        # 替换为你自己的第4通道读取/计算逻辑
        fourth_ch = np.zeros((x_rgb.shape[0], img_h, img_w, 1), dtype=np.float32)
        x_4ch = np.concatenate([x_rgb, fourth_ch], axis=-1)
        yield x_4ch, y

train_gen = four_channel_wrapper(base_train_gen)

场景C：4个通道分别存储在4个结构一致的目录中

如果每个样本的4个通道对应4个独立文件，分别存在4个根目录下（每个根目录下都有20个同名类文件夹），可以初始化4个基础生成器，固定相同随机种子保证样本顺序对齐，再拼接通道：

# 4个通道的数据集根目录
ch_dirs = ["./ch1_train", "./ch2_train", "./ch3_train", "./ch4_train"]
base_gens = []
for dir_path in ch_dirs:
    gen = train_datagen.flow_from_directory(
        directory=dir_path,
        target_size=(224,224),
        color_mode="grayscale", # 每个通道为单通道灰度图
        batch_size=32,
        class_mode="categorical",
        shuffle=True,
        seed=42 # 必须固定相同种子，保证4个生成器输出样本顺序一致
    )
    base_gens.append(gen)

def multi_ch_wrapper(gens):
    while True:
        batch_data = [next(g) for g in gens]
        x_list = [d[0] for d in batch_data]
        y = batch_data[0][1]
        # 校验标签完全一致，避免样本错配
        for d in batch_data[1:]:
            assert np.array_equal(y, d[1])
        x_4ch = np.concatenate(x_list, axis=-1)
        yield x_4ch, y

train_gen = multi_ch_wrapper(base_gens)

3. 训练调用 #

验证集按相同逻辑生成后，直接传入fit接口即可，注意步数要从基础生成器的samples属性计算，避免自定义生成器无法自动获取步数：

model.fit(
    train_gen,
    steps_per_epoch=base_train_gen.samples // base_train_gen.batch_size,
    epochs=30,
    validation_data=val_gen,
    validation_steps=base_val_gen.samples // base_val_gen.batch_size
)

注意事项 #

• 正式训练前先取1个批次打印张量形状，确认输入形状为(batch_size, img_h, img_w, 4)、标签形状为(batch_size, 20)，避免通道维度错误。
• 不要将所有数据一次性加载到内存拼接，上述生成器方案为按需读取，内存占用和原生ImageDataGenerator一致，可支持大规模数据集。
• 如果模型为4分支多输入结构（每个通道单独输入一个分支），只需要在自定义生成器的返回值中，将4个通道拆分后作为列表返回即可。

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