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

# 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]]]
# 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 6 surrounding positions
for i in range(3):
    for s in [slice(None, -2), slice(2, None)]:
        slices = tuple(slice(1, -1) if i != j else s for j in range(3))
        m_surround_z |= t_pad_z.__getitem__(slices)
# Surface points are positive points surrounded by some zero
m_surf = m_pos & m_surround_z
coords_surf = tf.where(m_surf)
# Find coordinates that are zero
coords_z = tf.where(~m_pos)
# 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.]]]

编辑：这是使用TensorFlow循环将距离计算划分为多个块的一种方法：

# Following from before
coords_surf = ...
coords_z = ...
CHUNK_SIZE = 1_000 # Choose chunk size
dtype = tf.float32
# If using TF 2.x you can know in advance the size of the tensor array
# (although the element shape will not be constant due to the last chunk)
num_z = tf.shape(coords_z)[0]
arr = tf.TensorArray(dtype, size=(num_z - 1) // CHUNK_SIZE + 1, element_shape=[None], infer_shape=False)
_, arr = tf.while_loop(lambda i, arr: i < num_z,
                       lambda i, arr: (i + CHUNK_SIZE, arr.write(i // CHUNK_SIZE,
                           tf.reduce_min(tf.linalg.norm(tf.dtypes.cast(
                               tf.reshape(coords_z[i:i + CHUNK_SIZE], [-1, 1, 3]) - coords_surf,
                           dtype), axis=-1), axis=-1))),
                       [tf.constant(0, tf.int32), arr])
min_dists = arr.concat()
out = tf.scatter_nd(coords_z, min_dists, [w, h, d])

计算张量流中3D张量的像素方向距离？

1 个答案: