如何解决tf.image.resize_with_crop_or_pad处理独热分割掩码补全不符合预期的问题

核心需求是将tf.image.resize_with_crop_or_pad默认的0填充，替换为指定背景类的独热编码，以下是两种成熟可行的方案：

方案1：手动实现裁剪+自定义padding（推荐，稳定性最高） #

完全复现原生方法的居中裁剪/填充逻辑，直接指定填充值为背景类独热编码，不会修改原掩码有效区域的数值，兼容性最好。

import tensorflow as tf

def resize_with_crop_or_pad_onehot(x, target_height, target_width, background_class=0):
    """
    x: 输入独热掩码，shape为 [batch_size, H, W, num_classes]
    target_height: 目标高度
    target_width: 目标宽度
    background_class: 填充区域对应的语义类别，默认0为背景
    """
    input_shape = tf.shape(x)
    H, W, C = input_shape[1], input_shape[2], input_shape[3]
    
    # 计算高度方向裁剪/填充量
    delta_h = target_height - H
    pad_top = tf.maximum(delta_h // 2, 0)
    pad_bottom = tf.maximum(delta_h - pad_top, 0)
    crop_top = tf.maximum(-delta_h // 2, 0)
    crop_bottom = H - tf.maximum(delta_h, 0) - crop_top
    
    # 计算宽度方向裁剪/填充量
    delta_w = target_width - W
    pad_left = tf.maximum(delta_w // 2, 0)
    pad_right = tf.maximum(delta_w - pad_left, 0)
    crop_left = tf.maximum(-delta_w // 2, 0)
    crop_right = W - tf.maximum(delta_w, 0) - crop_left
    
    # 先执行裁剪
    x_cropped = x[:, crop_top:H - crop_bottom, crop_left:W - crop_right, :]
    
    # 构造背景类独热填充值
    pad_value = tf.one_hot(background_class, depth=C, dtype=x.dtype)
    
    # 执行填充（batch和通道维度无需填充）
    x_padded = tf.pad(
        x_cropped,
        paddings=[[0, 0], [pad_top, pad_bottom], [pad_left, pad_right], [0, 0]],
        mode="CONSTANT",
        constant_values=pad_value
    )
    return x_padded

# 调用示例
x = get_one_hot_segmentation_mask()  # [batch_size, N, M, num_classes] 
x = resize_with_crop_or_pad_onehot(x, target_height=N//2, target_width=2*M)

方案2：修正默认方法的填充结果（代码更简洁） #

利用「严格生成的独热掩码像素不可能全为0」的特性，直接识别原生方法产生的全0填充区域，替换为目标背景值。

x = get_one_hot_segmentation_mask()  # [batch_size, N, M, num_classes] 
num_classes = tf.shape(x)[-1]
# 先调用原生方法做尺寸调整
x = tf.image.resize_with_crop_or_pad(x, target_height=N//2, target_width=2*M)
# 识别全0的填充像素
zero_pad_mask = tf.reduce_all(x == 0, axis=-1, keepdims=True)
# 替换为背景类独热值，这里假设背景为第0类
background_val = tf.one_hot(0, depth=num_classes, dtype=x.dtype)
x = tf.where(zero_pad_mask, background_val, x)

注意：该方案仅适用于输入独热掩码生成逻辑严格、不存在全0有效像素的场景。

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