def _check_images_and_flow_with_mask(self,
image1,
image2,
flow,
mask,
save_images=False,
plot_dir='/tmp/flow_images'):
self.assertGreaterEqual(np.min(image1), 0.)
self.assertLessEqual(np.max(image1), 1.)
self.assertGreaterEqual(np.min(mask), 0.)
self.assertLessEqual(np.max(mask), 1.)
# Check that the image2 warped by flow1 into image1 has lower pixelwise
# error than the unwarped image
mean_unwarped_diff = np.mean(mask * np.abs(image1 - image2))
warp = smurf_utils.flow_to_warp(flow)
image2_to_image1 = mask * smurf_utils.resample(image2, warp)
mean_warped_diff = np.mean(mask * np.abs(image2_to_image1 - image1))
if save_images:
plot_images(image1,
image2,
flow,
image2_to_image1,
plot_dir=plot_dir)
# check that the warped image has lower pixelwise error than the unwarped
self.assertLess(mean_warped_diff, mean_unwarped_diff)
def _check_images_and_flow(self, images, flow): # Check that the image2 warped by flow1 into image1 has lower pixelwise # error than the unwarped image image1, image2 = tf.unstack(images) image1 = tf.expand_dims(image1, axis=0) image2 = tf.expand_dims(image2, axis=0) flow = tf.expand_dims(flow, axis=0) mean_unwarped_diff = np.mean(np.abs(image1 - image2)) warp = smurf_utils.flow_to_warp(flow) image2_to_image1 = smurf_utils.resample(image2, warp) mean_warped_diff = np.mean(np.abs(image2_to_image1 - image1)) self.assertLess(mean_warped_diff, mean_unwarped_diff)
def complete_paper_plot(plot_dir,
index,
image1,
image2,
flow_uv,
ground_truth_flow_uv=None,
flow_valid_occ=None,
predicted_occlusion=None,
ground_truth_occlusion=None,
frame_skip=None):
def post_imshow(name, plot_dir):
plt.xticks([])
plt.yticks([])
if frame_skip is not None:
filename = str(index) + '_' + str(frame_skip) + '_' + name
plt.savefig(os.path.join(plot_dir, filename), bbox_inches='tight')
else:
filepath = str(index) + '_' + name
plt.savefig(os.path.join(plot_dir, filepath), bbox_inches='tight')
plt.clf()
warp = smurf_utils.flow_to_warp(tf.convert_to_tensor(flow_uv))
image1_reconstruction = smurf_utils.resample(
tf.expand_dims(image2, axis=0), tf.expand_dims(warp, axis=0))[0]
flow_uv = -flow_uv[:, :, ::-1]
if ground_truth_flow_uv is not None:
ground_truth_flow_uv = -ground_truth_flow_uv[:, :, ::-1]
plt.figure()
plt.clf()
plt.imshow(image1)
post_imshow('image1_rgb', plot_dir)
plt.imshow(image1_reconstruction)
post_imshow('image1_reconstruction_rgb', plot_dir)
plt.imshow(image1_reconstruction * predicted_occlusion)
post_imshow('image1_reconstruction_occlusions_rgb', plot_dir)
plt.imshow((image1 + image2) / 2.)
post_imshow('image_rgb', plot_dir)
plt.imshow(flow_to_rgb(flow_uv))
post_imshow('predicted_flow', plot_dir)
if ground_truth_flow_uv is not None and flow_valid_occ is not None:
plt.imshow(flow_to_rgb(ground_truth_flow_uv * flow_valid_occ))
post_imshow('ground_truth_flow', plot_dir)
endpoint_error = np.sum((ground_truth_flow_uv - flow_uv)**2,
axis=-1,
keepdims=True)**0.5
plt.imshow((endpoint_error * flow_valid_occ)[:, :, 0],
cmap='viridis',
vmin=0,
vmax=40)
post_imshow('flow_error', plot_dir)
if predicted_occlusion is not None:
plt.imshow((predicted_occlusion[:, :, 0]) * 255, cmap='Greys')
post_imshow('predicted_occlusion', plot_dir)
if ground_truth_occlusion is not None:
plt.imshow((ground_truth_occlusion[:, :, 0]) * 255, cmap='Greys')
post_imshow('ground_truth_occlusion', plot_dir)
plt.close('all')
def call(self, feature_dict, training=False, backward=False):
"""Run the model."""
context = None
flow = None
flow_up = None
context_up = None
flows = []
if backward:
feature_pyramid1 = feature_dict['features2']
feature_pyramid2 = feature_dict['features1']
else:
feature_pyramid1 = feature_dict['features1']
feature_pyramid2 = feature_dict['features2']
# Go top down through the levels to the second to last one to estimate flow.
for level, (features1, features2) in reversed(
list(enumerate(
zip(feature_pyramid1,
feature_pyramid2)))[self._output_flow_at_level:]):
# init flows with zeros for coarsest level if needed
if self._shared_flow_decoder and flow_up is None:
batch_size, height, width, _ = features1.shape.as_list()
flow_up = tf.zeros([batch_size, height, width, 2])
if self._num_context_up_channels:
num_channels = int(self._num_context_up_channels *
self._channel_multiplier)
context_up = tf.zeros(
[batch_size, height, width, num_channels])
# Warp features2 with upsampled flow from higher level.
if flow_up is None or not self._use_feature_warp:
warped2 = features2
else:
warp_up = smurf_utils.flow_to_warp(flow_up)
warped2 = smurf_utils.resample(features2, warp_up)
# Compute cost volume by comparing features1 and warped features2.
features1_normalized, warped2_normalized = normalize_features(
[features1, warped2],
normalize=self._normalize_before_cost_volume,
center=self._normalize_before_cost_volume,
moments_across_channels=True,
moments_across_images=True)
if self._use_cost_volume:
cost_volume = compute_cost_volume(features1_normalized,
warped2_normalized,
max_displacement=4)
else:
concat_features = Concatenate(axis=-1)(
[features1_normalized, warped2_normalized])
cost_volume = self._cost_volume_surrogate_convs[level](
concat_features)
cost_volume = LeakyReLU(alpha=self._leaky_relu_alpha)(cost_volume)
if self._shared_flow_decoder:
# this will ensure to work for arbitrary feature sizes per level
conv_1x1 = self._1x1_shared_decoder[level]
features1 = conv_1x1(features1)
# Compute context and flow from previous flow, cost volume, and features1.
if flow_up is None:
x_in = Concatenate(axis=-1)([cost_volume, features1])
else:
if context_up is None:
x_in = Concatenate(axis=-1)(
[flow_up, cost_volume, features1])
else:
x_in = Concatenate(axis=-1)(
[context_up, flow_up, cost_volume, features1])
# Use dense-net connections.
x_out = None
if self._shared_flow_decoder:
# reuse the same flow decoder on all levels
flow_layers = self._flow_layers
else:
flow_layers = self._flow_layers[level]
for layer in flow_layers[:-1]:
x_out = layer(x_in)
x_in = Concatenate(axis=-1)([x_in, x_out])
context = x_out
flow = flow_layers[-1](context)
if (training and self._drop_out_rate):
maybe_dropout = tf.cast(
tf.math.greater(tf.random.uniform([]),
self._drop_out_rate), tf.float32)
context *= maybe_dropout
flow *= maybe_dropout
if flow_up is not None and self._accumulate_flow:
flow += flow_up
# Upsample flow for the next lower level.
flow_up = upsample(flow, is_flow=True)
if self._num_context_up_channels:
context_up = self._context_up_layers[level](context)
# Append results to list.
flows.insert(0, flow)
# Refine flow at level '_output_flow_at_level'.
refinement = self._refine_model(context, flow)
if (training and self._drop_out_rate):
refinement *= tf.cast(
tf.math.greater(tf.random.uniform([]), self._drop_out_rate),
tf.float32)
refined_flow = flow + refinement
flows[0] = refined_flow
# Upsample flow to the highest available feature resolution.
for _ in range(self._output_flow_at_level):
upsampled_flow = upsample(flows[0], is_flow=True)
flows.insert(0, upsampled_flow)
# Upsample flow to the original input resolution.
upsampled_flow = upsample(flows[0], is_flow=True)
flows.insert(0, upsampled_flow)
return [tf.cast(flow, tf.float32) for flow in flows]