如何使flow_from_directory的color_mode不改变18通道输入形状？

我的图像由5个单通道和1个三通道拼接而成，保存为18通道的.npy文件；掩码为6个单通道，因无法在axis=2维度拼接，也转为18通道.npy文件。训练时出现ValueError，提示模型期望输入形状为(512, 512, 18)，但实际得到(512, 512, 3)。将image_color_mode改为rgba后，输入形状变为(512, 512, 4)，问题应该出在这里，请问如何让color_mode不修改输入形状？

错误信息 #

發生例外狀況: ValueError
Error when checking input: expected input_1 to have shape (512, 512, 18) but got array with shape (512, 512, 3)
File "C:\Labbb\testing\unet_mao0\main.py", line 96, in
train_history=model.fit_generator(myGene,steps_per_epoch=200,epochs=30,callbacks=[model_checkpoint])
ValueError: Error when checking input: expected input_1 to have shape (512, 512, 18) but got array with shape (512, 512, 3)

相关代码 #

def trainGenerator(batch_size, train_path, image_folder, mask_folder, aug_dict,image_color_mode='rgb', mask_color_mode='rgb', 
                   image_save_prefix="image", mask_save_prefix="mask", 
                   flag_multi_class=False, num_class=7, save_to_dir=None, target_size=(512, 512), seed=1, file_extension=".npy"):

    image_datagen = ImageDataGenerator(**aug_dict)
    mask_datagen = ImageDataGenerator(**aug_dict)
    image_generator = image_datagen.flow_from_directory(
        train_path,
        classes = image_folder,
        class_mode = None,
        color_mode = image_color_mode,
        target_size = target_size,
        batch_size = batch_size,
        save_to_dir = save_to_dir,
        save_prefix  = image_save_prefix,
        seed = seed)
    mask_generator = mask_datagen.flow_from_directory(
        train_path,
        classes = mask_folder,
        class_mode = None,
        color_mode = mask_color_mode,
        target_size = target_size,
        batch_size = batch_size,
        save_to_dir = save_to_dir,
        save_prefix  = mask_save_prefix,
        seed = seed)

核心问题是flow_from_directory本来就不是用来加载.npy文件的，它默认会把文件当作标准图像格式处理，color_mode参数只针对jpg/png这类图像生效，会强制修改通道数，直接破坏你的18通道数据。要解决这个问题，得自定义生成器来读取.npy文件：

1. 写一个自定义的.npy数据生成器
   用Keras的Sequence类实现生成器，完全保留原始通道数，还能同步做数据增强：

import numpy as np
import os
import cv2
from keras.utils import Sequence
from keras.preprocessing.image import ImageDataGenerator

class NPYDataGenerator(Sequence):
    def __init__(self, batch_size, train_path, image_folder, mask_folder, aug_dict, target_size=(512,512), seed=1):
        self.batch_size = batch_size
        self.image_dir = os.path.join(train_path, image_folder)
        self.mask_dir = os.path.join(train_path, mask_folder)
        self.augmentor = ImageDataGenerator(**aug_dict)
        self.target_size = target_size
        self.seed = seed
        # 取所有.npy文件名（确保图像和掩码文件名一一对应）
        self.filenames = [f for f in os.listdir(self.image_dir) if f.endswith('.npy')]
        
    def __len__(self):
        # 计算每个epoch的步数
        return int(np.ceil(len(self.filenames) / self.batch_size))
    
    def __getitem__(self, idx):
        batch_files = self.filenames[idx*self.batch_size : (idx+1)*self.batch_size]
        batch_imgs = []
        batch_masks = []
        
        for filename in batch_files:
            # 加载原始.npy文件，保留所有通道
            img = np.load(os.path.join(self.image_dir, filename))
            mask = np.load(os.path.join(self.mask_dir, filename))
            
            # 如果尺寸不符合目标大小，调整尺寸（保留通道维度）
            if img.shape[:2] != self.target_size:
                img = cv2.resize(img, self.target_size)
                mask = cv2.resize(mask, self.target_size)
            
            # 同步增强图像和掩码：先合并再增强，避免增强参数不一致
            combined = np.concatenate([img, mask], axis=-1)
            augmented = self.augmentor.random_transform(combined)
            img_aug = augmented[..., :img.shape[-1]]
            mask_aug = augmented[..., img.shape[-1]:]
            
            batch_imgs.append(img_aug)
            batch_masks.append(mask_aug)
        
        return np.array(batch_imgs), np.array(batch_masks)

2. 替换原生成器使用自定义类
   训练时直接用这个生成器替代原来的trainGenerator：

# 替换成你的参数
myGene = NPYDataGenerator(
    batch_size=8,
    train_path="你的训练数据根路径",
    image_folder="图像文件夹名",
    mask_folder="掩码文件夹名",
    aug_dict=你的数据增强字典,
    target_size=(512,512)
)

# 正常训练即可
train_history=model.fit_generator(myGene,steps_per_epoch=200,epochs=30,callbacks=[model_checkpoint])

3. 补充说明
   别再用flow_from_directory加载.npy文件了，这个函数的设计目标是处理标准图像格式，不管你设置什么color_mode，它都会强制将数据转换成3/4/1通道，完全忽略.npy的原始通道数，这就是你之前遇到形状不匹配的根本原因。

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