├── ARGAN ├── Discriminator.py ├── Eval.py ├── Input.py ├── Losses.py ├── Model.py ├── Test.py ├── Train.py ├── avg_psnr.py ├── results.txt ├── run_debug.sh ├── run_eval.sh ├── run_test.sh └── run_train.sh ├── README.md ├── baseline ├── Input.py ├── Losses.py ├── Losses2.py ├── Model.py ├── Test.py ├── Train.py ├── avg_psnr.py ├── avg_psnr_results.txt ├── results.txt ├── run.sh ├── run_debug.sh ├── run_eval.sh └── run_test.sh └── result.png /ARGAN/Discriminator.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | 3 | import tensorflow as tf 4 | import re 5 | 6 | alpha = 0.2 7 | TOWER_NAME ='tower' 8 | def _variable_on_cpu(name, shape, initializer): 9 | """Helper to create a Variable stored on CPU memory. 10 | 11 | Args: 12 | name: name of the variable 13 | shape: list of ints 14 | initializer: initializer for Variable 15 | 16 | Returns: 17 | Variable Tensor 18 | """ 19 | with tf.device('/cpu:0'): 20 | dtype = tf.float32# if FLAGS.use_fp16 else tf.float32 21 | var = tf.get_variable(name, shape, initializer=initializer, dtype=dtype) 22 | return var 23 | 24 | 25 | def _variable_with_weight_decay(name, shape, stddev=1e-3): 26 | """Helper to create an initialized Variable with weight decay. 27 | 28 | Note that the Variable is initialized with a truncated normal distribution. 29 | A weight decay is added only if one is specified. 30 | 31 | Args: 32 | name: name of the variable 33 | shape: list of ints 34 | stddev: standard deviation of a truncated Gaussian 35 | wd: add L2Loss weight decay multiplied by this float. If None, weight 36 | decay is not added for this Variable. 37 | 38 | Returns: 39 | Variable Tensor 40 | """ 41 | dtype = tf.float32# if FLAGS.use_fp16 else tf.float32 42 | var = _variable_on_cpu( 43 | name, 44 | shape, 45 | tf.truncated_normal_initializer(stddev=stddev, dtype=dtype)) 46 | 47 | return var 48 | 49 | 50 | 51 | def batch_norm(inputs, is_training, pop_mean, pop_var, beta, gamma, decay = 0.999): 52 | def train(): 53 | batch_mean, batch_var = tf.nn.moments(inputs,[0,1,2], name='moments') 54 | train_mean = tf.assign(pop_mean, pop_mean * decay + batch_mean * (1 - decay)) 55 | train_var = tf.assign(pop_var , pop_var * decay + batch_var * (1 - decay)) 56 | with tf.control_dependencies([train_mean, train_var]): 57 | return tf.nn.batch_normalization(inputs, batch_mean, batch_var, beta, gamma, variance_epsilon=1e-3) 58 | 59 | def testing(): 60 | return tf.nn.batch_normalization(inputs, pop_mean, pop_var, beta, gamma, variance_epsilon=1e-3) 61 | 62 | return tf.cond(is_training, train, testing) 63 | 64 | 65 | 66 | 67 | 68 | 69 | class Discriminator(object): 70 | def __init__(self,sess,optimizer, batch_size=32): 71 | self.sess = sess 72 | self.optimizer = optimizer 73 | self.batch_size = batch_size 74 | self.THETA = {} 75 | self.init_model() 76 | self.saver = tf.train.Saver(self.THETA, max_to_keep=100) 77 | 78 | 79 | def get_training(self): 80 | return self.training 81 | 82 | def set_training(self, training): 83 | self.training = tf.assign(self.training, tf.constant(training, shape=[])) 84 | 85 | def init_model(self): 86 | def init_bn_params(block_name,ind,size=64): 87 | self.THETA['{}beta{}'.format(block_name,ind)] = tf.Variable( 88 | tf.truncated_normal([size], mean=0.0, stddev=1e-3), name='{}beta{}'.format(block_name,ind), trainable=True) 89 | self.THETA['{}gamma{}'.format(block_name,ind)]= tf.Variable( 90 | tf.truncated_normal([size], mean=1.0, stddev=1e-3), name='{}gamma{}'.format(block_name,ind), trainable=True) 91 | self.THETA['{}pop_mean{}'.format(block_name,ind)] = tf.Variable(tf.zeros([size]), name = '{}pop_mean{}'.format(block_name,ind), trainable=False) 92 | self.THETA['{}pop_var{}'.format(block_name,ind)] = tf.Variable(tf.ones([size]) , name = '{}pop_var{}'.format(block_name,ind) , trainable=False) 93 | 94 | 95 | def init_block(block_name,in_channels,out_channels): 96 | self.THETA['{}w1'.format(block_name)] = _variable_with_weight_decay(shape=[3, 3, in_channels, out_channels], name='{}w1'.format(block_name)) 97 | self.THETA['{}b1'.format(block_name)] = _variable_with_weight_decay(shape=[ out_channels] , name='{}b1'.format(block_name)) 98 | 99 | init_bn_params('{}'.format(block_name),1,size=out_channels) 100 | 101 | self.THETA['{}w2'.format(block_name)] = _variable_with_weight_decay(shape=[3, 3, out_channels, out_channels], name='{}w2'.format(block_name)) 102 | self.THETA['{}b2'.format(block_name)] = _variable_with_weight_decay(shape=[ out_channels] , name='{}b2'.format(block_name)) 103 | 104 | init_bn_params('{}'.format(block_name),2,size=out_channels) 105 | 106 | init_block('block1', 3,64) 107 | init_block('block2', 64,128) 108 | init_block('block3',128,256) 109 | init_block('block4',256,512) 110 | 111 | 112 | 113 | self.THETA['dense1w1'] = _variable_with_weight_decay(shape=[8192, 1024], name='dense1w1') 114 | self.THETA['dense1b1'] = _variable_with_weight_decay(shape=[1024] , name='dense1b1') 115 | 116 | 117 | self.THETA['dense2w1'] = _variable_with_weight_decay(shape=[1024, 1], name='dense2w1') 118 | self.THETA['dense2b1'] = _variable_with_weight_decay(shape=[1] , name='dense2b1') 119 | 120 | 121 | 122 | 123 | 124 | def inference(self,inp,phase_train): 125 | 126 | def conv(blk,ind,inp, stride): 127 | x = tf.nn.conv2d(inp, self.THETA['{}w{}'.format(blk,ind)], strides=[1, stride, stride, 1], padding='SAME',name="{}w{}".format(blk,ind)) 128 | x = tf.nn.bias_add(x, self.THETA['{}b{}'.format(blk,ind)], name="{}b{}".format(blk,ind)) 129 | return x 130 | 131 | def bn(blk,ind,inp): 132 | x = batch_norm(inp, 133 | phase_train, 134 | self.THETA['{}pop_mean{}'.format(blk,ind)], 135 | self.THETA['{}pop_var{}'.format(blk,ind)], 136 | self.THETA['{}beta{}'.format(blk,ind)], 137 | self.THETA['{}gamma{}'.format(blk,ind)], 138 | ) 139 | return x 140 | 141 | 142 | def block_conv(inp, name): 143 | with tf.variable_scope(name): 144 | x = conv(name,1,inp,stride=1) 145 | x = bn(name,1,x) 146 | c1 = tf.maximum(alpha*x,x)#tf.nn.relu(x) 147 | 148 | x = conv(name ,2,c1,stride=2) 149 | x = bn(name,2,x) 150 | c2 = tf.maximum(alpha*x,x)#tf.nn.relu(x) 151 | 152 | return c2 153 | 154 | x = block_conv(inp,'block1') 155 | x = block_conv(x ,'block2') 156 | x = block_conv(x ,'block3') 157 | x = block_conv(x ,'block4') 158 | shape = x.get_shape().as_list() 159 | x = tf.reshape(x, [-1,shape[1]*shape[2]*shape[3]]) 160 | x = tf.matmul(x,self.THETA['dense1w1']) + self.THETA['dense1b1'] 161 | x = tf.maximum(alpha*x,x)#tf.nn.relu(x) 162 | 163 | x = tf.matmul(x,self.THETA['dense2w1']) #+ self.THETA['dense2b1'] 164 | # x = tf.nn.sigmoid(x) 165 | # so that it return logits not probabilties 166 | out = x 167 | return out 168 | 169 | 170 | def save_model(self,file_name): 171 | self.saver.save(self.sess,file_name) 172 | 173 | def load_model(self,file_name): 174 | self.saver.restore(self.sess, file_name) 175 | 176 | 177 | def get_gradients(self,total_loss): 178 | grads = self.optimizer.compute_gradients(total_loss) 179 | 180 | return grads 181 | 182 | def get_backprop(self,total_loss,global_step): 183 | solver = self.optimizer.minimize(total_loss,global_step=global_step) 184 | return solver 185 | -------------------------------------------------------------------------------- /ARGAN/Eval.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | parser = argparse.ArgumentParser(description='Expert denoisers training') 3 | parser.add_argument('--image-folder', type=str, default='', metavar='N', required=True, help='Image folder containing all the images') 4 | parser.add_argument('--image-filename', type=str, default='', metavar='N', required=True, help='Image file name') 5 | parser.add_argument('--quality', type=int, default='', metavar='N', required=True, help='quality') 6 | args = parser.parse_args() 7 | 8 | import skimage.measure as measure 9 | from PIL import Image 10 | import numpy as np 11 | import os 12 | import sys 13 | import math 14 | 15 | quality = args.quality 16 | image_folder = args.image_folder 17 | clean_image_folder = os.path.join(image_folder , "clean_images") 18 | noisy_image_folder = os.path.join(image_folder , "noise_std_" + str(quality)) 19 | recon_image_folder = os.path.join(image_folder , "recon_std_" + str(quality)) 20 | 21 | def mse(imageA, imageB): 22 | # the 'Mean Squared Error' between the two images is the 23 | # sum of the squared difference between the two images; 24 | # note: the two images must have the same dimension 25 | err = np.sum((imageA.astype("float") - imageB.astype("float")) ** 2) 26 | err /= float(imageA.shape[0] * imageA.shape[1]* imageA.shape[2]) 27 | 28 | # return the MSE, the lower the error, the more "similar" 29 | # the two images are 30 | return err 31 | 32 | 33 | 34 | def psnr(img1, img2): 35 | mse = np.mean( (img1 - img2) ** 2 ) 36 | if mse == 0: 37 | return 100 38 | PIXEL_MAX = 255.0 39 | return 20 * math.log10(PIXEL_MAX / math.sqrt(mse)) 40 | 41 | 42 | image_file = args.image_filename 43 | #print image_file 44 | with open(image_file,'r') as f: 45 | for l in f.readlines(): 46 | l = l.strip() 47 | l = l.rsplit('/',1)[1] 48 | 49 | Ic = np.asarray(Image.open(os.path.join(clean_image_folder,l)), dtype = 'float32') 50 | 51 | In = np.asarray(Image.open(os.path.join(noisy_image_folder,l)), dtype = 'float32') 52 | Ir = np.asarray(Image.open(os.path.join(recon_image_folder,l)), dtype = 'float32') 53 | 54 | print l+" "+str(psnr(Ic,In))+" "+str(psnr(Ic,Ir)) 55 | -------------------------------------------------------------------------------- /ARGAN/Input.py: -------------------------------------------------------------------------------- 1 | import os 2 | import glob 3 | import tensorflow as tf 4 | 5 | import random 6 | from tqdm import tqdm 7 | from itertools import product 8 | 9 | class BatchImageInput(): 10 | def __init__(self, file_paths, noise_levels, crop_size = [64, 64, 3], batch_size = 32): 11 | self.batch_size = batch_size 12 | self.noise_levels = noise_levels 13 | self.crop_size = crop_size 14 | self.file_paths = file_paths 15 | 16 | def read_one_file(self, filename_queue): 17 | print filename_queue 18 | image_reader = tf.WholeFileReader() 19 | _,I = image_reader.read(filename_queue) 20 | Id = tf.image.decode_image(I, channels=3) 21 | Id = tf.to_float(Id) 22 | Ic = tf.random_crop(Id, size = self.crop_size) 23 | 24 | tt = tf.cast(Ic,tf.uint8) 25 | tt = tf.image.encode_jpeg(tt,quality=20) # have to remove this hardcoding later 26 | In = tf.image.decode_jpeg(tt, channels=3) 27 | In = tf.reshape(In, self.crop_size) 28 | In = tf.to_float(In) 29 | 30 | return Ic,In 31 | 32 | def get_minibatch_tensors(self, num_epochs=None): 33 | input_queue = tf.train.string_input_producer(self.file_paths, num_epochs=num_epochs, shuffle=True) 34 | noises_ = tf.constant(self.noise_levels) 35 | shuff_noises_ = tf.random_shuffle(noises_) 36 | 37 | Ic, In = self.read_one_file(input_queue) 38 | # adding noise to a clean image 39 | # In = Ic + tf.random_normal(Ic.get_shape(), mean=0.0, stddev=shuff_noises_[0], dtype=tf.float32) 40 | 41 | Ic_batch, In_batch = tf.train.batch([Ic, In], batch_size=self.batch_size, capacity = 100) 42 | # (X,Y) 43 | return In_batch/255.0, Ic_batch/255.0 44 | 45 | if __name__ == '__main__': 46 | with open('imagenet_val_paths.txt', 'r') as f: 47 | file_paths = f.read().splitlines() 48 | 49 | EI = BatchImageInput(file_paths,[10.0,15.0],batch_size = 32) 50 | 51 | x_,y_ = EI.get_minibatch_tensors() 52 | 53 | config = tf.ConfigProto() 54 | config.gpu_options.allow_growth = True 55 | print "haha" 56 | with tf.Session(config=config) as sess: 57 | tf.global_variables_initializer() 58 | coord = tf.train.Coordinator() 59 | threads = tf.train.start_queue_runners(sess=sess, coord=coord) 60 | x,y = sess.run([x_,y_]) 61 | print x[0][:10,:10,0],y[0][:10,:10,0] 62 | print x[0].shape 63 | coord.request_stop() 64 | coord.join(threads) 65 | 66 | -------------------------------------------------------------------------------- /ARGAN/Losses.py: -------------------------------------------------------------------------------- 1 | from tensorflow.python.ops import math_ops 2 | import tensorflow as tf 3 | import numpy 4 | 5 | def lp_loss(y_true, y_pred): 6 | diff_sq = tf.squared_difference(y_true,y_pred) 7 | loss = tf.reduce_mean(diff_sq, axis=None) 8 | return loss 9 | 10 | 11 | 12 | def tv_loss(y_true,y_pred): 13 | def total_variation(images): 14 | pixel_dif1 = images[:, 1:, :, :] - images[:, :-1, :, :] 15 | pixel_dif2 = images[:, :, 1:, :] - images[:, :, :-1, :] 16 | sum_axis = [1, 2, 3] 17 | tot_var = math_ops.reduce_sum(math_ops.abs(pixel_dif1), axis=sum_axis) + \ 18 | math_ops.reduce_sum(math_ops.abs(pixel_dif2), axis=sum_axis) 19 | return tot_var 20 | 21 | loss = tf.reduce_mean(total_variation(y_pred)) 22 | return loss 23 | 24 | 25 | 26 | def perceptual_loss(y_true,y_pred, layer_list=['conv5_4'], alphas=[1.0]): 27 | def conv_layer(bottom, weight, bias=None, s=1, padding='SAME', relu=True, group=1): 28 | if group==1: 29 | conv = tf.nn.conv2d(bottom, weight, [1, s, s, 1], padding=padding) 30 | else: 31 | input_split = tf.split(bottom, group, 3) 32 | weight_split = tf.split(weight, group, 3) 33 | conv_1 = tf.nn.conv2d(input_split[0], weight_split[0], [1, s, s, 1], padding=padding) 34 | conv_2 = tf.nn.conv2d(input_split[1], weight_split[1], [1, s, s, 1], padding=padding) 35 | conv = tf.concat([conv_1, conv_2], 3) 36 | if bias is None: 37 | if relu: 38 | return tf.nn.relu(conv) 39 | else: 40 | return conv 41 | else: 42 | bias = tf.nn.bias_add(conv, bias) 43 | if relu: 44 | return tf.nn.relu(bias) 45 | else: 46 | return bias 47 | def max_pool(bottom, k=3, s=1, padding='SAME'): 48 | return tf.nn.max_pool(bottom, ksize=[1, k, k, 1], strides=[1, s, s, 1], padding=padding) 49 | def vgg_19(image, weights, biases, keep_prob=1.0): 50 | activations = {} 51 | shapes = image.get_shape().as_list() 52 | if shapes[3] == 1: 53 | rgb = tf.concat(axis = 3, values=[image,image,image]) 54 | else: 55 | rgb = image 56 | 57 | rgb_scaled = rgb*255.0 58 | red, green, blue = tf.split(axis=3, num_or_size_splits=3, value=rgb_scaled) 59 | VGG_MEAN = [103.939, 116.779, 123.68] 60 | bgr = tf.concat(axis=3, values=[ 61 | blue - VGG_MEAN[0], 62 | green - VGG_MEAN[1], 63 | red - VGG_MEAN[2], 64 | ]) 65 | 66 | image = bgr 67 | 68 | 69 | with tf.name_scope("conv1"): 70 | activations['conv1_1'] = conv_layer(image, weights['conv1_1'], biases['conv1_1']) 71 | activations['conv1_2'] = conv_layer(activations['conv1_1'], weights['conv1_2'], biases['conv1_2']) 72 | activations['pool1'] = max_pool(activations['conv1_2'], k=2, s=2) 73 | 74 | with tf.name_scope("conv2"): 75 | activations['conv2_1'] = conv_layer(activations['pool1'], weights['conv2_1'], biases['conv2_1']) 76 | activations['conv2_2'] = conv_layer(activations['conv2_1'], weights['conv2_2'], biases['conv2_2']) 77 | activations['pool2'] = max_pool(activations['conv2_2'], k=2, s=2) 78 | 79 | with tf.name_scope("conv3"): 80 | activations['conv3_1'] = conv_layer(activations['pool2'], weights['conv3_1'], biases['conv3_1']) 81 | activations['conv3_2'] = conv_layer(activations['conv3_1'], weights['conv3_2'], biases['conv3_2']) 82 | activations['conv3_3'] = conv_layer(activations['conv3_2'], weights['conv3_3'], biases['conv3_3']) 83 | activations['conv3_4'] = conv_layer(activations['conv3_3'], weights['conv3_4'], biases['conv3_3']) 84 | activations['pool3'] = max_pool(activations['conv3_4'], k=2, s=2) 85 | 86 | with tf.name_scope("conv4"): 87 | activations['conv4_1'] = conv_layer(activations['pool3'], weights['conv4_1'], biases['conv4_1']) 88 | activations['conv4_2'] = conv_layer(activations['conv4_1'], weights['conv4_2'], biases['conv4_2']) 89 | activations['conv4_3'] = conv_layer(activations['conv4_2'], weights['conv4_3'], biases['conv4_3']) 90 | activations['conv4_4'] = conv_layer(activations['conv4_3'], weights['conv4_4'], biases['conv4_4']) 91 | activations['pool4'] = max_pool(activations['conv4_4'], k=2, s=2) 92 | 93 | with tf.name_scope("conv5"): 94 | activations['conv5_1'] = conv_layer(activations['pool4'], weights['conv5_1'], biases['conv5_1']) 95 | activations['conv5_2'] = conv_layer(activations['conv5_1'], weights['conv5_2'], biases['conv5_2']) 96 | activations['conv5_3'] = conv_layer(activations['conv5_2'], weights['conv5_3'], biases['conv5_3']) 97 | activations['conv5_4'] = conv_layer(activations['conv5_3'], weights['conv5_4'], biases['conv5_4']) 98 | activations['pool5'] = max_pool(activations['conv5_4'], k=2, s=2) 99 | return activations 100 | 101 | def get_weights(): 102 | net = numpy.load('../Resources/VGG/vgg19.npy').item() 103 | weights = {} 104 | biases = {} 105 | for name in net.keys(): 106 | weights[name] = tf.Variable(tf.constant(net[name][0]), dtype='float32' ,name=name+"_weight", trainable=False) 107 | biases[name] = tf.Variable(tf.constant(net[name][1]), dtype='float32' ,name=name+"_bias", trainable=False) 108 | return weights,biases 109 | 110 | weights,biases = get_weights() 111 | 112 | y_true_activations = vgg_19(y_true,weights,biases) 113 | y_pred_activations = vgg_19(y_pred,weights,biases) 114 | 115 | for layer_name in layer_list: 116 | assert layer_name in y_true_activations.keys() 117 | 118 | loss = 0 119 | for layer_name,alpha in zip(layer_list,alphas): 120 | y = y_true_activations[layer_name]/255.0 121 | x = y_pred_activations[layer_name]/255.0 122 | loss += alpha* lp_loss(y,x) 123 | return loss 124 | 125 | 126 | def cross_entropy_sigmoid(logits,labels): 127 | loss = tf.nn.sigmoid_cross_entropy_with_logits(logits = logits,labels = labels) 128 | return tf.reduce_mean(loss) 129 | -------------------------------------------------------------------------------- /ARGAN/Model.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | 3 | import tensorflow as tf 4 | import re 5 | 6 | 7 | TOWER_NAME ='tower' 8 | def _variable_on_cpu(name, shape, initializer): 9 | """Helper to create a Variable stored on CPU memory. 10 | 11 | Args: 12 | name: name of the variable 13 | shape: list of ints 14 | initializer: initializer for Variable 15 | 16 | Returns: 17 | Variable Tensor 18 | """ 19 | with tf.device('/cpu:0'): 20 | dtype = tf.float32# if FLAGS.use_fp16 else tf.float32 21 | var = tf.get_variable(name, shape, initializer=initializer, dtype=dtype) 22 | return var 23 | 24 | 25 | def _variable_with_weight_decay(name, shape, stddev=5e-2): 26 | """Helper to create an initialized Variable with weight decay. 27 | 28 | Note that the Variable is initialized with a truncated normal distribution. 29 | A weight decay is added only if one is specified. 30 | 31 | Args: 32 | name: name of the variable 33 | shape: list of ints 34 | stddev: standard deviation of a truncated Gaussian 35 | wd: add L2Loss weight decay multiplied by this float. If None, weight 36 | decay is not added for this Variable. 37 | 38 | Returns: 39 | Variable Tensor 40 | """ 41 | dtype = tf.float32# if FLAGS.use_fp16 else tf.float32 42 | var = _variable_on_cpu( 43 | name, 44 | shape, 45 | tf.truncated_normal_initializer(stddev=stddev, dtype=dtype)) 46 | return var 47 | 48 | 49 | 50 | def batch_norm(inputs, is_training, pop_mean, pop_var, beta, gamma, decay = 0.999): 51 | def train(): 52 | batch_mean, batch_var = tf.nn.moments(inputs,[0,1,2], name='moments') 53 | train_mean = tf.assign(pop_mean, pop_mean * decay + batch_mean * (1 - decay)) 54 | train_var = tf.assign(pop_var , pop_var * decay + batch_var * (1 - decay)) 55 | with tf.control_dependencies([train_mean, train_var]): 56 | return tf.nn.batch_normalization(inputs, batch_mean, batch_var, beta, gamma, variance_epsilon=1e-3) 57 | 58 | def testing(): 59 | return tf.nn.batch_normalization(inputs, pop_mean, pop_var, beta, gamma, variance_epsilon=1e-3) 60 | 61 | return tf.cond(is_training, train, testing) 62 | 63 | 64 | 65 | 66 | 67 | 68 | class Model(object): 69 | def __init__(self,sess,optimizer, batch_size=32, tot_res_blocks=2): 70 | self.sess = sess 71 | self.optimizer = optimizer 72 | self.batch_size = batch_size 73 | self.THETA = {} 74 | self.tot_res_blocks = tot_res_blocks 75 | self.init_model() 76 | self.saver = tf.train.Saver(self.THETA, max_to_keep=100) 77 | 78 | 79 | def get_training(self): 80 | return self.training 81 | 82 | def set_training(self, training): 83 | self.training = tf.assign(self.training, tf.constant(training, shape=[])) 84 | 85 | def init_model(self): 86 | def init_bn_params(block_name,ind,size=64): 87 | self.THETA['{}beta{}'.format(block_name,ind)] = tf.Variable( 88 | tf.truncated_normal([size], mean=0.0, stddev=5e-2), name='{}beta{}'.format(block_name,ind), trainable=True) 89 | self.THETA['{}gamma{}'.format(block_name,ind)]= tf.Variable( 90 | tf.truncated_normal([size], mean=1.0, stddev=5e-2), name='{}gamma{}'.format(block_name,ind), trainable=True) 91 | self.THETA['{}pop_mean{}'.format(block_name,ind)] = tf.Variable(tf.zeros([size]), name = '{}pop_mean{}'.format(block_name,ind), trainable=False) 92 | self.THETA['{}pop_var{}'.format(block_name,ind)] = tf.Variable(tf.ones([size]) , name = '{}pop_var{}'.format(block_name,ind) , trainable=False) 93 | 94 | self.THETA['inpw1'] = _variable_with_weight_decay(shape=[3, 3, 3, 64], name='inw1') 95 | self.THETA['inpb1'] = _variable_with_weight_decay(shape=[ 64], name='inb1') 96 | 97 | for ind in range(1,self.tot_res_blocks+1): 98 | 99 | self.THETA['b{}w1'.format(ind)] = _variable_with_weight_decay(shape=[3, 3, 64, 64], name='b{}w1'.format(ind)) 100 | self.THETA['b{}b1'.format(ind)] = _variable_with_weight_decay(shape=[ 64] , name='b{}b1'.format(ind)) 101 | 102 | init_bn_params('b{}'.format(ind),1) 103 | 104 | self.THETA['b{}w2'.format(ind)] = _variable_with_weight_decay(shape=[3, 3, 64, 64], name='b{}w2'.format(ind)) 105 | self.THETA['b{}b2'.format(ind)] = _variable_with_weight_decay(shape=[ 64] , name='b{}b2'.format(ind)) 106 | 107 | init_bn_params('b{}'.format(ind),2) 108 | 109 | 110 | self.THETA['recon0w1'] = _variable_with_weight_decay(shape=[3, 3, 64, 64], name='recon0w1') 111 | self.THETA['recon0b1'] = _variable_with_weight_decay(shape=[64] , name='recon0b1') 112 | init_bn_params('recon0',1) 113 | 114 | self.THETA['recon1w1'] = _variable_with_weight_decay(shape=[3, 3, 64, 256], name='recon1w1') 115 | self.THETA['recon1b1'] = _variable_with_weight_decay(shape=[256] , name='recon1b1') 116 | 117 | self.THETA['recon2w1'] = _variable_with_weight_decay(shape=[3, 3, 256, 256], name='recon2w1') 118 | self.THETA['recon2b1'] = _variable_with_weight_decay(shape=[256] , name='recon2b1') 119 | 120 | self.THETA['recon3w1'] = _variable_with_weight_decay(shape=[3, 3, 256, 3], name='recon3w1') 121 | self.THETA['recon3b1'] = _variable_with_weight_decay(shape=[3] , name='recon3b1') 122 | 123 | 124 | 125 | 126 | def inference(self,inp, phase_train): 127 | 128 | def conv(blk,ind,inp): 129 | x = tf.nn.conv2d(inp, self.THETA['{}w{}'.format(blk,ind)], strides=[1, 1, 1, 1], padding='SAME',name="{}w{}".format(blk,ind)) 130 | x = tf.nn.bias_add(x, self.THETA['{}b{}'.format(blk,ind)], name="{}b{}".format(blk,ind)) 131 | return x 132 | 133 | def bn(blk,ind,inp): 134 | x = batch_norm(inp, 135 | phase_train, 136 | self.THETA['{}pop_mean{}'.format(blk,ind)], 137 | self.THETA['{}pop_var{}'.format(blk,ind)], 138 | self.THETA['{}beta{}'.format(blk,ind)], 139 | self.THETA['{}gamma{}'.format(blk,ind)], 140 | ) 141 | return x 142 | 143 | 144 | def block_conv(inp, name): 145 | with tf.variable_scope(name): 146 | x = conv(name,1,inp) 147 | x = bn(name,1,x) 148 | c1 = tf.nn.relu(x) 149 | 150 | x = conv(name,2,c1) 151 | x = bn(name,2,x) 152 | 153 | x = tf.add(inp,x) 154 | return x 155 | 156 | 157 | x = conv('inp',1,inp) 158 | o0 = tf.nn.relu(x) 159 | o = o0 160 | 161 | # all residual blocks 162 | for ind in range(1,self.tot_res_blocks+1): 163 | o = block_conv(o,"b{}".format(ind)) 164 | 165 | x = conv('recon0',1,o) 166 | x = bn('recon0',1,x) 167 | x = tf.add(x,o0) 168 | 169 | x = conv('recon1',1,x) 170 | x = tf.nn.relu(x) 171 | 172 | x = conv('recon2',1,x) 173 | x = tf.nn.relu(x) 174 | 175 | x = conv('recon3',1,x) 176 | 177 | out = x 178 | return out 179 | 180 | 181 | def save_model(self,file_name): 182 | self.saver.save(self.sess,file_name) 183 | 184 | def load_model(self,file_name): 185 | self.saver.restore(self.sess, file_name) 186 | 187 | 188 | def get_gradients(self,total_loss): 189 | grads = self.optimizer.compute_gradients(total_loss) 190 | 191 | return grads 192 | 193 | def get_backprop(self,total_loss,global_step): 194 | solver = self.optimizer.minimize(total_loss,global_step=global_step) 195 | return solver 196 | -------------------------------------------------------------------------------- /ARGAN/Test.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | parser = argparse.ArgumentParser(description='Expert denoisers training') 3 | parser.add_argument('--batch-size', type=int, default=32, metavar='N', required=True, help='input batch size for training (default: 32)') 4 | parser.add_argument('--stride', type=int, default=16, metavar='N', required=True, help='input batch size for training (default: 32)') 5 | #parser.add_argument('--noise-level', type=str, default=5.0, metavar='N', help='No of iterations') 6 | parser.add_argument('--model-snapshot-file', type=str, required=True, help='input batch size for training (default: 32)') 7 | parser.add_argument('--image-dir', type=str, required=True, help='input batch size for training (default: 32)') 8 | parser.add_argument('--test-data-list-file', type=str, metavar='N', required=True, help='test base folder is required') 9 | parser.add_argument('--save-images', type=str, metavar='N', required=True, help='save images') 10 | parser.add_argument('--quality', type=int, metavar='N', required=True, help='compression quality') 11 | args = parser.parse_args() 12 | 13 | 14 | import numpy as np 15 | import tensorflow as tf 16 | import os 17 | import numpy 18 | import glob 19 | from itertools import product 20 | from PIL import Image 21 | import skimage 22 | from skimage.util import view_as_windows 23 | import json 24 | import skimage.measure as measure 25 | import sys 26 | import importlib 27 | import cPickle as pickle 28 | from tqdm import tqdm 29 | #import matplotlib.pyplot as plt 30 | from Model import Model 31 | 32 | #config = tf.ConfigProto(gpu_options=tf.GPUOptions(per_process_gpu_memory_fraction=0.5)) 33 | config = tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True)) 34 | try: 35 | tf.Session(config=config) 36 | except: pass 37 | 38 | 39 | 40 | def make_dirs(dirs): 41 | for _dir in dirs: 42 | if not os.path.isdir(_dir): 43 | os.mkdir(_dir) 44 | 45 | def reconstruct_image_from_patches(Pp_p,s, In_p): 46 | Op = numpy.zeros(s) 47 | i=0 48 | for x in range(Op.shape[0]): 49 | for y in range(Op.shape[1]): 50 | Op[x,y,:] = Pp_p[i].reshape((s[3],s[4],3)) 51 | i = i+1 52 | 53 | Ip_p = numpy.zeros_like(In_p) 54 | Ip_w = numpy.zeros_like(In_p) 55 | 56 | for x in range(Op.shape[0]): 57 | for y in range(Op.shape[1]): 58 | Ip_p[x*stride[0]:(x)*stride[0]+64,y*stride[1]:(y)*stride[1]+64,:] += Op[x,y,0] 59 | Ip_w[x*stride[0]:(x)*stride[0]+64,y*stride[1]:(y)*stride[1]+64,:] += 1.0 60 | #psnr[x*stride:(x+1)*stride,y*stride:(y+1)*stride] = 2*sqrt(mean_squared_error(Op[x,y], Pc[x,y])) 61 | Ip_w[Ip_w==0.0] = 1.0 62 | Ip_p = Ip_p/Ip_w 63 | 64 | return Ip_p[32:-32,32:-32] # removing padding 65 | 66 | 67 | def denoise_image(In, stride): 68 | In_p = skimage.util.pad(In,32, 'reflect') 69 | In_p = In_p[:,:,32:(2*32+3)-32] 70 | Pn_p = view_as_windows(In_p[:,:,:],window_shape=(64,64,3),step=stride) 71 | 72 | 73 | ss = Pn_p.shape 74 | Pn_p = Pn_p.reshape(ss[0]*ss[1],ss[3],ss[4],3)/255.0 75 | Pp_p = numpy.zeros_like(Pn_p) 76 | 77 | for ind in range(0,ss[0]*ss[1],batch_size): 78 | #print ind,ind+batch_size 79 | mini_batch = Pn_p[ind:ind+batch_size] 80 | # print type(mini_batch) 81 | # print mini_batch.shape 82 | # print "{} = {}".format("max of input",np.max(mini_batch)) 83 | predk = E.sess.run([y_],feed_dict={x_:mini_batch, phase_train:False}) 84 | # print type(predk[0]) 85 | # print predk[0].shape 86 | # print "{} ={}".format("max of output",np.max(predk[0])) 87 | # print len(predk) 88 | # sys.exit(0) 89 | Pp_p[ind:ind+batch_size,:,:,:] = predk[0] 90 | Pp_p = Pp_p*255.0 91 | Ip = reconstruct_image_from_patches(Pp_p,ss, In_p) 92 | return post_process(Ip) 93 | 94 | def post_process(I): 95 | # clip values outside range 96 | I[I>255.0]=255.0 97 | I[I< 0.0]= 0.0 98 | return I 99 | 100 | def lp_loss(y_true, y_pred): 101 | loss = tf.reduce_mean(tf.losses.mean_squared_error(predictions=y_pred, labels=y_true)) 102 | return loss 103 | 104 | def read_image(file_name): 105 | return Image.open(file_name) 106 | 107 | def jpeg_compress(I, quality,file_name): 108 | img_path = file_name[:-4] + '_20.jpg' 109 | I.save(".temp","JPEG",quality = quality) 110 | return numpy.ascontiguousarray( Image.open(".temp").convert('RGB'), dtype = 'float32') 111 | 112 | def img_generator(quality, test_data_file): 113 | f = open(test_data_file,'r') 114 | for l in f.readlines(): 115 | l = l.strip() 116 | I = read_image(l) 117 | #print l 118 | In = jpeg_compress(I, quality,l) 119 | I = numpy.asarray(I, dtype='float32') 120 | yield I, In, l, l 121 | 122 | generator = img_generator(args.quality, args.test_data_list_file) 123 | batch_size = args.batch_size 124 | stride = (args.stride,args.stride,args.stride) 125 | 126 | LG ='E2' 127 | 128 | log_dir = os.path.join(LG, 'LOG') 129 | clean_images_save_folder = os.path.join(LG, 'save', 'clean_images') 130 | noisy_images_save_folder = os.path.join(LG, 'save', 'noise_std_{}'.format(args.quality)) 131 | reconstructed_images_save_folder = os.path.join(LG, 'save', 'recon_std_{}'.format(args.quality)) 132 | 133 | # make_dirs([args.exp_code, log_file, os.path.join(args.exp_code, 'save'), 134 | # clean_images_save_folder, noisy_images_save_folder]) 135 | 136 | make_dirs([LG, log_dir, os.path.join(LG, 'save'), clean_images_save_folder]) 137 | make_dirs([LG, log_dir, os.path.join(LG, 'save'), noisy_images_save_folder]) 138 | make_dirs([LG, log_dir, os.path.join(LG, 'save'), reconstructed_images_save_folder]) 139 | 140 | with tf.Session(config=config) as sess: 141 | optimizer = tf.train.AdamOptimizer(1e-3) 142 | E = Model(sess,optimizer=optimizer, batch_size = None, tot_res_blocks=5) 143 | x_ = tf.placeholder(tf.float32,[None,64,64,3],name="noisy_images") 144 | phase_train = tf.placeholder(tf.bool, [], name='phase_train') 145 | y_ = E.inference(x_, phase_train) 146 | 147 | init = tf.global_variables_initializer() 148 | E.sess.run(init) 149 | 150 | E.load_model(args.model_snapshot_file) 151 | 152 | for Ic, In, clean_image_name, noisy_image_name in tqdm(generator): 153 | In = post_process(In) 154 | Ir = denoise_image(In, stride) 155 | if args.save_images == 'True': 156 | # this will currently give an error 157 | Image.fromarray(Ir.astype(numpy.uint8)).save( 158 | os.path.join(reconstructed_images_save_folder, 159 | os.path.basename(clean_image_name)) 160 | ) 161 | Image.fromarray(Ic.astype(numpy.uint8)).save( 162 | os.path.join(clean_images_save_folder, 163 | os.path.basename(clean_image_name)) 164 | ) 165 | Image.fromarray(In.astype(numpy.uint8)).save( 166 | os.path.join(noisy_images_save_folder, 167 | os.path.basename(clean_image_name)) 168 | ) 169 | 170 | 171 | print "Done!" 172 | import argparse 173 | parser = argparse.ArgumentParser(description='Expert denoisers training') 174 | -------------------------------------------------------------------------------- /ARGAN/Train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import numpy 3 | import numpy as np 4 | numpy.random.seed(1) 5 | import tensorflow as tf 6 | tf.set_random_seed(1) 7 | import sys 8 | from random import shuffle 9 | import os 10 | import glob 11 | from PIL import Image 12 | import matplotlib.pyplot as plt 13 | import matplotlib 14 | #matplotlib.use('Agg') 15 | from Model import Model 16 | from Discriminator import Discriminator 17 | from Input import BatchImageInput 18 | import Losses 19 | import time 20 | 21 | #global config 22 | config = tf.ConfigProto(gpu_options=tf.GPUOptions(per_process_gpu_memory_fraction=0.5)) 23 | #config = tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True)) 24 | try: 25 | tf.Session(config=config) 26 | except: pass 27 | 28 | 29 | parser = argparse.ArgumentParser(description='Expert denoisers training') 30 | 31 | parser.add_argument('--lr1', type=float, metavar='N', default=1e-5, required=True, help='learning rate is required') 32 | parser.add_argument('--lr2', type=float, metavar='N', default=1e-5, required=True, help='learning rate is required') 33 | parser.add_argument('--batch-size', type=int, default=32, metavar='N', required=True, help='input batch size for training (default: 32)') 34 | parser.add_argument('--tot-grad-updates', type=int, default=10000, metavar='N', help='No of iterations') 35 | parser.add_argument('--start-step', type=int, default=1, metavar='N', help='training start step') 36 | parser.add_argument('--log-file', type=str, default=None, metavar='N', help='No of iterations') 37 | parser.add_argument('--train-dir', type=str, default=None, metavar='N', help='No of iterations') 38 | parser.add_argument('--valdn-dir', type=str, default=None, metavar='N', help='No of iterations') 39 | parser.add_argument('--log-interval', type=int, default=10000, metavar='N', required=True, help='frequency of logging') 40 | parser.add_argument('--snapshot-interval', type=int, default=10000, metavar='N', required=True, help='frequency of snapshotting') 41 | parser.add_argument('--exp-code', type=str, default=10000, metavar='N', required=True, help='experiment code') 42 | 43 | args = parser.parse_args() 44 | 45 | LGXY = "{}".format(args.exp_code) 46 | 47 | def make_dirs(dirs): 48 | for _dir in dirs: 49 | if not os.path.isdir(_dir): 50 | os.mkdir(_dir) 51 | 52 | make_dirs([LGXY, os.path.join(LGXY, 'LOG'), os.path.join(LGXY, 'checkpoints')]) 53 | 54 | 55 | # redirect error messages to std out or file 56 | import sys 57 | 58 | ave_out = sys.stdout 59 | if args.log_file is not None: 60 | f = open(os.path.join(LGXY,args.log_file),'a') 61 | sys.stdout = f 62 | 63 | 64 | with tf.Session(config=config) as sess: 65 | step_ = tf.Variable(0, trainable=False) 66 | #lr_ = tf.train.exponential_decay(args.lr, step_, 100, 0.999,staircase=True) 67 | #tf.summary.scalar("lr", lr_) 68 | tf.summary.scalar("step", step_) 69 | optimizer1 = tf.train.AdamOptimizer(args.lr1) 70 | optimizer2 = tf.train.GradientDescentOptimizer(args.lr2) 71 | 72 | # model 73 | E = Model(sess,optimizer=optimizer1, batch_size = args.batch_size, tot_res_blocks=5) 74 | D = Discriminator(sess,optimizer=optimizer2, batch_size = args.batch_size) 75 | 76 | # --beg-- train data 77 | with open('../Resources/imagenet_val_paths.txt', 'r') as f: 78 | file_paths = [os.path.join(args.train_dir,fname) for fname in f.read().splitlines()] 79 | Train_Data = BatchImageInput(file_paths,[20],batch_size = args.batch_size) 80 | 81 | # --beg-- validation data 82 | with open('../Resources/bsds500_test_paths.txt', 'r') as f: 83 | file_paths = [os.path.join(args.valdn_dir,fname) for fname in f.read().splitlines()] 84 | Valdn_Data = BatchImageInput(file_paths,[20],batch_size = args.batch_size) 85 | # --end-- validation data 86 | 87 | # computation 88 | x_ = tf.placeholder(tf.float32,[args.batch_size,64,64,3],name="noisy_images") 89 | y_ = tf.placeholder(tf.float32,[args.batch_size,64,64,3],name="clean_images") 90 | a_ = tf.placeholder(tf.float32,[], name="alpha") 91 | phase_train = tf.placeholder(tf.bool, [], name='is_training') 92 | 93 | # z_ = E.inference(x_, phase_train) 94 | t_ = D.inference(y_, phase_train) # clean sample logits 95 | f_ = D.inference(x_, phase_train) # reconstructed sample logits 96 | 97 | prob_t_ = tf.reduce_mean(tf.nn.sigmoid(t_)) 98 | prob_f_ = tf.reduce_mean(tf.nn.sigmoid(f_)) 99 | 100 | Loss_list = { 101 | 'lp':(Losses.lp_loss,0.9), 102 | 'pl':(Losses.perceptual_loss,0.1), 103 | # 'tv':Losses.tv_loss, 104 | } 105 | 106 | Loss_ops = {} 107 | Backprop_ops = {} 108 | LOSS = 0 109 | for loss_name in Loss_list: 110 | loss_name 111 | # compute the loss 112 | Loss_ops[loss_name] = Loss_list[loss_name][0](y_,x_) 113 | # weighted sum 114 | LOSS = LOSS+ (Loss_ops[loss_name]*Loss_list[loss_name][1]) 115 | # add o summary 116 | tf.summary.scalar(loss_name+'_loss', Loss_ops[loss_name]) 117 | 118 | Loss_ops['total'] = LOSS 119 | tf.summary.scalar('loss_recon', LOSS) 120 | 121 | def return_random (min_val , max_val): 122 | # return tf.random_uniform([1], minval = min_val , maxval = max_val , dtype = tf.float32) 123 | return tf.constant(1.0) 124 | 125 | pos_loss = -tf.reduce_mean((tf.constant(1.0)*tf.log(prob_t_))) #+ ((1.0 - return_random(0.8,1.2))*tf.log(1.0-prob_t_))) 126 | neg_loss = -tf.reduce_mean((1.0 - tf.constant(0.0))*tf.log(1.0-prob_f_)) 127 | 128 | disc_true_loss = pos_loss 129 | disc_false_loss = neg_loss 130 | disc_loss = pos_loss + neg_loss 131 | Loss_ops['pos_loss'] = pos_loss 132 | Loss_ops['neg_loss'] = neg_loss 133 | Loss_ops['disc_loss'] = disc_loss 134 | Loss_ops['disc_true_loss']= disc_true_loss 135 | Loss_ops['disc_false_loss'] = disc_false_loss 136 | 137 | 138 | BACK_PROP_OP_adv = E.get_backprop( 139 | genr_loss, 140 | step_ 141 | ) 142 | 143 | BACK_PROP_OP_dis = D.get_backprop( 144 | disc_loss, 145 | step_ 146 | ) 147 | BACK_PROP_OP_true_dis = D.get_backprop( 148 | disc_true_loss, 149 | step_ 150 | ) 151 | BACK_PROP_OP_false_dis = D.get_backprop( 152 | disc_false_loss, 153 | step_ 154 | ) 155 | 156 | # tf.summary.image("predk_image",z_) 157 | tf.summary.image("clean_image",y_) 158 | tf.summary.image("noisy_image",x_) 159 | 160 | 161 | # DATA FLOW 162 | 163 | X_train_,Y_train_ = Train_Data.get_minibatch_tensors() 164 | X_valdn_,Y_valdn_ = Valdn_Data.get_minibatch_tensors() 165 | t_coord = tf.train.Coordinator() 166 | t_threads = tf.train.start_queue_runners(coord=t_coord) 167 | 168 | '''If you want to save all the variables and visualize them in tensorboard then un comment this''' 169 | ## Add histograms for trainable variables. 170 | # for var in tf.trainable_variables(): 171 | # tf.summary.histogram(var.op.name, var) 172 | # for v_,g_ in GRADS_: 173 | # if v_ is not None: 174 | # tf.summary.histogram(v_.op.name,g_) 175 | 176 | summary_op = tf.summary.merge_all() 177 | 178 | writer = tf.summary.FileWriter(os.path.join(LGXY, 'LOG'), graph=E.sess.graph) 179 | init = tf.global_variables_initializer() 180 | # E.sess.run(init) 181 | D.sess.run(init) 182 | # E.load_model('../Exp000906/E2/checkpoints/E2.100000.ckpt') 183 | 184 | step = args.start_step 185 | stop_step = args.tot_grad_updates - args.start_step 186 | i=0 187 | j=0 188 | while step <= stop_step: 189 | X_train_batch, Y_train_batch = E.sess.run([X_train_,Y_train_]) 190 | 191 | summary = E.sess.run( 192 | [summary_op], 193 | feed_dict={x_:X_train_batch,y_:Y_train_batch, phase_train:True} 194 | ) 195 | 196 | prob_t, prob_f, pos_l,neg_l,total_disc_loss,tl,fl = D.sess.run( 197 | [prob_t_,prob_f_,pos_loss,neg_loss,disc_loss,t_,f_], 198 | feed_dict={ 199 | x_:X_train_batch, 200 | y_:Y_train_batch, 201 | phase_train:True 202 | } 203 | ) 204 | if( True ): 205 | 206 | print step 207 | print "{} = {}".format("Total disc loss",total_disc_loss) 208 | print "{} = {}".format("Positive loss",pos_l) 209 | print "{} = {}".format("Negative loss",neg_l) 210 | print "{} = {} {} = {}".format("Prob_t",prob_t,"Prob_f",prob_f) 211 | print "{} = {}".format("True logit",np.mean(tl)) 212 | print "{} = {}".format("False logit",np.mean(fl)) 213 | print "{} = {}".format("Mean of x train batch",np.mean(X_train_batch)) 214 | print "{} = {}".format("Mean of y train batch",np.mean(Y_train_batch)) 215 | print "{} = {}".format("Mean squarred error between actual pair",((Y_train_batch - X_train_batch)**2).mean(axis =None)) 216 | print "************************" 217 | 218 | if(False): 219 | sys.exit(0) 220 | 221 | j=j+1 222 | _ = D.sess.run( 223 | [BACK_PROP_OP_dis], 224 | feed_dict={ 225 | x_:X_train_batch, 226 | y_:Y_train_batch, 227 | phase_train:True 228 | } 229 | ) 230 | 231 | 232 | writer.add_summary(summary[0], step) 233 | 234 | if step % args.log_interval == 0 or step == stop_step: 235 | 236 | X_valdn_batch, Y_valdn_batch = D.sess.run([X_valdn_,Y_valdn_]) 237 | 238 | 239 | # getting all the losses 240 | prob_tv, prob_fv, pos_lv,neg_lv,total_disc_lossv = D.sess.run( 241 | [prob_t_,prob_f_,pos_loss,neg_loss,disc_loss], 242 | feed_dict={ 243 | x_:X_train_batch, 244 | y_:Y_train_batch, 245 | phase_train:False 246 | } 247 | ) 248 | 249 | LOSSES = {} 250 | for loss_op in Loss_ops: 251 | op = Loss_ops[loss_op] 252 | out= D.sess.run( 253 | [op], 254 | feed_dict={ 255 | x_:X_train_batch, 256 | y_:Y_train_batch, 257 | phase_train:True 258 | } 259 | ) 260 | LOSSES[loss_op+'_train'] = out[0] 261 | out= D.sess.run( 262 | [op], 263 | feed_dict={ 264 | x_:X_valdn_batch, 265 | y_:Y_valdn_batch, 266 | phase_train:True 267 | } 268 | ) 269 | LOSSES[loss_op+'_valdn'] = out[0] 270 | 271 | # Printing differrent losses 272 | loss_str_list = [] 273 | for key in LOSSES: 274 | if key=='disc_loss_train' or key=='disc_true_loss_train' or key=='disc_false_loss_train': 275 | 276 | loss_str_list.append(" {} = {:e}".format(key,LOSSES[key])) 277 | 278 | prob_t, prob_f = D.sess.run([prob_t_, prob_f_], feed_dict = {x_: X_valdn_batch, y_:Y_valdn_batch, phase_train:False}) 279 | 280 | 281 | sys.stdout.flush() 282 | if step % args.snapshot_interval == 0 or step == stop_step: 283 | E.save_model(os.path.join(LGXY, 'checkpoints', "{}.{:05d}.ckpt".format(LGXY,step))) 284 | 285 | step += 1 286 | 287 | print("Done") 288 | 289 | 290 | 291 | try: 292 | f.close() 293 | except: 294 | pass 295 | sys.stdout = _save_out 296 | -------------------------------------------------------------------------------- /ARGAN/avg_psnr.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | parser = argparse.ArgumentParser(description='Expert denoisers training') 3 | parser.add_argument('--image-filename', type=str, default='', metavar='N', required=True, help='File containing PSNR values of LIVE1 dataset') 4 | args = parser.parse_args() 5 | 6 | import sys 7 | 8 | 9 | noisy_psnr = [] 10 | recon_psnr = [] 11 | diff_psnr = [] 12 | i=0 13 | with open(args.image_filename,'r') as f: 14 | for l in f.readlines(): 15 | i=i+1 16 | l = l.strip() 17 | l = l.split() 18 | noisy_psnr.append(float(l[1])) 19 | recon_psnr.append(float(l[2])) 20 | diff_psnr.append( float(l[2]) - float(l[1])) 21 | 22 | 23 | avg_noisy_psnr = sum(noisy_psnr)/len(noisy_psnr) 24 | avg_recon_psnr = sum(recon_psnr)/len(recon_psnr) 25 | avg_diff_psnr = sum(diff_psnr)/len(diff_psnr) 26 | print "Average noisy PSNR = "+str(avg_noisy_psnr) 27 | print "Average recon PSNR = "+str(avg_recon_psnr) 28 | print "Average diffr PSNR = "+str(avg_diff_psnr) 29 | 30 | -------------------------------------------------------------------------------- /ARGAN/results.txt: -------------------------------------------------------------------------------- 1 | bikes.bmp 26.3445868553 27.1688501387 2 | building2.bmp 23.9975501141 24.5423449921 3 | buildings.bmp 26.0241046434 26.5020789427 4 | caps.bmp 31.4807162018 31.288299572 5 | carnivaldolls.bmp 27.6292937529 28.039722361 6 | cemetry.bmp 26.1793601901 26.4032215497 7 | churchandcapitol.bmp 27.2096019982 27.7525397636 8 | coinsinfountain.bmp 27.4504342959 27.8406877285 9 | dancers.bmp 25.0964979096 25.5820944186 10 | flowersonih35.bmp 24.0114022058 24.4477142501 11 | house.bmp 29.0179141483 29.2473426855 12 | lighthouse2.bmp 28.5988659411 28.8435702711 13 | lighthouse3.bmp 29.0727291461 29.4213411952 14 | manfishing.bmp 28.2856791518 28.5661033944 15 | monarch.bmp 29.7104554821 30.3399322702 16 | ocean.bmp 30.4518323677 30.5628900881 17 | paintedhouse.bmp 26.9858234226 27.2606274756 18 | parrots.bmp 32.0304203593 31.4072441612 19 | plane.bmp 30.6925221939 31.08430449 20 | rapids.bmp 27.5289651742 28.1551730198 21 | sailing1.bmp 28.1411129006 28.2693431985 22 | sailing2.bmp 31.0939653815 31.1845742387 23 | sailing3.bmp 30.6184272341 30.684935991 24 | sailing4.bmp 28.7732202545 28.99005838 25 | statue.bmp 30.3289649006 30.6846973551 26 | stream.bmp 24.6731855956 25.108156921 27 | studentsculpture.bmp 25.1515873305 25.4700957536 28 | woman.bmp 27.0525324249 27.5226662447 29 | womanhat.bmp 30.1287130633 30.0531796482 30 | -------------------------------------------------------------------------------- /ARGAN/run_debug.sh: -------------------------------------------------------------------------------- 1 | CUDA_VISIBLE_DEVICES=2 \ 2 | python Train.py \ 3 | --lr1=1e-5 \ 4 | --lr2=1e-4 \ 5 | --batch-size=32 \ 6 | --log-interval=10 \ 7 | --snapshot-interval=10000 \ 8 | --tot-grad-updates=110000 \ 9 | --train-dir='../Resources/Datasets/IMAGENET/val' \ 10 | --valdn-dir='../Resources/Datasets/BSDS/BSR/BSDS500/data/images/test' \ 11 | --exp-code=E2 \ 12 | #--log-file=log.txt \ 13 | -------------------------------------------------------------------------------- /ARGAN/run_eval.sh: -------------------------------------------------------------------------------- 1 | python Eval.py \ 2 | --image-folder E2/save \ 3 | --image-filename ../Resources/live1_list.txt \ 4 | --quality 20 5 | 6 | -------------------------------------------------------------------------------- /ARGAN/run_test.sh: -------------------------------------------------------------------------------- 1 | CUDA_VISIBLE_DEVICES=1 \ 2 | python Test.py \ 3 | --batch-size 32 \ 4 | --stride 16 \ 5 | --quality 20 \ 6 | --model-snapshot-file E2a/checkpoints/E2a.40000.ckpt \ 7 | --image-dir='../Resources/Datasets/LIVE1/' \ 8 | --test-data-list-file ../Resources/live1_list.txt \ 9 | --save-image True \ 10 | -------------------------------------------------------------------------------- /ARGAN/run_train.sh: -------------------------------------------------------------------------------- 1 | CUDA_VISIBLE_DEVICES=2 \ 2 | python Train.py \ 3 | --lr1=1e-4 \ 4 | --lr2=1e-4 \ 5 | --batch-size=32 \ 6 | --log-interval=10 \ 7 | --snapshot-interval=10000 \ 8 | --tot-grad-updates=110000 \ 9 | --train-dir='../Resources/Datasets/IMAGENET/val' \ 10 | --valdn-dir='../Resources/Datasets/BSDS/BSR/BSDS500/data/images/test' \ 11 | --exp-code=E2a \ 12 | --log-file=log.txt \ 13 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Compression artifacts reduction 2 | This repository contains the codes for different techniques for reducing JPEG compression artifacts in images 3 | 4 | ## ARGAN 5 | * ARGAN folder contains the code which employs adversarial training for reducing compression artifacts. The generator network consists of a residual network with 5 residual blocks. 6 | * The discriminator network consists of 8 convolutional layers and 2 fully connected layers (no pooling layers). 7 | * Some of the hacks articulated in [Gan hacks](https://github.com/soumith/ganhacks) have been used. 8 | * The overall architecture has been inspired from [SRGAN paper](https://arxiv.org/abs/1609.04802) 9 | 10 | ## Baseline 11 | * baseline folder contains the code for several baseline experiments without using any adversarial training. 12 | * We use a generator similar to the ARGAN folder (5 residual blocks), which aims to reduce the JPEG compression artifacts conditioned on the input compressed image. 13 | * Variety of loss functions can be used; we report the model's performances on a) only L2 loss b) L2 + Total variation (TV) c) L2 + perceptual loss. 14 | * Perceptual loss has inspired from this [paper](https://arxiv.org/abs/1609.04802), for we’ve used a pre-trained VGG 15 | network and compared the activation of 4 th conv layer in the 5th block for the ground truth and the 16 | reconstructed image. 17 | 18 | ## Dependencies 19 | ``` 20 | Python 2.7, Tensorflow 1.0.1 21 | ``` 22 | ## Performance 23 | 24 | The following are reported after 100000 iterations (which takes about 2 days when run on Nvidia GeForce GTX Titan X) and tested as average on 29 images of [LIVE1 dataset](http://live.ece.utexas.edu/research/quality/subjective.htm) 25 | * **ARGAN:** 26 | Noisy PNSR = 28.06 dB | Reconstructed PNSR = 29.02 dB | Improvement = 0.96 dB 27 | * **Baseline 1 (only L2 loss):** 28 | Noisy PSNR = 28.06 dB | Reconstructed PSNR = 28.82 dB | Improvement = 0.76 dB 29 | * **Baseline 2 (L2 + TV):** 30 | Noisy PSNR = 28.06 dB | Reconstructed PSNR = 23.54 dB | Improvement = -4.52 dB 31 | * **Baseline 3 (L2 + perceptual loss):** 32 | Noisy PSNR = 28.06 dB | Reconstructed PSNR = 28.77 dB | Improvement = 0.71 dB 33 | * **Baseline 4 (0.9xL2 + 0.1xperceptual loss):** 34 | Noisy PSNR = 28.06 dB | Reconstructed PSNR = 28.83 dB | Improvement = 0.77 dB 35 | 36 | ![](result.png) 37 | 38 | 39 | -------------------------------------------------------------------------------- /baseline/Input.py: -------------------------------------------------------------------------------- 1 | import os 2 | import glob 3 | import tensorflow as tf 4 | 5 | import random 6 | from tqdm import tqdm 7 | from itertools import product 8 | 9 | class BatchImageInput(): 10 | def __init__(self, file_paths, noise_levels, crop_size = [64, 64, 3], batch_size = 32): 11 | self.batch_size = batch_size 12 | self.noise_levels = noise_levels 13 | self.crop_size = crop_size 14 | self.file_paths = file_paths 15 | 16 | def read_one_file(self, filename_queue): 17 | print filename_queue 18 | image_reader = tf.WholeFileReader() 19 | _,I = image_reader.read(filename_queue) 20 | Id = tf.image.decode_image(I, channels=3) 21 | Id = tf.to_float(Id) 22 | Ic = tf.random_crop(Id, size = self.crop_size) 23 | 24 | tt = tf.cast(Ic,tf.uint8) 25 | tt = tf.image.encode_jpeg(tt,quality=20) # have to remove this hardcoding later 26 | In = tf.image.decode_jpeg(tt, channels=3) 27 | In = tf.reshape(In, self.crop_size) 28 | In = tf.to_float(In) 29 | 30 | return Ic,In 31 | 32 | def get_minibatch_tensors(self, num_epochs=None): 33 | input_queue = tf.train.string_input_producer(self.file_paths, num_epochs=num_epochs, shuffle=True) 34 | noises_ = tf.constant(self.noise_levels) 35 | shuff_noises_ = tf.random_shuffle(noises_) 36 | 37 | Ic, In = self.read_one_file(input_queue) 38 | # adding noise to a clean image 39 | # In = Ic + tf.random_normal(Ic.get_shape(), mean=0.0, stddev=shuff_noises_[0], dtype=tf.float32) 40 | 41 | Ic_batch, In_batch = tf.train.batch([Ic, In], batch_size=self.batch_size, capacity = 100) 42 | # (X,Y) 43 | return In_batch/255.0, Ic_batch/255.0 44 | 45 | if __name__ == '__main__': 46 | with open('imagenet_val_paths.txt', 'r') as f: 47 | file_paths = f.read().splitlines() 48 | 49 | EI = BatchImageInput(file_paths,[10.0,15.0],batch_size = 32) 50 | 51 | x_,y_ = EI.get_minibatch_tensors() 52 | 53 | config = tf.ConfigProto() 54 | config.gpu_options.allow_growth = True 55 | print "haha" 56 | with tf.Session(config=config) as sess: 57 | tf.global_variables_initializer() 58 | coord = tf.train.Coordinator() 59 | threads = tf.train.start_queue_runners(sess=sess, coord=coord) 60 | x,y = sess.run([x_,y_]) 61 | print x[0][:10,:10,0],y[0][:10,:10,0] 62 | print x[0].shape 63 | coord.request_stop() 64 | coord.join(threads) 65 | 66 | -------------------------------------------------------------------------------- /baseline/Losses.py: -------------------------------------------------------------------------------- 1 | from tensorflow.python.ops import math_ops 2 | import tensorflow as tf 3 | import numpy 4 | 5 | def lp_loss(y_true, y_pred): 6 | diff_sq = tf.squared_difference(y_true,y_pred) 7 | loss = tf.reduce_mean(diff_sq, axis=None) 8 | return loss 9 | 10 | 11 | 12 | def tv_loss(y_true,y_pred): 13 | def total_variation(images): 14 | pixel_dif1 = images[:, 1:, :, :] - images[:, :-1, :, :] 15 | pixel_dif2 = images[:, :, 1:, :] - images[:, :, :-1, :] 16 | sum_axis = [1, 2, 3] 17 | tot_var = math_ops.reduce_sum(math_ops.abs(pixel_dif1), axis=sum_axis) + \ 18 | math_ops.reduce_sum(math_ops.abs(pixel_dif2), axis=sum_axis) 19 | return tot_var 20 | 21 | loss = tf.reduce_sum(total_variation(y_pred)) 22 | return loss 23 | 24 | 25 | 26 | def perceptual_loss(y_true,y_pred, layer_list=['conv5_4'], alphas=[1.0]): 27 | def conv_layer(bottom, weight, bias=None, s=1, padding='SAME', relu=True, group=1): 28 | if group==1: 29 | conv = tf.nn.conv2d(bottom, weight, [1, s, s, 1], padding=padding) 30 | else: 31 | input_split = tf.split(bottom, group, 3) 32 | weight_split = tf.split(weight, group, 3) 33 | conv_1 = tf.nn.conv2d(input_split[0], weight_split[0], [1, s, s, 1], padding=padding) 34 | conv_2 = tf.nn.conv2d(input_split[1], weight_split[1], [1, s, s, 1], padding=padding) 35 | conv = tf.concat([conv_1, conv_2], 3) 36 | if bias is None: 37 | if relu: 38 | return tf.nn.relu(conv) 39 | else: 40 | return conv 41 | else: 42 | bias = tf.nn.bias_add(conv, bias) 43 | if relu: 44 | return tf.nn.relu(bias) 45 | else: 46 | return bias 47 | def max_pool(bottom, k=3, s=1, padding='SAME'): 48 | return tf.nn.max_pool(bottom, ksize=[1, k, k, 1], strides=[1, s, s, 1], padding=padding) 49 | def vgg_19(image, weights, biases, keep_prob=1.0): 50 | activations = {} 51 | shapes = image.get_shape().as_list() 52 | if shapes[3] == 1: 53 | rgb = tf.concat(axis = 3, values=[image,image,image]) 54 | else: 55 | rgb = image 56 | 57 | rgb_scaled = rgb*255.0 58 | red, green, blue = tf.split(axis=3, num_or_size_splits=3, value=rgb_scaled) 59 | VGG_MEAN = [103.939, 116.779, 123.68] 60 | bgr = tf.concat(axis=3, values=[ 61 | blue - VGG_MEAN[0], 62 | green - VGG_MEAN[1], 63 | red - VGG_MEAN[2], 64 | ]) 65 | 66 | image = bgr 67 | 68 | 69 | with tf.name_scope("conv1"): 70 | activations['conv1_1'] = conv_layer(image, weights['conv1_1'], biases['conv1_1']) 71 | activations['conv1_2'] = conv_layer(activations['conv1_1'], weights['conv1_2'], biases['conv1_2']) 72 | activations['pool1'] = max_pool(activations['conv1_2'], k=2, s=2) 73 | 74 | with tf.name_scope("conv2"): 75 | activations['conv2_1'] = conv_layer(activations['pool1'], weights['conv2_1'], biases['conv2_1']) 76 | activations['conv2_2'] = conv_layer(activations['conv2_1'], weights['conv2_2'], biases['conv2_2']) 77 | activations['pool2'] = max_pool(activations['conv2_2'], k=2, s=2) 78 | 79 | with tf.name_scope("conv3"): 80 | activations['conv3_1'] = conv_layer(activations['pool2'], weights['conv3_1'], biases['conv3_1']) 81 | activations['conv3_2'] = conv_layer(activations['conv3_1'], weights['conv3_2'], biases['conv3_2']) 82 | activations['conv3_3'] = conv_layer(activations['conv3_2'], weights['conv3_3'], biases['conv3_3']) 83 | activations['conv3_4'] = conv_layer(activations['conv3_3'], weights['conv3_4'], biases['conv3_3']) 84 | activations['pool3'] = max_pool(activations['conv3_4'], k=2, s=2) 85 | 86 | with tf.name_scope("conv4"): 87 | activations['conv4_1'] = conv_layer(activations['pool3'], weights['conv4_1'], biases['conv4_1']) 88 | activations['conv4_2'] = conv_layer(activations['conv4_1'], weights['conv4_2'], biases['conv4_2']) 89 | activations['conv4_3'] = conv_layer(activations['conv4_2'], weights['conv4_3'], biases['conv4_3']) 90 | activations['conv4_4'] = conv_layer(activations['conv4_3'], weights['conv4_4'], biases['conv4_4']) 91 | activations['pool4'] = max_pool(activations['conv4_4'], k=2, s=2) 92 | 93 | with tf.name_scope("conv5"): 94 | activations['conv5_1'] = conv_layer(activations['pool4'], weights['conv5_1'], biases['conv5_1']) 95 | activations['conv5_2'] = conv_layer(activations['conv5_1'], weights['conv5_2'], biases['conv5_2']) 96 | activations['conv5_3'] = conv_layer(activations['conv5_2'], weights['conv5_3'], biases['conv5_3']) 97 | activations['conv5_4'] = conv_layer(activations['conv5_3'], weights['conv5_4'], biases['conv5_4']) 98 | activations['pool5'] = max_pool(activations['conv5_4'], k=2, s=2) 99 | return activations 100 | 101 | def get_weights(): 102 | net = numpy.load('../Resources/VGG/vgg19.npy').item() 103 | weights = {} 104 | biases = {} 105 | for name in net.keys(): 106 | weights[name] = tf.Variable(tf.constant(net[name][0]), dtype='float32' ,name=name+"_weight", trainable=False) 107 | biases[name] = tf.Variable(tf.constant(net[name][1]), dtype='float32' ,name=name+"_bias", trainable=False) 108 | return weights,biases 109 | 110 | weights,biases = get_weights() 111 | 112 | y_true_activations = vgg_19(y_true,weights,biases) 113 | y_pred_activations = vgg_19(y_pred,weights,biases) 114 | 115 | for layer_name in layer_list: 116 | assert layer_name in y_true_activations.keys() 117 | 118 | loss = 0 119 | for layer_name,alpha in zip(layer_list,alphas): 120 | y = y_true_activations[layer_name]/255.0 121 | x = y_pred_activations[layer_name]/255.0 122 | loss += alpha* lp_loss(y,x) 123 | print loss 124 | return loss 125 | -------------------------------------------------------------------------------- /baseline/Losses2.py: -------------------------------------------------------------------------------- 1 | from tensorflow.python.ops import math_ops 2 | import tensorflow as tf 3 | import numpy 4 | 5 | def lp_loss(y_true, y_pred): 6 | loss = tf.reduce_mean(tf.losses.mean_squared_error(predictions=y_pred, labels=y_true)) 7 | return loss 8 | 9 | 10 | 11 | def tv_loss(y_true,y_pred): 12 | def total_variation(images): 13 | pixel_dif1 = images[:, 1:, :, :] - images[:, :-1, :, :] 14 | pixel_dif2 = images[:, :, 1:, :] - images[:, :, :-1, :] 15 | sum_axis = [1, 2, 3] 16 | tot_var = math_ops.reduce_sum(math_ops.abs(pixel_dif1), axis=sum_axis) + \ 17 | math_ops.reduce_sum(math_ops.abs(pixel_dif2), axis=sum_axis) 18 | return tot_var 19 | 20 | loss = tf.reduce_sum(total_variation(y_pred)) 21 | return loss 22 | 23 | 24 | 25 | def perceptual_loss(y_true,y_pred, layer_list=['pool5'], alphas=[1.0]): 26 | def conv_layer(bottom, weight, bias=None, s=1, padding='SAME', relu=True, group=1): 27 | if group==1: 28 | conv = tf.nn.conv2d(bottom, weight, [1, s, s, 1], padding=padding) 29 | else: 30 | input_split = tf.split(bottom, group, 3) 31 | weight_split = tf.split(weight, group, 3) 32 | conv_1 = tf.nn.conv2d(input_split[0], weight_split[0], [1, s, s, 1], padding=padding) 33 | conv_2 = tf.nn.conv2d(input_split[1], weight_split[1], [1, s, s, 1], padding=padding) 34 | conv = tf.concat([conv_1, conv_2], 3) 35 | if bias is None: 36 | if relu: 37 | return tf.nn.relu(conv) 38 | else: 39 | return conv 40 | else: 41 | bias = tf.nn.bias_add(conv, bias) 42 | if relu: 43 | return tf.nn.relu(bias) 44 | else: 45 | return bias 46 | def max_pool(bottom, k=3, s=1, padding='SAME'): 47 | return tf.nn.max_pool(bottom, ksize=[1, k, k, 1], strides=[1, s, s, 1], padding=padding) 48 | def vgg_16(image, weights, biases, keep_prob=1.0): 49 | activations = {} 50 | shapes = image.get_shape().as_list() 51 | if shapes[3] == 1: 52 | rgb = tf.concat(axis = 3, values=[image,image,image]) 53 | else: 54 | rgb = image 55 | 56 | rgb_scaled = rgb*255.0 57 | red, green, blue = tf.split(axis=3, num_or_size_splits=3, value=rgb_scaled) 58 | VGG_MEAN = [103.939, 116.779, 123.68] 59 | bgr = tf.concat(axis=3, values=[ 60 | blue - VGG_MEAN[0], 61 | green - VGG_MEAN[1], 62 | red - VGG_MEAN[2], 63 | ]) 64 | 65 | image = bgr 66 | 67 | 68 | with tf.name_scope("conv1"): 69 | activations['conv1_1'] = conv_layer(image, weights['conv1_1'], biases['conv1_1']) 70 | activations['conv1_2'] = conv_layer(activations['conv1_1'], weights['conv1_2'], biases['conv1_2']) 71 | activations['pool1'] = max_pool(activations['conv1_2'], k=2, s=2) 72 | 73 | with tf.name_scope("conv2"): 74 | activations['conv2_1'] = conv_layer(activations['pool1'], weights['conv2_1'], biases['conv2_1']) 75 | activations['conv2_2'] = conv_layer(activations['conv2_1'], weights['conv2_2'], biases['conv2_2']) 76 | activations['pool2'] = max_pool(activations['conv2_2'], k=2, s=2) 77 | 78 | with tf.name_scope("conv3"): 79 | activations['conv3_1'] = conv_layer(activations['pool2'], weights['conv3_1'], biases['conv3_1']) 80 | activations['conv3_2'] = conv_layer(activations['conv3_1'], weights['conv3_2'], biases['conv3_2']) 81 | activations['conv3_3'] = conv_layer(activations['conv3_2'], weights['conv3_3'], biases['conv3_3']) 82 | activations['pool3'] = max_pool(activations['conv3_3'], k=2, s=2) 83 | 84 | with tf.name_scope("conv4"): 85 | activations['conv4_1'] = conv_layer(activations['pool3'], weights['conv4_1'], biases['conv4_1']) 86 | activations['conv4_2'] = conv_layer(activations['conv4_1'], weights['conv4_2'], biases['conv4_2']) 87 | activations['conv4_3'] = conv_layer(activations['conv4_2'], weights['conv4_3'], biases['conv4_3']) 88 | activations['pool4'] = max_pool(activations['conv4_3'], k=2, s=2) 89 | 90 | with tf.name_scope("conv5"): 91 | activations['conv5_1'] = conv_layer(activations['pool4'], weights['conv5_1'], biases['conv5_1']) 92 | activations['conv5_2'] = conv_layer(activations['conv5_1'], weights['conv5_2'], biases['conv5_2']) 93 | activations['conv5_3'] = conv_layer(activations['conv5_2'], weights['conv5_3'], biases['conv5_3']) 94 | activations['pool5'] = max_pool(activations['conv5_3'], k=2, s=2) 95 | return activations 96 | 97 | def get_weights(): 98 | net = numpy.load('../Resources/VGG/vgg16.npy').item() 99 | weights = {} 100 | biases = {} 101 | for name in net.keys(): 102 | weights[name] = tf.Variable(tf.constant(net[name][0]), dtype='float32' ,name=name+"_weight", trainable=False) 103 | biases[name] = tf.Variable(tf.constant(net[name][1]), dtype='float32' ,name=name+"_bias", trainable=False) 104 | return weights,biases 105 | 106 | weights,biases = get_weights() 107 | 108 | y_true_activations = vgg_16(y_true,weights,biases) 109 | y_pred_activations = vgg_16(y_pred,weights,biases) 110 | 111 | for layer_name in layer_list: 112 | assert layer_name in y_true_activations.keys() 113 | 114 | loss = 0 115 | for layer_name,alpha in zip(layer_list,alphas): 116 | y = y_true_activations[layer_name] 117 | x = y_pred_activations[layer_name] 118 | loss += alpha*lp_loss(y,x) 119 | print loss 120 | return loss 121 | -------------------------------------------------------------------------------- /baseline/Model.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | 3 | import tensorflow as tf 4 | import re 5 | 6 | 7 | TOWER_NAME ='tower' 8 | def _variable_on_cpu(name, shape, initializer): 9 | """Helper to create a Variable stored on CPU memory. 10 | 11 | Args: 12 | name: name of the variable 13 | shape: list of ints 14 | initializer: initializer for Variable 15 | 16 | Returns: 17 | Variable Tensor 18 | """ 19 | with tf.device('/cpu:0'): 20 | dtype = tf.float32# if FLAGS.use_fp16 else tf.float32 21 | var = tf.get_variable(name, shape, initializer=initializer, dtype=dtype) 22 | return var 23 | 24 | 25 | def _variable_with_weight_decay(name, shape, stddev=5e-2): 26 | """Helper to create an initialized Variable with weight decay. 27 | 28 | Note that the Variable is initialized with a truncated normal distribution. 29 | A weight decay is added only if one is specified. 30 | 31 | Args: 32 | name: name of the variable 33 | shape: list of ints 34 | stddev: standard deviation of a truncated Gaussian 35 | wd: add L2Loss weight decay multiplied by this float. If None, weight 36 | decay is not added for this Variable. 37 | 38 | Returns: 39 | Variable Tensor 40 | """ 41 | dtype = tf.float32# if FLAGS.use_fp16 else tf.float32 42 | var = _variable_on_cpu( 43 | name, 44 | shape, 45 | tf.truncated_normal_initializer(stddev=stddev, dtype=dtype)) 46 | return var 47 | 48 | 49 | 50 | def batch_norm(inputs, is_training, pop_mean, pop_var, beta, gamma, decay = 0.999): 51 | def train(): 52 | batch_mean, batch_var = tf.nn.moments(inputs,[0,1,2], name='moments') 53 | train_mean = tf.assign(pop_mean, pop_mean * decay + batch_mean * (1 - decay)) 54 | train_var = tf.assign(pop_var , pop_var * decay + batch_var * (1 - decay)) 55 | with tf.control_dependencies([train_mean, train_var]): 56 | return tf.nn.batch_normalization(inputs, batch_mean, batch_var, beta, gamma, variance_epsilon=1e-3) 57 | 58 | def testing(): 59 | return tf.nn.batch_normalization(inputs, pop_mean, pop_var, beta, gamma, variance_epsilon=1e-3) 60 | 61 | return tf.cond(is_training, train, testing) 62 | 63 | 64 | 65 | 66 | 67 | 68 | class Model(object): 69 | def __init__(self,sess,optimizer, batch_size=32, tot_res_blocks=2): 70 | self.sess = sess 71 | self.optimizer = optimizer 72 | self.batch_size = batch_size 73 | self.THETA = {} 74 | self.tot_res_blocks = tot_res_blocks 75 | self.init_model() 76 | self.saver = tf.train.Saver(self.THETA, max_to_keep=100) 77 | 78 | 79 | def get_training(self): 80 | return self.training 81 | 82 | def set_training(self, training): 83 | self.training = tf.assign(self.training, tf.constant(training, shape=[])) 84 | 85 | def init_model(self): 86 | def init_bn_params(block_name,ind,size=64): 87 | self.THETA['{}beta{}'.format(block_name,ind)] = tf.Variable( 88 | tf.truncated_normal([size], mean=0.0, stddev=5e-2), name='{}beta{}'.format(block_name,ind), trainable=True) 89 | self.THETA['{}gamma{}'.format(block_name,ind)]= tf.Variable( 90 | tf.truncated_normal([size], mean=1.0, stddev=5e-2), name='{}gamma{}'.format(block_name,ind), trainable=True) 91 | self.THETA['{}pop_mean{}'.format(block_name,ind)] = tf.Variable(tf.zeros([size]), name = '{}pop_mean{}'.format(block_name,ind), trainable=False) 92 | self.THETA['{}pop_var{}'.format(block_name,ind)] = tf.Variable(tf.ones([size]) , name = '{}pop_var{}'.format(block_name,ind) , trainable=False) 93 | 94 | self.THETA['inpw1'] = _variable_with_weight_decay(shape=[3, 3, 3, 64], name='inw1') 95 | self.THETA['inpb1'] = _variable_with_weight_decay(shape=[ 64], name='inb1') 96 | 97 | for ind in range(1,self.tot_res_blocks+1): 98 | 99 | self.THETA['b{}w1'.format(ind)] = _variable_with_weight_decay(shape=[3, 3, 64, 64], name='b{}w1'.format(ind)) 100 | self.THETA['b{}b1'.format(ind)] = _variable_with_weight_decay(shape=[ 64] , name='b{}b1'.format(ind)) 101 | 102 | init_bn_params('b{}'.format(ind),1) 103 | 104 | self.THETA['b{}w2'.format(ind)] = _variable_with_weight_decay(shape=[3, 3, 64, 64], name='b{}w2'.format(ind)) 105 | self.THETA['b{}b2'.format(ind)] = _variable_with_weight_decay(shape=[ 64] , name='b{}b2'.format(ind)) 106 | 107 | init_bn_params('b{}'.format(ind),2) 108 | 109 | 110 | self.THETA['recon0w1'] = _variable_with_weight_decay(shape=[3, 3, 64, 64], name='recon0w1') 111 | self.THETA['recon0b1'] = _variable_with_weight_decay(shape=[64] , name='recon0b1') 112 | init_bn_params('recon0',1) 113 | 114 | self.THETA['recon1w1'] = _variable_with_weight_decay(shape=[3, 3, 64, 256], name='recon1w1') 115 | self.THETA['recon1b1'] = _variable_with_weight_decay(shape=[256] , name='recon1b1') 116 | 117 | self.THETA['recon2w1'] = _variable_with_weight_decay(shape=[3, 3, 256, 256], name='recon2w1') 118 | self.THETA['recon2b1'] = _variable_with_weight_decay(shape=[256] , name='recon2b1') 119 | 120 | self.THETA['recon3w1'] = _variable_with_weight_decay(shape=[3, 3, 256, 3], name='recon3w1') 121 | self.THETA['recon3b1'] = _variable_with_weight_decay(shape=[3] , name='recon3b1') 122 | 123 | 124 | 125 | 126 | def inference(self,inp, phase_train): 127 | 128 | def conv(blk,ind,inp): 129 | x = tf.nn.conv2d(inp, self.THETA['{}w{}'.format(blk,ind)], strides=[1, 1, 1, 1], padding='SAME',name="{}w{}".format(blk,ind)) 130 | x = tf.nn.bias_add(x, self.THETA['{}b{}'.format(blk,ind)], name="{}b{}".format(blk,ind)) 131 | return x 132 | 133 | def bn(blk,ind,inp): 134 | x = batch_norm(inp, 135 | phase_train, 136 | self.THETA['{}pop_mean{}'.format(blk,ind)], 137 | self.THETA['{}pop_var{}'.format(blk,ind)], 138 | self.THETA['{}beta{}'.format(blk,ind)], 139 | self.THETA['{}gamma{}'.format(blk,ind)], 140 | ) 141 | return x 142 | 143 | 144 | def block_conv(inp, name): 145 | with tf.variable_scope(name): 146 | x = conv(name,1,inp) 147 | x = bn(name,1,x) 148 | c1 = tf.nn.relu(x) 149 | 150 | x = conv(name,2,c1) 151 | x = bn(name,2,x) 152 | 153 | x = tf.add(inp,x) 154 | return x 155 | 156 | 157 | x = conv('inp',1,inp) 158 | o0 = tf.nn.relu(x) 159 | o = o0 160 | 161 | # all residual blocks 162 | for ind in range(1,self.tot_res_blocks+1): 163 | o = block_conv(o,"b{}".format(ind)) 164 | 165 | x = conv('recon0',1,o) 166 | x = bn('recon0',1,x) 167 | x = tf.add(x,o0) 168 | 169 | x = conv('recon1',1,x) 170 | x = tf.nn.relu(x) 171 | 172 | x = conv('recon2',1,x) 173 | x = tf.nn.relu(x) 174 | 175 | x = conv('recon3',1,x) 176 | 177 | out = x 178 | return out 179 | 180 | 181 | def save_model(self,file_name): 182 | self.saver.save(self.sess,file_name) 183 | 184 | def load_model(self,file_name): 185 | self.saver.restore(self.sess, file_name) 186 | 187 | 188 | def get_gradients(self,total_loss): 189 | grads = self.optimizer.compute_gradients(total_loss) 190 | 191 | return grads 192 | 193 | def get_backprop(self,total_loss,global_step): 194 | solver = self.optimizer.minimize(total_loss,global_step=global_step) 195 | return solver 196 | -------------------------------------------------------------------------------- /baseline/Test.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | parser = argparse.ArgumentParser(description='Expert denoisers training') 3 | parser.add_argument('--batch-size', type=int, default=32, metavar='N', required=True, help='input batch size for training (default: 32)') 4 | parser.add_argument('--stride', type=int, default=16, metavar='N', required=True, help='input batch size for training (default: 32)') 5 | #parser.add_argument('--noise-level', type=str, default=5.0, metavar='N', help='No of iterations') 6 | parser.add_argument('--model-snapshot-file', type=str, required=True, help='input batch size for training (default: 32)') 7 | parser.add_argument('--image-dir', type=str, required=True, help='input batch size for training (default: 32)') 8 | parser.add_argument('--test-data-list-file', type=str, metavar='N', required=True, help='test base folder is required') 9 | parser.add_argument('--save-images', type=str, metavar='N', required=True, help='save images') 10 | parser.add_argument('--quality', type=int, metavar='N', required=True, help='compression quality') 11 | args = parser.parse_args() 12 | 13 | 14 | import numpy as np 15 | import tensorflow as tf 16 | import os 17 | import numpy 18 | import glob 19 | from itertools import product 20 | from PIL import Image 21 | import skimage 22 | from skimage.util import view_as_windows 23 | import json 24 | import skimage.measure as measure 25 | import sys 26 | import importlib 27 | import cPickle as pickle 28 | from tqdm import tqdm 29 | #import matplotlib.pyplot as plt 30 | from Model import Model 31 | 32 | #config = tf.ConfigProto(gpu_options=tf.GPUOptions(per_process_gpu_memory_fraction=0.5)) 33 | config = tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True)) 34 | try: 35 | tf.Session(config=config) 36 | except: pass 37 | 38 | 39 | 40 | def make_dirs(dirs): 41 | for _dir in dirs: 42 | if not os.path.isdir(_dir): 43 | os.mkdir(_dir) 44 | 45 | def reconstruct_image_from_patches(Pp_p,s, In_p): 46 | Op = numpy.zeros(s) 47 | i=0 48 | for x in range(Op.shape[0]): 49 | for y in range(Op.shape[1]): 50 | Op[x,y,:] = Pp_p[i].reshape((s[3],s[4],3)) 51 | i = i+1 52 | 53 | Ip_p = numpy.zeros_like(In_p) 54 | Ip_w = numpy.zeros_like(In_p) 55 | 56 | for x in range(Op.shape[0]): 57 | for y in range(Op.shape[1]): 58 | Ip_p[x*stride[0]:(x)*stride[0]+64,y*stride[1]:(y)*stride[1]+64,:] += Op[x,y,0] 59 | Ip_w[x*stride[0]:(x)*stride[0]+64,y*stride[1]:(y)*stride[1]+64,:] += 1.0 60 | #psnr[x*stride:(x+1)*stride,y*stride:(y+1)*stride] = 2*sqrt(mean_squared_error(Op[x,y], Pc[x,y])) 61 | Ip_w[Ip_w==0.0] = 1.0 62 | Ip_p = Ip_p/Ip_w 63 | 64 | return Ip_p[32:-32,32:-32] # removing padding 65 | 66 | #r=0 67 | def denoise_image(In, stride): 68 | In_p = skimage.util.pad(In,32, 'reflect') 69 | In_p = In_p[:,:,32:(2*32+3)-32] 70 | Pn_p = view_as_windows(In_p[:,:,:],window_shape=(64,64,3),step=stride) 71 | 72 | 73 | ss = Pn_p.shape 74 | Pn_p = Pn_p.reshape(ss[0]*ss[1],ss[3],ss[4],3)/255.0 75 | Pp_p = numpy.zeros_like(Pn_p) 76 | # r=0 77 | for ind in range(0,ss[0]*ss[1],batch_size): 78 | # print ind,ind+batch_size 79 | mini_batch = Pn_p[ind:ind+batch_size] 80 | # print "{} = {}".format("Input shape",mini_batch[0].shape) 81 | # print mini_batch[0] 82 | # np.save('ex1'+str(r+1)+'.npy',mini_batch) 83 | 84 | predk = E.sess.run([y_],feed_dict={x_:mini_batch, phase_train:False}) 85 | # np.save('ez1'+str(r+1)+'.npy',predk) 86 | # print "{} = {}".format("Shape of the batch",predk[0].shape) 87 | # print "{} = {}".format("Max of the batch",np.max(predk[0])) 88 | # print "{} = {}".format("Shape of the input",mini_batch.shape) 89 | # print "{} = {}".format("Max of the input",np.max(mini_batch)) 90 | # r = r + 1 91 | # print "{} ={}".format("Shape of the output",predk[0][0].shape) 92 | # print predk[0][0] 93 | # r=r+1 94 | # print "{} = {}".format("MSE of inpit and output",((predk[0][0] - mini_batch[0])**2).mean(axis=None)) 95 | # print "{} = {}".format("PSNR of input and output",measure.compare_psnr(predk[0][0],mini_batch[0], dynamic_range = 1.0)) 96 | # if( r==2): 97 | # sys.exit(0) 98 | # print "{} = {}".format("MSE OF INPUT AND RECONSTRUCTED",((predk[0] - mini_batch)**2).mean(axis = None)) 99 | Pp_p[ind:ind+batch_size,:,:,:] = predk[0] 100 | Pp_p = Pp_p*255.0 101 | Ip = reconstruct_image_from_patches(Pp_p,ss, In_p) 102 | return post_process(Ip) 103 | 104 | def post_process(I): 105 | # clip values outside range 106 | I[I>255.0]=255.0 107 | I[I< 0.0]= 0.0 108 | return I 109 | 110 | def lp_loss(y_true, y_pred): 111 | loss = tf.reduce_mean(tf.losses.mean_squared_error(predictions=y_pred, labels=y_true)) 112 | return loss 113 | 114 | def read_image(file_name): 115 | return Image.open(file_name) 116 | 117 | def jpeg_compress(I, quality,file_name): 118 | img_path = file_name[:-4] + '_20.jpg' 119 | I.save(".temp","JPEG",quality = quality) 120 | return numpy.ascontiguousarray( Image.open(".temp").convert('RGB'), dtype = 'float32') 121 | 122 | def img_generator(quality, test_data_file): 123 | f = open(test_data_file,'r') 124 | for l in f.readlines(): 125 | l = l.strip() 126 | I = read_image(l) 127 | #print l 128 | In = jpeg_compress(I, quality,l) 129 | I = numpy.asarray(I, dtype='float32') 130 | yield I, In, l, l 131 | 132 | 133 | generator = img_generator(args.quality, args.test_data_list_file) 134 | batch_size = args.batch_size 135 | stride = (args.stride,args.stride,args.stride) 136 | 137 | LG ='E2' 138 | 139 | log_dir = os.path.join(LG, 'LOG') 140 | clean_images_save_folder = os.path.join(LG, 'save', 'clean_images') 141 | noisy_images_save_folder = os.path.join(LG, 'save', 'noise_std_{}'.format(args.quality)) 142 | reconstructed_images_save_folder = os.path.join(LG, 'save', 'recon_std_{}'.format(args.quality)) 143 | 144 | # make_dirs([args.exp_code, log_file, os.path.join(args.exp_code, 'save'), 145 | # clean_images_save_folder, noisy_images_save_folder]) 146 | 147 | make_dirs([LG, log_dir, os.path.join(LG, 'save'), clean_images_save_folder]) 148 | make_dirs([LG, log_dir, os.path.join(LG, 'save'), noisy_images_save_folder]) 149 | make_dirs([LG, log_dir, os.path.join(LG, 'save'), reconstructed_images_save_folder]) 150 | 151 | with tf.Session(config=config) as sess: 152 | 153 | optimizer = tf.train.AdamOptimizer(1e-3) 154 | E = Model(sess,optimizer=optimizer, batch_size = None, tot_res_blocks=5) 155 | x_ = tf.placeholder(tf.float32,[None,64,64,3],name="noisy_images") 156 | phase_train = tf.placeholder(tf.bool, [], name='phase_train') 157 | y_ = E.inference(x_, phase_train) 158 | 159 | init = tf.global_variables_initializer() 160 | E.sess.run(init) 161 | 162 | E.load_model(args.model_snapshot_file) 163 | 164 | for Ic, In, clean_image_name, noisy_image_name in tqdm(generator): 165 | In = post_process(In) 166 | Ir = denoise_image(In, stride) 167 | if args.save_images == 'True': 168 | # this will currently give an error 169 | Image.fromarray(Ir.astype(numpy.uint8)).save( 170 | os.path.join(reconstructed_images_save_folder, 171 | os.path.basename(clean_image_name)) 172 | ) 173 | Image.fromarray(Ic.astype(numpy.uint8)).save( 174 | os.path.join(clean_images_save_folder, 175 | os.path.basename(clean_image_name)) 176 | ) 177 | Image.fromarray(In.astype(numpy.uint8)).save( 178 | os.path.join(noisy_images_save_folder, 179 | os.path.basename(clean_image_name)) 180 | ) 181 | 182 | 183 | print "Done!" 184 | import argparse 185 | parser = argparse.ArgumentParser(description='Expert denoisers training') 186 | -------------------------------------------------------------------------------- /baseline/Train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import numpy 3 | import numpy as np 4 | numpy.random.seed(1) 5 | import tensorflow as tf 6 | tf.set_random_seed(1) 7 | 8 | from random import shuffle 9 | import os 10 | import glob 11 | from PIL import Image 12 | 13 | from Model import Model 14 | from Input import BatchImageInput 15 | import Losses 16 | import time 17 | 18 | #global config 19 | #config = tf.ConfigProto(gpu_options=tf.GPUOptions(per_process_gpu_memory_fraction=0.5)) 20 | config = tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True)) 21 | try: 22 | tf.Session(config=config) 23 | except: pass 24 | 25 | 26 | parser = argparse.ArgumentParser(description='Expert denoisers training') 27 | 28 | parser.add_argument('--lr', type=float, metavar='N', default=1e-5, required=True, help='learning rate is required') 29 | parser.add_argument('--batch-size', type=int, default=32, metavar='N', required=True, help='input batch size for training (default: 32)') 30 | parser.add_argument('--tot-grad-updates', type=int, default=10000, metavar='N', help='No of iterations') 31 | parser.add_argument('--start-step', type=int, default=1, metavar='N', help='training start step') 32 | parser.add_argument('--log-file', type=str, default=None, metavar='N', help='No of iterations') 33 | parser.add_argument('--train-dir', type=str, default=None, metavar='N', help='No of iterations') 34 | parser.add_argument('--valdn-dir', type=str, default=None, metavar='N', help='No of iterations') 35 | parser.add_argument('--log-interval', type=int, default=10000, metavar='N', required=True, help='frequency of logging') 36 | parser.add_argument('--snapshot-interval', type=int, default=10000, metavar='N', required=True, help='frequency of snapshotting') 37 | parser.add_argument('--exp-code', type=str, default=10000, metavar='N', required=True, help='experiment code') 38 | 39 | args = parser.parse_args() 40 | 41 | LGXY = "{}".format(args.exp_code) 42 | 43 | def make_dirs(dirs): 44 | for _dir in dirs: 45 | if not os.path.isdir(_dir): 46 | os.mkdir(_dir) 47 | 48 | make_dirs([LGXY, os.path.join(LGXY, 'LOG'), os.path.join(LGXY, 'checkpoints')]) 49 | 50 | 51 | # redirect error messages to std out or file 52 | import sys 53 | _save_out = sys.stdout 54 | if args.log_file is not None: 55 | f = open(os.path.join(LGXY,args.log_file),'a') 56 | sys.stdout = f 57 | 58 | with tf.Session(config=config) as sess: 59 | step_ = tf.Variable(0, trainable=False) 60 | lr_ = tf.train.exponential_decay(args.lr, step_, 100, 0.999,staircase=True) 61 | tf.summary.scalar("lr", lr_) 62 | tf.summary.scalar("step", step_) 63 | optimizer = tf.train.AdamOptimizer(lr_) 64 | 65 | # model 66 | E = Model(sess,optimizer=optimizer, batch_size = args.batch_size, tot_res_blocks=5) 67 | 68 | # --beg-- train data 69 | with open('../Resources/imagenet_val_paths.txt', 'r') as f: 70 | file_paths = [os.path.join(args.train_dir,fname) for fname in f.read().splitlines()] 71 | Train_Data = BatchImageInput(file_paths,[20],batch_size = args.batch_size) 72 | 73 | # --beg-- validation data 74 | with open('../Resources/bsds500_test_paths.txt', 'r') as f: 75 | file_paths = [os.path.join(args.valdn_dir,fname) for fname in f.read().splitlines()] 76 | Valdn_Data = BatchImageInput(file_paths,[20],batch_size = args.batch_size) 77 | # --end-- validation data 78 | 79 | # computation 80 | x_ = tf.placeholder(tf.float32,[args.batch_size,64,64,3],name="noisy_images") 81 | y_ = tf.placeholder(tf.float32,[args.batch_size,64,64,3],name="clean_images") 82 | a_ = tf.placeholder(tf.float32,[], name="alpha") 83 | phase_train = tf.placeholder(tf.bool, [], name='is_training') 84 | 85 | z_ = E.inference(x_, phase_train) 86 | Loss_list = { 87 | 'lp':(Losses.lp_loss,0.9), 88 | 'pl':(Losses.perceptual_loss,0.1), 89 | # 'tv':Losses.tv_loss, 90 | } 91 | 92 | Loss_ops = {} 93 | Backprop_ops = {} 94 | LOSS = 0 95 | for loss_name in Loss_list: 96 | print loss_name 97 | # compute the loss 98 | Loss_ops[loss_name] = Loss_list[loss_name][0](y_,z_) 99 | # weighted sum 100 | LOSS = LOSS+ (Loss_ops[loss_name]*Loss_list[loss_name][1]) 101 | # add o summary 102 | tf.summary.scalar(loss_name+'_loss', Loss_ops[loss_name]) 103 | 104 | Loss_ops['total'] = LOSS 105 | tf.summary.scalar('total_loss', LOSS) 106 | 107 | BACK_PROP_OP = E.get_backprop( 108 | LOSS, 109 | step_ 110 | ) 111 | 112 | tf.summary.image("predk_image",z_) 113 | tf.summary.image("clean_image",y_) 114 | tf.summary.image("noisy_image",x_) 115 | 116 | 117 | # DATA FLOW 118 | 119 | X_train_,Y_train_ = Train_Data.get_minibatch_tensors() 120 | X_valdn_,Y_valdn_ = Valdn_Data.get_minibatch_tensors() 121 | t_coord = tf.train.Coordinator() 122 | t_threads = tf.train.start_queue_runners(coord=t_coord) 123 | 124 | 125 | ## Add histograms for trainable variables. 126 | #for var in tf.trainable_variables(): 127 | # tf.summary.histogram(var.op.name, var) 128 | #for v_,g_ in GRADS_: 129 | # if v_ is not None: 130 | # tf.summary.histogram(v_.op.name,g_) 131 | 132 | summary_op = tf.summary.merge_all() 133 | 134 | writer = tf.summary.FileWriter(os.path.join(LGXY, 'LOG'), graph=E.sess.graph) 135 | init = tf.global_variables_initializer() 136 | E.sess.run(init) 137 | 138 | step = args.start_step 139 | stop_step = args.tot_grad_updates - args.start_step 140 | 141 | while step <= stop_step: 142 | # get train batch 143 | X_train_batch, Y_train_batch = E.sess.run([X_train_,Y_train_]) 144 | 145 | # get op summary 146 | summary = E.sess.run( 147 | [summary_op], 148 | feed_dict={x_:X_train_batch,y_:Y_train_batch, phase_train:False} 149 | ) 150 | 151 | # run one grad update 152 | _ = E.sess.run( 153 | [BACK_PROP_OP], 154 | feed_dict={ 155 | x_:X_train_batch, 156 | y_:Y_train_batch, 157 | phase_train:True 158 | } 159 | ) 160 | writer.add_summary(summary[0], step) 161 | 162 | 163 | if step % args.log_interval == 0 or step == stop_step: 164 | 165 | X_valdn_batch, Y_valdn_batch = E.sess.run([X_valdn_,Y_valdn_]) 166 | 167 | # getting all the losses 168 | LOSSES = {} 169 | for loss_op in Loss_ops: 170 | op = Loss_ops[loss_op] 171 | out= E.sess.run( 172 | [op], 173 | feed_dict={ 174 | x_:X_train_batch, 175 | y_:Y_train_batch, 176 | phase_train:False 177 | } 178 | ) 179 | LOSSES[loss_op+'_train'] = out[0] 180 | out= E.sess.run( 181 | [op], 182 | feed_dict={ 183 | x_:X_valdn_batch, 184 | y_:Y_valdn_batch, 185 | phase_train:False 186 | } 187 | ) 188 | LOSSES[loss_op+'_valdn'] = out[0] 189 | 190 | 191 | # printing all losses 192 | loss_str_list = [] 193 | for key in LOSSES: 194 | loss_str_list.append(" {} = {:e}".format(key,LOSSES[key])) 195 | print("[{}] Step = {:5d} {}".format(time.strftime("%A:%B-%d-%Y %H:%M:%S"),step, "".join(loss_str_list))) 196 | 197 | sys.stdout.flush() 198 | if step % args.snapshot_interval == 0 or step == stop_step: 199 | E.save_model(os.path.join(LGXY, 'checkpoints', "{}.{:05d}.ckpt".format(LGXY,step))) 200 | 201 | step += 1 202 | 203 | print("Done") 204 | #tf.reset_default_graph() 205 | 206 | 207 | try: 208 | f.close() 209 | except: 210 | pass 211 | sys.stdout = _save_out 212 | -------------------------------------------------------------------------------- /baseline/avg_psnr.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | parser = argparse.ArgumentParser(description='Expert denoisers training') 3 | parser.add_argument('--image-filename', type=str, default='', metavar='N', required=True, help='File containing PSNR values of LIVE1 dataset') 4 | args = parser.parse_args() 5 | 6 | import sys 7 | 8 | 9 | noisy_psnr = [] 10 | recon_psnr = [] 11 | diff_psnr = [] 12 | i=0 13 | with open(args.image_filename,'r') as f: 14 | for l in f.readlines(): 15 | i=i+1 16 | l = l.strip() 17 | l = l.split() 18 | noisy_psnr.append(float(l[1])) 19 | recon_psnr.append(float(l[2])) 20 | diff_psnr.append( float(l[2]) - float(l[1])) 21 | 22 | 23 | avg_noisy_psnr = sum(noisy_psnr)/len(noisy_psnr) 24 | avg_recon_psnr = sum(recon_psnr)/len(recon_psnr) 25 | avg_diff_psnr = sum(diff_psnr)/len(diff_psnr) 26 | print "Average noisy PSNR = "+str(avg_noisy_psnr) 27 | print "Average recon PSNR = "+str(avg_recon_psnr) 28 | print "Average diffr PSNR = "+str(avg_diff_psnr) 29 | 30 | -------------------------------------------------------------------------------- /baseline/avg_psnr_results.txt: -------------------------------------------------------------------------------- 1 | Average noisy PSNR = 28.0607057901 2 | Average recon PSNR = 28.8379574926 3 | Average diffr PSNR = 0.777251702476 4 | -------------------------------------------------------------------------------- /baseline/results.txt: -------------------------------------------------------------------------------- 1 | bikes.bmp 26.3445871238 26.5181686657 2 | building2.bmp 23.9975502847 24.0337194875 3 | buildings.bmp 26.0241052781 25.9818589208 4 | caps.bmp 31.480716799 31.5081896704 5 | carnivaldolls.bmp 27.6292930857 27.2967063542 6 | cemetry.bmp 26.1793606878 26.2413333081 7 | churchandcapitol.bmp 27.2096019756 27.1292305676 8 | coinsinfountain.bmp 27.4504358068 27.2594348475 9 | dancers.bmp 25.0964985704 25.0167865545 10 | flowersonih35.bmp 24.0114013845 24.1183093935 11 | house.bmp 29.0179136968 29.1437620744 12 | lighthouse2.bmp 28.5988660231 28.7307600157 13 | lighthouse3.bmp 29.0727283528 29.0244882376 14 | manfishing.bmp 28.2856803354 27.6391749862 15 | monarch.bmp 29.7104547936 30.0136115837 16 | ocean.bmp 30.4518313939 30.4051665653 17 | paintedhouse.bmp 26.9858236489 26.9219885502 18 | parrots.bmp 32.0304208111 31.8605091725 19 | plane.bmp 30.6925222603 29.3694228625 20 | rapids.bmp 27.5289660075 27.6915463343 21 | sailing1.bmp 28.1411130482 28.1597126421 22 | sailing2.bmp 31.0939656612 31.1013001004 23 | sailing3.bmp 30.6184287948 30.5732574031 24 | sailing4.bmp 28.7732208948 28.7162033074 25 | statue.bmp 30.3289651333 29.7194091231 26 | stream.bmp 24.6731856621 24.8868319738 27 | studentsculpture.bmp 25.151586129 25.1442424697 28 | woman.bmp 27.0525324742 27.1237802689 29 | womanhat.bmp 30.1287117948 30.1159388864 30 | -------------------------------------------------------------------------------- /baseline/run.sh: -------------------------------------------------------------------------------- 1 | CUDA_VISIBLE_DEVICES=2 \ 2 | python Train.py \ 3 | --lr=1e-5 \ 4 | --batch-size=32 \ 5 | --log-interval=1000 \ 6 | --snapshot-interval=10000 \ 7 | --tot-grad-updates=110000 \ 8 | --train-dir='../Resources/Datasets/IMAGENET/val' \ 9 | --valdn-dir='../Resources/Datasets/BSDS/BSR/BSDS500/data/images/test' \ 10 | --exp-code=E2 \ 11 | --log-file=log.txt \ 12 | -------------------------------------------------------------------------------- /baseline/run_debug.sh: -------------------------------------------------------------------------------- 1 | CUDA_VISIBLE_DEVICES=2 \ 2 | python Train.py \ 3 | --lr=1e-5 \ 4 | --batch-size=32 \ 5 | --log-interval=10 \ 6 | --snapshot-interval=10000 \ 7 | --tot-grad-updates=110000 \ 8 | --train-dir='../Resources/Datasets/IMAGENET/val' \ 9 | --valdn-dir='../Resources/Datasets/BSDS/BSR/BSDS500/data/images/test' \ 10 | --exp-code=E2 \ 11 | #--log-file=log.txt \ 12 | -------------------------------------------------------------------------------- /baseline/run_eval.sh: -------------------------------------------------------------------------------- 1 | python Eval.py \ 2 | --image-folder E2/save \ 3 | --image-filename ../Resources/live1_list.txt \ 4 | --quality 20 5 | 6 | -------------------------------------------------------------------------------- /baseline/run_test.sh: -------------------------------------------------------------------------------- 1 | CUDA_VISIBLE_DEVICES=2 \ 2 | python Test.py \ 3 | --batch-size 32 \ 4 | --stride 16 \ 5 | --quality 20 \ 6 | --model-snapshot-file E2/checkpoints/E2.20000.ckpt \ 7 | --image-dir='../Resources/Datasets/LIVE1/' \ 8 | --test-data-list-file ../Resources/live1_list.txt \ 9 | --save-image True \ 10 | -------------------------------------------------------------------------------- /result.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/utkarshojha/compression-artifacts-reduction/6ebf11ff813e4bb64ab8437a56c6ffd2f99b1f7a/result.png --------------------------------------------------------------------------------