注意:tf.image.non_max_suppression
没有做我正在寻找的事情!
我试图执行类似于Canny edge detector的非最大抑制(NMS)。具体来说,如果2D数组中的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步(即没有下采样)的最大池和窗口大小,然后使用{{1检查该值是否等于最大池值,如果不是则设置为0。还有更好的方法吗?答案 0 :(得分:3)
回答我自己的问题(虽然对更好的解决方案持开放态度):
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()