基于TensorFlow实现图像缩放、图像与掩码均Resize的差异化预处理

我拥有一个包含图像和掩码的数据集，已通过以下代码加载并合并：

images_p = tf.keras.utils.image_dataset_from_directory(
        path_imgs, batch_size=None, shuffle=False, label_mode=None)
    
masks_p = tf.keras.utils.image_dataset_from_directory(
        path_masks, batch_size=None, shuffle=False, label_mode=None,
        color_mode='grayscale')
    
dataset = tf.data.Dataset.zip((images_p, masks_p))

我需要对其进行预处理，但图像和掩码的预处理逻辑不同，例如掩码不应被缩放。目前我有如下代码：

def resize_and_rescale(image, mask=False):
    image = tf.image.resize(image, (IMG_HEIGHT, IMG_WIDTH))
    if mask is False:
        image = image / 255.0
    return image
    
def prepare(ds, shuffle=False, augment=False):
    
    # Resize and rescale 
    ds = ds.map(lambda x, y: (resize_and_rescale(x), y), 
                num_parallel_calls=AUTOTUNE)
    
    
    if shuffle:
        ds = ds.shuffle(buffer_size=1000)
    
    # Batch dataset
    ds = ds.batch(BATCH_SIZE)
    
    # Use data augmentation only on the training set
    if augment:
      ds = ds.map(lambda x, y: (transform(x), y), 
                  num_parallel_calls=AUTOTUNE)
    
    return ds.prefetch(buffer_size=AUTOTUNE)

我已将数据集拆分为训练集、验证集和测试集：

train_ds, val_ds, test_ds = split_train_test_val(dataset,
                                                 SEED,
                                                 train_split=0.8, 
                                                 val_split=0.1, 
                                                 test_split=0.1,
                                                 shuffle=True, 
                                                 shuffle_size=1000)

并计划通过以下方式应用prepare函数：

train_ds = prepare(train_ds, shuffle=True, augment=True)
val_ds = prepare(val_ds)
test_ds = prepare(test_ds)

请问如何修改代码，实现对图像和掩码均执行Resize操作，仅对图像执行缩放（除以255）的预处理？

你可以通过拆分预处理逻辑、同步处理图像和掩码的方式实现需求，具体修改如下：

1. 拆分图像与掩码的预处理函数 #

把原来的单一函数拆分成两个独立逻辑，分别对应图像和掩码的处理：

def preprocess_image(image):
    # 图像执行resize+像素缩放
    image = tf.image.resize(image, (IMG_HEIGHT, IMG_WIDTH))
    image = image / 255.0
    return image

def preprocess_mask(mask):
    # 掩码仅执行resize，不做缩放
    # 注意：如果掩码是类别标签（如分割任务的像素类别ID），必须用最近邻插值避免类别值模糊
    mask = tf.image.resize(mask, (IMG_HEIGHT, IMG_WIDTH), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
    return mask

2. 修改prepare函数的映射逻辑 #

在prepare函数中，同时对图像和掩码应用各自的预处理：

def prepare(ds, shuffle=False, augment=False):
    # 同步处理图像和掩码
    ds = ds.map(lambda x, y: (preprocess_image(x), preprocess_mask(y)), 
                num_parallel_calls=AUTOTUNE)
    
    if shuffle:
        ds = ds.shuffle(buffer_size=1000)
    
    ds = ds.batch(BATCH_SIZE)
    
    if augment:
        # 补充提示：如果是空间类增强（如翻转、平移），需要同步变换图像和掩码
        # 若只是像素级增强（如亮度调整），则仅处理图像即可
        ds = ds.map(lambda x, y: (transform(x), y), 
                    num_parallel_calls=AUTOTUNE)
    
    return ds.prefetch(buffer_size=AUTOTUNE)

关键提醒 #

• 如果你的掩码是分割任务的类别标签，一定要用NEAREST_NEIGHBOR插值，否则插值后会产生非整数的类别值，导致后续损失计算出错。
• 若使用空间类数据增强（比如随机翻转、平移），必须保证图像和掩码的变换完全同步，否则会出现图像与掩码错位的问题，这种情况下需要编写同时处理两者的增强函数。

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