如何构建TensorFlow的(image,label)数据集及基于增强分割图的权重映射

首先，我会先拆解你的两个核心需求的实现逻辑，再针对性修复代码里的问题：

一、构建基础的(image, label)分割数据集 #

你当前用process_path读取图像和分割图的思路是对的，标准实现流程如下：

1. 准备好图像路径与对应分割图路径的配对列表（也就是你的train_im_files和train_seg_files）
2. 用tf.data.Dataset.from_tensor_slices生成初始数据集
3. 通过map调用路径处理函数，把路径转为预处理后的图像与分割标签对，示例代码：

def process_path(image_path, seg_path):
    # 读取并预处理图像
    img = tf.io.read_file(image_path)
    img = tf.image.decode_jpeg(img, channels=3)
    img = tf.image.convert_image_dtype(img, tf.float32)
    # 读取并预处理分割图
    seg = tf.io.read_file(seg_path)
    seg = tf.image.decode_png(seg, channels=1)
    seg = tf.cast(seg, tf.int32)
    return img, seg

执行后就能得到每个元素为(图像, 分割标签)的标准数据集。

二、类别权重的两种实现方式 #

在分割任务中，类别权重主要有两种使用场景，你可以根据需求选择：

1. 全局类别权重（适配class_weight参数） #

如果是针对整个数据集的类别不平衡，计算全局的类别权重字典（比如{0:1.0, 1:4.2}），可以直接传给model.fit的class_weight参数。计算逻辑是：统计所有分割图中每个类别的像素总数，用「总像素数 ÷（类别数 × 该类别像素数）」得到对应权重。

2. 逐样本权重映射（Weight Map） #

这是你代码中尝试实现的方式：给每张增强后的分割图生成对应的像素级权重图，让稀有类别像素获得更高权重。这种情况下，权重图必须和增强后的图像、分割图一一绑定，不能分开生成独立数据集。

三、修复你的代码问题：让权重映射基于增强后的分割图 #

你的代码核心问题在于：把数据增强和权重生成拆成了两个独立的map操作，且训练时没有将权重图与(image, seg)配对，还错误地把make_weight_map_batch函数直接传给了需要字典的class_weight参数。

正确的做法是把数据增强和权重生成合并到同一个处理步骤，让增强后的分割图立刻生成对应的权重图，最终数据集的每个元素变为(图像, (分割图, 权重图))，或者(图像, 分割图, 权重图)，再在训练时将权重图传入损失函数。

修正后的完整代码示例 #

# 1. 构建基础数据集
tf_dataset = tf.data.Dataset.from_tensor_slices((train_im_files, train_seg_files))
tf_dataset = tf_dataset.map(process_path, num_parallel_calls=tf.data.AUTOTUNE, deterministic=True)

# 2. 合并增强与权重生成逻辑（注意这里处理单样本，而非batch）
def augment_and_gen_weights(image, seg_mask):
    # 执行数据增强（把你的augment_batch改成单样本处理的augment_sample）
    aug_img, aug_seg = augment_sample(image, seg_mask)
    # 基于增强后的分割图生成权重图（同理，把make_weight_map_batch改成单样本版make_weight_map）
    weight_map = make_weight_map(aug_seg)
    # 返回绑定后的三元组或二元组（二元组的话把seg和weight打包）
    return aug_img, (aug_seg, weight_map)

# 应用合并后的map操作
tf_augmented_with_weights = tf_dataset.map(augment_and_gen_weights, num_parallel_calls=tf.data.AUTOTUNE, deterministic=True)

# 3. 批量处理
batch_dataset = tf_augmented_with_weights.batch(5)

# 4. 自定义带权重的损失函数
def weighted_segmentation_loss(y_true, y_pred):
    # 从y_true中解包分割标签和权重图
    seg_label, weight_map = y_true
    # 基础损失（这里用稀疏分类交叉熵，根据你的任务调整）
    base_loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)(seg_label, y_pred)
    # 应用权重图，求加权平均损失
    return tf.reduce_mean(base_loss * weight_map)

# 编译模型并训练
model.compile(optimizer='adam', loss=weighted_segmentation_loss)
model.fit(batch_dataset, epochs=100)

关键注意事项 #

• 你的augment_batch和make_weight_map_batch是处理批量数据的函数，但tf.data的map默认处理单样本（除非先做batch再map），所以要改成单样本处理逻辑，这样增强和权重生成的对应关系才不会乱。
• 绝对不能把权重生成和增强拆成两个独立数据集，因为tf.data的每个map是独立迭代的，会导致增强后的分割图和权重图完全不匹配。
• class_weight参数只适用于全局类别权重，如果你用的是逐样本weight map，必须把权重图作为数据集的一部分传入，在损失函数中使用。

补充：全局class_weight的生成方法 #

如果你实际需要的是全局类别权重字典，可以用以下代码计算：

def calc_global_class_weights(seg_dataset, num_classes):
    class_counts = tf.zeros(num_classes, dtype=tf.float32)
    total_pixels = 0
    # 遍历所有分割图统计类别像素数
    for seg_mask in seg_dataset:
        counts = tf.math.bincount(tf.reshape(seg_mask, [-1]), minlength=num_classes)
        class_counts += counts
        total_pixels += tf.size(seg_mask)
    # 计算权重：总像素数 / (类别数 × 每个类别的像素数)
    class_weights = total_pixels / (num_classes * class_counts)
    # 转为字典格式
    return {i: class_weights[i].numpy() for i in range(num_classes)}

# 提取分割图数据集用于计算
seg_dataset = tf_dataset.map(lambda x, y: y)
class_weights = calc_global_class_weights(seg_dataset, num_classes=你的类别数)

# 训练时传入class_weight参数
model.fit(tf_augmented.batch(5), epochs=100, class_weight=class_weights)

内容的提问来源于stack exchange，提问作者Tolga Dinçer