├── .gitignore ├── Evaluation ├── NTIRE_PeakSNR_imgs.m └── psnr_first_100.m ├── MODEL_ADASR.py ├── MODEL_EDSR.py ├── Test.py ├── Train.py ├── Train2.py ├── Variations ├── MODEL_EDSR.py ├── MODEL_EDSR_2x2_kernel.py ├── MODEL_EDSR_Flatten.py ├── MODEL_EDSR_extra_smoothing_layer.py ├── MODEL_EDSR_transposed_convolution.py └── MODEL_EDSR_x2x2x2.py └── train_data_generation.m /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | -------------------------------------------------------------------------------- /Evaluation/NTIRE_PeakSNR_imgs.m: -------------------------------------------------------------------------------- 1 | function res = NTIRE_PeakSNR_imgs(F, G, scale) 2 | % NTIRE 2017 image super-resolution challenge scoring function 3 | % 4 | % F - original image 5 | % G - distorted image 6 | % scale factor - determins the number of boundary pixels to ignore (6+scale) 7 | % 8 | % returns res, the PSNR over all pixel values 9 | 10 | if ischar(F) 11 | F = imread(F); 12 | end 13 | if ischar(G) 14 | G = imread(G); 15 | end 16 | 17 | img_size = size(F); 18 | F = F(1:img_size(1)- mod(img_size(1), scale), 1:img_size(2) - mod(img_size(2), scale), :); 19 | 20 | boundarypixels = 0; 21 | if exist('scale','var') 22 | boundarypixels = 6+scale; 23 | F = F(boundarypixels+1:end-boundarypixels,boundarypixels+1:end-boundarypixels,:); 24 | G = G(boundarypixels+1:end-boundarypixels,boundarypixels+1:end-boundarypixels,:); 25 | end 26 | 27 | if max(F(:)) > 1 28 | F = im2double(F); 29 | end 30 | if max(G(:)) > 1 31 | G = im2double(G); 32 | end 33 | E = F - G; % error signal 34 | N = numel(E); % Assume the original signal is at peak (|F|=1) 35 | res = 10*log10( N / sum(E(:).^2) ); -------------------------------------------------------------------------------- /Evaluation/psnr_first_100.m: -------------------------------------------------------------------------------- 1 | originalDir = 'DIV2K_valid_HR'; 2 | outputDir = 'EDSR_9/output'; 3 | count = 0; 4 | psnrv = []; 5 | for i = 801:900 6 | fnameOriginal = strcat(sprintf('%04d', i), '.png'); 7 | fnameOutput = strcat(sprintf('%04d', i), 'x8.png'); 8 | f_path = fullfile(originalDir, fnameOriginal); 9 | img_original = imread(f_path); 10 | f_path = fullfile(outputDir, fnameOutput); 11 | img_output = imread(f_path); 12 | img_size = size(img_original); 13 | %img_output = imresize(img_output, [img_size(1), img_size(2)]); 14 | psnrv = [psnrv, NTIRE_PeakSNR_imgs(img_original, img_output, 8)]; 15 | end 16 | display(psnrv) 17 | display(mean(psnrv)) 18 | -------------------------------------------------------------------------------- /MODEL_ADASR.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | def resBlock(x, channels=256, kernel_size=[3,3], scale=1): 5 | tmp = tf.layers.conv2d(x, channels, kernel_size, activation=None, padding='SAME') 6 | tmp = tf.nn.relu(tmp) 7 | tmp = tf.layers.conv2d(tmp, channels, kernel_size, activation=None, padding='SAME') 8 | tmp *= scale 9 | return x + tmp 10 | 11 | def _phase_shift(I, r): 12 | return tf.depth_to_space(I, r) 13 | 14 | def PS(X, r, color=False): 15 | if color: 16 | Xc = tf.split(X, 3, 3) 17 | X = tf.concat([_phase_shift(x, r) for x in Xc],3) 18 | else: 19 | X = _phase_shift(X, r) 20 | return X 21 | 22 | def upsample(x, scale=8, features=256, activation=tf.nn.relu): 23 | assert scale in [2, 3, 4, 8, 16, 32] 24 | x = tf.layers.conv2d(x, features, [3,3], activation=activation, padding='SAME') 25 | if scale == 2: 26 | ps_features = 3*(scale**2) 27 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') #Increase channel depth 28 | #x = slim.conv2d_transpose(x,ps_features,6,stride=1,activation_fn=activation) 29 | x = PS(x, 2, color=True) 30 | elif scale == 3: 31 | ps_features =3*(scale**2) 32 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') 33 | #x = slim.conv2d_transpose(x,ps_features,9,stride=1,activation_fn=activation) 34 | x = PS(x, 3, color=True) 35 | elif scale == 4: 36 | ps_features = 3*(2**2) 37 | for i in range(2): 38 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') 39 | #x = slim.conv2d_transpose(x,ps_features,6,stride=1,activation_fn=activation) 40 | x = PS(x, 2, color=True) 41 | elif scale == 8: 42 | ps_features = 3*(8*8) 43 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') 44 | #x = tf.layers.conv2d_transpose(x,3,kernel_size=(3,3),strides=2,activation=activation, padding='SAME') 45 | x = PS(x, 8, color=True) 46 | elif scale == 16: 47 | ps_features = 3*(8*8) 48 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') 49 | x = PS(x, 8, color=True) 50 | ps_features = 3*(2*2) 51 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') 52 | x = PS(x, 2, color=True) 53 | elif scale == 32: 54 | ps_features = 3*(8*8) 55 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') 56 | x = PS(x, 8, color=True) 57 | ps_features = 3*(2**2) 58 | for i in range(2): 59 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') 60 | #x = slim.conv2d_transpose(x,ps_features,6,stride=1,activation_fn=activation) 61 | x = PS(x, 2, color=True) 62 | return x 63 | 64 | def EDSR_Block(tensor, scale=8, feature_size=256, num_layers=32, SCALING_FACTOR=0.1): 65 | ''' 66 | First convolution layer, Feature extraction 67 | ''' 68 | tensor = tf.layers.conv2d(tensor, feature_size, [3,3], padding='SAME') #Linear activation by default 69 | 70 | conv_1 = tensor #backup 71 | 72 | ''' 73 | Building resBlocks 74 | ''' 75 | for i in range(num_layers): 76 | tensor = resBlock(tensor, feature_size, scale=SCALING_FACTOR) 77 | 78 | ''' 79 | Add back the conv_1 tensor 80 | ''' 81 | tensor = tf.layers.conv2d(tensor, feature_size, [3,3], padding='SAME') 82 | tensor += conv_1 83 | 84 | ''' 85 | Upsampling 86 | ''' 87 | tensor = upsample(tensor, scale, feature_size, activation=None) 88 | 89 | return tensor 90 | 91 | 92 | def model(input_tensor, scale=8, feature_size=256, num_layers=32): 93 | with tf.device("/gpu:2"): 94 | 95 | SCALING_FACTOR = 0.1 96 | 97 | ''' 98 | Preprocessing by subtracting the batch mean 99 | ''' 100 | input_mean = tf.reduce_mean(input_tensor, 2, keep_dims=True) 101 | input_mean = tf.reduce_mean(input_mean, 1, keep_dims=True) 102 | tensor = input_tensor - input_mean 103 | 104 | ''' 105 | Downsample 106 | ''' 107 | tensor_down_2 = tf.layers.average_pooling2d(tensor, pool_size=[2,2], strides=2, padding='VALID') 108 | tensor_down_4 = tf.layers.average_pooling2d(tensor_down_2, pool_size=[2,2], strides=2, padding='VALID') 109 | 110 | ''' 111 | Superresolution 112 | ''' 113 | tensor_output_down_4 = EDSR_Block(tensor_down_4, 2, feature_size, num_layers, SCALING_FACTOR=SCALING_FACTOR) 114 | 115 | tensor_output_down_2 = EDSR_Block(tf.concat([tensor_output_down_4, tensor_down_2], axis=3), 2, feature_size, num_layers, SCALING_FACTOR=SCALING_FACTOR) 116 | 117 | tensor_output = EDSR_Block(tf.concat([tensor_output_down_2, tensor], axis=3), scale, feature_size, num_layers, SCALING_FACTOR=SCALING_FACTOR) 118 | 119 | ''' 120 | Downsample 121 | ''' 122 | tensor_final_output = tf.layers.conv2d(tensor_output, 3, [3,3], padding='SAME') 123 | 124 | tensor_final_output = tf.clip_by_value(tensor_final_output+input_mean, 0.0, 255.0) 125 | 126 | return tensor_final_output 127 | -------------------------------------------------------------------------------- /MODEL_EDSR.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | SCALING_FACTOR = 0.1 5 | 6 | def resBlock(x, channels=256, kernel_size=[3,3], scale=1): 7 | tmp = tf.layers.conv2d(x, channels, kernel_size, activation=None, padding='SAME') 8 | tmp = tf.nn.relu(tmp) 9 | tmp = tf.layers.conv2d(tmp, channels, kernel_size, activation=None, padding='SAME') 10 | tmp *= scale 11 | return x + tmp 12 | 13 | def _phase_shift(I, r): 14 | return tf.depth_to_space(I, r) 15 | 16 | def PS(X, r, color=False): 17 | if color: 18 | Xc = tf.split(X, 3, 3) 19 | X = tf.concat([_phase_shift(x, r) for x in Xc],3) 20 | else: 21 | X = _phase_shift(X, r) 22 | return X 23 | 24 | def upsample(x, scale=8, features=256, activation=tf.nn.relu): 25 | assert scale in [2, 3, 4, 8] 26 | x = tf.layers.conv2d(x, features, [3,3], activation=activation, padding='SAME') 27 | if scale == 2: 28 | ps_features = 3*(scale**2) 29 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') #Increase channel depth 30 | #x = slim.conv2d_transpose(x,ps_features,6,stride=1,activation_fn=activation) 31 | x = PS(x, 2, color=True) 32 | elif scale == 3: 33 | ps_features =3*(scale**2) 34 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') 35 | #x = slim.conv2d_transpose(x,ps_features,9,stride=1,activation_fn=activation) 36 | x = PS(x, 3, color=True) 37 | elif scale == 4: 38 | ps_features = 3*(2**2) 39 | for i in range(2): 40 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') 41 | #x = slim.conv2d_transpose(x,ps_features,6,stride=1,activation_fn=activation) 42 | x = PS(x, 2, color=True) 43 | elif scale == 8: 44 | ps_features = 3*(8*8) 45 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') 46 | #x = tf.layers.conv2d_transpose(x,3,kernel_size=(3,3),strides=2,activation=activation, padding='SAME') 47 | x = PS(x, 8, color=True) 48 | return x 49 | 50 | def model(input_tensor, scale=8, feature_size=256, num_layers=32): 51 | with tf.device("/gpu:0"): 52 | ''' 53 | Preprocessing by subtracting the batch mean 54 | ''' 55 | #input_mean = tf.reduce_mean(input_tensor) 56 | input_mean = tf.reduce_mean(input_tensor, 2, keep_dims=True) 57 | input_mean = tf.reduce_mean(input_mean, 1, keep_dims=True) 58 | tensor = input_tensor - input_mean 59 | 60 | ''' 61 | First convolution layer 62 | ''' 63 | tensor = tf.layers.conv2d(tensor, feature_size, [3,3], padding='SAME') #Linear activation by default 64 | 65 | conv_1 = tensor #backup 66 | 67 | ''' 68 | Building resBlocks 69 | ''' 70 | for i in range(num_layers): 71 | tensor = resBlock(tensor, feature_size, scale=SCALING_FACTOR) 72 | 73 | ''' 74 | Add back the conv_1 tensor 75 | ''' 76 | tensor = tf.layers.conv2d(tensor, feature_size, [3,3], padding='SAME') 77 | tensor += conv_1 78 | 79 | ''' 80 | Upsampling 81 | ''' 82 | tensor = upsample(tensor, scale, feature_size, activation=None) 83 | 84 | tensor = tf.clip_by_value(tensor+input_mean, 0.0, 255.0) 85 | 86 | return tensor 87 | -------------------------------------------------------------------------------- /Test.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from scipy import misc 3 | from PIL import Image 4 | import tensorflow as tf 5 | import glob, os, re 6 | import scipy.io 7 | from MODEL_EDSR import model 8 | from skimage import color 9 | #from MODEL_FACTORIZED import model_factorized 10 | import time 11 | 12 | input_dir = '../DIV2K_valid_LR_x8' 13 | output_dir = './output2' 14 | 15 | model_path = './checkpoints/EDSR_const_clip_0.01_epoch_012.ckpt-404627' 16 | 17 | def get_img_list(): 18 | files = [x for x in os.listdir(input_dir) if x.endswith('.png') and not x.startswith('.')] 19 | return sorted(files) 20 | 21 | def get_test_image(img_list, index): 22 | filename = img_list[index] 23 | path = os.path.join(input_dir, filename) 24 | img = Image.open(path) 25 | return np.array(img) 26 | 27 | def run_sess(sess, img): 28 | shape = img.shape 29 | target = np.zeros((img.shape[0]*8,img.shape[1]*8, 3)) 30 | img = sess.run(train_output, feed_dict={input_tensor: img.reshape(1, img.shape[0], img.shape[1], 3), target_tensor: target.reshape(1, target.shape[0], target.shape[1], 3)}) 31 | img = img.reshape(shape[0]*8, shape[1]*8, 3) 32 | return img 33 | 34 | 35 | def test_VDSR_with_sess(): 36 | saver.restore(sess, model_path) 37 | img_list = get_img_list() 38 | for i in range(len(img_list)): 39 | img = get_test_image(img_list, i) 40 | 41 | img_1 = np.rot90(img) 42 | img_2 = np.rot90(img_1) 43 | img_3 = np.rot90(img_2) 44 | img_4 = np.flipud(img) 45 | img_5 = np.fliplr(img) 46 | img_6 = np.roll(img, 1, axis=2) 47 | img_7 = np.roll(img, 2, axis=2) 48 | 49 | 50 | output_0 = run_sess(sess, img) 51 | output_1 = np.rot90(run_sess(sess, img_1),3) 52 | output_2 = np.rot90(run_sess(sess, img_2),2) 53 | output_3 = np.rot90(run_sess(sess, img_3),1) 54 | output_4 = np.flipud(run_sess(sess, img_4)) 55 | output_5 = np.fliplr(run_sess(sess, img_5)) 56 | output_6 = np.roll(run_sess(sess, img_6), 2, axis=2) 57 | output_7 = np.roll(run_sess(sess, img_7), 1, axis=2) 58 | 59 | img = (output_0 + output_1 + output_2 + output_3 + output_4 + output_5 + output_6 + output_7)/8 60 | 61 | img = Image.fromarray(img.astype('uint8')) 62 | img.save(os.path.join(output_dir, img_list[i]), compress_level=0) 63 | 64 | 65 | with tf.Session() as sess: 66 | input_tensor = tf.placeholder(tf.float32, shape=(1, None, None, 3)) 67 | target_tensor = tf.placeholder(tf.float32, shape=(1, None, None, 3)) 68 | shared_model = tf.make_template('shared_model', model) 69 | train_output, output_edge, target_edge = shared_model(input_tensor, target_tensor, scale=8, feature_size=256, num_layers=32) 70 | saver = tf.train.Saver() 71 | initializer = tf.global_variables_initializer() 72 | sess.run(initializer) 73 | test_VDSR_with_sess() 74 | -------------------------------------------------------------------------------- /Train.py: -------------------------------------------------------------------------------- 1 | import os, glob, re, signal, sys, argparse, threading, time 2 | from random import shuffle 3 | import random 4 | import tensorflow as tf 5 | from PIL import Image 6 | import numpy as np 7 | import scipy.io 8 | from MODEL_EDSR import model 9 | 10 | DATA_PATH = "/usr/project/xtmp/yb44/EDSR_HR_train_40" 11 | #DATA_PATH = "data/train" 12 | INPUT_IMG_SIZE = (40,40) 13 | OUTPUT_IMG_SIZE = (320, 320) 14 | BATCH_SIZE = 20 15 | BASE_LR = 0.0004 16 | DECAY_RATE = 0.96 17 | DECAY_STEP = 10000 18 | MAX_EPOCH = 120 19 | 20 | parser = argparse.ArgumentParser() 21 | parser.add_argument("--model_path") 22 | args = parser.parse_args() 23 | model_path = args.model_path 24 | 25 | def log10(x): 26 | numerator = tf.log(x) 27 | denominator = tf.log(tf.constant(10, dtype=numerator.dtype)) 28 | return numerator / denominator 29 | 30 | def get_train_list(data_path): 31 | l = glob.glob(os.path.join(data_path,"*")) 32 | l = [f for f in l if os.path.basename(f).endswith('.mat')] 33 | print len(l) 34 | return l 35 | 36 | def get_image_batch(train_list,offset,batch_size): 37 | target_list = train_list[offset:offset+batch_size] 38 | input_list = [] 39 | gt_list = [] 40 | for mat in target_list: 41 | data = scipy.io.loadmat(mat) 42 | input_img = data['input_patch'] 43 | gt_img = data['target_patch'] 44 | input_list.append(input_img) 45 | gt_list.append(gt_img) 46 | input_list = np.array(input_list) 47 | input_list.resize([BATCH_SIZE, INPUT_IMG_SIZE[1], INPUT_IMG_SIZE[0], 3]) 48 | gt_list = np.array(gt_list) 49 | gt_list.resize([BATCH_SIZE, OUTPUT_IMG_SIZE[1], OUTPUT_IMG_SIZE[0], 3]) 50 | return input_list, gt_list 51 | 52 | if __name__ == '__main__': 53 | train_list = get_train_list(DATA_PATH) 54 | 55 | train_input = tf.placeholder(tf.float32, shape=(BATCH_SIZE, INPUT_IMG_SIZE[0], INPUT_IMG_SIZE[1], 3)) 56 | train_gt = tf.placeholder(tf.float32, shape=(BATCH_SIZE, OUTPUT_IMG_SIZE[0], OUTPUT_IMG_SIZE[1], 3)) 57 | #output_edges = tf.placeholder(tf.float32, shape=(BATCH_SIZE, OUTPUT_IMG_SIZE[0], OUTPUT_IMG_SIZE[1], 3)) 58 | #target_edges = tf.placeholder(tf.float32, shape=(BATCH_SIZE, OUTPUT_IMG_SIZE[0], OUTPUT_IMG_SIZE[1], 3)) 59 | 60 | shared_model = tf.make_template('shared_model', model) 61 | #train_output, output_edges, target_edges = shared_model(train_input, train_gt, scale=8, feature_size=256, num_layers=32) 62 | train_output = shared_model(train_input, scale=8, feature_size=256, num_layers=32) 63 | #loss = tf.reduce_sum(tf.nn.l2_loss(tf.subtract(train_output, train_gt))) 64 | loss_1 = tf.reduce_mean(tf.losses.absolute_difference(train_output, train_gt)) 65 | #loss_2 = tf.reduce_mean(tf.losses.absolute_difference(output_edges, target_edges)) 66 | loss = loss_1 67 | tf.summary.scalar("loss", loss) 68 | 69 | mse = tf.reduce_mean(tf.squared_difference(train_output, train_gt)) 70 | PSNR = tf.constant(255**2,dtype=tf.float32)/mse 71 | PSNR = tf.constant(10,dtype=tf.float32)*log10(PSNR) 72 | 73 | global_step = tf.Variable(0, trainable=False) 74 | learning_rate = tf.train.exponential_decay(BASE_LR, global_step*BATCH_SIZE, DECAY_STEP, DECAY_RATE, staircase=True) 75 | tf.summary.scalar("learning rate", learning_rate) 76 | 77 | optimizer = tf.train.AdamOptimizer(learning_rate)#tf.train.MomentumOptimizer(learning_rate, 0.9) 78 | opt = optimizer.minimize(loss, global_step=global_step) 79 | 80 | saver = tf.train.Saver() 81 | 82 | shuffle(train_list) 83 | config = tf.ConfigProto() 84 | config.gpu_options.allow_growth = True 85 | 86 | with tf.Session(config=config) as sess: 87 | if not os.path.exists('logs'): 88 | os.mkdir('logs') 89 | merged = tf.summary.merge_all() 90 | file_writer = tf.summary.FileWriter('logs', sess.graph) 91 | 92 | initializer = tf.global_variables_initializer() 93 | sess.run(initializer) 94 | 95 | if model_path: 96 | print "restore model..." 97 | saver.restore(sess, model_path) 98 | assign_op = global_step.assign_op(0) 99 | sess.run(assign_op) 100 | print "Done" 101 | 102 | saved = False 103 | log_file = open('log.csv', 'w') 104 | for epoch in xrange(0, MAX_EPOCH): 105 | for step in range(len(train_list)//BATCH_SIZE): 106 | offset = step*BATCH_SIZE 107 | input_data, gt_data = get_image_batch(train_list, offset, BATCH_SIZE) 108 | feed_dict = {train_input: input_data, train_gt: gt_data} 109 | _,l,output,lr, g_step, mse_val, psnr_val = sess.run([opt, loss, train_output, learning_rate, global_step, mse, PSNR], feed_dict=feed_dict) 110 | print "[epoch %2.4f] loss %.4f\tmse %.4f\tpsnr %.4f\tlr %.5f"%(epoch+(float(step)*BATCH_SIZE/len(train_list)), np.sum(l)/BATCH_SIZE, mse_val, psnr_val, lr) 111 | log_file.write('%2.4f,%.4f,%.4f' % (epoch+(float(step)*BATCH_SIZE/len(train_list)), np.sum(l)/BATCH_SIZE, psnr_val) + '\n') 112 | log_file.flush() 113 | del input_data, gt_data 114 | if not saved: 115 | saver.save(sess, "/usr/project/xtmp/webster/EDSR_15_checkpoints/initial.ckpt") 116 | saved = True 117 | 118 | saver.save(sess, "/usr/project/xtmp/webster/EDSR_15_checkpoints/EDSR_const_clip_0.01_epoch_%03d.ckpt" % epoch ,global_step=global_step) 119 | -------------------------------------------------------------------------------- /Train2.py: -------------------------------------------------------------------------------- 1 | import os, glob, re, signal, sys, argparse, threading, time 2 | from random import shuffle 3 | import random 4 | import tensorflow as tf 5 | from PIL import Image 6 | import numpy as np 7 | import scipy.io 8 | from MODEL_ADASR import model 9 | 10 | DATA_PATH = "/usr/project/xtmp/webster/training_data" 11 | #DATA_PATH = "data/train" 12 | INPUT_IMG_SIZE = (40,40) 13 | OUTPUT_IMG_SIZE = (320, 320) 14 | BATCH_SIZE = 25 15 | BASE_LR = 0.0004 16 | DECAY_RATE = 0.96 17 | DECAY_STEP = 60000 18 | MAX_EPOCH = 120 19 | 20 | parser = argparse.ArgumentParser() 21 | parser.add_argument("--model_path") 22 | args = parser.parse_args() 23 | model_path = args.model_path 24 | 25 | def log10(x): 26 | numerator = tf.log(x) 27 | denominator = tf.log(tf.constant(10, dtype=numerator.dtype)) 28 | return numerator / denominator 29 | 30 | def get_train_list(data_path): 31 | l = glob.glob(os.path.join(data_path,"*")) 32 | l = [f for f in l if os.path.basename(f).endswith('.mat')] 33 | print len(l) 34 | return l 35 | 36 | def get_image_batch(train_list,offset,batch_size): 37 | target_list = train_list[offset:offset+batch_size] 38 | input_list = [] 39 | gt_list = [] 40 | for mat in target_list: 41 | data = scipy.io.loadmat(mat) 42 | input_img = data['input_patch'] 43 | gt_img = data['target_patch'] 44 | input_list.append(input_img) 45 | gt_list.append(gt_img) 46 | input_list = np.array(input_list) 47 | input_list.resize([BATCH_SIZE, INPUT_IMG_SIZE[1], INPUT_IMG_SIZE[0], 3]) 48 | gt_list = np.array(gt_list) 49 | gt_list.resize([BATCH_SIZE, OUTPUT_IMG_SIZE[1], OUTPUT_IMG_SIZE[0], 3]) 50 | return input_list, gt_list 51 | 52 | if __name__ == '__main__': 53 | train_list = get_train_list(DATA_PATH) 54 | 55 | train_input = tf.placeholder(tf.float32, shape=(BATCH_SIZE, INPUT_IMG_SIZE[0], INPUT_IMG_SIZE[1], 3)) 56 | train_gt = tf.placeholder(tf.float32, shape=(BATCH_SIZE, OUTPUT_IMG_SIZE[0], OUTPUT_IMG_SIZE[1], 3)) 57 | 58 | shared_model = tf.make_template('shared_model', model) 59 | 60 | train_output = shared_model(train_input, scale=8, feature_size=256, num_layers=32) 61 | #loss = tf.reduce_sum(tf.nn.l2_loss(tf.subtract(train_output, train_gt))) 62 | loss = tf.reduce_mean(tf.losses.absolute_difference(train_output, train_gt)) 63 | #loss = -tf.reduce_mean(ssim(train_output, train_gt, max_val=255)) 64 | tf.summary.scalar("loss", loss) 65 | 66 | mse = tf.reduce_mean(tf.squared_difference(train_output, train_gt)) 67 | PSNR = tf.constant(255**2,dtype=tf.float32)/mse 68 | PSNR = tf.constant(10,dtype=tf.float32)*log10(PSNR) 69 | 70 | global_step = tf.Variable(0, trainable=False) 71 | learning_rate = tf.train.exponential_decay(BASE_LR, global_step*BATCH_SIZE, DECAY_STEP, DECAY_RATE, staircase=True) 72 | tf.summary.scalar("learning rate", learning_rate) 73 | 74 | optimizer = tf.train.AdamOptimizer(learning_rate)#tf.train.MomentumOptimizer(learning_rate, 0.9) 75 | opt = optimizer.minimize(loss, global_step=global_step) 76 | #opt_2 = optimizer.minimize(loss_2, global_step=global_step) 77 | 78 | saver = tf.train.Saver() 79 | 80 | shuffle(train_list) 81 | config = tf.ConfigProto() 82 | config.gpu_options.allow_growth = True 83 | 84 | with tf.Session(config=config) as sess: 85 | if not os.path.exists('logs'): 86 | os.mkdir('logs') 87 | merged = tf.summary.merge_all() 88 | file_writer = tf.summary.FileWriter('logs', sess.graph) 89 | 90 | initializer = tf.global_variables_initializer() 91 | sess.run(initializer) 92 | 93 | if model_path: 94 | print "restore model..." 95 | saver.restore(sess, model_path) 96 | assign_op = global_step.assign(0) 97 | sess.run(assign_op) 98 | print "Done" 99 | 100 | saved = False 101 | log_file = open('log.csv', 'w') 102 | counter = 0 103 | for epoch in xrange(0, MAX_EPOCH): 104 | for step in range(len(train_list)//BATCH_SIZE): 105 | offset = step*BATCH_SIZE 106 | input_data, gt_data = get_image_batch(train_list, offset, BATCH_SIZE) 107 | feed_dict = {train_input: input_data, train_gt: gt_data} 108 | _,l,output,lr, g_step, mse_val, psnr_val = sess.run([opt, loss, train_output, learning_rate, global_step, mse, PSNR], feed_dict=feed_dict) 109 | print "[epoch %2.4f] loss %.4f\tmse %.4f\tpsnr %.4f\tlr %.5f"%(epoch+(float(step)*BATCH_SIZE/len(train_list)), np.sum(l)/BATCH_SIZE, mse_val, psnr_val, lr) 110 | log_file.write('%2.4f,%.4f,%.4f' % (epoch+(float(step)*BATCH_SIZE/len(train_list)), np.sum(l)/BATCH_SIZE, psnr_val) + '\n') 111 | log_file.flush() 112 | del input_data, gt_data 113 | if not saved: 114 | saver.save(sess, "/usr/project/xtmp/webster/ADASR_checkpoints/initial.ckpt") 115 | saved = True 116 | if epoch+(float(step)*BATCH_SIZE/len(train_list))/0.2 > counter: 117 | saver.save(sess, "/usr/project/xtmp/webster/ADASR_checkpoints/ADASR_const_clip_0.01_epoch_%03d.ckpt" % epoch ,global_step=global_step) 118 | counter += 1 119 | -------------------------------------------------------------------------------- /Variations/MODEL_EDSR.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | SCALING_FACTOR = 0.1 5 | 6 | def resBlock(x, channels=256, kernel_size=[3,3], scale=1): 7 | tmp = tf.layers.conv2d(x, channels, kernel_size, activation=None, padding='SAME') 8 | tmp = tf.nn.relu(tmp) 9 | tmp = tf.layers.conv2d(tmp, channels, kernel_size, activation=None, padding='SAME') 10 | tmp *= scale 11 | return x + tmp 12 | 13 | def _phase_shift(I, r): 14 | return tf.depth_to_space(I, r) 15 | 16 | def PS(X, r, color=False): 17 | if color: 18 | Xc = tf.split(X, 3, 3) 19 | X = tf.concat([_phase_shift(x, r) for x in Xc],3) 20 | else: 21 | X = _phase_shift(X, r) 22 | return X 23 | 24 | def upsample(x, scale=8, features=256, activation=tf.nn.relu): 25 | assert scale in [2, 3, 4, 8] 26 | x = tf.layers.conv2d(x, features, [3,3], activation=activation, padding='SAME') 27 | if scale == 2: 28 | ps_features = 3*(scale**2) 29 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') #Increase channel depth 30 | #x = slim.conv2d_transpose(x,ps_features,6,stride=1,activation_fn=activation) 31 | x = PS(x, 2, color=True) 32 | elif scale == 3: 33 | ps_features =3*(scale**2) 34 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') 35 | #x = slim.conv2d_transpose(x,ps_features,9,stride=1,activation_fn=activation) 36 | x = PS(x, 3, color=True) 37 | elif scale == 4: 38 | ps_features = 3*(2**2) 39 | for i in range(2): 40 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') 41 | #x = slim.conv2d_transpose(x,ps_features,6,stride=1,activation_fn=activation) 42 | x = PS(x, 2, color=True) 43 | elif scale == 8: 44 | ps_features = 3*(8*8) 45 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') 46 | #x = tf.layers.conv2d_transpose(x,3,kernel_size=(3,3),strides=2,activation=activation, padding='SAME') 47 | x = PS(x, 8, color=True) 48 | return x 49 | 50 | def model(input_tensor, target_tensor, scale=8, feature_size=256, num_layers=32): 51 | with tf.device("/gpu:0"): 52 | ''' 53 | Preprocessing by subtracting the batch mean 54 | ''' 55 | input_mean = tf.reduce_mean(input_tensor) 56 | tensor = input_tensor - input_mean 57 | 58 | ''' 59 | First convolution layer 60 | ''' 61 | tensor = tf.layers.conv2d(tensor, feature_size, [3,3], padding='SAME') #Linear activation by default 62 | 63 | conv_1 = tensor #backup 64 | 65 | ''' 66 | Building resBlocks 67 | ''' 68 | for i in range(num_layers): 69 | tensor = resBlock(tensor, feature_size, scale=SCALING_FACTOR) 70 | 71 | ''' 72 | Add back the conv_1 tensor 73 | ''' 74 | tensor = tf.layers.conv2d(tensor, feature_size, [3,3], padding='SAME') 75 | tensor += conv_1 76 | 77 | ''' 78 | Upsampling 79 | ''' 80 | tensor = upsample(tensor, scale, feature_size, activation=None) 81 | 82 | tensor = tf.clip_by_value(tensor+input_mean, 0.0, 255.0) 83 | 84 | sobel = tf.constant([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], tf.float32) 85 | sobel_3 = tf.stack([sobel, sobel, sobel], axis=2) 86 | sobel_filter = tf.reshape(sobel_3, [1, 3, 3, 3]) 87 | output_edges = tf.nn.conv2d(tensor, sobel_filter, strides=[1,1,1,1], padding='SAME') 88 | output_clipped_edges = tf.clip_by_value(output_edges, 50.0, 255.0) 89 | 90 | target_edges = tf.nn.conv2d(target_tensor, sobel_filter, strides=[1,1,1,1], padding='SAME') 91 | target_clipped_edges = tf.clip_by_value(target_edges, 50.0, 255.0) 92 | 93 | return tensor, output_clipped_edges, target_clipped_edges 94 | -------------------------------------------------------------------------------- /Variations/MODEL_EDSR_2x2_kernel.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | SCALING_FACTOR = 0.1 5 | 6 | def resBlock(x, channels=256, kernel_size=[3,3], scale=0.1): 7 | tmp = tf.layers.conv2d(x, channels, kernel_size, activation=None, padding='same') 8 | tmp = tf.nn.relu(tmp) 9 | tmp = tf.layers.conv2d(tmp, channels, kernel_size, activation=None, padding='same') 10 | tmp *= scale 11 | return x + tmp 12 | 13 | def _phase_shift(I, r): 14 | return tf.depth_to_space(I, r) 15 | 16 | def PS(X, r, color=False): 17 | if color: 18 | Xc = tf.split(X, 3, 3) 19 | X = tf.concat([_phase_shift(x, r) for x in Xc],3) 20 | else: 21 | X = _phase_shift(X, r) 22 | return X 23 | 24 | def upsample(x, scale=8, features=256, activation=tf.nn.relu): 25 | assert scale in [2, 3, 4, 8] 26 | x = tf.layers.conv2d(x, features, [3,3], activation=activation, padding='same') 27 | if scale == 2: 28 | ps_features = 3*(scale**2) 29 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='same') #Increase channel depth 30 | #x = slim.conv2d_transpose(x,ps_features,6,stride=1,activation_fn=activation) 31 | x = PS(x, 2, color=True) 32 | elif scale == 3: 33 | ps_features =3*(scale**2) 34 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='same') 35 | #x = slim.conv2d_transpose(x,ps_features,9,stride=1,activation_fn=activation) 36 | x = PS(x, 3, color=True) 37 | elif scale == 4: 38 | ps_features = 3*(2**2) 39 | for i in range(2): 40 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='same') 41 | #x = slim.conv2d_transpose(x,ps_features,6,stride=1,activation_fn=activation) 42 | x = PS(x, 2, color=True) 43 | elif scale == 8: 44 | ps_features = 3*(2*2) 45 | for i in range(3): 46 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='same') 47 | #x = slim.conv2d_transpose(x,ps_features,6,stride=1,activation_fn=activation) 48 | x = PS(x, 2, color=True) 49 | return x 50 | 51 | def model(input_tensor, scale=8, feature_size=256, num_layers=32): 52 | with tf.device("/gpu:0"): 53 | ''' 54 | Preprocessing by subtracting the batch mean 55 | ''' 56 | input_mean = tf.reduce_mean(input_tensor) 57 | tensor = input_tensor - input_mean 58 | 59 | ''' 60 | First convolution layer 61 | ''' 62 | tensor = tf.layers.conv2d(tensor, feature_size, [3,3], padding='same') #Linear activation by default 63 | 64 | conv_1 = tensor #backup 65 | 66 | ''' 67 | Building resBlocks 68 | ''' 69 | for i in range(num_layers): 70 | #tensor_kernel_3 = resBlock(tensor, feature_size, kernel_size=[3,3], scale=SCALING_FACTOR) 71 | #tensor_kernel_4 = resBlock(tensor, feature_size, kernel_size=[4,4], scale=SCALING_FACTOR) 72 | #tensor_kernel_5 = resBlock(tensor, feature_size, kernel_size=[5,5], scale=SCALING_FACTOR) 73 | #tensor = tf.concat([tensor_kernel_3, tensor_kernel_4, tensor_kernel_5], axis=3) #similar to Inception model 74 | #tensor = tf.layers.conv2d(tensor, feature_size, [1,1], activation=None, padding='same') 75 | tensor = resBlock(tensor, feature_size, kernel_size=[2,2], scale=SCALING_FACTOR) 76 | 77 | ''' 78 | Add back the conv_1 tensor 79 | ''' 80 | tensor = tf.layers.conv2d(tensor, feature_size, [3,3], padding='same') 81 | tensor += conv_1 82 | 83 | ''' 84 | Upsampling 85 | ''' 86 | tensor = upsample(tensor, scale, feature_size, activation=None) 87 | 88 | tensor = tf.clip_by_value(tensor+input_mean, 0.0, 255.0) 89 | return tensor 90 | -------------------------------------------------------------------------------- /Variations/MODEL_EDSR_Flatten.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | SCALING_FACTOR = 0.1 5 | 6 | def resBlock(x, channels=256, kernel_size=[3,3], scale=1): 7 | tmp = tf.layers.conv2d(x, channels, kernel_size, activation=None, padding='SAME') 8 | tmp = tf.nn.relu(tmp) 9 | tmp = tf.layers.conv2d(tmp, channels, kernel_size, activation=None, padding='SAME') 10 | tmp *= scale 11 | return x + tmp 12 | 13 | def _phase_shift(I, r): 14 | return tf.depth_to_space(I, r) 15 | 16 | def PS(X, r, color=False): 17 | if color: 18 | Xc = tf.split(X, 3, 3) 19 | X = tf.concat([_phase_shift(x, r) for x in Xc],3) 20 | else: 21 | X = _phase_shift(X, r) 22 | return X 23 | 24 | def upsample(x, scale=8, features=256, activation=tf.nn.relu): 25 | assert scale in [2, 3, 4, 8] 26 | x = tf.layers.conv2d(x, features, [3,3], activation=activation, padding='SAME') 27 | if scale == 2: 28 | ps_features = 3*(scale**2) 29 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') #Increase channel depth 30 | #x = slim.conv2d_transpose(x,ps_features,6,stride=1,activation_fn=activation) 31 | x = PS(x, 2, color=True) 32 | elif scale == 3: 33 | ps_features =3*(scale**2) 34 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') 35 | #x = slim.conv2d_transpose(x,ps_features,9,stride=1,activation_fn=activation) 36 | x = PS(x, 3, color=True) 37 | elif scale == 4: 38 | ps_features = 3*(2**2) 39 | for i in range(2): 40 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') 41 | #x = slim.conv2d_transpose(x,ps_features,6,stride=1,activation_fn=activation) 42 | x = PS(x, 2, color=True) 43 | elif scale == 8: 44 | ps_features = 3*(8*8) 45 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') 46 | #x = tf.layers.conv2d_transpose(x,3,kernel_size=(3,3),strides=2,activation=activation, padding='SAME') 47 | x = PS(x, 8, color=True) 48 | return x 49 | 50 | def flatten(input_tensor): # Duplicate the G Layer, then depth to space mapping 51 | rgb_layers = tf.split(input_tensor, [1,1,1], axis=3) 52 | rggb_layers = [rgb_layers[0], rgb_layers[1], rgb_layers[1], rgb_layers[2]] 53 | stacked = tf.concat(rggb_layers, axis=3) 54 | flattened = tf.depth_to_space(stacked, 2) 55 | return flattened 56 | 57 | 58 | def model(input_tensor, scale=8, feature_size=256, num_layers=32): 59 | with tf.device("/gpu:0"): 60 | scale = scale / 2 61 | 62 | ''' 63 | Preprocessing by subtracting the batch mean 64 | ''' 65 | #input_mean = tf.reduce_mean(input_tensor) 66 | input_mean = tf.reduce_mean(input_tensor, 2, keep_dims=True) 67 | input_mean = tf.reduce_mean(input_mean, 1, keep_dims=True) 68 | tensor = input_tensor - input_mean 69 | 70 | ''' 71 | Preprocessing by flattening the image 72 | ''' 73 | tensor = flatten(tensor) 74 | 75 | ''' 76 | First convolution layer 77 | ''' 78 | tensor = tf.layers.conv2d(tensor, feature_size, [6,6], padding='SAME') #Linear activation by default 79 | 80 | conv_1 = tensor #backup 81 | 82 | ''' 83 | Building resBlocks 84 | ''' 85 | for i in range(num_layers): 86 | tensor = resBlock(tensor, feature_size, scale=SCALING_FACTOR) 87 | 88 | ''' 89 | Add back the conv_1 tensor 90 | ''' 91 | tensor = tf.layers.conv2d(tensor, feature_size, [3,3], padding='SAME') 92 | tensor += conv_1 93 | 94 | ''' 95 | Upsampling 96 | ''' 97 | tensor = upsample(tensor, scale, feature_size, activation=None) 98 | 99 | tensor = tf.clip_by_value(tensor+input_mean, 0.0, 255.0) 100 | 101 | return tensor 102 | -------------------------------------------------------------------------------- /Variations/MODEL_EDSR_extra_smoothing_layer.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | def resBlock(x, channels=256, kernel_size=[3,3], scale=0.1): 5 | tmp = tf.layers.conv2d(x, channels, kernel_size, activation=None, padding='same') 6 | tmp = tf.nn.relu(tmp) 7 | tmp = tf.layers.conv2d(tmp, channels, kernel_size, activation=None, padding='same') 8 | tmp *= scale 9 | return x + tmp 10 | 11 | def _phase_shift(I, r): 12 | return tf.depth_to_space(I, r) 13 | 14 | def PS(X, r, color=False): 15 | if color: 16 | Xc = tf.split(X, 3, 3) 17 | X = tf.concat([_phase_shift(x, r) for x in Xc],3) 18 | else: 19 | X = _phase_shift(X, r) 20 | return X 21 | 22 | def upsample(x, scale=8, features=256, activation=tf.nn.relu): 23 | assert scale in [2, 3, 4, 8] 24 | x = tf.layers.conv2d(x, features, [3,3], activation=activation, padding='same') 25 | if scale == 2: 26 | ps_features = 3*(scale**2) 27 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='same') #Increase channel depth 28 | #x = slim.conv2d_transpose(x,ps_features,6,stride=1,activation_fn=activation) 29 | x = PS(x, 2, color=True) 30 | elif scale == 3: 31 | ps_features =3*(scale**2) 32 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='same') 33 | #x = slim.conv2d_transpose(x,ps_features,9,stride=1,activation_fn=activation) 34 | x = PS(x, 3, color=True) 35 | elif scale == 4: 36 | ps_features = 3*(2**2) 37 | for i in range(2): 38 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='same') 39 | #x = tf.layers.conv2d_transpose(x,ps_features,6,stride=1,activation_fn=activation) 40 | x = PS(x, 2, color=True) 41 | elif scale == 8: 42 | ps_features = 3*(2*2) 43 | for i in range(3): 44 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='same') 45 | #x = tf.layers.conv2d_transpose(x,3,kernel_size=(3,3),strides=2,activation=activation, padding='same') 46 | x = PS(x, 2, color=True) 47 | return x 48 | 49 | def smoothing(X, feature_size=256, activation=tf.nn.relu): 50 | return tf.layers.conv2d(X, feature_size, [3,3], activation=activation, padding='same') 51 | 52 | def model(input_tensor, scale=8, feature_size=256, num_layers=32): 53 | with tf.device("/gpu:3"): 54 | ''' 55 | Preprocessing by subtracting the batch mean 56 | ''' 57 | input_mean = tf.reduce_mean(input_tensor) 58 | tensor = input_tensor - input_mean 59 | 60 | ''' 61 | First convolution layer 62 | ''' 63 | tensor = tf.layers.conv2d(tensor, feature_size, [3,3], padding='same') #Linear activation by default 64 | 65 | conv_1 = tensor #backup 66 | 67 | SCALING_FACTOR = 0.1 68 | 69 | ''' 70 | Building resBlocks 71 | ''' 72 | for i in range(num_layers): 73 | tensor = resBlock(tensor, feature_size, kernel_size=[3,3], scale=SCALING_FACTOR) 74 | 75 | ''' 76 | Add back the conv_1 tensor 77 | ''' 78 | tensor = tf.layers.conv2d(tensor, feature_size, [3,3], padding='same') 79 | tensor += conv_1 80 | 81 | ''' 82 | Upsampling 83 | ''' 84 | tensor = upsample(tensor, scale, feature_size, activation=None) 85 | 86 | ''' 87 | More convolution for smoothing and reducing ghosting 88 | ''' 89 | for i in range(2): 90 | tensor = smoothing(tensor, feature_size=3, activation=tf.nn.relu) 91 | 92 | tensor = tf.layers.conv2d(tensor, 3, [3,3], activation=None, padding='same') 93 | 94 | tensor = tf.clip_by_value(tensor+input_mean, 0.0, 255.0) 95 | return tensor 96 | -------------------------------------------------------------------------------- /Variations/MODEL_EDSR_transposed_convolution.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | SCALING_FACTOR = 0.1 5 | 6 | def resBlock(x, channels=256, kernel_size=[3,3], scale=1): 7 | tmp = tf.layers.conv2d(x, channels, kernel_size, activation=None, padding='SAME') 8 | tmp = tf.nn.relu(tmp) 9 | tmp = tf.layers.conv2d(tmp, channels, kernel_size, activation=None, padding='SAME') 10 | tmp *= scale 11 | return x + tmp 12 | 13 | def _phase_shift(I, r): 14 | return tf.depth_to_space(I, r) 15 | 16 | def PS(X, r, color=False): 17 | if color: 18 | Xc = tf.split(X, 3, 3) 19 | X = tf.concat([_phase_shift(x, r) for x in Xc],3) 20 | else: 21 | X = _phase_shift(X, r) 22 | return X 23 | 24 | def upsample(x, scale=8, features=256, activation=tf.nn.relu): 25 | assert scale in [2, 3, 4, 8] 26 | x = tf.layers.conv2d(x, features, [3,3], activation=activation, padding='SAME') 27 | if scale == 2: 28 | ps_features = 3*(scale**2) 29 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') #Increase channel depth 30 | #x = slim.conv2d_transpose(x,ps_features,6,stride=1,activation_fn=activation) 31 | x = PS(x, 2, color=True) 32 | elif scale == 3: 33 | ps_features =3*(scale**2) 34 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') 35 | #x = slim.conv2d_transpose(x,ps_features,9,stride=1,activation_fn=activation) 36 | x = PS(x, 3, color=True) 37 | elif scale == 4: 38 | ps_features = 3*(2**2) 39 | for i in range(2): 40 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') 41 | #x = slim.conv2d_transpose(x,ps_features,6,stride=1,activation_fn=activation) 42 | x = PS(x, 2, color=True) 43 | elif scale == 8: 44 | ps_features = 3*(2*2) 45 | for i in range(3): 46 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') 47 | x = tf.layers.conv2d_transpose(x,3,kernel_size=(3,3),strides=2,activation=activation, padding='SAME') 48 | #x = PS(x, 2, color=True) 49 | return x 50 | 51 | def model(input_tensor, target_tensor, scale=8, feature_size=256, num_layers=32): 52 | with tf.device("/gpu:0"): 53 | ''' 54 | Preprocessing by subtracting the batch mean 55 | ''' 56 | input_mean = tf.reduce_mean(input_tensor) 57 | tensor = input_tensor - input_mean 58 | 59 | ''' 60 | First convolution layer 61 | ''' 62 | tensor = tf.layers.conv2d(tensor, feature_size, [3,3], padding='SAME') #Linear activation by default 63 | 64 | conv_1 = tensor #backup 65 | 66 | ''' 67 | Building resBlocks 68 | ''' 69 | for i in range(num_layers): 70 | tensor = resBlock(tensor, feature_size, scale=SCALING_FACTOR) 71 | 72 | ''' 73 | Add back the conv_1 tensor 74 | ''' 75 | tensor = tf.layers.conv2d(tensor, feature_size, [3,3], padding='SAME') 76 | tensor += conv_1 77 | 78 | ''' 79 | Upsampling 80 | ''' 81 | tensor = upsample(tensor, scale, feature_size, activation=None) 82 | 83 | tensor = tf.clip_by_value(tensor+input_mean, 0.0, 255.0) 84 | 85 | sobel = tf.constant([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], tf.float32) 86 | sobel_3 = tf.stack([sobel, sobel, sobel], axis=2) 87 | sobel_filter = tf.reshape(sobel_3, [1, 3, 3, 3]) 88 | output_edges = tf.nn.conv2d(tensor, sobel_filter, strides=[1,1,1,1], padding='SAME') 89 | output_clipped_edges = tf.clip_by_value(output_edges, 50.0, 255.0) 90 | 91 | target_edges = tf.nn.conv2d(target_tensor, sobel_filter, strides=[1,1,1,1], padding='SAME') 92 | target_clipped_edges = tf.clip_by_value(target_edges, 50.0, 255.0) 93 | 94 | return tensor, output_clipped_edges, target_clipped_edges 95 | -------------------------------------------------------------------------------- /Variations/MODEL_EDSR_x2x2x2.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | SCALING_FACTOR = 0.1 5 | 6 | def resBlock(x, channels=256, kernel_size=[3,3], scale=1): 7 | tmp = tf.layers.conv2d(x, channels, kernel_size, activation=None, padding='SAME') 8 | tmp = tf.nn.relu(tmp) 9 | tmp = tf.layers.conv2d(tmp, channels, kernel_size, activation=None, padding='SAME') 10 | tmp *= scale 11 | return x + tmp 12 | 13 | def _phase_shift(I, r): 14 | return tf.depth_to_space(I, r) 15 | 16 | def PS(X, r, color=False): 17 | if color: 18 | Xc = tf.split(X, 3, 3) 19 | X = tf.concat([_phase_shift(x, r) for x in Xc],3) 20 | else: 21 | X = _phase_shift(X, r) 22 | return X 23 | 24 | def upsample(x, scale=8, features=256, activation=tf.nn.relu): 25 | assert scale in [2, 3, 4, 8] 26 | x = tf.layers.conv2d(x, features, [3,3], activation=activation, padding='SAME') 27 | if scale == 2: 28 | ps_features = 3*(scale**2) 29 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') #Increase channel depth 30 | #x = slim.conv2d_transpose(x,ps_features,6,stride=1,activation_fn=activation) 31 | x = PS(x, 2, color=True) 32 | elif scale == 3: 33 | ps_features =3*(scale**2) 34 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') 35 | #x = slim.conv2d_transpose(x,ps_features,9,stride=1,activation_fn=activation) 36 | x = PS(x, 3, color=True) 37 | elif scale == 4: 38 | ps_features = 3*(2**2) 39 | for i in range(2): 40 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') 41 | #x = slim.conv2d_transpose(x,ps_features,6,stride=1,activation_fn=activation) 42 | x = PS(x, 2, color=True) 43 | elif scale == 8: 44 | ps_features = 3*(2*2) 45 | for i in range(3): 46 | x = tf.layers.conv2d(x, ps_features, [3,3], activation=activation, padding='SAME') 47 | x = tf.layers.conv2d_transpose(x,3,kernel_size=(3,3),strides=2,activation=activation, padding='SAME') 48 | #x = PS(x, 2, color=True) 49 | return x 50 | 51 | def model(input_tensor, target_tensor, scale=8, feature_size=256, num_layers=32): 52 | with tf.device("/gpu:0"): 53 | ''' 54 | Preprocessing by subtracting the batch mean 55 | ''' 56 | input_mean = tf.reduce_mean(input_tensor) 57 | tensor = input_tensor - input_mean 58 | 59 | ''' 60 | First convolution layer 61 | ''' 62 | tensor = tf.layers.conv2d(tensor, feature_size, [3,3], padding='SAME') #Linear activation by default 63 | 64 | conv_1 = tensor #backup 65 | 66 | ''' 67 | Building resBlocks 68 | ''' 69 | for i in range(num_layers): 70 | tensor = resBlock(tensor, feature_size, scale=SCALING_FACTOR) 71 | 72 | ''' 73 | Add back the conv_1 tensor 74 | ''' 75 | tensor = tf.layers.conv2d(tensor, feature_size, [3,3], padding='SAME') 76 | tensor += conv_1 77 | 78 | ''' 79 | Upsampling 80 | ''' 81 | tensor = upsample(tensor, scale, feature_size, activation=None) 82 | 83 | tensor = tf.clip_by_value(tensor+input_mean, 0.0, 255.0) 84 | 85 | sobel = tf.constant([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], tf.float32) 86 | sobel_3 = tf.stack([sobel, sobel, sobel], axis=2) 87 | sobel_filter = tf.reshape(sobel_3, [1, 3, 3, 3]) 88 | output_edges = tf.nn.conv2d(tensor, sobel_filter, strides=[1,1,1,1], padding='SAME') 89 | output_clipped_edges = tf.clip_by_value(output_edges, 50.0, 255.0) 90 | 91 | target_edges = tf.nn.conv2d(target_tensor, sobel_filter, strides=[1,1,1,1], padding='SAME') 92 | target_clipped_edges = tf.clip_by_value(target_edges, 50.0, 255.0) 93 | 94 | return tensor, output_clipped_edges, target_clipped_edges 95 | -------------------------------------------------------------------------------- /train_data_generation.m: -------------------------------------------------------------------------------- 1 | target_dir = '/usr/project/xtmp/DIV2K_train_HR'; 2 | count = 0; 3 | factor = 8; 4 | patch_size = 400; 5 | f_lst = []; 6 | f_lst = [f_lst; dir(fullfile(target_dir, '*.png'))]; 7 | for f = 1:numel(f_lst) 8 | %fname = sprintf('%04d.png', i); 9 | file = f_lst(f); 10 | fname = file.name; 11 | if fname(1)=='.' 12 | continue; 13 | end 14 | target_path = fullfile(target_dir, fname); 15 | 16 | target_img = imread(target_path); 17 | 18 | img_size = size(target_img); 19 | 20 | target_img = target_img(1:img_size(1)- mod(img_size(1), factor), 1:img_size(2) - mod(img_size(2), factor), :); 21 | 22 | img_size = size(target_img); 23 | 24 | x_max = img_size(1) - patch_size; 25 | y_max = img_size(2) - patch_size; 26 | 27 | x = uint32(rand(1, 100) * x_max) + 1; 28 | y = uint32(rand(1, 100) * y_max) + 1; 29 | 30 | for i = 1:length(x) 31 | x_coord = x(i); 32 | y_coord = y(i); 33 | 34 | original_target_patch = target_img(x_coord:x_coord+patch_size-1, y_coord:y_coord+patch_size-1, :); 35 | original_input_patch = imresize(original_target_patch, 1/factor, 'bicubic'); 36 | 37 | target_patch = original_target_patch; 38 | input_patch = original_input_patch; 39 | patch_name = sprintf('/usr/project/xtmp/EDSR_HR_train/%d', count); 40 | save(patch_name, 'target_patch', 'input_patch', '-v6'); 41 | count = count + 1; 42 | 43 | target_patch = imrotate(original_target_patch, 90); 44 | input_patch = imrotate(original_input_patch, 90); 45 | patch_name = sprintf('/usr/project/xtmp/EDSR_HR_train/%d', count); 46 | save(patch_name, 'target_patch', 'input_patch', '-v6'); 47 | count = count + 1; 48 | 49 | target_patch = imrotate(original_target_patch, 180); 50 | input_patch = imrotate(original_input_patch, 180); 51 | patch_name = sprintf('/usr/project/xtmp/EDSR_HR_train/%d', count); 52 | save(patch_name, 'target_patch', 'input_patch', '-v6'); 53 | count = count + 1; 54 | 55 | target_patch = imrotate(original_target_patch, 270); 56 | input_patch = imrotate(original_input_patch, 270); 57 | patch_name = sprintf('/usr/project/xtmp/EDSR_HR_train/%d', count); 58 | save(patch_name, 'target_patch', 'input_patch', '-v6'); 59 | count = count + 1; 60 | 61 | target_patch = flipdim(original_target_patch, 1); 62 | input_patch = flipdim(original_input_patch, 1); 63 | patch_name = sprintf('/usr/project/xtmp/EDSR_HR_train/%d', count); 64 | save(patch_name, 'target_patch', 'input_patch', '-v6'); 65 | count = count + 1; 66 | 67 | target_patch = flipdim(original_target_patch, 2); 68 | input_patch = flipdim(original_input_patch, 2); 69 | patch_name = sprintf('/usr/project/xtmp/EDSR_HR_train/%d', count); 70 | save(patch_name, 'target_patch', 'input_patch', '-v6'); 71 | count = count + 1; 72 | 73 | display(count); 74 | end 75 | end 76 | --------------------------------------------------------------------------------