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)