def residual_block(inputs, output_channel=64, stride=1, scope='res_block'):
with tf.variable_scope(scope):
net = ops.conv2(
inputs, 3, output_channel, stride, use_bias=True, scope='conv_1')
net = tf.nn.relu(net)
net = ops.conv2(
net, 3, output_channel, stride, use_bias=True, scope='conv_2')
net = net + inputs
return net
def down_block(inputs, output_channel=64, stride=1, scope='down_block'):
with tf.variable_scope(scope):
net = ops.conv2(
inputs, 3, output_channel, stride, use_bias=True, scope='conv_1')
net = ops.lrelu(net, 0.2)
net = ops.conv2(
net, 3, output_channel, stride, use_bias=True, scope='conv_2')
net = ops.lrelu(net, 0.2)
net = ops.maxpool(net)
return net
def up_block(inputs, output_channel=64, stride=1, scope='up_block'):
with tf.variable_scope(scope):
net = ops.conv2(
inputs, 3, output_channel, stride, use_bias=True, scope='conv_1')
net = ops.lrelu(net, 0.2)
net = ops.conv2(
net, 3, output_channel, stride, use_bias=True, scope='conv_2')
net = ops.lrelu(net, 0.2)
new_shape = tf.shape(net)[1:-1] * 2
net = tf2.image.resize(net, new_shape)
return net
def generator_f_encoder(gen_inputs, num_resblock=10, reuse=False):
"""Generator function encoder."""
# The Bx residual blocks
def residual_block(inputs, output_channel=64, stride=1, scope='res_block'):
with tf.variable_scope(scope):
net = ops.conv2(
inputs, 3, output_channel, stride, use_bias=True, scope='conv_1')
net = tf.nn.relu(net)
net = ops.conv2(
net, 3, output_channel, stride, use_bias=True, scope='conv_2')
net = net + inputs
return net
with tf.variable_scope('generator_unit', reuse=reuse):
# The input layer
with tf.variable_scope('input_stage'):
net = ops.conv2(gen_inputs, 3, 64, 1, scope='conv')
stage1_output = tf.nn.relu(net)
net = stage1_output
# The residual block parts
for i in range(1, num_resblock + 1,
1): # should be 16 for TecoGAN, and 10 for TecoGANmini
name_scope = 'resblock_%d' % (i)
net = residual_block(net, 64, 1, name_scope)
return net
def generator_f_decoder(net,
gen_inputs,
gen_output_channels,
vsr_scale,
reuse=False):
"""Generator function decoder."""
with tf.variable_scope('generator_unit', reuse=reuse):
with tf.variable_scope('conv_tran2highres'):
if vsr_scale == 2:
net = ops.conv2_tran(
net, kernel=3, output_channel=64, stride=2, scope='conv_tran1')
net = tf.nn.relu(net)
if vsr_scale == 4:
net = ops.conv2_tran(net, 3, 64, 2, scope='conv_tran1')
net = tf.nn.relu(net)
net = ops.conv2_tran(net, 3, 64, 2, scope='conv_tran2')
net = tf.nn.relu(net)
with tf.variable_scope('output_stage'):
net = ops.conv2(net, 3, gen_output_channels, 1, scope='conv')
low_res_in = gen_inputs[:, :, :, 0:3] # ignore warped pre high res
bicubic_hi = ops.bicubic_x(low_res_in, scale=vsr_scale) # can put on GPU
net = net + bicubic_hi
net = ops.preprocess(net)
return net
def fnet(fnet_input, reuse=False):
"""Flow net."""
def down_block(inputs, output_channel=64, stride=1, scope='down_block'):
with tf.variable_scope(scope):
net = ops.conv2(
inputs, 3, output_channel, stride, use_bias=True, scope='conv_1')
net = ops.lrelu(net, 0.2)
net = ops.conv2(
net, 3, output_channel, stride, use_bias=True, scope='conv_2')
net = ops.lrelu(net, 0.2)
net = ops.maxpool(net)
return net
def up_block(inputs, output_channel=64, stride=1, scope='up_block'):
with tf.variable_scope(scope):
net = ops.conv2(
inputs, 3, output_channel, stride, use_bias=True, scope='conv_1')
net = ops.lrelu(net, 0.2)
net = ops.conv2(
net, 3, output_channel, stride, use_bias=True, scope='conv_2')
net = ops.lrelu(net, 0.2)
new_shape = tf.shape(net)[1:-1] * 2
net = tf2.image.resize(net, new_shape)
return net
with tf.variable_scope('autoencode_unit', reuse=reuse):
net = down_block(fnet_input, 32, scope='encoder_1')
net = down_block(net, 64, scope='encoder_2')
net = down_block(net, 128, scope='encoder_3')
net = up_block(net, 256, scope='decoder_1')
net = up_block(net, 128, scope='decoder_2')
net1 = up_block(net, 64, scope='decoder_3')
with tf.variable_scope('output_stage'):
net = ops.conv2(net1, 3, 32, 1, scope='conv1')
net = ops.lrelu(net, 0.2)
net2 = ops.conv2(net, 3, 2, 1, scope='conv2')
net = tf.tanh(net2) * 24.0
# the 24.0 is the max Velocity, details can be found in TecoGAN paper
return net
def inference(
input_lr_dir,
input_hr_dir,
input_dir_len,
num_resblock,
vsr_scale,
checkpoint_path,
output_dir,
output_pre,
output_name,
output_ext,
):
"""Main inference function."""
if checkpoint_path is None:
raise ValueError('The checkpoint file is needed to performing the test.')
# Declare the test data reader
inference_data = inference_data_loader(input_lr_dir, input_hr_dir,
input_dir_len)
input_shape = [
1,
] + list(inference_data.inputs[0].shape)
output_shape = [1, input_shape[1] * vsr_scale, input_shape[2] * vsr_scale, 3]
oh = input_shape[1] - input_shape[1] // 8 * 8
ow = input_shape[2] - input_shape[2] // 8 * 8
paddings = tf.constant([[0, 0], [0, oh], [0, ow], [0, 0]])
print('input shape:', input_shape)
print('output shape:', output_shape)
# build the graph
inputs_raw = tf.placeholder(tf.float32, shape=input_shape, name='inputs_raw')
pre_inputs = tf.Variable(
tf.zeros(input_shape), trainable=False, name='pre_inputs')
pre_gen = tf.Variable(tf.zeros(output_shape), trainable=False, name='pre_gen')
pre_warp = tf.Variable(
tf.zeros(output_shape), trainable=False, name='pre_warp')
transpose_pre = tf.space_to_depth(pre_warp, vsr_scale)
inputs_all = tf.concat((inputs_raw, transpose_pre), axis=-1)
with tf.variable_scope('generator'):
gen_output = generator_f(
inputs_all, 3, num_resblock, vsr_scale, reuse=False)
# Deprocess the images outputed from the model, and assign things for next
# frame
with tf.control_dependencies([tf.assign(pre_inputs, inputs_raw)]):
outputs = tf.assign(pre_gen, ops.deprocess(gen_output))
inputs_frames = tf.concat((pre_inputs, inputs_raw), axis=-1)
with tf.variable_scope('fnet'):
gen_flow_lr = fnet(inputs_frames, reuse=False)
gen_flow_lr = tf.pad(gen_flow_lr, paddings, 'SYMMETRIC')
deconv_flow = gen_flow_lr
deconv_flow = ops.conv2_tran(
deconv_flow, 3, 64, 2, scope='deconv_flow_tran1')
deconv_flow = tf.nn.relu(deconv_flow)
deconv_flow = ops.conv2_tran(
deconv_flow, 3, 64, 2, scope='deconv_flow_tran2')
deconv_flow = tf.nn.relu(deconv_flow)
deconv_flow = ops.conv2(deconv_flow, 3, 2, 1, scope='deconv_flow_conv')
gen_flow = ops.upscale_x(gen_flow_lr * 4.0, scale=vsr_scale)
gen_flow = deconv_flow + gen_flow
gen_flow.set_shape(output_shape[:-1] + [2])
pre_warp_hi = tfa.image.dense_image_warp(pre_gen, gen_flow)
pre_warp_hi = pre_warp_hi + extract_detail_ops(pre_warp_hi)
before_ops = tf.assign(pre_warp, pre_warp_hi)
print('Finish building the network')
if FLAGS.use_ema:
moving_average_decay = 0.99
global_step = tf.train.get_or_create_global_step()
ema = tf.train.ExponentialMovingAverage(moving_average_decay, global_step)
ema_vars = _get_ema_vars()
# In inference time, we only need to restore the weight of the generator
var_list = tf.trainable_variables()
restore_vars_dict = {}
if FLAGS.use_ema:
for v in var_list:
if re.match(v.name, '.*global_step.*'):
restore_vars_dict[v.name[:-2]] = v
else:
restore_vars_dict[v.name[:-2] + '/ExponentialMovingAverage'] = v
else:
restore_vars_dict = var_list
weight_initiallizer = tf.train.Saver(restore_vars_dict)
# Define the initialization operation
init_op = tf.global_variables_initializer()
local_init_op = tf.local_variables_initializer()
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
if not gfile.exists(output_dir):
gfile.mkdir(output_dir)
if not output_pre:
image_dir = output_dir
else:
image_dir = os.path.join(output_dir, output_pre)
if not gfile.exists(image_dir):
gfile.mkdir(image_dir)
with tf.Session(config=config) as sess:
# Load the pretrained model
sess.run(init_op)
sess.run(local_init_op)
print('Loading weights from ckpt model')
weight_initiallizer.restore(sess, checkpoint_path)
max_iter = len(inference_data.inputs)
srtime = 0
print('Frame evaluation starts!!')
for i in range(max_iter):
input_im = np.array([inference_data.inputs[i]]).astype(np.float32)
feed_dict = {inputs_raw: input_im}
t0 = time.time()
if i != 0:
sess.run(before_ops, feed_dict=feed_dict)
output_frame = sess.run(outputs, feed_dict=feed_dict)
srtime += time.time() - t0
if i >= 5:
name, _ = os.path.splitext(
os.path.basename(str(inference_data.paths_LR[i])))
filename = output_name + '_' + name
out_path = os.path.join(image_dir, '%s.%s' % (filename, output_ext))
print('saving image %s' % out_path)
with tf.gfile.Open(out_path, 'wb') as image_file:
img = np.clip(output_frame[0] * 255.0, 0, 255).astype(np.uint8)
_, buff = cv2.imencode('.png', img[:, :, ::-1])
image_file.write(buff.tostring())
else:
print('Warming up %d' % (5 - i))
tf.reset_default_graph()
print('total time ' + str(srtime) + ', frame number ' + str(max_iter))
