如何在Keras/TensorFlow自定义指标中对y_true进行膨胀操作？

直接用TensorFlow原生算子实现膨胀操作，避免张量转numpy的问题，以下是具体方案：

核心实现：用tf.nn.dilation2d实现等价于cv2.dilate的张量操作 #

形态学膨胀的本质是用结构元素遍历图像，取覆盖区域内的最大值，TensorFlow的tf.nn.dilation2d就是专门做这个的，完全兼容计算图模式，适合在自定义指标中使用。

膨胀函数实现 #

import tensorflow as tf

def tf_dilate(y_true, kernel_size=(3,3)):
    # y_true 形状为 (batch, height, width, channels)
    channels = tf.shape(y_true)[-1]
    # 创建与cv2.dilate等价的全1结构元素
    kernel = tf.ones(
        shape=kernel_size + (channels, channels),
        dtype=y_true.dtype
    )
    # 执行膨胀操作，参数对应cv2.dilate的默认行为
    dilated = tf.nn.dilation2d(
        input=y_true,
        filters=kernel,
        strides=[1, 1, 1, 1],  # 步长1，对应cv2.dilate的默认滑动方式
        padding='SAME',         # 边缘填充0，与cv2.dilate默认padding一致
        data_format='NHWC'      # 匹配TensorFlow默认的通道最后格式
    )
    return dilated

在U-net自定义指标中使用 #

以自定义基于膨胀标签的交并比指标为例：

from tensorflow.keras import metrics

class DilatedIoU(metrics.Metric):
    def __init__(self, name='dilated_iou', kernel_size=(3,3), **kwargs):
        super().__init__(name=name, **kwargs)
        self.kernel_size = kernel_size
        self.total_iou = self.add_weight(name='total', initializer='zeros')
        self.sample_count = self.add_weight(name='count', initializer='zeros')
    
    def update_state(self, y_true, y_pred, sample_weight=None):
        # 对真实标签执行膨胀
        y_true_dilated = tf_dilate(y_true, self.kernel_size)
        
        # 计算交并比
        intersection = tf.reduce_sum(y_pred * y_true_dilated, axis=[1,2,3])
        union = tf.reduce_sum(y_pred, axis=[1,2,3]) + tf.reduce_sum(y_true_dilated, axis=[1,2,3]) - intersection
        iou = intersection / (union + tf.keras.backend.epsilon())
        
        # 处理样本权重
        if sample_weight is not None:
            sample_weight = tf.cast(sample_weight, dtype=self.dtype)
            iou = iou * sample_weight
            self.sample_count.assign_add(tf.reduce_sum(sample_weight))
        else:
            self.sample_count.assign_add(tf.cast(tf.shape(y_true)[0], self.dtype))
        
        self.total_iou.assign_add(tf.reduce_sum(iou))
    
    def result(self):
        return self.total_iou / (self.sample_count + tf.keras.backend.epsilon())
    
    def reset_state(self):
        self.total_iou.assign(0.)
        self.sample_count.assign(0.)

关键说明 #

• 效果对齐：只要结构元素大小、padding方式和cv2.dilate一致，tf_dilate的输出和cv2.dilate处理numpy数组的结果完全等价
• 计算图兼容：全程使用TensorFlow原生算子，不会打破计算图，支持分布式训练、模型导出等场景
• 避免转numpy：无需将张量转numpy数组，解决了cv2.dilate无法直接处理张量的问题

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