def tv_loss(prediction, batch_size): batch_shape = (batch_size, PATCH_WIDTH, PATCH_HEIGHT, 3) tv_y_size = utils._tensor_size(prediction[:, 1:, :, :]) tv_x_size = utils._tensor_size(prediction[:, :, 1:, :]) y_tv = tf.nn.l2_loss(prediction[:, 1:, :, :] - prediction[:, :batch_shape[1] - 1, :, :]) x_tv = tf.nn.l2_loss(prediction[:, :, 1:, :] - prediction[:, :, :batch_shape[2] - 1, :]) loss_tv = 2 * (x_tv / tv_x_size + y_tv / tv_y_size) / batch_size return loss_tv
def content_loss(target, prediction, batch_size): CONTENT_LAYER = 'relu5_4' vgg_dir = '/gdata/huangjie/hdrnet/vgg_pretrained/imagenet-vgg-verydeep-19.mat' enhanced_vgg = vgg.net(vgg_dir, vgg.preprocess(prediction * 255)) dslr_vgg = vgg.net(vgg_dir, vgg.preprocess(target * 255)) content_size = utils._tensor_size(dslr_vgg[CONTENT_LAYER]) * batch_size loss_content = 2 * tf.nn.l2_loss(enhanced_vgg[CONTENT_LAYER] - dslr_vgg[CONTENT_LAYER]) / content_size return tf.reduce_mean(loss_content)
def content_loss(target, prediction, batch_size): CONTENT_LAYER = 'relu5_4' CONTENT_LAYER1 = 'relu3_4' CONTENT_LAYER2 = 'relu1_2' vgg_dir = './vgg_pretrained/imagenet-vgg-verydeep-19.mat' enhanced_vgg = vgg.net(vgg_dir, vgg.preprocess(prediction * 255)) dslr_vgg = vgg.net(vgg_dir, vgg.preprocess(target * 255)) content_size = utils._tensor_size(dslr_vgg[CONTENT_LAYER]) * batch_size content_size1 = utils._tensor_size(dslr_vgg[CONTENT_LAYER1]) * batch_size content_size2 = utils._tensor_size(dslr_vgg[CONTENT_LAYER2]) * batch_size loss_content = 2 * tf.nn.l2_loss(enhanced_vgg[CONTENT_LAYER] - dslr_vgg[CONTENT_LAYER]) / content_size loss_content1 = 2 * tf.nn.l2_loss(enhanced_vgg[CONTENT_LAYER1] - dslr_vgg[CONTENT_LAYER1]) / content_size1 loss_content2 = 2 * tf.nn.l2_loss(enhanced_vgg[CONTENT_LAYER2] - dslr_vgg[CONTENT_LAYER2]) / content_size2 return (tf.reduce_mean(loss_content) + tf.reduce_mean(loss_content1) + tf.reduce_mean(loss_content2)) / 3
loss_content = loss.content_loss(dslr_image, enhanced, batch_size) # 3) color loss enhanced_blur = lutils.blur(enhanced) dslr_blur = lutils.blur(dslr_image) #loss_color = tf.reduce_sum(tf.pow(dslr_blur - enhanced_blur, 2))/(2 * batch_size) loss_color = tf.reduce_sum( tf.abs(dslr_image - enhanced)) / (2 * batch_size) #loss_color = loss.color_loss(dslr_image, enhanced, batch_size) # 4) total variation loss batch_shape = (batch_size, PATCH_WIDTH, PATCH_HEIGHT, 3) tv_y_size = lutils._tensor_size(enhanced[:, 1:, :, :]) tv_x_size = lutils._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 # 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
# PSNR loss loss_psnr = 20 * utils.log10(1.0 / tf.sqrt(loss_mse)) # SSIM loss loss_ssim = tf.reduce_mean(tf.image.ssim(enhanced, dslr_, 1.0)) # MS-SSIM loss loss_ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(enhanced, dslr_, 1.0)) # Content loss CONTENT_LAYER = 'relu5_4' enhanced_vgg = vgg.net(vgg_dir, vgg.preprocess(enhanced * 255)) dslr_vgg = vgg.net(vgg_dir, vgg.preprocess(dslr_ * 255)) content_size = utils._tensor_size(dslr_vgg[CONTENT_LAYER]) * batch_size loss_content = 2 * tf.nn.l2_loss(enhanced_vgg[CONTENT_LAYER] - dslr_vgg[CONTENT_LAYER]) / content_size # Final loss function loss_generator = loss_mse * 20 + loss_content + (1 - loss_ssim) * 20 # Optimize network parameters generator_vars = [v for v in tf.compat.v1.global_variables() if v.name.startswith("generator")] train_step_gen = tf.compat.v1.train.AdamOptimizer(learning_rate).minimize(loss_generator, var_list=generator_vars) # Initialize and restore the variables print("Initializing variables...") sess.run(tf.compat.v1.global_variables_initializer()) saver = tf.compat.v1.train.Saver(var_list=generator_vars, max_to_keep=100)
discrim_target = tf.concat([adv_, 1 - adv_], 1) loss_discrim = -tf.reduce_sum(discrim_target * tf.log(tf.clip_by_value(discrim_predictions, 1e-10, 1.0))) loss_texture = -loss_discrim correct_predictions = tf.equal(tf.argmax(discrim_predictions, 1), tf.argmax(discrim_target, 1)) discim_accuracy = tf.reduce_mean(tf.cast(correct_predictions, tf.float32)) # 2) content loss CONTENT_LAYER = 'relu5_4' 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(dslr_vgg[CONTENT_LAYER]) * batch_size loss_content = 2 * tf.nn.l2_loss(enhanced_vgg[CONTENT_LAYER] - dslr_vgg[CONTENT_LAYER]) / content_size # content_size = num * height * weidth (of feature map) # tf.nn.l2_loss(a) = sum(a**2)/2 # 3) color loss enhanced_blur = utils.blur(enhanced) dslr_blur = utils.blur(dslr_image) loss_color = tf.reduce_sum(tf.pow(dslr_blur - enhanced_blur, 2)) / (2 * batch_size) # 4) total variation loss batch_shape = (batch_size, PATCH_WIDTH, PATCH_HEIGHT, 3) tv_y_size = utils._tensor_size(enhanced[:, 1:, :, :])
def perceptual_loss(gen_hq_vgg, hq_vgg, VGG_LAYER, function, batch_size): content_size = utils._tensor_size(hq_vgg[VGG_LAYER]) * batch_size loss_vgg = tf.reduce_sum( tf.pow(gen_hq_vgg[VGG_LAYER] - hq_vgg[VGG_LAYER], 2)) / content_size return loss_vgg
# content loss loss_content = ssim_loss + cx_loss # 3) color loss enhanced_blur = utils.blur(enhanced) dslr_blur = utils.blur(dslr_image) loss_color = tf.reduce_sum(tf.pow(dslr_blur - enhanced_blur, 2)) / (2 * batch_size) # 4) total variation loss batch_shape = (batch_size, PATCH_WIDTH, PATCH_HEIGHT, 3) tv_y_size = utils._tensor_size(enhanced[:, 1:, :, :]) # H tv_x_size = utils._tensor_size(enhanced[:, :, 1:, :]) # W 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 evaluation enhanced_flat = tf.reshape(enhanced, [-1, PATCH_SIZE])
loss_ssim = tf.reduce_mean(tf.image.ssim(bokeh_img, target_, 1.0)) # MS-SSIM loss loss_ms_ssim = tf.reduce_mean( tf.image.ssim_multiscale(bokeh_img, target_, 1.0)) # L1 loss loss_l1 = tf.compat.v1.losses.absolute_difference(bokeh_img, target_) # Content loss CONTENT_LAYER = 'relu5_4' bokeh_img_vgg = vgg.net(vgg_dir, vgg.preprocess(bokeh_img * 255)) target_vgg = vgg.net(vgg_dir, vgg.preprocess(target_ * 255)) content_size = utils._tensor_size(target_vgg[CONTENT_LAYER]) * batch_size loss_content = 2 * tf.nn.l2_loss(bokeh_img_vgg[CONTENT_LAYER] - target_vgg[CONTENT_LAYER]) / content_size # Final loss function if LEVEL > 1: loss_generator = loss_l1 * 100 else: loss_generator = loss_l1 * 10 + loss_content * 0.1 + (1 - loss_ssim) * 10 # Optimize network parameters generator_vars = [ v for v in tf.compat.v1.global_variables()