def setUp(self):
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
self.sess = tf.Session(config=config)
# Load from files.
s = 0.5 # Scale.
self.flow_ab = cv2.resize(
read_flow_file('common/utils/test_data/flow_ab.flo'), (0, 0),
fx=s,
fy=s) * s
self.img_a = cv2.resize(read_image(
'common/utils/test_data/image_a.png', as_float=True), (0, 0),
fx=s,
fy=s)
self.img_b = cv2.resize(read_image(
'common/utils/test_data/image_b.png', as_float=True), (0, 0),
fx=s,
fy=s)
self.flow_cd = cv2.resize(
read_flow_file('common/utils/test_data/flow_cd.flo'), (0, 0),
fx=s,
fy=s) * s
self.img_c = cv2.resize(read_image(
'common/utils/test_data/image_c.png', as_float=True), (0, 0),
fx=s,
fy=s)
self.img_d = cv2.resize(read_image(
'common/utils/test_data/image_d.png', as_float=True), (0, 0),
fx=s,
fy=s)
self.tf_true = tf.constant(True)
self.tf_false = tf.constant(False)
def test_tensorflow_vis(self):
"""
Tests the tensorflow implementation and the batching against the numpy version.
"""
flow_image = read_flow_file('common/utils/test_data/frame_0001.flo')
flow_image_2 = read_flow_file('common/utils/test_data/frame_0011.flo')
flow_ph = tf.placeholder(shape=[
None, flow_image.shape[0], flow_image.shape[1], flow_image.shape[2]
],
dtype=tf.float32)
visualization_tensor = get_tf_flow_visualization(flow_ph)
feed_dict = {flow_ph: np.stack([flow_image, flow_image_2], axis=0)}
visualization = self.sess.run(visualization_tensor,
feed_dict=feed_dict)
target_visualization = get_flow_visualization(flow_image) / 255.0
target_visualization_2 = get_flow_visualization(flow_image_2) / 255.0
self.assertTrue(
np.allclose(target_visualization, visualization[0], atol=0.06))
self.assertTrue(
np.allclose(target_visualization_2, visualization[1], atol=0.06))
if SHOW_FLOW_TEST_IMAGES:
show_image(visualization[0])
show_image(visualization[1])
def test_write(self):
flow_image_expected = read_flow_file(
'common/utils/test_data/frame_0001.flo')
write_flow_file(self.out_file, flow_image_expected)
flow_image = read_flow_file(self.out_file)
self.assertTrue(np.allclose(flow_image_expected, flow_image))
if SHOW_FLOW_TEST_IMAGES:
show_flow_image(flow_image)
def test_flow_vis(self):
flow_image = read_flow_file('common/utils/test_data/frame_0001.flo')
self.assertTrue(flow_image is not None)
expected_shape = (436, 1024, 2)
self.assertTupleEqual((436, 1024, 2), flow_image.shape)
if SHOW_FLOW_TEST_IMAGES:
show_flow_image(flow_image)
flow_image = read_flow_file('common/utils/test_data/frame_0011.flo')
self.assertTrue(flow_image is not None)
self.assertTupleEqual(expected_shape, flow_image.shape)
if SHOW_FLOW_TEST_IMAGES:
show_flow_image(flow_image)
def test_optical_flow_warp_sintel(self):
"""
Runs warp test against Sintel ground truth and checks for <2% average pixel error.
Also masks out the occluded warp regions when computing the error.
"""
# Load from files.
flow_ab = read_flow_file('pwcnet/warp/test_data/flow_ab.flo')
img_a = read_image('pwcnet/warp/test_data/image_a.png', as_float=True)
img_b = read_image('pwcnet/warp/test_data/image_b.png', as_float=True)
flow_cd = read_flow_file('pwcnet/warp/test_data/flow_cd.flo')
img_c = read_image('pwcnet/warp/test_data/image_c.png', as_float=True)
img_d = read_image('pwcnet/warp/test_data/image_d.png', as_float=True)
mask_ab = np.ones(shape=img_a.shape)
mask_cd = np.ones(shape=img_a.shape)
H = img_a.shape[0]
W = img_a.shape[1]
C = img_a.shape[2]
# Create the graph.
img_shape = [None, H, W, C]
flow_shape = [None, H, W, 2]
input = tf.placeholder(shape=img_shape, dtype=tf.float32)
flow_tensor = tf.placeholder(shape=flow_shape, dtype=tf.float32)
warped_tensor = backward_warp(input, flow_tensor)
# Run.
warped_image = self.sess.run(warped_tensor,
feed_dict={
input:
[img_b, img_d, mask_ab, mask_cd],
flow_tensor:
[flow_ab, flow_cd, flow_ab, flow_cd]
})
# Get the masked errors.
error_ab_img = np.abs(warped_image[0] - img_a) * warped_image[2]
error_cd_img = np.abs(warped_image[1] - img_c) * warped_image[3]
error_ab = np.mean(error_ab_img)
error_cd = np.mean(error_cd_img)
# Assert a < 1.3% average error.
self.assertLess(error_ab, 0.013)
self.assertLess(error_cd, 0.013)
if SHOW_WARPED_IMAGES:
show_image(warped_image[0])
show_image(warped_image[1])
show_image(error_ab_img)
show_image(error_cd_img)
def test_gradients(self):
"""
Test to see if the gradients flow.
"""
# Load from files.
flow_ab = read_flow_file('pwcnet/warp/test_data/flow_ab.flo')
img_b = read_image('pwcnet/warp/test_data/image_b.png', as_float=True)
H = img_b.shape[0]
W = img_b.shape[1]
C = img_b.shape[2]
img_shape = [None, H, W, C]
flow_shape = [None, H, W, 2]
input = tf.placeholder(shape=img_shape, dtype=tf.float32)
flow_tensor = tf.placeholder(shape=flow_shape, dtype=tf.float32)
warped_tensor = backward_warp(input, flow_tensor)
grad_op = tf.gradients(warped_tensor, [input, flow_tensor])
grads = self.sess.run(grad_op,
feed_dict={
input: [img_b],
flow_tensor: [flow_ab]
})
for gradient in grads:
self.assertNotAlmostEqual(np.sum(gradient), 0.0)
def _write_shard(shard_id, shard_range, image_a_paths, image_b_paths,
flow_paths, filename, directory, verbose, max_flow):
"""
:param shard_id: Index of the shard.
:param shard_range: Iteration range of the shard.
:param image_paths: Path of all images.
:param flow_paths: Path of all flows.
:param filename: Base name of the output shard.
:param directory: Output directory.
:param verbose: Whether to print to console.
:param max_flow: Float. Maximum flow magnitude of the flow image. Any examples with flow magnitude greater than this
will be ignored.
:return: Nothing.
"""
if verbose and len(shard_range) > 0:
print('Writing to shard', shard_id, 'data points', shard_range[0],
'to', shard_range[-1])
record_name = os.path.join(directory, str(shard_id) + '_' + filename)
options = tf.python_io.TFRecordOptions(
tf.python_io.TFRecordCompressionType.GZIP)
writer = tf.python_io.TFRecordWriter(record_name, options=options)
num_examples_written = 0
for i in shard_range:
# Read from file.
flow = read_flow_file(flow_paths[i])
if np.amax(np.linalg.norm(flow, axis=-1)) > max_flow:
if verbose:
print(flow_paths[i], 'has a flow magnitude greater than',
max_flow)
continue
# Read and decode images as bytes to save memory.
image_a = read_image(image_a_paths[i], as_float=False)
image_b = read_image(image_b_paths[i], as_float=False)
# Write to tf record.
H = image_a.shape[0]
W = image_a.shape[1]
image_a_raw = image_a.tostring()
image_b_raw = image_b.tostring()
flow_raw = flow.tostring()
example = tf.train.Example(features=tf.train.Features(
feature={
FlowDataSet.HEIGHT: tf_int64_feature(H),
FlowDataSet.WIDTH: tf_int64_feature(W),
FlowDataSet.IMAGE_A_RAW: tf_bytes_feature(image_a_raw),
FlowDataSet.IMAGE_B_RAW: tf_bytes_feature(image_b_raw),
FlowDataSet.FLOW_RAW: tf_bytes_feature(flow_raw)
}))
writer.write(example.SerializeToString())
num_examples_written += 1
writer.close()
if num_examples_written == 0:
# Delete the file if nothing was written to it.
if verbose:
print(record_name, 'is empty')
silently_remove_file(record_name)
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 single_warp_test_helper(self, flow_ab_path, img_a_path, img_b_path,
tolerance):
"""
Runs warp test for a set of 2 images and a flow between them.
Also masks out the occluded warp regions when computing the error.
:param flow_ab_path: Str.
:param img_a_path: Str.
:param img_b_path: Str.
:param tolerance: Float. Usually a percentage like 0.03.
"""
# Load from files.
flow_ab = read_flow_file(flow_ab_path)
img_a = read_image(img_a_path, as_float=True)
img_b = read_image(img_b_path, as_float=True)
mask_ab = np.ones(shape=img_a.shape)
H = img_a.shape[0]
W = img_a.shape[1]
C = img_a.shape[2]
# Create the graph.
img_shape = [None, H, W, C]
flow_shape = [None, H, W, 2]
input = tf.placeholder(shape=img_shape, dtype=tf.float32)
flow_tensor = tf.placeholder(shape=flow_shape, dtype=tf.float32)
warped_tensor = backward_warp(input, flow_tensor)
# Run.
warped_image = self.sess.run(warped_tensor,
feed_dict={
input: [img_b, mask_ab],
flow_tensor: [flow_ab, flow_ab]
})
# Get the masked errors.
error_ab_img = np.abs(warped_image[0] - img_a) * warped_image[1]
error_ab = np.mean(error_ab_img)
# Assert a < tolerance average error.
self.assertLess(error_ab, tolerance)
if SHOW_WARPED_IMAGES:
show_image(warped_image[0])
show_image(error_ab_img)
def test_pfm_read(self):
flow_image = read_flow_file('common/utils/test_data/flow.pfm')
self.assertTrue(flow_image is not None)
self.assertTupleEqual((540, 960, 2), flow_image.shape)
if SHOW_FLOW_TEST_IMAGES:
show_flow_image(flow_image)
