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))