计算张量流中3D张量的像素间距？

Question

我正在尝试在张量流中创建3d距离图（尺寸：W * H * D），以用于损失函数进行训练。我有一个地面真相（大小为W * H * D的二进制体积），将用于创建距离图，即，距离图中每个​​像素的值将是该像素到正值的最小距离（即像素= 1）在地面真实情况下的形状。L2.NORM等3d形状问题会导致轴减少为2D形状，并使该问题完全可微。任何建议或指示将不胜感激。

Answer 1

如果我理解正确，则需要计算从体积中的每个位置到给定类的最接近位置的距离。为简单起见，我将假设有趣的类标记为1，但希望您可以根据情况修改它。该代码用于TensorFlow 2.0，但对于1.x应该相同。

最简单的方法是使用1计算体积中所有坐标与每个坐标之间的距离，然后从那里选择最小距离。您可以这样做：

import tensorflow as tf

# Make input data
w, h, d = 10, 20, 30
w, h, d = 2, 3, 4
t = tf.random.stateless_uniform([w, h, d], (0, 0), 0, 2, tf.int32)
print(t.numpy())
# [[[0 1 0 0]
#   [0 0 0 0]
#   [1 1 0 1]]
#
#  [[1 0 0 0]
#   [0 0 0 0]
#   [1 1 0 0]]]
# Make coordinates
coords = tf.meshgrid(tf.range(w), tf.range(h), tf.range(d), indexing='ij')
coords = tf.stack(coords, axis=-1)
# Find coordinates that are positive
m = t > 0
coords_pos = tf.boolean_mask(coords, m)
# Find every pairwise distance
vec_d = tf.reshape(coords, [-1, 1, 3]) - coords_pos
# You may choose a difference precision type here
dists = tf.linalg.norm(tf.dtypes.cast(vec_d, tf.float32), axis=-1)
# Find minimum distances
min_dists = tf.reduce_min(dists, axis=-1)
# Reshape
out = tf.reshape(min_dists, [w, h, d])
print(out.numpy().round(3))
# [[[1.    0.    1.    2.   ]
#   [1.    1.    1.414 1.   ]
#   [0.    0.    1.    0.   ]]
#
#  [[0.    1.    1.414 2.236]
#   [1.    1.    1.414 1.414]
#   [0.    0.    1.    1.   ]]]

尽管这可能不是最有效的解决方案，但它可能对您足够好。最明智的做法是在每个位置的相邻区域中搜索最接近的正位置，但这通常很难有效地进行，而且在TensorFlow中以向量化的方式更是如此。但是，我们可以通过以下两种方法来改进上面的代码。一方面，我们知道1的位置将始终为零距离，因此无需进行计算。另一方面，如果3D体积中的1类表示某种密集的形状，则仅计算相对于该形状表面的距离，便可以节省一些时间。所有其他正位置将必须与形状外部的位置具有更大的距离。因此，我们可以做与以前相同的事情，但是只计算从非正位置到正表面位置的距离。您可以这样做：

import tensorflow as tf
from itertools import product

# Make input data
w, h, d = 10, 20, 30
w, h, d = 2, 3, 4
t = tf.dtypes.cast(tf.random.stateless_uniform([w, h, d], (0, 0)) > .15, tf.int32)
print(t.numpy())
# [[[1 1 1 1]
#   [1 1 1 1]
#   [1 1 0 0]]
# 
#  [[1 1 1 1]
#   [1 1 1 1]
#   [1 1 1 1]]]
# Make coordinates
coords = tf.meshgrid(tf.range(w), tf.range(h), tf.range(d), indexing='ij')
coords = tf.stack(coords, axis=-1)
# Find coordinates that are positive and on the surface
# (surrounded but at least one 0)
t_pad_z = tf.pad(t, [(1, 1), (1, 1), (1, 1)]) <= 0
m_pos = t > 0
m_surround_z = tf.zeros_like(m_pos)
# Go through the 26 surrounding positions
for s1, s2, s3 in product([slice(None, -2), slice(1, -1), slice(2, None)], repeat=3):
    if s1 == s2 == s3 == slice(1, -1): continue
    m_surround_z |= t_pad_z[s1, s2, s3]
m_surf = m_pos & m_surround_z
coords_surf = tf.boolean_mask(coords, m_surf)
# Find coordinates that are zero
coords_z = tf.boolean_mask(coords, ~m_surf)
# Find every pairwise distance
vec_d = tf.reshape(coords_z, [-1, 1, 3]) - coords_surf
dists = tf.linalg.norm(tf.dtypes.cast(vec_d, tf.float32), axis=-1)
# Find minimum distances
min_dists = tf.reduce_min(dists, axis=-1)
# Put minimum distances in output array
out = tf.scatter_nd(coords_z, min_dists, [w, h, d])
print(out.numpy().round(3))
# [[[0. 0. 0. 0.]
#   [0. 0. 0. 0.]
#   [0. 0. 1. 1.]]
#
#  [[0. 0. 0. 0.]
#   [0. 0. 0. 0.]
#   [0. 0. 0. 0.]]]