def psnr(self, input_HR, input_LR):
# Calculating Peak Signal-to-noise-ratio
# Using equations from here:
# https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio
mse = tf.reduce_mean(
tf.squared_difference(input_HR, self.gen_G(input_LR)))
PSNR = tf.constant(255**2, dtype=tf.float32) / mse
PSNR = tf.constant(10, dtype=tf.float32) * utils.log10(PSNR)
return PSNR
def __init__(self, img_size=50, num_layers=32, feature_size=256, scale=2, output_channels=3):
print("Building EDSR...")
with tf.name_scope(name='image_input'):
self.input = x = tf.placeholder(tf.float32, [None, img_size, img_size, output_channels])
mean_x = tf.reduce_mean(x)
image_input = x - mean_x
with tf.name_scope(name='target_input'):
self.target = y = tf.placeholder(tf.float32, [None, img_size*scale, img_size*scale, output_channels])
mean_y = tf.reduce_mean(y)
image_target = y - mean_x
resnet_out = self.resnet(image_input, feature_size, num_layers, reuse=False, scope='_g')
self.debug = self.resnet(image_target, feature_size, num_layers, reuse=True, scope='_g')
g_output = self.upconv(resnet_out, scale, feature_size)
self.g_ini_out = output = g_output # slim.conv2d(x,output_channels,[3,3])
self.g_ini_loss = g_ini_loss = tf.reduce_mean(tf.losses.absolute_difference(image_target, g_output)) # L1 loss
mse = tf.reduce_mean(tf.squared_difference(image_target, g_output))
self.PSNR = PSNR = tf.constant(10,dtype=tf.float32)*utils.log10(tf.constant(255**2,dtype=tf.float32)/mse)
# discriminator
d_resnet_fake = self.resnet(g_output, feature_size, num_layers, reuse=True, scope='_g')
d_logits_fake = self.classification(d_resnet_fake, feature_size, reuse=False)
# d_resnet_real.shape = (batchsize, 100, 100, 128)
self.d_resnet_real = d_resnet_real = self.resnet(image_target, feature_size, num_layers, reuse=True, scope='_g')
self.d_logits_real = d_logits_real = self.classification(d_resnet_real, feature_size, reuse=True)
d_loss1 = tf.losses.sigmoid_cross_entropy(d_logits_real, tf.ones_like(d_logits_real), scope='d1')
d_loss2 = tf.losses.sigmoid_cross_entropy(d_logits_fake, tf.zeros_like(d_logits_fake), scope='d2')
self.d_loss = d_loss1 + d_loss2
# generator
g_res_loss = 2e-6 * tf.reduce_mean(tf.squared_difference(d_resnet_real, d_resnet_fake))
g_gan_loss = 1e-3 * tf.losses.sigmoid_cross_entropy(d_logits_fake, tf.ones_like(d_logits_fake), scope='g')
self.g_loss = g_ini_loss + g_gan_loss + g_res_loss
self.g_ini_var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'resnet_g')+tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'upconv')
self.d_var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'resnet_g')+tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'classification')
self.g_var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'resnet_g')+tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'upconv')
tf.summary.scalar("loss",self.g_ini_loss)
tf.summary.scalar("PSNR",PSNR)
tf.summary.image("input_image",self.input+mean_x)
tf.summary.image("target_image",self.target+mean_y)
tf.summary.image("output_image",self.g_ini_out+mean_x)
self.sess = tf.Session()
self.saver = tf.train.Saver()
print("Done building!")
def build_recon_model(self,blu,QP):
self.images = tf.placeholder(tf.float32, [None, None, None, self.c_dim], name='images')
self.labels = tf.placeholder(tf.float32, [None, None, None, self.c_dim], name='labels')
self.images_norm=self.images/255-128/255#tf.divide(self.images,255)
self.labels_norm=self.labels/255-128/255#tf.divide(self.labels,255)
self.weights = {
'w1': tf.get_variable('w1',[5, 5, 1, 64], initializer=variance_scaling_initializer()),#initializer=tf.random_normal_initializer(stddev=np.sqrt(2.0/25))),
'w2_1': tf.get_variable('w2_1',[3, 3, 64, 32],initializer=variance_scaling_initializer()),#initializer=tf.random_normal_initializer(stddev=np.sqrt(2.0/9/64))),
'w2_2': tf.get_variable('w2_2',[5, 5, 64, 16], initializer=variance_scaling_initializer()),#initializer=tf.random_normal_initializer(stddev=np.sqrt(2.0/25/64))),
'w3_1': tf.get_variable('w3_1',[3, 3, 48, 16],initializer=variance_scaling_initializer() ),#initializer=tf.random_normal_initializer(stddev=np.sqrt(2.0/9/48))),
'w3_2': tf.get_variable('w3_2',[1, 1, 48, 32], initializer=variance_scaling_initializer()),#initializer=tf.random_normal_initializer(stddev=np.sqrt(2.0/1/48))),
'w4': tf.get_variable('w4',[3, 3, 48, 1], initializer=variance_scaling_initializer())#initializer=tf.random_normal_initializer(stddev=np.sqrt(2.0/25/48)))
}
self.biases = {
'b1': tf.get_variable('b1',[64],initializer=tf.constant_initializer(0)),
'b2_1': tf.get_variable('b2_1',[32], initializer=tf.constant_initializer(0)),
'b2_2': tf.get_variable('b2_2', [16], initializer=tf.constant_initializer(0)),
'b3_1': tf.get_variable('b3_1',[16], initializer=tf.constant_initializer(0)),
'b3_2': tf.get_variable('b3_2', [32], initializer=tf.constant_initializer(0)),
'b4': tf.get_variable('b3', [1], initializer=tf.constant_initializer(0))
}
stepw, self.blu_ub, ratio = quantization.loadQpara(QP)
self.stepw_list=stepw
self.stepw = {'w1': stepw[0], 'w2_1': stepw[1], 'w2_2': stepw[2], 'w3_1': stepw[3], 'w3_2': stepw[4],
'w4': stepw[5]}
self.stepw_b = {'b1': stepw[0], 'b2_1': stepw[1], 'b2_2': stepw[2], 'b3_1': stepw[3], 'b3_2': stepw[4],
'b4': stepw[5]}
self.ratio = {'b1': ratio[0], 'b2_1': ratio[1], 'b2_2': ratio[2], 'b3_1': ratio[3], 'b3_2': ratio[4],
'b4': ratio[5]}
if blu:
self.residual=self.model_blu()
else:
self.residual=self.model()
self.pred = tf.add(self.residual,self.images_norm,name='pred')
# Loss function (MSE)
self.loss = tf.nn.l2_loss(tf.subtract(self.labels_norm, self.pred))#tf.reduce_mean(tf.square(self.labels - self.pred))
self.gloal_step=tf.Variable(0,trainable=False,name='global_step')
tf.summary.scalar("loss", self.loss)
mse = tf.reduce_mean(tf.squared_difference(self.labels, self.pred * 255))
PSNR = tf.constant(255 ** 2, dtype=tf.float32) / mse
PSNR = tf.constant(10, dtype=tf.float32) * utils.log10(PSNR)
tf.summary.scalar("PSNR", PSNR)
tf.summary.image("input_image", tf.cast(self.images * 255, tf.uint8))
tf.summary.image("target_image", tf.cast(self.labels * 255, tf.uint8))
tf.summary.image("output_image", tf.cast(self.pred * 255, tf.uint8))
self.saver=tf.train.Saver(max_to_keep=30)
def evaluate(input_op, clean_imgs, eval_type):
batch_size = input_op.get_shape().as_list()[0]
input_op = tf.reshape(input_op, [batch_size, -1])
clean_imgs = tf.reshape(clean_imgs, [batch_size, -1])
if eval_type == 'mse':
result = tf.reduce_mean(tf.squared_difference(input_op, clean_imgs),
axis=0)
elif eval_type == 'psnr':
mse = tf.reduce_mean(tf.squared_difference(input_op, clean_imgs),
axis=0)
psnr = tf.constant(255**2, dtype=tf.float32) / mse
result = tf.constant(10, dtype=tf.float32) * log10(psnr)
elif eval_type == 'ssim':
pass
else:
raise NotImplementedError
return result
def discriminator_loss(y_true, y_pred):
return log10(y_true + 1) + log10(y_pred + 1)
def adversary_loss(y_true, y_pred):
return log10(y_pred + 1)
def __init__(self):
print("Building MYSR...")
self.is_continued = False
self.imgsize = config.TRAIN.imgsize
self.output_channels = config.TRAIN.output_channels
self.scale = config.TRAIN.scale
self.epoch = config.TRAIN.n_epoch
self.batch_size = config.TRAIN.batch_size
self.save_model_dir = config.TRAIN.save_model_dir
# Placeholder for image inputs
# self.input = x = tf.placeholder(tf.float32, [None, self.imgsize, self.imgsize, self.output_channels])
self.input = x = tf.placeholder(
tf.float32, [None, None, None, self.output_channels])
# Placeholder for upscaled image ground-truth
# self.target = y = tf.placeholder(tf.float32, [None, self.imgsize*self.scale, self.imgsize*self.scale, self.output_channels])
self.target = y = tf.placeholder(
tf.float32, [None, None, None, self.output_channels])
# 验证集Placeholder
self.valid_input = []
self.valid_target = []
# 输入预处理
# result = result / (255. / 2.)
# TODO: 后边有relu层，注意将输入图像收缩至[-1, 1]区间是否合适
# result = result - 1.
# image_input = x - (255. / 2.)
# image_input = image_input - 1
# image_target = y - (255. / 2.)
# image_target = image_target - 1
image_input = x / (255. / 2.)
image_input = image_input - 1
image_target = y / (255. / 2.)
image_target = image_target - 1
# 貌似即使收缩至[-1, 1]区间，卷积层依旧可以有效适应，只是注意最后一层不能使用relu，b毕竟relu值域在[0, x]
# ENCODER
# 入口
# One convolution before res blocks and to convert to required feature depth
x = slim.conv2d(image_input, 256, [3, 3])
conv_1 = x
# scaling_factor = 0.1
scaling_factor = 0.1
# Add the residual blocks to the model
for i in range(12):
x = utils.resBlock(x, 256, scale=scaling_factor)
# One more convolution, and then we add the output of our first conv layer
x = slim.conv2d(x, 256, [3, 3])
x += conv_1
# Upsample output of the convolution
# x = utils.upsample(x, self.scale, 256, None)
# TODO:试试新的上采样
x = tf.depth_to_space(x, self.scale)
x = slim.conv2d(x,
self.output_channels, [3, 3],
activation_fn=tf.nn.tanh)
# One final convolution on the upsampling output
output = x # slim.conv2d(x,output_channels,[3,3])
# 结果 注意预处理的值
# self.out = tf.clip_by_value(output+(255. / 2.), 0.0, 255.0)
self.out = tf.clip_by_value((output + 1) * (255. / 2.), 0.0, 255.0)
self.loss = loss = tf.reduce_mean(
tf.losses.absolute_difference(image_target, output))
# Calculating Peak Signal-to-noise-ratio
# Using equations from here: https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio
mse = tf.reduce_mean(
tf.squared_difference((image_target + 1) * (255. / 2.),
tf.clip_by_value((output + 1) * (255. / 2.),
0.0, 255.0)))
PSNR = tf.constant(255**2, dtype=tf.float32) / mse
PSNR = tf.constant(10, dtype=tf.float32) * utils.log10(PSNR)
# Scalar to keep track for loss
tf.summary.scalar("loss", self.loss)
tf.summary.scalar("PSNR", PSNR)
# Image summaries for input, target, and output
tf.summary.image("input_image", tf.cast(self.input, tf.uint8))
tf.summary.image("target_image", tf.cast(self.target, tf.uint8))
tf.summary.image("output_image", tf.cast(self.out, tf.uint8))
# Tensorflow graph setup... session, saver, etc.
config_tf = tf.ConfigProto()
config_tf.gpu_options.allow_growth = True
self.sess = tf.Session(config=config_tf)
self.saver = tf.train.Saver()
print("Done building!")
def __init__(self,
img_size=32,
global_layers=16,
local_layers=8,
feature_size=64,
scale=2,
output_channels=3):
print("Building RDN...")
self.img_size = img_size
self.scale = scale
self.output_channels = output_channels
#Placeholder for image inputs
self.input = x = tf.placeholder(tf.float32,
[None, None, None, output_channels])
#Placeholder for upscaled image ground-truth
self.target = y = tf.placeholder(tf.float32,
[None, None, None, output_channels])
self.is_training = tf.placeholder(tf.bool, name='is_training')
"""
Preprocessing as mentioned in the paper, by subtracting the mean
However, the subtract the mean of the entire dataset they use. As of
now, I am subtracting the mean of each batch
"""
mean_x = 127
image_input = x - mean_x
mean_y = 127
image_target = y - mean_y
scaling_factor = 0.1
x1 = slim.conv2d(image_input, feature_size, [3, 3])
x = slim.conv2d(x1, feature_size, [3, 3])
outputs = []
for i in range(global_layers):
x = utils.resDenseBlock(x,
feature_size,
layers=local_layers,
scale=scaling_factor)
outputs.append(x)
x = tf.concat(outputs, 3)
x = slim.conv2d(x, feature_size, [1, 1], activation_fn=None)
x = slim.conv2d(x, feature_size, [3, 3])
x = x + x1
x = utils.upsample(x, scale, feature_size)
#output = slim.conv2d(x,output_channels,[3,3])
output = tf.layers.conv2d(x,
output_channels, (3, 3),
padding='same',
use_bias=False)
#l1 loss
self.loss = tf.reduce_mean(
tf.losses.absolute_difference(image_target, output))
self.out = tf.clip_by_value(output + mean_x, 0.0, 255.0)
#Calculating Peak Signal-to-noise-ratio
#Using equations from here: https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio
mse = tf.reduce_mean(tf.squared_difference(image_target, output))
PSNR = tf.constant(255**2, dtype=tf.float32) / mse
self.PSNR = tf.constant(10, dtype=tf.float32) * utils.log10(PSNR)
#Scalar to keep track for loss
tf.summary.scalar("loss", self.loss)
tf.summary.scalar("PSNR", self.PSNR)
#Image summaries for input, target, and output
tf.summary.image("input_image", tf.cast(self.input, tf.uint8))
tf.summary.image("target_image", tf.cast(self.target, tf.uint8))
tf.summary.image("output_image", tf.cast(self.out, tf.uint8))
#Tensorflow graph setup... session, saver, etc.
self.sess = tf.Session()
self.saver = tf.train.Saver()
print("Done building!")
def build_EDSR_WGAN_att(self): ###
"""
Build SRCNN model
"""
# Define input and label images
self.input = tf.placeholder(
tf.float32,
[None, self.image_size * 2, self.image_size * 2, self.color_dim],
name='images')
self.image_target = tf.placeholder(
tf.float32,
[None, self.image_size * 2, self.image_size * 2, self.color_dim],
name='labels')
self.curr_batch_size = tf.placeholder(tf.int32, shape=[])
self.dropout = tf.placeholder(tf.float32, name='dropout')
self.lr = tf.placeholder(tf.float32, name='learning_rate')
self.image_input = self.input / 255.
self.target = target = self.image_target / 255.
# Initial model_zoo
mz = model_zoo.model_zoo(self.image_input, self.dropout, self.is_train,
self.model_ticket)
### Build model
gen_f = mz.build_model({
"d_inputs": None,
"scale": self.scale,
"feature_size": 64,
"reuse": False,
"is_training": True,
"net": "Gen"
})
dis_t, lp_t = mz.build_model({
"d_inputs": self.target,
"scale": self.scale,
"feature_size": 64,
"reuse": False,
"is_training": True,
"net": "Dis",
"d_model": "PatchWGAN_GP"
})
dis_f, lp_f = mz.build_model({
"d_inputs": gen_f,
"scale": self.scale,
"feature_size": 64,
"reuse": True,
"is_training": True,
"net": "Dis",
"d_model": "PatchWGAN_GP"
})
#### WGAN-GP ####
# Calculate gradient penalty
self.epsilon = epsilon = tf.random_uniform(
[self.curr_batch_size, 1, 1, 1], 0.0, 1.0)
x_hat = epsilon * self.target + (1. - epsilon) * (gen_f)
d_hat = mz.build_model({
"d_inputs": x_hat,
"scale": self.scale,
"feature_size": 64,
"reuse": True,
"is_training": True,
"net": "Dis",
"d_model": "PatchWGAN_GP"
})
d_gp = tf.gradients(d_hat, [x_hat])[0]
d_gp = tf.sqrt(tf.reduce_sum(tf.square(d_gp), axis=[1, 2, 3]))
d_gp = tf.reduce_mean((d_gp - 1.0)**2) * 10
self.disc_ture_loss = disc_ture_loss = tf.reduce_mean(dis_t)
self.disc_fake_loss = disc_fake_loss = tf.reduce_mean(dis_f)
perceptual_loss = tf.pow(lp_t - lp_f, 2)
reconstucted_weight = 25.0
lp_weight = 0.25
self.d_loss = (disc_fake_loss - disc_ture_loss) + d_gp
self.g_l1loss = tf.reduce_mean(tf.pow(target - gen_f, 2))
self.g_loss = -1.0 * disc_fake_loss + reconstucted_weight * self.g_l1loss + lp_weight * perceptual_loss
train_variables = tf.trainable_variables()
generator_variables = [
v for v in train_variables if v.name.startswith("EDSR_gen")
]
discriminator_variables = [
v for v in train_variables if v.name.startswith("EDSR_dis")
]
self.train_d = tf.train.AdamOptimizer(
self.lr, beta1=0.5,
beta2=0.9).minimize(self.d_loss, var_list=discriminator_variables)
self.train_g = tf.train.AdamOptimizer(self.lr, beta1=0.5,
beta2=0.9).minimize(
self.g_loss,
var_list=generator_variables)
self.g_output = gen_f
mse = tf.reduce_mean(tf.squared_difference(target * 255.,
gen_f * 255.))
PSNR = tf.constant(255**2, dtype=tf.float32) / mse
PSNR = tf.constant(10, dtype=tf.float32) * log10(PSNR)
mse_ref = tf.reduce_mean(
tf.squared_difference(target * 255., self.image_input * 255.))
PSNR_ref = tf.constant(255**2, dtype=tf.float32) / mse_ref
PSNR_ref = tf.constant(10, dtype=tf.float32) * log10(PSNR_ref)
with tf.name_scope('train_summary'):
tf.summary.scalar("l1_loss", self.g_l1loss, collections=['train'])
tf.summary.scalar("d_loss", self.d_loss, collections=['train'])
tf.summary.scalar("d_true_loss",
disc_ture_loss,
collections=['train'])
tf.summary.scalar("d_fake_loss",
disc_fake_loss,
collections=['train'])
tf.summary.scalar("d_lp", perceptual_loss,
collections=['train']) # for tuning
tf.summary.scalar("MSE", mse, collections=['train'])
tf.summary.scalar("PSNR", PSNR, collections=['train'])
tf.summary.image("input_image", self.input, collections=['train'])
tf.summary.image("target_image",
target * 255,
collections=['train'])
tf.summary.image("output_image",
gen_f * 255,
collections=['train'])
tf.summary.histogram("d_false", dis_f, collections=['train'])
tf.summary.histogram("d_true", dis_t, collections=['train'])
self.merged_summary_train = tf.summary.merge_all('train')
with tf.name_scope('test_summary'):
tf.summary.scalar("loss", self.g_l1loss, collections=['test'])
tf.summary.scalar("d_loss", self.d_loss, collections=['test'])
tf.summary.scalar("g_loss", disc_fake_loss, collections=['test'])
tf.summary.scalar("d_lp", perceptual_loss,
collections=['test']) # for tuning
tf.summary.scalar("MSE", mse, collections=['test'])
tf.summary.scalar("PSNR", PSNR, collections=['test'])
tf.summary.scalar("PSNR_ref", PSNR_ref, collections=['test'])
tf.summary.image("input_image", self.input, collections=['test'])
tf.summary.image("target_image",
target * 255,
collections=['test'])
tf.summary.image("output_image", gen_f * 255, collections=['test'])
self.merged_summary_test = tf.summary.merge_all('test')
self.saver = tf.train.Saver()
self.best_saver = tf.train.Saver()
loss_tv = 2 * (x_tv / tv_x_size + y_tv / tv_y_size) / batch_size
# 5) ssim
loss_ssim = loss.Mssim_loss(dslr_image, enhanced)
# final loss
loss_generator = w_content * loss_content + w_texture * loss_texture + w_color * loss_color + w_tv * loss_tv + w_ssim * loss_ssim
# psnr loss
enhanced_flat = tf.reshape(enhanced, [-1, PATCH_SIZE])
loss_mse = tf.reduce_sum(tf.pow(dslr_ - enhanced_flat,
2)) / (PATCH_SIZE * batch_size)
loss_psnr = 20 * lutils.log10(1.0 / tf.sqrt(loss_mse))
# optimize parameters of image enhancement (generator) and discriminator networks
generator_vars = [
v for v in tf.global_variables() if v.name.startswith("generator")
]
discriminator_vars = [
v for v in tf.global_variables() if v.name.startswith("discriminator")
]
train_step_gen = tf.train.AdamOptimizer(learning_rate).minimize(
loss_generator, var_list=generator_vars)
train_step_disc = tf.train.AdamOptimizer(learning_rate).minimize(
loss_discrim, var_list=discriminator_vars)
def __init__(self):
print("Building MYSR...")
self.is_continued = False
self.imgsize = config.TRAIN.imgsize
self.output_channels = config.TRAIN.output_channels
self.scale = config.TRAIN.scale
self.epoch = config.TRAIN.n_epoch
self.batch_size = config.TRAIN.batch_size
self.save_model_dir = config.TRAIN.save_model_dir
self.test_input = config.TEST.hr_img_path
self.test_output = config.TEST.sr_img_path
# 初始化log文件
utils.log_message(config.VALID.log_file, "w+", "START")
utils.log_message(config.TEST.log_file, "w+", "START")
# Placeholder for image inputs
# self.input = x = tf.placeholder(tf.float32, [None, self.imgsize, self.imgsize, self.output_channels])
self.input = x = tf.placeholder(
tf.float32, [None, None, None, self.output_channels])
self.input_bicubic = tf.placeholder(
tf.float32, [None, None, None, self.output_channels])
# Placeholder for upscaled image ground-truth
# self.target = y = tf.placeholder(tf.float32, [None, self.imgsize*self.scale, self.imgsize*self.scale, self.output_channels])
self.target = y = tf.placeholder(
tf.float32, [None, None, None, self.output_channels])
# 输入预处理
# TODO: 后边有relu层，注意将输入图像收缩至[-1, 1]区间是否合适
image_input = x / (255. / 2.)
image_input -= 1
image_target = y / (255. / 2.)
image_target -= 1
image_input_bicubic = self.input_bicubic / (255. / 2.)
image_input_bicubic -= 1
# TODO:加载模型
output = modellib.BICUBIC(self, image_input, image_input_bicubic)
# output = modellib.SR_CNN(self, image_input, 3)
# output = modellib.SR_CNN(self, image_input, 27)
# output = modellib.VDSR_v1(self, image_input, image_input_bicubic, 64, 18)
# output = modellib.VDSR_v1_b(self, image_input, image_input_bicubic, 64, 18)
# output = modellib.SRDenseNetX2(self, image_input, 16, 8)
# output = modellib.SRDenseNetX4(self, image_input, 16, 8)
# output = modellib.MYSR_v6(self, image_input, 64, 16)
# output = modellib.MYSR_v5_Dense1(self, image_input, 64, 4, 16)
# output = modellib.MYSR_v5_Dense2(self, image_input, 64, 4, 16)
# output = modellib.MYSR_v5_3(self, image_input, image_input_bicubic, 64, 16)
# output = modellib.MYSR_v5_0(self, image_input, image_input_bicubic, 64, 16)
# output = modellib.MYSR_v5_0_b(self, image_input, image_input_bicubic, 64, 16)
# output = modellib.MYSR_v5_1(self, image_input, image_input_bicubic, 64, 16)
# output = modellib.MYSR_v5(self, image_input, 64, 16)
# output = modellib.MYSR_v5_b(self, image_input, image_input_bicubic, 64, 16)
# output = modellib.MYSR_v4(self, image_input, image_input_bicubic)
# output = modellib.EDSR_v1(self, image_input, 128, 16)
# output = modellib.EDSR_v1_b(self, image_input, 128, 16)
# output = modellib.WDSR_v1(self, image_input, 64, 16)
# 结果 注意预处理的值
# self.out = tf.clip_by_value(output+(255. / 2.), 0.0, 255.0)
self.out = tf.clip_by_value((output + 1) * (255. / 2.), 0.0, 255.0)
self.loss = loss = tf.reduce_mean(
tf.losses.absolute_difference(image_target, output))
# Calculating Peak Signal-to-noise-ratio
# Using equations from here: https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio
mse = tf.reduce_mean(
tf.squared_difference((image_target + 1) * (255. / 2.),
tf.clip_by_value((output + 1) * (255. / 2.),
0.0, 255.0)))
PSNR = tf.constant(255**2, dtype=tf.float32) / mse
self.PSNR = PSNR = tf.constant(10,
dtype=tf.float32) * utils.log10(PSNR)
# self.PSNR = tf.image.psnr((image_target+1)*(255. / 2.), self.out, max_val=255.)
SSIM = tf.image.ssim((image_target + 1) * (255. / 2.),
self.out,
max_val=255.)
self.SSIM = tf.reduce_mean(SSIM)
# Scalar to keep track for loss
tf.summary.scalar("loss", self.loss)
tf.summary.scalar("PSNR", PSNR)
tf.summary.scalar("SSIM", self.SSIM)
# Image summaries for input, target, and output
tf.summary.image("input_image", tf.cast(self.input, tf.uint8))
tf.summary.image("target_image", tf.cast(self.target, tf.uint8))
tf.summary.image("output_image", tf.cast(self.out, tf.uint8))
# Tensorflow graph setup... session, saver, etc.
config_tf = tf.ConfigProto()
config_tf.gpu_options.allow_growth = True
self.sess = tf.Session(config=config_tf)
self.saver = tf.train.Saver()
print("Done building!")
def __init__(self,img_size=32,num_layers=32,feature_size=128,scale=2,output_channels=1):
print("Building EDSR...")
self.img_size = img_size
self.scale = scale
self.output_channels = output_channels
#Placeholder for image inputs
self.input = x = tf.placeholder(tf.float32,[None,img_size,img_size,output_channels])
#Placeholder for upscaled image ground-truth
self.target = y = tf.placeholder(tf.float32,[None,img_size*scale,img_size*scale,output_channels])
"""
Preprocessing as mentioned in the paper, by subtracting the mean
However, the subtract the mean of the entire dataset they use. As of
now, I am subtracting the mean of each batch
"""
mean_x = tf.reduce_mean(self.input)
image_input =x- mean_x
mean_y = tf.reduce_mean(self.target)
image_target =y- mean_y
#One convolution before res blocks and to convert to required feature depth
x = slim.conv2d(image_input,feature_size,[3,3])
#Store the output of the first convolution to add later
conv_1 = x
"""
This creates `num_layers` number of resBlocks
a resBlock is defined in the paper as
(excuse the ugly ASCII graph)
x
|\
| \
| conv2d
| relu
| conv2d
| /
|/
+ (addition here)
|
result
"""
"""
Doing scaling here as mentioned in the paper:
`we found that increasing the number of feature
maps above a certain level would make the training procedure
numerically unstable. A similar phenomenon was
reported by Szegedy et al. We resolve this issue by
adopting the residual scaling with factor 0.1. In each
residual block, constant scaling layers are placed after the
last convolution layers. These modules stabilize the training
procedure greatly when using a large number of filters.
In the test phase, this layer can be integrated into the previous
convolution layer for the computational efficiency.'
"""
scaling_factor = 0.1
#Add the residual blocks to the model
for i in range(num_layers):
x = utils.resBlock(x,feature_size,scale=scaling_factor)
#One more convolution, and then we add the output of our first conv layer
x = slim.conv2d(x,feature_size,[3,3])
x += conv_1
#Upsample output of the convolution
x = utils.upsample(x,scale,feature_size,None)
#One final convolution on the upsampling output
output =x# slim.conv2d(x,output_channels,[3,3])
self.out = tf.clip_by_value(output+mean_x,0.0,255.0)
self.loss = loss = tf.reduce_mean(tf.losses.absolute_difference(image_target,output))
#Calculating Peak Signal-to-noise-ratio
#Using equations from here: https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio
mse = tf.reduce_mean(tf.squared_difference(image_target,output))
PSNR = tf.constant(255**2,dtype=tf.float32)/mse
PSNR = tf.constant(10,dtype=tf.float32)*utils.log10(PSNR)
#Scalar to keep track for loss
tf.summary.scalar("loss",self.loss)
tf.summary.scalar("PSNR",PSNR)
#Image summaries for input, target, and output
tf.summary.image("input_image",tf.cast(self.input,tf.uint8))
tf.summary.image("target_image",tf.cast(self.target,tf.uint8))
tf.summary.image("output_image",tf.cast(self.out,tf.uint8))
#Tensorflow graph setup... session, saver, etc.
self.sess = tf.Session()
self.saver = tf.train.Saver()
print("Done building!")
x_tv = tf.nn.l2_loss(enhanced[:,:,1:,:] - enhanced[:,:,:batch_shape[2]-1,:])
loss_tv = 2 * (x_tv/tv_x_size + y_tv/tv_y_size) / batch_size
else:
loss_tv = zero_
# final loss
loss_generator = w_content * loss_content + w_gray * loss_texture + w_color * loss_color + w_gradient * loss_gradient + w_tv * loss_tv + w_laplacian * loss_laplacian
loss_discrim = loss_discrim_color * w_color + loss_discrim_texture * w_gray + loss_discrim_gradient * w_gradient + loss_discrim_laplacian * w_laplacian
#loss for pre-training identity mapping
loss_generator_identity = tf.nn.l2_loss(phone_image - enhanced) / PATCH_SIZE + tf.nn.l2_loss(rec_image - enhanced) / PATCH_SIZE
# psnr score
enhanced_flat = tf.reshape(enhanced, [-1, PATCH_SIZE])
loss_mse = tf.reduce_sum(tf.pow(dslr_ - enhanced_flat, 2))/(PATCH_SIZE * batch_size)
psnr_score = 20 * utils.log10(1.0 / tf.sqrt(loss_mse))
# lpips score
perceptual_score= tf.reduce_mean(lpips(enhanced, dslr_image, model='net-lin', net='alex'))
loss_perceptual_color = tf.reduce_mean(lpips(enhanced_blur, dslr_blur, model='net-lin', net='alex'))
loss_perceptual_texture = tf.reduce_mean(lpips(tf.tile(enhanced_gray, [1, 1, 1, 3]), tf.tile(dslr_gray, [1, 1, 1, 3]), model='net-lin', net='alex'))
# VGG score
enhanced_vgg = vgg.net(vgg_dir, vgg.preprocess(enhanced * 255))
dslr_vgg = vgg.net(vgg_dir, vgg.preprocess(dslr_image * 255))
content_size = utils._tensor_size(enhanced_vgg[CONTENT_LAYER]) * batch_size
VGG_loss = 2 * tf.nn.l2_loss(enhanced_vgg[CONTENT_LAYER] - dslr_vgg[CONTENT_LAYER]) / content_size
# optimize parameters of image enhancement (generator) and discriminator networks
generator_vars = [v for v in tf.global_variables() if v.name.startswith("generator")]
score = 0.0
file_path = "./vaeganCeleba2/sample_psnr/"
ori_list, gen_list = read_image_list(file_path)
ori_list.sort(compare)
gen_list.sort(compare)
print len(ori_list)
for i in range(len(ori_list)):
with tf.gfile.FastGFile(ori_list[i]) as image_file:
img1_str = image_file.read()
with tf.gfile.FastGFile(gen_list[i]) as image_file:
img2_str = image_file.read()
input_img = tf.placeholder(tf.string)
decoded_image = tf.expand_dims(tf.image.decode_png(input_img, channels=3),
0)
with tf.Session() as sess:
img1 = sess.run(decoded_image, feed_dict={input_img: img1_str})
img2 = sess.run(decoded_image, feed_dict={input_img: img2_str})
mse_loss = tf.reduce_mean(tf.cast((img1 - img2)**2, tf.float32))
print "mse_loss", sess.run(mse_loss)
print "psnr", sess.run(20 * log10(255 / tf.sqrt(mse_loss)))
score += sess.run(20 * log10(255 / tf.sqrt(mse_loss)))
print "result", score / len(ori_list)
def __init__(self,
img_size=32,
num_layers=32,
feature_size=256,
scale=2,
output_channels=3):
print("Building EDSR...")
self.img_size = img_size
self.scale = scale
self.output_channels = output_channels
#Placeholder for image inputs
self.input = x = tf.placeholder(tf.float32,
[None, None, None, output_channels],
name='crzin')
#Placeholder for upscaled image ground-truth
self.target = y = tf.placeholder(tf.float32,
[None, None, None, output_channels])
"""
Preprocessing as mentioned in the paper, by subtracting the mean
However, the subtract the mean of the entire dataset they use. As of
now, I am subtracting the mean of each batch
"""
#mean_x = 127#tf.reduce_mean(self.input)
image_input = x
#mean_y = 127#tf.reduce_mean(self.target)
image_target = y
x = slim.conv2d(image_input, 64, [3, 3], padding="SAME")
x = slim.conv2d(x, 64, [3, 3], padding="SAME")
x = slim.conv2d(x, 32, [3, 3], padding="SAME")
x = slim.conv2d_transpose(x, 3, [3, 3], 2, padding="SAME")
#x = slim.conv2d(x,64,[5,5],stride=1,activation_fn=tf.nn.tanh,padding = "SAME")
#print("x.shap",x.get_shape())
#x = slim.conv2d(x,64,[3,3],stride=1,activation_fn=tf.nn.tanh,padding = "SAME")
#print("x.shap",x.get_shape())
#x = slim.conv2d(x,32,[3,3],stride=1,activation_fn=tf.nn.tanh,padding = "SAME")
#print("x.shap",x.get_shape())
#x = slim.conv2d_transpose(x,3,[4,4],stride=2,activation_fn=tf.nn.tanh,padding = "SAME")
#print("x.shap",x.get_shape())
output = x #slim.conv2d(x,output_channels,[3,3])
#self.out = tf.clip_by_value(output+mean_x,0.0,255.0)
self.out = tf.clip_by_value(output, 0.0, 255.0, name='crzout')
#self.loss = loss = tf.reduce_mean(tf.losses.absolute_difference(image_target,output))
self.loss = loss = tf.reduce_mean(tf.square(image_target - output))
#Calculating Peak Signal-to-noise-ratio
#Using equations from here: https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio
mse = tf.reduce_mean(tf.squared_difference(image_target, output))
PSNR = tf.constant(255**2, dtype=tf.float32) / mse
PSNR = tf.constant(10, dtype=tf.float32) * utils.log10(PSNR)
#Scalar to keep track for loss
tf.summary.scalar("loss", self.loss)
tf.summary.scalar("PSNR", PSNR)
#Image summaries for input, target, and output
tf.summary.image("input_image", tf.cast(self.input, tf.uint8))
tf.summary.image("target_image", tf.cast(self.target, tf.uint8))
tf.summary.image("output_image", tf.cast(self.out, tf.uint8))
#Tensorflow graph setup... session, saver, etc.
self.sess = tf.Session()
self.saver = tf.train.Saver()
print("Done building!")
def __init__(self):
print("Building MYSR...")
self.imgsize = config.TRAIN.imgsize
self.output_channels = config.TRAIN.output_channels
self.scale = config.TRAIN.scale
# Placeholder for image inputs
self.input = x = tf.placeholder(
tf.float32, [None, None, None, self.output_channels])
# Placeholder for upscaled image ground-truth
self.target = y = tf.placeholder(
tf.float32, [None, None, None, self.output_channels])
# 输入预处理
# result = result / (255. / 2.)
# TODO: 后边有relu层，注意将输入图像收缩至[-1, 1]区间是否合适
# result = result - 1.
image_input = x / (255. / 2.)
image_input = image_input - 1
image_target = y / (255. / 2.)
image_target = image_target - 1
# 貌似即使收缩至[-1, 1]区间，卷积层依旧可以有效适应，只是注意最后一层不能使用relu，b毕竟relu值域在[0, x]
# ENCODER
# 入口
x = slim.conv2d(image_input, 64, [5, 5]) # 入口适当大点？
conv_1 = x
# ENCODER-resBlock-64
scaling_factor = 1
for i in range(3):
x = utils.resBlock(x, 64, scale=scaling_factor)
x = slim.conv2d(image_input, 128, [3, 3])
# ENCODER-resBlock-128
scaling_factor = 1
for i in range(4):
x = utils.resBlock(x, 128, scale=scaling_factor)
x = slim.conv2d(image_input, 256, [3, 3])
# ENCODER-resBlock-256
scaling_factor = 1
for i in range(5):
x = utils.resBlock(x, 256, scale=scaling_factor)
# Upsample output of the convolution
x = utils.upsample(x, self.scale, 128, None)
# DECODER-resBlock-64
scaling_factor = 0.1
for i in range(4):
x = utils.resBlock(x, 64, scale=scaling_factor)
# DECODER-resBlock-32
scaling_factor = 0.1
for i in range(3):
x = utils.resBlock(x, 32, scale=scaling_factor)
# DECODER-resBlock-16
scaling_factor = 0.1
for i in range(2):
x = utils.resBlock(x, 16, scale=scaling_factor)
# DECODER-resBlock-8
scaling_factor = 0.1
for i in range(1):
x = utils.resBlock(x, 8, scale=scaling_factor)
# DECODER-输出
# TODO: 貌似这里直接使用conv会破坏逐步修复结构？反而会因为缺少feature_map而导致精细度降低？
x = slim.conv2d(x, self.output_channels, [3, 3])
output = x
# 结果
self.out = tf.clip_by_value((output + 1) * (255. / 2.), 0.0, 255.0)
self.loss = loss = tf.reduce_mean(
tf.losses.absolute_difference(image_target, output))
# Calculating Peak Signal-to-noise-ratio
# Using equations from here: https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio
mse = tf.reduce_mean(tf.squared_difference(image_target, output))
PSNR = tf.constant(255**2, dtype=tf.float32) / mse
PSNR = tf.constant(10, dtype=tf.float32) * utils.log10(PSNR)
# Scalar to keep track for loss
tf.summary.scalar("loss", self.loss)
tf.summary.scalar("PSNR", PSNR)
# Image summaries for input, target, and output
tf.summary.image("input_image", tf.cast(self.input, tf.uint8))
tf.summary.image("target_image", tf.cast(self.target, tf.uint8))
tf.summary.image("output_image", tf.cast(self.out, tf.uint8))
# Tensorflow graph setup... session, saver, etc.
self.sess = tf.Session()
self.saver = tf.train.Saver()
print("Done building!")
tv_y_size = utils._tensor_size(enhanced[:,1:,:,:])
tv_x_size = utils._tensor_size(enhanced[:,:,1:,:])
y_tv = tf.nn.l2_loss(enhanced[:,1:,:,:] - enhanced[:,:batch_shape[1]-1,:,:])
x_tv = tf.nn.l2_loss(enhanced[:,:,1:,:] - enhanced[:,:,:batch_shape[2]-1,:])
loss_tv = 2 * (x_tv/tv_x_size + y_tv/tv_y_size) / batch_size
# final loss
loss_generator = w_content * loss_content + w_texture * loss_texture + w_color * loss_color + w_tv * loss_tv
# psnr loss
enhanced_flat = tf.reshape(enhanced, [-1, PATCH_SIZE])
loss_mse = tf.reduce_mean(tf.squared_difference(dslr_,enhanced_flat))
loss_psnr = tf.constant(1.0,dtype=tf.float32)/loss_mse
loss_psnr = tf.constant(10,dtype=tf.float32)*utils.log10(loss_psnr)#loss_mse = tf.reduce_sum(tf.pow(dslr_ - enhanced_flat, 2))/(PATCH_SIZE * batch_size)
#loss_psnr = 20 * utils.log10(1.0 / tf.sqrt(loss_mse))
#Scalar to keep track for loss
tf.summary.scalar("loss_content",loss_content)
tf.summary.scalar("loss_texture",loss_texture)
tf.summary.scalar("loss_tv",loss_tv)
tf.summary.scalar("PSNR",loss_psnr)
# optimize parameters of image enhancement (generator) and discriminator networks
generator_vars = [v for v in tf.global_variables() if v.name.startswith("generator")]
discriminator_vars = [v for v in tf.global_variables() if v.name.startswith("discriminator")]
train_step_gen = tf.train.AdamOptimizer(learning_rate).minimize(loss_generator, var_list=generator_vars)
def __init__(self,img_size=32,num_layers=32,feature_size=256,scale=2,output_channels=3):
print("Building EDSR...")
#Placeholder for image inputs
self.input = x = tf.placeholder(tf.float32,[None,img_size,img_size,output_channels])
#Placeholder for upscaled image ground-truth
self.target = y = tf.placeholder(tf.float32,[None,img_size*scale,img_size*scale,output_channels])
"""
Preprocessing as mentioned in the paper, by subtracting the mean
However, the subtract the mean of the entire dataset they use. As of
now, I am subtracting the mean of each batch
"""
mean_x = tf.reduce_mean(self.input)
image_input =x- mean_x
mean_y = tf.reduce_mean(self.target)
image_target =y- mean_y
#One convolution before res blocks and to convert to required feature depth
x = slim.conv2d(image_input,feature_size,[3,3])
#Store the output of the first convolution to add later
conv_1 = x
"""
This creates `num_layers` number of resBlocks
a resBlock is defined in the paper as
(excuse the ugly ASCII graph)
x
|\
| \
| conv2d
| relu
| conv2d
| /
|/
+ (addition here)
|
result
"""
"""
Doing scaling here as mentioned in the paper:
`we found that increasing the number of feature
maps above a certain level would make the training procedure
numerically unstable. A similar phenomenon was
reported by Szegedy et al. We resolve this issue by
adopting the residual scaling with factor 0.1. In each
residual block, constant scaling layers are placed after the
last convolution layers. These modules stabilize the training
procedure greatly when using a large number of filters.
In the test phase, this layer can be integrated into the previous
convolution layer for the computational efficiency.'
"""
scaling_factor = 0.1
#Add the residual blocks to the model
for i in range(num_layers):
x = utils.resBlock(x,feature_size,scale=scaling_factor)
#One more convolution, and then we add the output of our first conv layer
x = slim.conv2d(x,feature_size,[3,3])
x += conv_1
#Upsample output of the convolution
x = utils.upsample(x,scale,feature_size,None)
#One final convolution on the upsampling output
output =x# slim.conv2d(x,output_channels,[3,3])
self.out = tf.clip_by_value(output+mean_x,0.0,255.0)
self.loss = loss = tf.reduce_mean(tf.losses.absolute_difference(image_target,output))
#Calculating Peak Signal-to-noise-ratio
#Using equations from here: https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio
mse = tf.reduce_mean(tf.squared_difference(image_target,output))
PSNR = tf.constant(255**2,dtype=tf.float32)/mse
PSNR = tf.constant(10,dtype=tf.float32)*utils.log10(PSNR)
#Scalar to keep track for loss
tf.summary.scalar("loss",self.loss)
tf.summary.scalar("PSNR",PSNR)
#Image summaries for input, target, and output
tf.summary.image("input_image",tf.cast(self.input,tf.uint8))
tf.summary.image("target_image",tf.cast(self.target,tf.uint8))
tf.summary.image("output_image",tf.cast(self.out,tf.uint8))
#Tensorflow graph setup... session, saver, etc.
self.sess = tf.Session()
self.saver = tf.train.Saver()
print("Done building!")
def build_EDSR_WGAN(self): ###
"""
Build SRCNN model
"""
# Define input and label images
self.input = tf.placeholder(
tf.float32,
[None, self.image_size * 2, self.image_size * 2, self.color_dim],
name='images')
self.image_target = tf.placeholder(
tf.float32,
[None, self.image_size * 2, self.image_size * 2, self.color_dim],
name='labels')
self.curr_batch_size = tf.placeholder(tf.int32, shape=[])
self.dropout = tf.placeholder(tf.float32, name='dropout')
self.lr = tf.placeholder(tf.float32, name='learning_rate')
self.image_input = self.input / 255.
self.target = target = self.image_target / 255.
# Initial model_zoo
mz = model_zoo.model_zoo(self.image_input, self.dropout, self.is_train,
self.model_ticket)
### Build model
gen_f = mz.build_model({
"d_inputs": None,
"d_target": self.target,
"scale": self.scale,
"feature_size": 64,
"reuse": False,
"is_training": True,
"net": "Gen"
})
dis_t = mz.build_model({
"d_inputs": self.target,
"d_target": self.target,
"scale": self.scale,
"feature_size": 64,
"reuse": False,
"is_training": True,
"net": "Dis",
"d_model": "PatchWGAN"
})
dis_f = mz.build_model({
"d_inputs": gen_f,
"d_target": self.target,
"scale": self.scale,
"feature_size": 64,
"reuse": True,
"is_training": True,
"net": "Dis",
"d_model": "PatchWGAN"
})
"""
#### WGAN-GP ####
# Calculate gradient penalty
self.epsilon = epsilon = tf.random_uniform([self.curr_batch_size, 1, 1, 1], 0.0, 1.0)
x_hat = epsilon * self.target + (1. - epsilon) * (gen_f)
d_hat = mz.build_model({"d_inputs":x_hat,"d_target":self.target,"scale":self.scale,"feature_size" :64, "reuse":True, "is_training":True, "net":"Dis", "d_model":"PatchWGAN_GP"})
d_gp = tf.gradients(d_hat, [x_hat])[0]
d_gp = tf.sqrt(tf.reduce_sum(tf.square(d_gp), axis=[1,2,3]))
d_gp = tf.reduce_mean((d_gp - 1.0)**2) * 10
self.disc_ture_loss = disc_ture_loss = tf.reduce_mean(dis_t)
self.disc_fake_loss = disc_fake_loss = tf.reduce_mean(dis_f)
# W(P_data, P_G) = min{ E_x~P_G[D(x)] - E_x~P_data[D(x)] + lamda*E_x~P_penalty[(D'(x)-1)^2] } ~ max{ V(G,D) }
self.d_loss = disc_fake_loss - disc_ture_loss + d_gp
# Generator loss
reconstucted_weight = 1.0
self.g_l1loss = tf.reduce_mean(tf.losses.absolute_difference(target, gen_f))
self.g_loss = -1.0*disc_fake_loss + reconstucted_weight*self.g_l1loss
train_variables = tf.trainable_variables()
generator_variables = [v for v in train_variables if v.name.startswith("EDSR_gen")]
discriminator_variables = [v for v in train_variables if v.name.startswith("EDSR_dis")]
self.train_d = tf.train.AdamOptimizer(self.lr, beta1=0.5, beta2=0.9).minimize(self.d_loss, var_list=discriminator_variables)
self.train_g = tf.train.AdamOptimizer(self.lr, beta1=0.5, beta2=0.9).minimize(self.g_loss, var_list=generator_variables)
"""
######## WGAN #######
self.disc_ture_loss = disc_ture_loss = tf.reduce_mean(dis_t)
disc_fake_loss = tf.reduce_mean(dis_f)
# reconstucted_weight = 1.0 #StarGAN is 10 v3
reconstucted_weight = 50 #StarGAN is 10
self.d_loss = disc_fake_loss - disc_ture_loss
self.g_l1loss = tf.reduce_mean(
tf.losses.absolute_difference(target, gen_f))
self.g_loss = -1.0 * disc_fake_loss + reconstucted_weight * self.g_l1loss
train_variables = tf.trainable_variables()
generator_variables = [
v for v in train_variables if v.name.startswith("EDSR_gen")
]
discriminator_variables = [
v for v in train_variables if v.name.startswith("EDSR_dis")
]
self.clip_discriminator_var_op = [
var.assign(tf.clip_by_value(var, -0.01, 0.01))
for var in discriminator_variables
]
alpha = 0.00005
optimizer = tf.train.RMSPropOptimizer(self.lr)
gvs_d = optimizer.compute_gradients(self.d_loss,
var_list=discriminator_variables)
gvs_g = optimizer.compute_gradients(self.g_loss,
var_list=generator_variables)
gvs_g_l1 = optimizer.compute_gradients(self.g_l1loss,
var_list=generator_variables)
gvs_g_dis = optimizer.compute_gradients(-1.0 * disc_fake_loss,
var_list=generator_variables)
wgvs_d = [(grad * alpha, var) for grad, var in gvs_d]
wgvs_g = [(grad * alpha, var) for grad, var in gvs_g]
self.train_d = optimizer.apply_gradients(wgvs_d)
self.train_g = optimizer.apply_gradients(wgvs_g)
self.g_output = gen_f
mse = tf.reduce_mean(tf.squared_difference(target * 255.,
gen_f * 255.))
PSNR = tf.constant(255**2, dtype=tf.float32) / mse
PSNR = tf.constant(10, dtype=tf.float32) * log10(PSNR)
mse_ref = tf.reduce_mean(
tf.squared_difference(target * 255., self.image_input * 255.))
PSNR_ref = tf.constant(255**2, dtype=tf.float32) / mse_ref
PSNR_ref = tf.constant(10, dtype=tf.float32) * log10(PSNR_ref)
with tf.name_scope('train_summary'):
tf.summary.scalar("l1_loss", self.g_l1loss, collections=['train'])
tf.summary.scalar("d_loss", self.d_loss, collections=['train'])
tf.summary.scalar("d_true_loss",
disc_ture_loss,
collections=['train'])
tf.summary.scalar("d_fake_loss",
disc_fake_loss,
collections=['train'])
# tf.summary.scalar("grad_loss", d_gp, collections=['train'])
tf.summary.scalar("dis_f_mean",
tf.reduce_mean(dis_f),
collections=['train'])
tf.summary.scalar("dis_t_mean",
tf.reduce_mean(dis_t),
collections=['train'])
tf.summary.scalar("MSE", mse, collections=['train'])
tf.summary.scalar("PSNR", PSNR, collections=['train'])
tf.summary.image("input_image", self.input, collections=['train'])
tf.summary.image("target_image",
target * 255,
collections=['train'])
tf.summary.image("output_image",
gen_f * 255,
collections=['train'])
# tf.summary.image("enhence_img",(2.0*gen_f-target)*255, collections=['train'])
# tf.summary.image("dis_f_img",100*dis_f*255, collections=['train'])
# tf.summary.image("dis_t_img",100*dis_t*255, collections=['train'])
# tf.summary.image("dis_diff",tf.abs(dis_f-dis_t)*255, collections=['train'])
tf.summary.histogram("d_false", dis_f, collections=['train'])
tf.summary.histogram("d_true", dis_t, collections=['train'])
idx = 0
for grad, var in gvs_g_l1:
if idx < 3:
tf.summary.histogram(var.name + '/gradient_l1',
grad,
collections=['train'])
idx += 1
else:
break
idx = 0
for grad, var in gvs_g_dis:
if idx < 3:
tf.summary.histogram(var.name + '/gradient_dis',
grad,
collections=['train'])
idx += 1
else:
break
self.merged_summary_train = tf.summary.merge_all('train')
with tf.name_scope('test_summary'):
tf.summary.scalar("loss", self.g_l1loss, collections=['test'])
tf.summary.scalar("d_loss", self.d_loss, collections=['test'])
tf.summary.scalar("g_loss", disc_fake_loss, collections=['test'])
tf.summary.scalar("MSE", mse, collections=['test'])
tf.summary.scalar("PSNR", PSNR, collections=['test'])
tf.summary.scalar("PSNR_ref", PSNR_ref, collections=['test'])
tf.summary.image("input_image", self.input, collections=['test'])
tf.summary.image("target_image",
target * 255,
collections=['test'])
tf.summary.image("output_image", gen_f * 255, collections=['test'])
self.merged_summary_test = tf.summary.merge_all('test')
self.saver = tf.train.Saver()
