## -*- coding: utf-8 -*- import tensorflow as tf import numpy as np import time import glob import scipy import argparse import os import subprocess from random import choice from PIL import Image from utils import LoadImage, DownSample, AVG_PSNR, depth_to_space_3D, DynFilter3D, LoadParams from nets import FR_16L, FR_28L, FR_52L parser = argparse.ArgumentParser() parser.add_argument('--L', metavar='L', type=int, default=choice([16, 28, 52]), help='Network depth: One of 16, 28, 52') parser.add_argument('--T', metavar='T', default='L', help='Input type: L(Low-resolution) or G(Ground-truth)') parser.add_argument('--dataset', default=None, help='Name of dataset') parser.add_argument('--in_dir', metavar='in_dir', default=None, help='Directory to process') parser.add_argument('--out_dir', metavar='out_dir', default='./results', help='Directory to output to') parser.add_argument('--network_dir', default='.', help='Path to networks') parser.add_argument('--mov_from_dirs', action='store_true') args = parser.parse_args() # Size of input temporal radius T_in = 7 # Upscaling factor R = 4 # Selecting filters and residual generating network if args.L == 16: FR = FR_16L elif args.L == 28: FR = FR_28L elif args.L == 52: FR = FR_52L else: print('Invalid network depth: {} (Must be one of 16, 28, 52)'.format(args.L)) exit(1) if not(args.T == 'L' or args.T =='G'): print('Invalid input type: {} (Must be L(Low-resolution) or G(Ground-truth))'.format(args.T)) exit(1) def process_dir(dir): dir_partz = dir.split('/') dataset = args.dataset = dir_partz[-2] part = dir_partz[-1] tag = '_'.join([dataset, str(args.L) + 'L', part]) out_path = os.path.join(args.out_dir, dataset, part) render_path = os.path.join(args.out_dir, 'renders') os.makedirs(out_path, exist_ok=True) if args.mov_from_dirs: os.makedirs(render_path, exist_ok=True) dir_frames = sorted(glob.glob(os.path.join(dir, '*.png'))) # print(dir_frames) frames = [] for f in dir_frames: frames.append(LoadImage(f)) frames = np.asarray(frames) frames_padded = np.lib.pad(frames, pad_width=((T_in//2,T_in//2),(0,0),(0,0),(0,0)), mode='constant') for i in range(frames.shape[0]): print('Scene {}: Frame {}/{} processing'.format(tag, i+1, frames.shape[0])) in_L = frames_padded[i:i+T_in] # select T_in frames in_L = in_L[np.newaxis,:,:,:,:] out_H = sess.run(GH, feed_dict={L: in_L, is_train: False}) out_H = np.clip(out_H, 0, 1) Image.fromarray(np.around(out_H[0,0]*255).astype(np.uint8)).save('{}/frame_{:05d}.png'.format(out_path, i+1)) if args.mov_from_dirs: subprocess.call([ 'ffmpeg', '-hide_banner', '-i', os.path.join(out_path, 'frame_%05d.png'), '-y', '-c:v', 'libx264', '-preset', 'slow', '-crf', '19', '-vf', 'fps=25', '-pix_fmt', 'yuv420p', os.path.join(render_path, tag + '.mp4') ]) def G(x, is_train): # shape of x: [B,T_in,H,W,C] # Generate filters and residual # Fx: [B,1,H,W,1*5*5,R*R] # Rx: [B,1,H,W,3*R*R] Fx, Rx = FR(x, is_train) x_c = [] for c in range(3): t = DynFilter3D(x[:,T_in//2:T_in//2+1,:,:,c], Fx[:,0,:,:,:,:], [1,5,5]) # [B,H,W,R*R] t = tf.depth_to_space(t, R) # [B,H*R,W*R,1] x_c += [t] x = tf.concat(x_c, axis=3) # [B,H*R,W*R,3] x = tf.expand_dims(x, axis=1) Rx = depth_to_space_3D(Rx, R) # [B,1,H*R,W*R,3] x += Rx return x # Gaussian kernel for downsampling def gkern(kernlen=13, nsig=1.6): import scipy.ndimage.filters as fi # create nxn zeros inp = np.zeros((kernlen, kernlen)) # set element at the middle to one, a dirac delta inp[kernlen//2, kernlen//2] = 1 # gaussian-smooth the dirac, resulting in a gaussian filter mask return fi.gaussian_filter(inp, nsig) h = gkern(13, 1.6) # 13 and 1.6 for x4 h = h[:,:,np.newaxis,np.newaxis].astype(np.float32) # Network H = tf.placeholder(tf.float32, shape=[None, T_in, None, None, 3]) L_ = DownSample(H, h, R) L = L_[:,:,2:-2,2:-2,:] # To minimize boundary artifact is_train = tf.placeholder(tf.bool, shape=[]) # Phase ,scalar with tf.variable_scope('G') as scope: GH = G(L, is_train) params_G = [v for v in tf.global_variables() if v.name.startswith('G/')] # Session config = tf.ConfigProto() config.gpu_options.allow_growth=True with tf.Session(config=config) as sess: tf.global_variables_initializer().run() # Load parameters LoadParams(sess, [params_G], in_file=os.path.join(args.network_dir, 'params_{}L_x{}.h5'.format(args.L, R))) if args.T == 'L': # Test using Low-resolution videos if args.in_dir: for dir in sorted(glob.glob(os.path.join(args.in_dir, '*'))): process_dir(dir) else: for dir in sorted(glob.glob('./inputs/L/*')): process_dir(dir) elif args.T == 'G': # Test using GT videos avg_psnrs = [] dir_inputs = glob.glob('./inputs/G/*') for v in dir_inputs: scene_name = v.split('/')[-1] os.makedirs('./results/{}L/G/{}/'.format(args.L, scene_name), exist_ok=True) dir_frames = glob.glob(v + '/*.png') dir_frames.sort() frames = [] for f in dir_frames: frames.append(LoadImage(f)) frames = np.asarray(frames) frames_padded = np.lib.pad(frames, pad_width=((T_in//2,T_in//2),(0,0),(0,0),(0,0)), mode='constant') frames_padded = np.lib.pad(frames_padded, pad_width=((0,0),(8,8),(8,8),(0,0)), mode='reflect') out_Hs = [] for i in range(frames.shape[0]): print('Scene {}: Frame {}/{} processing'.format(scene_name, i+1, frames.shape[0])) in_H = frames_padded[i:i+T_in] # select T_in frames in_H = in_H[np.newaxis,:,:,:,:] out_H = sess.run(GH, feed_dict={H: in_H, is_train: False}) out_H = np.clip(out_H, 0, 1) Image.fromarray(np.around(out_H[0,0]*255).astype(np.uint8)).save('./results/{}L/G/{}/frame_{:05d}.png'.format(args.L, scene_name, i+1)) out_Hs.append(out_H[0, 0]) out_Hs = np.asarray(out_Hs) avg_psnr = AVG_PSNR(((frames)*255).astype(np.uint8)/255.0, ((out_Hs)*255).astype(np.uint8)/255.0, vmin=0, vmax=1, t_border=2, sp_border=8) avg_psnrs.append(avg_psnr) print('Scene {}: PSNR {}'.format(scene_name, avg_psnr))