如何将图像列表与分割掩码列表转换为TensorFlow Dataset？

问题背景 #

我有100张形状为(128, 128, 3)的图像列表，以及100张形状为(128, 128, 1)的分割掩码列表，尝试将它们转换为keras.Model可用的Dataset时遇到困难。试过slice、TFRecord等方法均无效，训练时loss为nan，精度不足1%；但使用tensorflow-dataset的load函数加载数据则运行正常。以下是我的代码：

import cv2
from pycocotools import mask as maskUtils
from keras.callbacks import ModelCheckpoint
from pycocotools.coco import COCO
import numpy as np

from unet import *

coco_images_dir = '../dataset/images/'
coco_json_path = '../dataset/COCO_Football Pixel.json'
checkpoint_path = "unet3plus_ckpt.weights.h5"
export_path = "unet3plus_export.weights.h5"
input_size = [128, 128]
batch_size = 30

# 加载JSON文件
coco = COCO(coco_json_path)
img_ids = coco.getImgIds()
cat_ids = coco.getCatIds()
images = []
segmentation_masks = []
for img_id in img_ids:
    img = coco.loadImgs(img_id)[0]
    file_name = coco_images_dir + img['file_name']
    image = cv2.cvtColor(cv2.imread(file_name), cv2.COLOR_BGR2RGB)
    image = image.astype(np.float32) / 255.0
    image = tf.convert_to_tensor(image)
    image = tf.image.resize(image, input_size)
    images.append(image)
    annIds = coco.getAnnIds(imgIds=img_id, iscrowd=None)
    anns = coco.loadAnns(annIds)
    segmentation_mask = np.zeros((input_size[0], input_size[1], 1))
    for ann in anns:
        if 'segmentation' in ann:
            rle = coco.annToRLE(ann)
            mask = maskUtils.decode(rle)
            mask = tf.expand_dims(mask, axis=-1)
            mask = tf.convert_to_tensor(mask)
            mask = tf.image.convert_image_dtype(mask, tf.uint8)
            mask = tf.image.resize(mask, input_size)
            mask *= (cat_ids.index(ann['category_id']) + 1)
            segmentation_mask += mask
    segmentation_masks.append(segmentation_mask)

##################################################################
# What should I do with images and segmentation_masks?
##################################################################
dataset = ...

model = unet((input_size[0], input_size[1], 3), len(coco.getCatIds()))
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])
model.fit(dataset, epochs=100,
          callbacks=ModelCheckpoint(filepath=checkpoint_path, save_weights_only=True, save_best_only=False))
model.save(export_path)

问题根源分析 #

1. 掩码数据类型错误：sparse_categorical_crossentropy要求标签为整数类型，但当前代码中掩码叠加后是float类型，且可能出现非整数数值（比如多个掩码重叠区域叠加），导致损失计算异常出现nan。
2. Dataset构建不规范：直接用列表转Dataset时未做数据类型对齐、shuffle和batch处理，模型无法有效学习。

解决方案 #

步骤1：修正掩码处理逻辑 #

确保掩码最终是整数类型，且重叠区域取正确的类别值（避免叠加导致的非整数）：

# 替换原掩码循环部分
segmentation_mask = np.zeros((input_size[0], input_size[1], 1), dtype=np.int32)  # 改为int32类型
for ann in anns:
    if 'segmentation' in ann:
        rle = coco.annToRLE(ann)
        mask = maskUtils.decode(rle)
        # 先调整掩码尺寸再处理，避免尺寸变换后的精度问题
        mask = cv2.resize(mask, (input_size[1], input_size[0]), interpolation=cv2.INTER_NEAREST)
        mask = np.expand_dims(mask, axis=-1)
        class_idx = cat_ids.index(ann['category_id']) + 1  # 类别索引从1开始（0为背景）
        # 将当前类别的掩码区域赋值，重叠区域以后续类别覆盖（按需调整）
        segmentation_mask[mask == 1] = class_idx
segmentation_masks.append(segmentation_mask)

步骤2：构建规范的TF Dataset #

将图像和掩码列表转换为numpy数组，再用tf.data.Dataset.from_tensor_slices构建，添加必要的预处理步骤：

##################################################################
# 处理images和segmentation_masks
##################################################################
# 将列表转换为numpy数组
images_np = np.array([img.numpy() for img in images], dtype=np.float32)
masks_np = np.array(segmentation_masks, dtype=np.int32)

# 构建Dataset
dataset = tf.data.Dataset.from_tensor_slices((images_np, masks_np))

# 添加shuffle、batch、prefetch，提升训练效率
dataset = dataset.shuffle(buffer_size=len(images_np)) \
                 .batch(batch_size) \
                 .prefetch(tf.data.AUTOTUNE)

完整修改后的关键代码片段 #

# 加载JSON文件部分修改掩码处理
coco = COCO(coco_json_path)
img_ids = coco.getImgIds()
cat_ids = coco.getCatIds()
images = []
segmentation_masks = []
for img_id in img_ids:
    img = coco.loadImgs(img_id)[0]
    file_name = coco_images_dir + img['file_name']
    image = cv2.cvtColor(cv2.imread(file_name), cv2.COLOR_BGR2RGB)
    image = image.astype(np.float32) / 255.0
    image = tf.convert_to_tensor(image)
    image = tf.image.resize(image, input_size)
    images.append(image)
    
    annIds = coco.getAnnIds(imgIds=img_id, iscrowd=None)
    anns = coco.loadAnns(annIds)
    segmentation_mask = np.zeros((input_size[0], input_size[1], 1), dtype=np.int32)
    for ann in anns:
        if 'segmentation' in ann:
            rle = coco.annToRLE(ann)
            mask = maskUtils.decode(rle)
            mask = cv2.resize(mask, (input_size[1], input_size[0]), interpolation=cv2.INTER_NEAREST)
            mask = np.expand_dims(mask, axis=-1)
            class_idx = cat_ids.index(ann['category_id']) + 1
            segmentation_mask[mask == 1] = class_idx
    segmentation_masks.append(segmentation_mask)

# 构建Dataset
images_np = np.array([img.numpy() for img in images], dtype=np.float32)
masks_np = np.array(segmentation_masks, dtype=np.int32)
dataset = tf.data.Dataset.from_tensor_slices((images_np, masks_np))
dataset = dataset.shuffle(len(images_np)).batch(batch_size).prefetch(tf.data.AUTOTUNE)

# 后续模型编译和训练不变
model = unet((input_size[0], input_size[1], 3), len(coco.getCatIds()))
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])
model.fit(dataset, epochs=100,
          callbacks=ModelCheckpoint(filepath=checkpoint_path, save_weights_only=True, save_best_only=False))
model.save(export_path)

额外注意事项 #

• 类别索引从1开始是因为0默认为背景类，确保与sparse_categorical_crossentropy的要求匹配（标签范围为0到num_classes-1或1到num_classes，需对应模型输出的类别数）。
• 掩码 resize 时使用INTER_NEAREST插值，避免出现模糊的非整数像素值，保证掩码的类别准确性。
• 如果数据集较大，建议使用生成器或TFRecord来避免内存占用过高，但当前100张图片用numpy数组完全没问题。

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