python - Tensorflow 非极大值抑制

Question

注:tf.image.non_max_suppression不符合我的要求!

我正在尝试执行类似于 Canny edge detector 的非最大抑制 (NMS) .具体来说，二维数组上的 NMS 将保留一个值，如果它是窗口内的最大值，否则抑制它(设置为 0)。

例如，考虑矩阵

[[3 2 1 4 2 3]
[1 4 2 1 5 2]
[2 2 3 2 1 3]]


如果我们考虑 3 x 3 的窗口大小，那么结果应该是

[[0 0 0 0 0 0]
[0 4 0 0 5 0]
[0 0 0 0 0 0]]


我四处搜索，但在 tf.image 和 tf.nn 中找不到执行此操作的任何内容。某处是否有执行 NMS 的代码？如果没有，如何在 Tensorflow (Python) 中高效地实现 NMS？

谢谢!

编辑:我想出了一种方法来解决这个问题，但我不确定是否有更好的方法:采用 1 个步幅(即没有下采样)和窗口大小的最大池，然后使用 tf. where 检查值是否等于最大池值，如果不等于则设置为 0。有没有更好的办法？

Answer 1

回答我自己的问题(尽管对更好的解决方案持开放态度):

import tensorflow as tf
import numpy as np

def non_max_suppression(input, window_size):
    # input: B x W x H x C
    pooled = tf.nn.max_pool(input, ksize=[1, window_size, window_size, 1], strides=[1,1,1,1], padding='SAME')
    output = tf.where(tf.equal(input, pooled), input, tf.zeros_like(input))

    # NOTE: if input has negative values, the suppressed values can be higher than original
    return output # output: B X W X H x C

sess = tf.InteractiveSession()

x = np.array([[3,2,1,4,2,3],[1,4,2,1,5,2],[2,2,3,2,1,3]], dtype=np.float32).reshape([1,3,6,1])
inp = tf.Variable(x)
out = non_max_suppression(inp, 3)

sess.run(tf.global_variables_initializer())
print out.eval().reshape([3,6])
'''
[[ 0.  0.  0.  0.  0.  0.]
 [ 0.  4.  0.  0.  5.  0.]
 [ 0.  0.  0.  0.  0.  0.]]
'''

sess.close()