def test_visualization(self):
if not VISUALIZE:
return
cur_dir = os.path.dirname(os.path.abspath(__file__))
root_dir = os.path.join(cur_dir, '..')
flow_path = os.path.join(root_dir, 'pwcnet', 'warp', 'test_data',
'flow_ab.flo')
image_path = os.path.join(root_dir, 'pwcnet', 'warp', 'test_data',
'image_a.png')
flow_ab = [read_flow_file(flow_path)]
img_a = [read_image(image_path, as_float=True)]
t_tensor = tf.placeholder(tf.float32, None)
flow_ab_tensor = tf.placeholder(tf.float32, np.shape(flow_ab))
img_a_tensor = tf.placeholder(tf.float32, np.shape(img_a))
warp_tensor = forward_warp(img_a_tensor, t_tensor * flow_ab_tensor)
warp = self.sess.run(warp_tensor,
feed_dict={
flow_ab_tensor: flow_ab,
img_a_tensor: img_a,
t_tensor: 1.0
})
warp = np.clip(warp[0], 0.0, 1.0)
show_image(warp)
# For writing to video.
if WRITE_TO_VIDEO:
import cv2
import mpimg
height = img_a[0].shape[0]
width = img_a[0].shape[1]
writer = cv2.VideoWriter(cur_dir + '/outputs/warped.avi',
cv2.VideoWriter_fourcc(*'MJPG'), 20,
(width, height))
steps = 60
for i in range(steps):
print('Writing video at step %d' % i)
t = i * (1.0 / float(steps))
warped = self.sess.run(warp_tensor,
feed_dict={
flow_ab_tensor: flow_ab,
img_a_tensor: img_a,
t_tensor: t
})
warped = warped[0]
warped = np.clip(warped, 0.0, 1.0)
output_path = cur_dir + '/outputs/out-%.2f.png' % t
mpimg.imsave(output_path, warped)
writer.write(cv2.imread(output_path))
writer.release()
def get_forward(self, image_a, image_b, t, reuse_variables=tf.AUTO_REUSE):
"""
:param image_a: Tensor of shape [batch_size, H, W, 3].
:param image_b: Tensor of shape [batch_size, H, W, 3].
:param t: Float. Specifies the interpolation point (i.e 0 for image_a, 1 for image_b).
:return: interpolated: The interpolated image. Tensor of shape [batch_size, H, W, 3].
warped_a_b: Image and features from a forward-flowed towards b, before synthesis.
The first 3 channels are the image.
warped_b_a: Image and features from b forward-flowed towards a, before synthesis.
The first 3 channels are the image.
flow_a_b: Flow from a to b (centered at a).
flow_b_a: Flow from b to a (centered at b).
"""
self.enclosing_scope = tf.get_variable_scope()
with tf.variable_scope(self.name, reuse=reuse_variables):
batch_size = tf.shape(image_a)[0]
from_frames = tf.concat([image_a, image_b], axis=0)
to_frames = tf.concat([image_b, image_a], axis=0)
all_contexts = self.feature_extractor.get_context_features(
from_frames)
# TODO: Add instance normalization. Described in 3.3 of https://arxiv.org/pdf/1803.10967.pdf.
# Get a->b and b->a flows from PWCNet.
# TODO: Migrate to pwcnet.get_bidirectional.
all_flows, _ = self.pwcnet.get_forward(
from_frames, to_frames, reuse_variables=reuse_variables)
flow_a_b = all_flows[:batch_size]
flow_b_a = all_flows[batch_size:]
features_a = tf.concat([image_a, all_contexts[:batch_size]],
axis=-1)
features_b = tf.concat([image_b, all_contexts[batch_size:]],
axis=-1)
all_features = tf.concat([features_a, features_b], axis=0)
all_warp_flows = tf.concat([t * flow_a_b, (1.0 - t) * flow_b_a],
axis=0)
# Warp images and their contexts from a->b and from b->a.
all_warped = forward_warp(all_features, all_warp_flows)
warped_a_b = tf.stop_gradient(all_warped[:batch_size])
warped_b_a = tf.stop_gradient(all_warped[batch_size:])
# Feed into GridNet for final synthesis.
warped_combined = tf.concat([warped_a_b, warped_b_a], axis=-1)
synthesized, _, _, _ = self.gridnet.get_forward(warped_combined,
training=True)
return synthesized, warped_a_b, warped_b_a, flow_a_b, flow_b_a
def test_forward_warp_partial_1(self):
height = 2
width = 2
# Flow is in (x, y) order.
# Splats the top-left pixel right in the center.
flow = [[[[0.5, 0.5], [0, 0]], [[0, 0], [0, 0]]]]
features = [[[[4, 0], [0, 0]], [[1, 1], [0, 0]]]]
expected_warp = [[[[1, 0], [1, 0]], [[2, 1], [1, 0]]]]
flow_tensor = tf.placeholder(tf.float32, (1, height, width, 2))
features_tensor = tf.placeholder(tf.float32, (1, height, width, 2))
warp_tensor = forward_warp(features_tensor, flow_tensor)
warp = self.sess.run(warp_tensor,
feed_dict={
flow_tensor: flow,
features_tensor: features
})
self.assertEqual(warp.tolist(), expected_warp)
def test_forward_warp_batch(self):
height = 2
width = 2
# Flow is in (x, y) order.
# Splats the top-left pixel right in the center.
flow = [[[[0.5, 0.5], [0, 0]], [[0, 0], [0, 0]]],
[[[1, 1], [0, 0]], [[0, 0], [0, 0]]]]
features = [[[[4, 0], [0, 0]], [[1, 1], [0, 0]]],
[[[100, 0], [0, 0]], [[1, 1], [0, 0]]]]
expected_warp = [[[[1, 0], [1, 0]], [[2, 1], [1, 0]]],
[[[0, 0], [0, 0]], [[1, 1], [100, 0]]]]
flow_tensor = tf.placeholder(tf.float32, (2, height, width, 2))
features_tensor = tf.placeholder(tf.float32, (2, height, width, 2))
warp_tensor = forward_warp(features_tensor, flow_tensor)
warp = self.sess.run(warp_tensor,
feed_dict={
flow_tensor: flow,
features_tensor: features
})
self.assertEqual(warp[0].tolist(), expected_warp[0])
self.assertEqual(warp[1].tolist(), expected_warp[1])
# Check for gradients.
grads_tensor = tf.gradients(warp_tensor[0][0][0],
[flow_tensor, features_tensor])
for grad_tensor in grads_tensor:
self.assertNotEqual(grad_tensor, None)
grads = self.sess.run(grads_tensor,
feed_dict={
flow_tensor: flow,
features_tensor: features
})
flow_grads, feature_grads = grads[0][0], grads[1][0]
self.assertNotEqual(np.sum(flow_grads), 0.0)
self.assertNotEqual(np.sum(feature_grads), 0.0)
def test_forward_warp_oob(self):
"""
Note that oob == out of bounds.
"""
height = 3
width = 2
# Flow is in (x, y) order.
# Splats the top-left pixel right in the center.
flow = [[[[1.5, 0], [0, 0]], [[0, 0], [0, 0]], [[0, 0], [-10, -10]]]]
features = [[[[4, 0], [0, 0]], [[1, 1], [0, 0]], [[0, 0], [-4, -4]]]]
expected_warp = [[[[0, 0], [2, 0]], [[1, 1], [0, 0]], [[0, 0], [0,
0]]]]
flow_tensor = tf.placeholder(tf.float32, (1, height, width, 2))
features_tensor = tf.placeholder(tf.float32, (1, height, width, 2))
warp_tensor = forward_warp(features_tensor, flow_tensor)
warp = self.sess.run(warp_tensor,
feed_dict={
flow_tensor: flow,
features_tensor: features
})
self.assertEqual(warp.tolist(), expected_warp)