时间戳

首页 > 解决方案 > TensorFlow 2.2 中的自定义二值化（阈值化）图像层

image-processing - TensorFlow 2.2 中的自定义二值化（阈值化）图像层

问题描述

我正在 Tensorflow 中构建一个 CNN，其输出需要是二进制图像 (0, 1)。我想用一个可调变量（阈值）实现一个自定义层，它将格式化最终图像，例如：

>>> a
array([[ 0.02 ,  0.833,  0.778,  0.87 ],
       [ 0.979,  0.799,  0.461,  0.781],
       [ 0.118,  0.64 ,  0.143,  0.945],
       [ 0.522,  0.415,  0.265,  0.774]])

>>> np.where(a > threshold, 1, 0) 
array([[0, 1, 1, 1],
       [1, 1, 0, 1],
       [0, 1, 0, 1],
       [1, 0, 0, 1]])

我试图将其实现为：

class BinarizeLayer(keras.layers.Layer):
    def __init__(self, **kwargs):
        super(BinarizeLayer, self).__init__(**kwargs)
        self.threshold = tf.Variable(initial_value=0.5, dtype=tf.float32, name='BinaryThreshold')

    def call(self, x):
        cond = tf.math.less(x, tf.fill(tf.shape(x), self.threshold))
        out = tf.where(cond, tf.zeros(tf.shape(x)), tf.ones(tf.shape(x)))
        return out

但是，这会导致ValueError: An operation has 'None' for gradient.. 有没有解决的办法？我尝试将 tf.stop_gradient() 添加到“调用”的所有部分，但错误仍然存​​在。

标签： image-processingtensorflow2.0

解决方案

无需实现自定义层，tensorflow 已经支持 https://www.tensorflow.org/api_docs/python/tf/where

tf.where(a > threshold, 1, 0)

推荐阅读