基于Keras的U-Net多标签分割：双掩码输入的数据处理错误修正

错误原因 #

1. Y数组初始化形状错误：你初始化Y为(len(ids), im_height, im_width, 1)，但多标签任务需要每个样本对应2个通道的掩码，形状应为(len(ids), 224, 224, 2)。
2. 掩码堆叠维度冗余：mask1和mask2的形状是(224,224,1)，使用np.stack([mask1,mask2], axis=-1)会得到(224,224,1,2)的形状，多了一个不必要的维度，无法匹配目标数组的形状。

修正后的数据准备代码 #

import os
import numpy as np
from tensorflow.keras.preprocessing.image import load_img, img_to_array
from skimage.transform import resize
from tqdm import tqdm
from sklearn.model_selection import train_test_split

im_height, im_width = 224, 224
ids = next(os.walk("data/train/image"))[2] 
print("No. of images = ", len(ids))

# 修正X和Y的形状：Y改为2个通道对应两个标签
X = np.zeros((len(ids), im_height, im_width, 3), dtype=np.float32) 
Y = np.zeros((len(ids), im_height, im_width, 2), dtype=np.float32) 

for n, id_ in tqdm(enumerate(ids), total=len(ids)):
    # 加载并处理输入图像
    img = load_img(f"data/train/image/{id_}", color_mode="rgb")
    x_img = img_to_array(img)
    x_img = resize(x_img, (im_height, im_width, 3), 
                   mode='constant', preserve_range=True)
    
    # 加载并处理两个掩码
    mask1 = img_to_array(load_img(f"data/train/label1/{id_}", color_mode="grayscale"))
    mask2 = img_to_array(load_img(f"data/train/label2/{id_}", color_mode="grayscale"))
    
    # 调整掩码形状时去掉多余的通道维度，变成(224,224)
    mask1 = resize(mask1, (im_height, im_width), 
                  mode='constant', preserve_range=True)
    mask2 = resize(mask2, (im_height, im_width), 
                  mode='constant', preserve_range=True)
    
    # 在最后一个维度堆叠，得到(224,224,2)的掩码
    mask = np.stack([mask1, mask2], axis=-1)
    
    # 归一化并赋值
    X[n] = x_img / 255.0
    Y[n] = mask / 255.0

# 划分训练验证集
X_tr, X_val, Y_tr, Y_val = train_test_split(X, Y, test_size=0.3, random_state=42) 

关键修改点说明 #

• Y数组初始化：将Y的形状从(len(ids), im_height, im_width, 1)改为(len(ids), im_height, im_width, 2)，对应2个标签的输出通道。
• 掩码形状调整：调整resize的目标形状为(im_height, im_width)，去掉灰度图默认的1通道，得到形状为(224,224)的掩码。
• 堆叠方式：使用np.stack([mask1, mask2], axis=-1)将两个单通道掩码堆叠为(224,224,2)的多通道掩码，完美匹配Y数组的形状。

验证模型输入输出匹配 #

你的U-Net模型输出形状为(None, 224, 224, 2)（从模型层形状最后一行(None, None, None, 2)可推断，输入固定224×224时输出对应尺寸），修正后的Y数组形状(样本数,224,224,2)完全匹配模型的输出要求，训练时不会再出现维度不匹配的错误。

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