def test_collapse_dim(self):
# A tensor of shape [2, 3, 4].
tensor = tf.constant(
[[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]],
[[13, 14, 15, 16], [17, 18, 19, 20], [21, 22, 23, 24]]],
dtype=tf.uint8)
# Collapse using positive axis number
self.assertAllEqual(tf.reshape(tensor, [6, 4]),
utils.collapse_dim(tensor, 1))
self.assertAllEqual(tf.reshape(tensor, [2, 12]),
utils.collapse_dim(tensor, 2))
# Collapse using negative axis number
self.assertAllEqual(tf.reshape(tensor, [2, 12]),
utils.collapse_dim(tensor, -1))
self.assertAllEqual(tf.reshape(tensor, [6, 4]),
utils.collapse_dim(tensor, -2))
def apply_homography(homography, coords):
"""Transform grid of (x,y) texture coordinates by a homography.
Args:
homography: [..., 3, 3]
coords: [..., H, W, 2] (x,y) texture coordinates
Returns:
[..., H, W, 2] transformed coordinates.
"""
height = tf.shape(coords)[-3]
coords = homogenize(utils.collapse_dim(coords, -2)) # [..., H*W, 3]
# Instead of transposing the coords, transpose the homography and
# swap the order of multiplication.
coords = broadcasting_matmul(coords, homography, transpose_b=True)
# coords is now [..., H*W, 3]
return utils.split_dim(dehomogenize(coords), -2, height)