def pairs(tensor1, tensor2, name="pairs"):
"""Pairwise combination of elements from the two tensors.
Example::
t1 = [[0],[1]]
t2 = [2,3,4]
t12 = [[0,2],[1,2],[0,3],[1,3],[0,4],[1,4]]
tf.reduce_all(tf.equal(pairs(t1,t2),t12))
Args:
tensor1(``Tensor``): a tensor, python list, or numpy array
tensor2(``Tensor``): a tensor, python list, or numpy array
name: name for this operation (optional)
Returns:
``Tensor``: a ``Tensor`` of rank 2
"""
with ops.name_scope(name):
tensor1 = ops.convert_to_tensor(tensor1)
tensor2 = ops.convert_to_tensor(tensor2)
x, y = array_ops.meshgrid(tensor1, tensor2)
result = array_ops.stack([x, y], axis=-1)
result = array_ops.reshape(result, [-1, 2])
return result
def image_warp3D(image, flow, interp_method, name='dense_image_warp'):
with ops.name_scope(name):
batch_size, height, width, length, channels = (array_ops.shape(image)[0],
array_ops.shape(image)[1],
array_ops.shape(image)[2],
array_ops.shape(image)[3],
array_ops.shape(image)[4])
# The flow is defined on the image grid. Turn the flow into a list of query
# points in the grid space.
grid_x, grid_y, grid_z = array_ops.meshgrid(
math_ops.range(height), math_ops.range(width), math_ops.range(length))
stacked_grid = math_ops.cast(
array_ops.stack([grid_z, grid_y, grid_x], axis=3), flow.dtype)
batched_grid = array_ops.expand_dims(stacked_grid, axis=0)
query_points_on_grid = batched_grid - flow
print(image.get_shape().as_list())
#
query_points_flattened = array_ops.reshape(query_points_on_grid,
[batch_size, height * width * length, 3])
# Compute values at the query points, then reshape the result back to the
# image grid.
if interp_method == 'bilinear' or interp_method == 'Bilinear':
interpolated = _interpolate_bilinear3D(image, query_points_flattened)
elif interp_method == 'nearest_neighbor' or interp_method == 'NN':
interpolated = _interpolate_nearest3D(image, query_points_flattened)
else:
print('Running on bi-linear interpolation!')
interpolated = _interpolate_bilinear3D(image, query_points_flattened)
interpolated = array_ops.reshape(interpolated, [batch_size, height, width, length, channels])
return interpolated
def _compareDiff(self, x, y, use_gpu):
for index in ("ij", "xy"):
numpy_out = np.meshgrid(x, y, indexing=index)
tf_out = array_ops.meshgrid(x, y, indexing=index)
with self.test_session(use_gpu=use_gpu):
for xx, yy in zip(numpy_out, tf_out):
self.assertAllEqual(xx, yy.eval())
def _compareDiff(self, x, y, use_gpu):
for index in ("ij", "xy"):
numpy_out = np.meshgrid(x, y, indexing=index)
tf_out = array_ops.meshgrid(x, y, indexing=index)
with self.test_session(use_gpu=use_gpu):
for xx, yy in zip(numpy_out, tf_out):
self.assertAllEqual(xx, yy.eval())
def random_deformation_linear(shape, std, distance):
r"""Create a random deformation.
Parameters
----------
shape : sequence of 3 ints
Batch, height and width.
std : float
Correlation distance for the linear deformations.
distance : float
Expected total effective distance for the deformation.
Notes
-----
``distance`` must be significantly smaller than ``std`` to guarantee that
the deformation is smooth.
"""
grid_x, grid_y = array_ops.meshgrid(math_ops.range(shape[2]),
math_ops.range(shape[1]))
grid_x = tf.cast(grid_x[None, ..., None], 'float32')
grid_y = tf.cast(grid_y[None, ..., None], 'float32')
# Create mask to stop movement at edges
mask = (tf.cos(
(grid_x - shape[2] / 2 + 1) * np.pi / (shape[2] + 2)) * tf.cos(
(grid_y - shape[1] / 2 + 1) * np.pi / (shape[1] + 2)))**(0.25)
# Scale to get std 1 after smoothing
C = np.sqrt(2 * np.pi) * std
# Multiply by dt here to keep values small-ish for numerical purposes
momenta = distance * C * tf.random_normal(shape=[*shape, 2])
v = mask * gausssmooth(momenta, std)
return v
def _get_grid_locations(image_height, image_width):
"""Wrapper for array_ops.meshgrid."""
y_range = math_ops.linspace(0.0, math_ops.to_float(image_height) - 1,
image_height)
x_range = math_ops.linspace(0.0, math_ops.to_float(image_width) - 1,
image_width)
y_grid, x_grid = array_ops.meshgrid(y_range, x_range, indexing='ij')
return array_ops.stack((y_grid, x_grid), -1)
def dense_image_warp(image, flow, name='dense_image_warp'):
"""Image warping using per-pixel flow vectors.
Apply a non-linear warp to the image, where the warp is specified by a dense
flow field of offset vectors that define the correspondences of pixel values
in the output image back to locations in the source image. Specifically, the
pixel value at output[b, j, i, c] is
images[b, j - flow[b, j, i, 0], i - flow[b, j, i, 1], c].
The locations specified by this formula do not necessarily map to an int
index. Therefore, the pixel value is obtained by bilinear
interpolation of the 4 nearest pixels around
(b, j - flow[b, j, i, 0], i - flow[b, j, i, 1]). For locations outside
of the image, we use the nearest pixel values at the image boundary.
Args:
image: 4-D float `Tensor` with shape `[batch, height, width, channels]`.
flow: A 4-D float `Tensor` with shape `[batch, height, width, 2]`.
name: A name for the operation (optional).
Note that image and flow can be of type tf.half, tf.float32, or tf.float64,
and do not necessarily have to be the same type.
Returns:
A 4-D float `Tensor` with shape`[batch, height, width, channels]`
and same type as input image.
Raises:
ValueError: if height < 2 or width < 2 or the inputs have the wrong number
of dimensions.
"""
with ops.name_scope(name):
# batch_size, height, width, channels = image.get_shape().as_list()
batch_size, height, width, channels = \
array_ops.shape(image)[0], \
array_ops.shape(image)[1], \
array_ops.shape(image)[2], \
array_ops.shape(image)[3]
# The flow is defined on the image grid. Turn the flow into a list of query
# points in the grid space.
grid_x, grid_y = array_ops.meshgrid(math_ops.range(width),
math_ops.range(height))
stacked_grid = math_ops.cast(array_ops.stack([grid_y, grid_x], axis=2),
flow.dtype)
batched_grid = array_ops.expand_dims(stacked_grid, axis=0)
query_points_on_grid = batched_grid - flow
query_points_flattened = array_ops.reshape(
query_points_on_grid, [batch_size, height * width, 2])
# Compute values at the query points, then reshape the result back to the
# image grid.
interpolated = _interpolate_bilinear(image, query_points_flattened)
interpolated = array_ops.reshape(interpolated,
[batch_size, height, width, channels])
return interpolated
def dense_image_warp(image, flow, name='dense_image_warp'):
"""Image warping using per-pixel flow vectors.
Apply a non-linear warp to the image, where the warp is specified by a dense
flow field of offset vectors that define the correspondences of pixel values
in the output image back to locations in the source image. Specifically, the
pixel value at output[b, j, i, c] is
images[b, j - flow[b, j, i, 0], i - flow[b, j, i, 1], c].
The locations specified by this formula do not necessarily map to an int
index. Therefore, the pixel value is obtained by bilinear
interpolation of the 4 nearest pixels around
(b, j - flow[b, j, i, 0], i - flow[b, j, i, 1]). For locations outside
of the image, we use the nearest pixel values at the image boundary.
Args:
image: 4-D float `Tensor` with shape `[batch, height, width, channels]`.
flow: A 4-D float `Tensor` with shape `[batch, height, width, 2]`.
name: A name for the operation (optional).
Note that image and flow can be of type tf.half, tf.float32, or tf.float64,
and do not necessarily have to be the same type.
Returns:
A 4-D float `Tensor` with shape`[batch, height, width, channels]`
and same type as input image.
Raises:
ValueError: if height < 2 or width < 2 or the inputs have the wrong number
of dimensions.
"""
with ops.name_scope(name):
batch_size, height, width, channels = (array_ops.shape(image)[0],
array_ops.shape(image)[1],
array_ops.shape(image)[2],
array_ops.shape(image)[3])
# The flow is defined on the image grid. Turn the flow into a list of query
# points in the grid space.
grid_x, grid_y = array_ops.meshgrid(
math_ops.range(width), math_ops.range(height))
stacked_grid = math_ops.cast(
array_ops.stack([grid_y, grid_x], axis=2), flow.dtype)
batched_grid = array_ops.expand_dims(stacked_grid, axis=0)
query_points_on_grid = batched_grid - flow
query_points_flattened = array_ops.reshape(query_points_on_grid,
[batch_size, height * width, 2])
# Compute values at the query points, then reshape the result back to the
# image grid.
interpolated = _interpolate_bilinear(image, query_points_flattened)
interpolated = array_ops.reshape(interpolated,
[batch_size, height, width, channels])
return interpolated
def coord_indices(tensor, indices, axis=None):
"""
Coordinate indices according to tensor, e.g.
>> a = tf.random.normal((2, 3))
tf.Tensor(
[[-0.02075689 0.7327415 -0.70488846]
[ 2.8714576 -0.2355751 0.33274603]], shape=(2, 3), dtype=float32)
>> indices = tf.argmax(a, axis=0)
tf.Tensor([1 0 1], shape=(3,), dtype=int64) # can not use in gather method
>> indices = coord_indices(a, indices, axis=0)
tf.Tensor(
[[1 0]
[0 1]
[1 2]], shape=(3, 2), dtype=int64) # index along axis=0 was filled
>> tf.gather_nd(a, indices)
tf.Tensor([2.8714576 0.7327415 0.33274603], shape=(3,), dtype=float32)
All we have to do is to calculate missing indices along dim: 0~axis, axis+1~rank
which has same shape as indices, we just stack them together, we then get
gather-style indices
:param tensor: Tensor has rank at least is 1
:param indices: indices(e.g. from tf.argmax) need to coordinate
which must has rank that less than tensor's rank
:param axis: indices along tensor's ? axis
:return: coordinated indices
"""
rank = ndim(tensor)
if rank == 1:
return indices
shape = int_shape(tensor)
ind_shape = int_shape(indices)
if axis is None or axis == -1:
axis = rank
elif axis > rank:
raise ValueError("Axis %d out of rank %d" % (axis, rank))
if len(ind_shape) > rank:
raise ValueError("Indices to coordinate must have rank"
" less than tensor's rank, but received:"
" %d vs %d" % (ind_shape[-1], rank))
if axis != rank:
shape = list(shape)
shape[-1], shape[axis] = shape[axis], shape[-1]
aux_indices = [
math_ops.range(0, shape[i], dtype=indices.dtype)
for i in range(rank - 2, -1, -1)
]
aux_indices = array_ops.meshgrid(*aux_indices)
aux_indices.reverse()
if axis != rank:
aux_indices.insert(axis, indices)
else:
aux_indices.append(indices)
print(aux_indices)
indices = array_ops.stack(aux_indices, axis=-1)
return indices
def _compareDiffType(self, n, np_dtype, use_gpu):
inputs = []
for index in ("ij", "xy"):
for i in range(n):
x = np.linspace(-10, 10, 5).astype(np_dtype)
if np_dtype in (np.complex64, np.complex128):
x += 1j
inputs.append(x)
numpy_out = np.meshgrid(*inputs, indexing=index)
with self.test_session(use_gpu=use_gpu):
tf_out = array_ops.meshgrid(*inputs, indexing=index)
for X, _X in zip(numpy_out, tf_out):
self.assertAllEqual(X, _X.eval())
def _compareDiffType(self, n, np_dtype, use_gpu):
inputs = []
for index in ("ij", "xy"):
for i in range(n):
x = np.linspace(-10, 10, 5).astype(np_dtype)
if np_dtype in (np.complex64, np.complex128):
x += 1j
inputs.append(x)
numpy_out = np.meshgrid(*inputs, indexing=index)
with self.test_session(use_gpu=use_gpu):
tf_out = array_ops.meshgrid(*inputs, indexing=index)
for X, _X in zip(numpy_out, tf_out):
self.assertAllEqual(X, _X.eval())
def _get_boundary_locations(image_height, image_width, num_points_per_edge):
"""Compute evenly-spaced indices along edge of image."""
image_height = math_ops.to_float(image_height)
image_width = math_ops.to_float(image_width)
y_range = math_ops.linspace(0.0, image_height - 1, num_points_per_edge + 2)
x_range = math_ops.linspace(0.0, image_width - 1, num_points_per_edge + 2)
ys, xs = array_ops.meshgrid(y_range, x_range, indexing='ij')
is_boundary = math_ops.logical_or(
math_ops.logical_or(math_ops.equal(xs, 0),
math_ops.equal(xs, image_width - 1)),
math_ops.logical_or(math_ops.equal(ys, 0),
math_ops.equal(ys, image_height - 1)))
return array_ops.stack([array_ops.boolean_mask(ys, is_boundary),
array_ops.boolean_mask(xs, is_boundary)], axis=-1)
def gausssmooth(img, std):
# Create gaussian
size = tf.cast(std * 5, 'int32')
size = (size // 2) * 2 + 1
size_f = tf.cast(size, 'float32')
grid_x, grid_y = array_ops.meshgrid(math_ops.range(size),
math_ops.range(size))
grid_x = tf.cast(grid_x[None, ..., None], 'float32')
grid_y = tf.cast(grid_y[None, ..., None], 'float32')
gaussian = tf.exp(-((grid_x - size_f / 2 - 0.5) ** 2 + (grid_y - size_f / 2 + 0.5) ** 2) / std ** 2)
gaussian = gaussian / tf.reduce_sum(gaussian)
return fftconv(img, gaussian, 'same')
def testCompare(self):
for t in (np.float16, np.float32, np.float64, np.int32, np.int64,
np.complex64, np.complex128):
# Don't test the one-dimensional case, as
# old numpy versions don't support it
self._compare(2, t, False)
self._compare(3, t, False)
self._compare(4, t, False)
self._compare(5, t, False)
# Test for inputs with rank not equal to 1
x = [[1, 1], [1, 1]]
with self.assertRaisesRegexp(errors.InvalidArgumentError,
"needs to have rank 1"):
with self.test_session():
X, _ = array_ops.meshgrid(x, x)
X.eval()
def testCompare(self):
for t in (np.float16, np.float32, np.float64, np.int32, np.int64,
np.complex64, np.complex128):
# Don't test the one-dimensional case, as
# old numpy versions don't support it
self._compare(2, t, False)
self._compare(3, t, False)
self._compare(4, t, False)
self._compare(5, t, False)
# Test for inputs with rank not equal to 1
x = [[1, 1], [1, 1]]
with self.assertRaisesRegexp(errors.InvalidArgumentError,
"needs to have rank 1"):
with self.test_session():
X, _ = array_ops.meshgrid(x, x)
X.eval()
def affine_grid(theta, size: (list, tuple), name='affine_grid'):
with ops.name_scope(name):
x = gen_math_ops.lin_space(-1., 1., size[1])
y = gen_math_ops.lin_space(-1., 1., size[2])
x_t, y_t = array_ops.meshgrid(x, y)
x_t = array_ops.reshape(x_t, shape=(-1, ))
y_t = array_ops.reshape(y_t, shape=(-1, ))
ones = array_ops.ones_like(x_t)
grids = array_ops.stack([x_t, y_t, ones])
grids = array_ops.expand_dims(grids, axis=0)
grids = array_ops.tile(grids,
multiples=array_ops.stack([size[0], 1, 1]))
grids = float32(grids)
theta = float32(theta)
grids = math_ops.matmul(theta, grids)
grids = array_ops.reshape(grids, shape=(size[0], 2, size[1], size[2]))
return grids
def meshgrid(*xi, **kwargs):
"""This currently requires copy=True and sparse=False."""
sparse = kwargs.get('sparse', False)
if sparse:
raise ValueError('meshgrid doesnt support returning sparse arrays yet')
copy = kwargs.get('copy', True)
if not copy:
raise ValueError('meshgrid only supports copy=True')
indexing = kwargs.get('indexing', 'xy')
xi = [np_array_ops.asarray(arg) for arg in xi]
kwargs = {'indexing': indexing}
outputs = array_ops.meshgrid(*xi, **kwargs)
return outputs
def dense_image_warp(image, flow, name='dense_image_warp'):
with ops.name_scope(name):
batch_size, height, width, channels = array_ops.unstack(
array_ops.shape(image))
# The flow is defined on the image grid. Turn the flow into a list of query
# points in the grid space.
grid_x, grid_y = array_ops.meshgrid(
math_ops.range(width), math_ops.range(height))
stacked_grid = math_ops.cast(
array_ops.stack([grid_y, grid_x], axis=2), flow.dtype)
batched_grid = array_ops.expand_dims(stacked_grid, axis=0)
query_points_on_grid = batched_grid - flow
query_points_flattened = array_ops.reshape(query_points_on_grid,
[batch_size, height * width, 2])
# Compute values at the query points, then reshape the result back to the
# image grid.
interpolated = _interpolate_bilinear(image, query_points_flattened)
interpolated = array_ops.reshape(interpolated,
[batch_size, height, width, channels])
return interpolated
def modified_dense_image_warp(image, flow, name='dense_image_warp'):
"""
Modified version of Tensorflow's dense_image_warp function.
This assumes that flow vectors are stored at source location, so the
pixel value at output[b, j, i, c] is
images[b, j + flow[b, j, i, 0], i + flow[b, j, i, 1], c]
"""
with ops.name_scope(name):
batch_size, height, width, channels = (array_ops.shape(image)[0],
array_ops.shape(image)[1],
array_ops.shape(image)[2],
array_ops.shape(image)[3])
# The flow is defined on the image grid. Turn the flow into a list of query
# points in the grid space.
grid_x, grid_y = array_ops.meshgrid(math_ops.range(width),
math_ops.range(height))
stacked_grid = math_ops.cast(array_ops.stack([grid_y, grid_x], axis=2),
flow.dtype)
batched_grid = array_ops.expand_dims(stacked_grid, axis=0)
# ToDo: this is a hack, change the quaery points instead.
# swap u and v in the flow tensor
flow = array_ops.reverse(flow, axis=[3])
query_points_on_grid = batched_grid + flow
query_points_flattened = array_ops.reshape(
query_points_on_grid, [batch_size, height * width, 2])
# Compute values at the query points, then reshape the result back to the
# image grid.
interpolated = _interpolate_bilinear(image, query_points_flattened)
interpolated = array_ops.reshape(interpolated,
[batch_size, height, width, channels])
return interpolated
def testEmptyMeshgrid(self):
self.assertEqual(array_ops.meshgrid(), [])
def random_deformation_momentum(shape, std, distance, stepsize=0.1):
r"""Create a random diffeomorphic deformation.
Parameters
----------
shape : sequence of 3 ints
Batch, height and width.
std : float
Correlation distance for the linear deformations.
distance : float
Expected total effective distance for the deformation.
stepsize : float
How large each step should be (as a propotion of ``std``).
Notes
-----
``distance`` should typically not be more than a small fraction of the
sidelength of the image.
The computational time is is propotional to
.. math::
\frac{distance}{std * stepsize}
"""
grid_x, grid_y = array_ops.meshgrid(math_ops.range(shape[2]),
math_ops.range(shape[1]))
grid_x = tf.cast(grid_x[None, ..., None], 'float32')
grid_y = tf.cast(grid_y[None, ..., None], 'float32')
base_coordinates = tf.concat([grid_y, grid_x], axis=-1)
coordinates = tf.identity(base_coordinates)
# Create mask to stop movement at edges
mask = (tf.cos(
(grid_x - shape[2] / 2 + 1) * np.pi / (shape[2] + 2)) * tf.cos(
(grid_y - shape[1] / 2 + 1) * np.pi / (shape[1] + 2)))**(0.25)
# Total distance is given by std * n_steps * dt, we use this
# to work out the exact numbers.
n_steps = tf.cast(tf.ceil(distance / (std * stepsize)), 'int32')
dt = distance / (tf.cast(n_steps, 'float32') * std)
# Scale to get std 1 after smoothing
C = np.sqrt(2 * np.pi) * std**2
# Multiply by dt here to keep values small-ish for numerical purposes
momenta = dt * C * tf.random_normal(shape=[*shape, 2])
# Using a while loop, generate the deformation step-by-step.
def cond(i, from_coordinates, momenta):
return i < n_steps
def body(i, from_coordinates, momenta):
v = mask * gausssmooth(momenta, std)
d1 = matmul_transposed(jacobian(momenta), v)
d2 = matmul(jacobian(v), momenta)
d3 = div(v) * momenta
momenta = momenta - dt * (d1 + d2 + d3)
v = dense_image_warp(v, from_coordinates - base_coordinates)
from_coordinates = dense_image_warp(from_coordinates, v)
return i + 1, from_coordinates, momenta
i = tf.constant(0, dtype=tf.int32)
i, from_coordinates, momenta = tf.while_loop(cond, body,
[i, coordinates, momenta])
from_total_offset = from_coordinates - base_coordinates
return from_total_offset
def spatial_softmax(features,
temperature=None,
name=None,
variables_collections=None,
trainable=True,
data_format='NHWC'):
"""Computes the spatial softmax of a convolutional feature map.
First computes the softmax over the spatial extent of each channel of a
convolutional feature map. Then computes the expected 2D position of the
points of maximal activation for each channel, resulting in a set of
feature keypoints [x1, y1, ... xN, yN] for all N channels.
Read more here:
"Learning visual feature spaces for robotic manipulation with
deep spatial autoencoders." Finn et al., http://arxiv.org/abs/1509.06113.
Args:
features: A `Tensor` of size [batch_size, W, H, num_channels]; the
convolutional feature map.
temperature: Softmax temperature (optional). If None, a learnable
temperature is created.
name: A name for this operation (optional).
variables_collections: Collections for the temperature variable.
trainable: If `True` also add variables to the graph collection
`GraphKeys.TRAINABLE_VARIABLES` (see `tf.Variable`).
data_format: A string. `NHWC` (default) and `NCHW` are supported.
Returns:
feature_keypoints: A `Tensor` with size [batch_size, num_channels * 2];
the expected 2D locations of each channel's feature keypoint (normalized
to the range (-1,1)). The inner dimension is arranged as
[x1, y1, ... xN, yN].
Raises:
ValueError: If unexpected data_format specified.
ValueError: If num_channels dimension is unspecified.
"""
shape = array_ops.shape(features)
static_shape = features.shape
height, width, num_channels = shape[1], shape[2], static_shape[3]
if num_channels.value is None:
raise ValueError('The num_channels dimension of the inputs to '
'`spatial_softmax` should be defined. Found `None`.')
with ops.name_scope(name, 'spatial_softmax', [features]) as name:
# Create tensors for x and y coordinate values, scaled to range [-1, 1].
pos_x, pos_y = array_ops.meshgrid(math_ops.lin_space(-1.,
1.,
num=height),
math_ops.lin_space(-1.,
1.,
num=width),
indexing='ij')
pos_x = array_ops.reshape(pos_x, [height * width])
pos_y = array_ops.reshape(pos_y, [height * width])
if temperature is None:
temperature_collections = utils.get_variable_collections(
variables_collections, name + 'temperature')
temperature = variables.model_variable(
name + 'temperature',
shape=(),
dtype=dtypes.float32,
initializer=init_ops.ones_initializer(),
collections=temperature_collections,
trainable=trainable)
# We assume all ops are [NBATCH, HEIGHT, WIDTH, CHANNELS] but this code
# does not! It will reorder them appropriately.
features = array_ops.reshape(
array_ops.transpose(features, [0, 3, 1, 2]),
[-1, height * width])
softmax_attention = nn.softmax(features / temperature)
expected_x = math_ops.reduce_sum(pos_x * softmax_attention, [1],
keep_dims=True)
expected_y = math_ops.reduce_sum(pos_y * softmax_attention, [1],
keep_dims=True)
expected_xy = array_ops.concat([expected_x, expected_y], 1)
feature_keypoints = array_ops.reshape(expected_xy,
[-1, num_channels.value * 2])
feature_keypoints.set_shape([None, num_channels.value * 2])
return feature_keypoints
def dense_image_warp_3D(tensors, name='dense_image_warp'):
"""Image warping using per-pixel flow vectors.
Apply a non-linear warp to the image, where the warp is specified by a dense
flow field of offset vectors that define the correspondences of pixel values
in the output image back to locations in the source image. Specifically, the
pixel value at output[b, j, i, k, c] is
images[b, j - flow[b, j, i, k, 0], i - flow[b, j, i, k, 1], k - flow[b, j, i, k, 2], c].
The locations specified by this formula do not necessarily map to an int
index. Therefore, the pixel value is obtained by trilinear
interpolation of the 8 nearest pixels around
(b, j - flow[b, j, i, k, 0], i - flow[b, j, i, k, 1], k - flow[b, j, i, k, 2]). For locations outside
of the image, we use the nearest pixel values at the image boundary.
Args:
image: 5-D float `Tensor` with shape `[batch, height, width, depth, channels]`.
flow: A 5-D float `Tensor` with shape `[batch, height, width, depth, 3]`.
name: A name for the operation (optional).
Note that image and flow can be of type tf.half, tf.float32, or tf.float64,
and do not necessarily have to be the same type.
Returns:
A 5-D float `Tensor` with shape`[batch, height, width, depth, channels]`
and same type as input image.
Raises:
ValueError: if height < 2 or width < 2 or the inputs have the wrong number
of dimensions.
"""
DEBUG = 0
image = tensors[0]
flow = tensors[1]
batch_size, height, width, depth, channels = (array_ops.shape(image)[0],
array_ops.shape(image)[1],
array_ops.shape(image)[2],
array_ops.shape(image)[3],
array_ops.shape(image)[4])
# The flow is defined on the image grid. Turn the flow into a list of query
# points in the grid space.
#grid_x, grid_y, grid_z = array_ops.meshgrid(math_ops.range(width), math_ops.range(height), math_ops.range(depth))
#stacked_grid = math_ops.cast(array_ops.stack([grid_y, grid_x, grid_z], axis=3), flow.dtype)
grid_i, grid_j, grid_k = array_ops.meshgrid(math_ops.range(height),
math_ops.range(width),
math_ops.range(depth),
indexing='ij')
stacked_grid = math_ops.cast(
array_ops.stack([grid_i, grid_j, grid_k], axis=3), flow.dtype)
batched_grid = array_ops.expand_dims(stacked_grid,
axis=0) #add the batch dim on axis 0
query_points_on_grid = batched_grid - flow
if DEBUG:
query_points_on_grid = K.print_tensor(
query_points_on_grid, message="query_points_on_grid is:")
query_points_flattened = array_ops.reshape(
query_points_on_grid, [batch_size, height * width * depth, 3])
if DEBUG:
query_points_flattened = K.print_tensor(
query_points_flattened, message="query_points_flattened is:")
# Compute values at the query points, then reshape the result back to the
# image grid.
interpolated = interpolate_trilinear(image,
query_points_flattened,
indexing='ijk')
if DEBUG:
interpolated = K.print_tensor(interpolated, message='interpolated is:')
interpolated = array_ops.reshape(
interpolated, [batch_size, height, width, depth, channels])
return interpolated
