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