在tensorflow中将二进制掩码转换为边界框

时间:2020-09-09 04:49:39

标签: tensorflow tensor

我在张量流中有一个二进制掩码作为张量。

如何使用张量流操作将此二进制掩码转换为边界框?

1 个答案:

答案 0 :(得分:0)

经过一些工作,我设法解决了它。请注意,给出的解决方案仅适用于单个对象,但是通过一些调整,您也可以将其应用于多个对象。

基本上,您要检查整个轴上是否有真实像素。从边缘开始,然后进一步向内移动,直到命中具有至少一个真实像素的轴。对左,右,顶部和底部执行此操作。

image = tf.io.read_file('mask.png')
image = tf.io.decode_png(image)
image = tf.image.resize(image, size=(300, 300), method='nearest')

rows = tf.math.count_nonzero(image, axis=0, keepdims=None, dtype=tf.bool) # return true if any pixels in the given row is true
rows = tf.squeeze(rows, axis=1) #make a scalar

columns = tf.math.count_nonzero(image, axis=1, keepdims=None, dtype=tf.bool)
columns = tf.squeeze(columns, axis=1)

def indicies_by_value(value): return tf.where(tf.equal(value, True))[:,-1] #return all the indices where mask is present along given axis

#coordinates
y_min = indicies_by_value(columns)[0] #first true pixel along axis
y_max = indicies_by_value(columns)[-1] #last true pixel along axis
x_min = indicies_by_value(rows)[0]
x_max = indicies_by_value(rows)[-1]

#apply the bounding box
image = tf.image.convert_image_dtype(image, dtype=tf.float32)
img = tf.expand_dims(image, axis=-0)
img = tf.reshape(img, shape=[1, 300, 300, 1])

box = tf.stack([y_min, x_min, y_max, x_max], axis=0)
box = tf.math.divide(box, 300)
box = box.numpy()
boxes = box.reshape([1,1,4])

colors = np.array([[0.5, 0.9, 0.5], [0.5, 0.9, 0.5]])
boundning_box = tf.image.draw_bounding_boxes(img, boxes, colors)

tf.keras.preprocessing.image.save_img('boxed.png', boundning_box.numpy()[0])