def test_inference(self):
"""Test that inference runs and produces output of the right size."""
image1 = np.random.randn(256, 256, 3).astype('float32')
image2 = np.random.randn(256, 256, 3).astype('float32')
uflow = UFlow()
flow = uflow.infer(image1, image2)
correct_shape = np.equal(flow.shape, [256, 256, 2]).all()
self.assertTrue(correct_shape)
def test_train_step(self):
"""Test a single training step."""
ds = tf.data.Dataset.from_tensor_slices(
(tf.zeros([1, 2, 256, 256, 3],
dtype=tf.float32), {})).repeat().batch(1)
it = tf.compat.v1.data.make_one_shot_iterator(ds)
uflow = UFlow()
log_update = uflow.train(it, num_steps=1)
self.assertNotEmpty(log_update)
def convert_graph(checkpoint_dir, height, width, output_row_col):
"""Converts uflow to a frozen model."""
print('Building the model.')
tf.compat.v1.reset_default_graph()
uflow = UFlow(checkpoint_dir=checkpoint_dir,
channel_multiplier=FLAGS.channel_multiplier,
num_levels=FLAGS.num_levels)
image1 = tf.compat.v1.placeholder(tf.float32, [height, width, 3],
name='first_image')
image2 = tf.compat.v1.placeholder(tf.float32, [height, width, 3],
name='second_image')
flow = uflow.infer_no_tf_function(image1,
image2,
input_height=FLAGS.resize_to_height,
input_width=FLAGS.resize_to_width)
if output_row_col:
flow = tf.identity(flow, 'flow')
else:
flow = tf.identity(flow[Ellipsis, ::-1], 'flow')
print('Loading the checkpoint.')
saver = tf.compat.v1.train.Saver()
sess = tf.compat.v1.Session()
sess.run(tf.compat.v1.global_variables_initializer())
sess.run(tf.compat.v1.local_variables_initializer())
# Apparently, uflow.restore() does not work here in graph mode.
saver.restore(sess, tf.train.latest_checkpoint(checkpoint_dir))
# Load images.
image1_np = cv2.imread(
FLAGS.image1_path)[:, :, ::-1].astype('float32') / 255.
image2_np = cv2.imread(
FLAGS.image2_path)[:, :, ::-1].astype('float32') / 255.
image1_np = cv2.resize(image1_np, (FLAGS.width, FLAGS.height))
image2_np = cv2.resize(image2_np, (FLAGS.width, FLAGS.height))
# Compute test flow and save.
flow_np = sess.run(flow, feed_dict={image1: image1_np, image2: image2_np})
uflow_plotting.plot_flow(image1_np, image2_np, flow_np, 'test_result.png',
checkpoint_dir)
print('Freezing and optimizing the graph.')
dg = tf.compat.v1.get_default_graph()
frozen_graph_def = tf.compat.v1.graph_util.convert_variables_to_constants(
sess, # The session is used to retrieve the weights.
dg.as_graph_def(), # The graph_def is used to retrieve the nodes.
['flow'] # The output node names are used to select the useful nodes.
)
print('Writing the result to', checkpoint_dir)
filename = os.path.join(checkpoint_dir, 'uflow.pb')
with tf.io.gfile.GFile(filename, 'wb') as f:
f.write(frozen_graph_def.SerializeToString())
def _export_frozen_graph():
# save weights
tempdir = tempfile.mkdtemp()
sess = tf.compat.v1.Session()
with sess.as_default():
uflow = UFlow(checkpoint_dir=tempdir)
image1 = tf.compat.v1.placeholder(tf.float32, [HEIGHT, WIDTH, 3],
name='first_image')
image2 = tf.compat.v1.placeholder(tf.float32, [HEIGHT, WIDTH, 3],
name='second_image')
flow = tf.identity(uflow.infer_no_tf_function(image1, image2), 'flow')
sess.run(tf.compat.v1.global_variables_initializer())
image1_np = np.random.randn(HEIGHT, WIDTH, 3)
image2_np = np.random.randn(HEIGHT, WIDTH, 3)
flow_output = sess.run(flow,
feed_dict={
image1: image1_np,
image2: image2_np
})
uflow.save()
# freeze model
freeze_model.convert_graph(tempdir, HEIGHT, WIDTH, output_row_col=True)
return tempdir, image1_np, image2_np, flow_output
def create_uflow():
"""Build the uflow model."""
build_selfsup_transformations = partial(
uflow_augmentation.build_selfsup_transformations,
crop_height=FLAGS.selfsup_crop_height,
crop_width=FLAGS.selfsup_crop_width,
max_shift_height=FLAGS.selfsup_max_shift,
max_shift_width=FLAGS.selfsup_max_shift,
resize=FLAGS.resize_selfsup)
# Define learning rate schedules [none, cosine, linear, expoential].
def learning_rate_fn():
step = tf.compat.v1.train.get_or_create_global_step()
effective_step = tf.maximum(step - FLAGS.lr_decay_after_num_steps + 1,
0)
lr_step_ratio = tf.cast(effective_step, 'float32') / float(
FLAGS.lr_decay_steps)
if FLAGS.lr_decay_type == 'none' or FLAGS.lr_decay_steps <= 0:
return FLAGS.gpu_learning_rate
elif FLAGS.lr_decay_type == 'cosine':
x = np.pi * tf.minimum(lr_step_ratio, 1.0)
return FLAGS.gpu_learning_rate * (tf.cos(x) + 1.0) / 2.0
elif FLAGS.lr_decay_type == 'linear':
return FLAGS.gpu_learning_rate * tf.maximum(
1.0 - lr_step_ratio, 0.0)
elif FLAGS.lr_decay_type == 'exponential':
return FLAGS.gpu_learning_rate * 0.5**lr_step_ratio
else:
raise ValueError('Unknown lr_decay_type', FLAGS.lr_decay_type)
occ_weights = {
'fb_abs': FLAGS.occ_weights_fb_abs,
'forward_collision': FLAGS.occ_weights_forward_collision,
'backward_zero': FLAGS.occ_weights_backward_zero,
}
# Switch off loss-terms that have weights < 1e-2.
occ_weights = {k: v for (k, v) in occ_weights.items() if v > 1e-2}
occ_thresholds = {
'fb_abs': FLAGS.occ_thresholds_fb_abs,
'forward_collision': FLAGS.occ_thresholds_forward_collision,
'backward_zero': FLAGS.occ_thresholds_backward_zero,
}
occ_clip_max = {
'fb_abs': FLAGS.occ_clip_max_fb_abs,
'forward_collision': FLAGS.occ_clip_max_forward_collision,
}
uflow = UFlow(
checkpoint_dir=FLAGS.checkpoint_dir,
optimizer=FLAGS.optimizer,
learning_rate=learning_rate_fn,
only_forward=FLAGS.only_forward,
level1_num_layers=FLAGS.level1_num_layers,
level1_num_filters=FLAGS.level1_num_filters,
level1_num_1x1=FLAGS.level1_num_1x1,
dropout_rate=FLAGS.dropout_rate,
build_selfsup_transformations=build_selfsup_transformations,
fb_sigma_teacher=FLAGS.fb_sigma_teacher,
fb_sigma_student=FLAGS.fb_sigma_student,
train_with_supervision=FLAGS.use_supervision,
train_with_gt_occlusions=FLAGS.use_gt_occlusions,
smoothness_edge_weighting=FLAGS.smoothness_edge_weighting,
teacher_image_version=FLAGS.teacher_image_version,
stop_gradient_mask=FLAGS.stop_gradient_mask,
selfsup_mask=FLAGS.selfsup_mask,
normalize_before_cost_volume=FLAGS.normalize_before_cost_volume,
original_layer_sizes=FLAGS.original_layer_sizes,
shared_flow_decoder=FLAGS.shared_flow_decoder,
channel_multiplier=FLAGS.channel_multiplier,
num_levels=FLAGS.num_levels,
use_cost_volume=FLAGS.use_cost_volume,
use_feature_warp=FLAGS.use_feature_warp,
accumulate_flow=FLAGS.accumulate_flow,
occlusion_estimation=FLAGS.occlusion_estimation,
occ_weights=occ_weights,
occ_thresholds=occ_thresholds,
occ_clip_max=occ_clip_max,
smoothness_at_level=FLAGS.smoothness_at_level,
)
return uflow