def main(): parser = argparse.ArgumentParser() parser.add_argument('--model', choices=['net-lin', 'net'], default='net-lin', help='net-lin or net') parser.add_argument('--net', choices=['squeeze', 'alex', 'vgg'], default='alex', help='squeeze, alex, or vgg') parser.add_argument('--version', type=str, default='0.1') args = parser.parse_args() ex_ref = load_image('./PerceptualSimilarity/imgs/ex_ref.png') ex_p0 = load_image('./PerceptualSimilarity/imgs/ex_p0.png') ex_p1 = load_image('./PerceptualSimilarity/imgs/ex_p1.png') session = tf.Session() image0_ph = tf.placeholder(tf.float32) image1_ph = tf.placeholder(tf.float32) lpips_fn = session.make_callable( lpips_tf.lpips(image0_ph, image1_ph, model=args.model, net=args.net, version=args.version), [image0_ph, image1_ph]) ex_d0 = lpips_fn(ex_ref, ex_p0) ex_d1 = lpips_fn(ex_ref, ex_p1) print('Distances: (%.3f, %.3f)' % (ex_d0, ex_d1))
def image_lpips(y_true, y_pred): y_true = join_reim_mag_output(y_true) y_pred = join_reim_mag_output(y_pred) y_true = K.tile(y_true, [1, 1, 1, 3]) y_pred = K.tile(y_pred, [1, 1, 1, 3]) return lpips_tf.lpips(y_true, y_pred, model='net-lin', net='alex')
def lpips(input0, input1): if input0.shape[-1].value == 1: input0 = tf.tile(input0, [1] * (input0.shape.ndims - 1) + [3]) if input1.shape[-1].value == 1: input1 = tf.tile(input1, [1] * (input1.shape.ndims - 1) + [3]) distance = lpips_tf.lpips(input0, input1) return -distance
def image_lpips(y_true, y_pred): y_true_image = utils.convert_tensor_to_image_domain(y_true) y_pred_image = utils.convert_tensor_to_image_domain(y_pred) y_true = join_reim_mag_output(y_true_image) y_pred = join_reim_mag_output(y_pred_image) y_true = K.tile(y_true, [1, 1, 1, 3]) y_pred = K.tile(y_pred, [1, 1, 1, 3]) return lpips_tf.lpips(y_true, y_pred, model='net-lin', net='alex')
def loss(self, y_true, y_pred): """Compute loss.""" y_true_rgb = tf.concat(3 * [y_true], axis=-1) y_pred_rgb = tf.concat(3 * [y_pred], axis=-1) return lpips_tf.lpips(y_true_rgb, y_pred_rgb, model=self.model, net=self.net)
def get_alexnet(inp, otp): # https://stackoverflow.com/questions/38376478/changing-the-scale-of-a-tensor-in-tensorflow with tf.name_scope('perceptual_loss'): axis = None inp_normalized = tf.div( tf.subtract(inp, tf.reduce_min(inp, axis=axis)), tf.subtract(tf.reduce_max(inp, axis=axis), tf.reduce_min(inp, axis=axis))) otp_normalized = tf.div( tf.subtract(otp, tf.reduce_min(otp, axis=axis)), tf.subtract(tf.reduce_max(otp, axis=axis), tf.reduce_min(otp, axis=axis))) permutation = [0, 2, 3, 1] inp_normalized = tf.transpose(inp_normalized, permutation) otp_normalized = tf.transpose(otp_normalized, permutation) # pad_val = (224-160)//2 # paddings = tf.constant([[0,0],[pad_val,pad_val],[pad_val,pad_val],[0,0]]) # ksizes = [1, 64, 64, 1] # strides = [1, 32, 32, 1] # rates = [1,1,1,1] # padding = 'SAME' # patches = tf.image.extract_image_patches(inp_normalized, ksizes, strides, rates, padding) # patches_shape = tf.shape(patches) # h = tf.shape(inp_normalized)[1] # w = tf.shape(inp_normalized)[2] # c = tf.shape(inp_normalized)[3] # patches = tf.reshape(patches, [tf.reduce_prod(patches_shape[0:3]), h, w, int(c)]) # inp_normalized = tf.pad(inp_normalized, paddings, 'CONSTANT') # otp_normalized = tf.pad(otp_normalized, paddings, 'CONSTANT') translator = model_translator(inp_normalized[:, :64, :64, :], otp_normalized[:, :64, :64, :], network='alexnet') translator.get_weights() # model_translator.get_alexnet_diff(inp_normalized[:,:64,:64,:], otp_normalized[:,:64,:64,:]) # img_content_normalized = (inp - np.min(inp)) / ( # np.max(inp) - np.min(inp)) # img_content_normalized = np.transpose(np.expand_dims(img_content_normalized, axis=0), # [0, 2, 3, 1]) # out_img_content_normalized = (otp - np.min(otp)) / ( # np.max(otp) - np.min(otp)) # print(out_img_content_normalized.shape) # out_img_content_normalized = np.transpose(out_img_content_normalized, # [0, 2, 3, 1]) distance_t = lpips_tf.lpips(inp_normalized, otp_normalized, model='net-lin', net='vgg', data_format='NCHW') return distance_t
def perceptual_loss_img(y_true, y_pred, model="net-lin", net="vgg"): """Compute the perceptual loss (PL) between two images. Parameters ---------- y_true : np.array Image 1. Either (h, w) of (N, h, w). If (N, h, w), the decorator `multiple_images_decorator` takes care of the sample dimension. y_pred : np.array Image 2. Either (h, w) of (N, h, w). If (N, h, w), the decorator `multiple_images_decorator` takes care of the sample dimension. model: str, {'net', 'net-lin'} Type of model (cf lpips_tf package). net: str, {'vgg', 'alex'} Type of network (cf lpips_tf package). Return ------ pl: float The Perceptual Loss (PL) metric. Loss metric, the lower the more similar the images are. Notes ----- We use the decorator just to make sure we do not run out of memory during a forward pass. Also, its fully convolutional but if the images are too small then might run into issues. """ # gray2rgb y_true = np.stack((y_true, ) * 3, axis=-1) y_pred = np.stack((y_pred, ) * 3, axis=-1) image0_ph = tf.placeholder(tf.float32) image1_ph = tf.placeholder(tf.float32) distance_t = lpips_tf.lpips(image0_ph, image1_ph, model=model, net=net) with tf.Session() as session: pl = session.run(distance_t, feed_dict={ image0_ph: y_true, image1_ph: y_pred }) tf.reset_default_graph() return pl.item()
def recontrruction_error(self, batch_size=32): image_shape = (batch_size, 64, 64, 3) image0 = np.random.random(image_shape) image1 = np.random.random(image_shape) image0_ph = tf.placeholder(tf.float32) image1_ph = tf.placeholder(tf.float32) distance_t = lpips_tf.lpips(image0_ph, image1_ph, model='net-lin', net='alex') to_batch = lambda x: np.concatenate( [np.tile(imgs, (1, 1, 1, 3)) for imgs in x], axis=0) total_distance = None for i in range(self.n_images): generated_imgs = to_batch([ self.converter.generator.predict( [self.content_codes[[j]], self.class_adain_params[[i]]])[0] for j in range(self.n_images) if i != j ]) same_content_imgs = to_batch( [self.curr_imgs[j] for j in range(self.n_images) if i != j]) with tf.Session() as session: distance = session.run( distance_t, feed_dict={ image0_ph: generated_imgs[idx * batch_size:(idx + 1) * batch_size], image1_ph: same_content_imgs[idx * batch_size:(idx + 1) * batch_size] }) total_distance = np.mean(distance) # for idx in range(int(np.ceil((self.n_images-1) / batch_size))): # distance = session.run(distance_t, feed_dict={image0_ph: generated_imgs[idx * batch_size:(idx + 1) * batch_size], # image1_ph: same_content_imgs[idx * batch_size:(idx + 1) * batch_size]}) # if total_distance is None: # total_distance = np.mean(distance) # else: # total_distance = ((total_distance*idx * batch_size) + np.sum(distance)) / ((idx * batch_size) + len(distance)) return total_distance
def lpips_distance(model, x, attack_name=None, adv_x=None, norm=2, distance_only=False, run_parallel=0): import tensorflow as tf if adv_x is None: adv_x = craft_attack(model, x, attack_name, norm) batch_size = 32 image_shape = (batch_size, x.shape[1], x.shape[2], x.shape[3]) image0_ph = tf.placeholder(tf.float32) image1_ph = tf.placeholder(tf.float32) distance_t = lpips_tf.lpips(image0_ph, image1_ph, model='net-lin', net='alex', model_dir="./metrics") def f(i): with tf.Session() as session: distance = session.run(distance_t, feed_dict={ image0_ph: x[i], image1_ph: adv_x[i] }) if distance_only: return distance else: return x[i], adv_x[i], distance if run_parallel > 0: with Pool(run_parallel) as p: return p.map(f, range(len(x))) else: for i in range(len(x)): yield f(i)
def main(): args, save_dir, load_dir = check_args(parse_arguments()) global BATCH_SIZE BATCH_SIZE = args.batch_size config_path = os.path.join(load_dir, 'params.pkl') if os.path.exists(config_path): config = pickle.load(open(config_path, 'rb')) output_width = config['output_width'] output_height = config['output_height'] resolution = output_height z_dim = config['z_dim'] else: output_width = output_height = 64 resolution = 64 z_dim = 100 ### open session run_config = tf.ConfigProto() run_config.gpu_options.allow_growth = True with tf.Session(config=run_config) as sess: dcgan = DCGAN(sess, output_width=output_width, output_height=output_height, batch_size=BATCH_SIZE, sample_num=BATCH_SIZE, z_dim=z_dim) load_success, load_counter = dcgan.load(load_dir) if not load_success: raise Exception("Checkpoint not found in " + load_dir) ### initialization init_val_ph = None init_val = {'pos': None, 'neg': None} if args.initialize_type == 'zero': z = tf.Variable(tf.zeros([BATCH_SIZE, z_dim], tf.float32), name='latent_z') elif args.initialize_type == 'random': np.random.seed(RANDOM_SEED) init_val_np = np.random.normal(size=(z_dim, )) init = np.tile(init_val_np, (BATCH_SIZE, 1)).astype(np.float32) z = tf.Variable(init, name='latent_z') elif args.initialize_type == 'nn': idx = 0 init_val['pos'] = np.load(os.path.join(args.nn_dir, 'pos_z.npy'))[:, idx, :] init_val['neg'] = np.load(os.path.join(args.nn_dir, 'neg_z.npy'))[:, idx, :] init_val_ph = tf.placeholder(dtype=tf.float32, name='init_ph', shape=(BATCH_SIZE, z_dim)) z = tf.Variable(init_val_ph, name='latent_z') else: raise NotImplementedError ### define variables x = tf.placeholder(tf.float32, shape=(BATCH_SIZE, resolution, resolution, 3)) x_hat = dcgan.generator(z, is_training=False) ### loss if args.distance == 'l2': print('use distance: l2') loss_l2 = tf.reduce_mean(tf.square(x_hat - x), axis=[1, 2, 3]) vec_loss = loss_l2 vec_losses = {'l2': loss_l2} elif args.distance == 'l2-lpips': print('use distance: lpips + l2') loss_l2 = tf.reduce_mean(tf.square(x_hat - x), axis=[1, 2, 3]) loss_lpips = lpips_tf.lpips(x_hat, x, normalize=False, model='net-lin', net='vgg', version='0.1') vec_losses = {'l2': loss_l2, 'lpips': loss_lpips} vec_loss = loss_l2 + LAMBDA2 * loss_lpips else: raise NotImplementedError ## regularizer norm = tf.reduce_sum(tf.square(z), axis=1) norm_penalty = (norm - z_dim)**2 if args.if_norm_reg: loss = tf.reduce_mean( vec_loss) + LAMBDA3 * tf.reduce_mean(norm_penalty) vec_losses['norm'] = norm_penalty else: loss = tf.reduce_mean(vec_loss) ### set up optimizer opt = tf.contrib.opt.ScipyOptimizerInterface( loss, var_list=[z], method='L-BFGS-B', options={'maxfun': args.maxfunc}) ### load query images pos_data_paths = get_filepaths_from_dir(args.pos_data_dir, ext='png')[:args.data_num] pos_query_imgs = np.array( [read_image(f, resolution) for f in pos_data_paths]) neg_data_paths = get_filepaths_from_dir(args.neg_data_dir, ext='png')[:args.data_num] neg_query_imgs = np.array( [read_image(f, resolution) for f in neg_data_paths]) ### run the optimization on query images query_loss, query_z, query_xhat = optimize_z( sess, z, x, x_hat, init_val_ph, init_val['pos'], pos_query_imgs, check_folder(os.path.join(save_dir, 'pos_results')), opt, vec_loss, vec_losses) save_files(save_dir, ['pos_loss'], [query_loss]) query_loss, query_z, query_xhat = optimize_z( sess, z, x, x_hat, init_val_ph, init_val['neg'], neg_query_imgs, check_folder(os.path.join(save_dir, 'neg_results')), opt, vec_loss, vec_losses) save_files(save_dir, ['neg_loss'], [query_loss])
import numpy as np import tensorflow as tf import lpips_tf, glob, os, scipy def read_color_img(filename): return scipy.misc.imread(filename).astype('float32') / 255 batch_size = 1 image_shape = (batch_size, None, None, 3) image0_ph = tf.placeholder(tf.float32) image1_ph = tf.placeholder(tf.float32) distance_t = lpips_tf.lpips(image0_ph, image1_ph, model='net-lin', net='alex') folder = 'E://DepthOfField//test_vis//sr_x4//rgb' gt_lists = glob.glob(os.path.join(folder, '*HR.png')) distance = [] with tf.Session() as sess: for i_file in range(len(gt_lists)): gt_name = gt_lists[i_file] img_name = gt_name.replace('HR.png', 'LR4_rcan_real.png') label = read_color_img(gt_name) output = read_color_img(img_name) d = sess.run(distance_t, feed_dict={ image0_ph: label, image1_ph: output }) distance.append(d)
def compute_metric(dir1, dir2, mode, mask=None, thre=0): if mode == 'perceptual': from models import dist_model as dm import torch from util import util global model if model is None: cwd = os.getcwd() os.chdir('../../PerceptualSimilarity') model = dm.DistModel() #model.initialize(model='net-lin',net='alex',use_gpu=True, spatial=True) model.initialize(model='net-lin', net='alex', use_gpu=True) print('Model [%s] initialized' % model.name()) os.chdir(cwd) if mode.startswith('perceptual_tf'): sys.path += ['../../lpips-tensorflow'] import lpips_tf import tensorflow as tf image0_ph = tf.placeholder(tf.float32, [1, None, None, 3]) image1_ph = tf.placeholder(tf.float32, [1, None, None, 3]) if mode == 'perceptual_tf': distance_t = lpips_tf.lpips(image0_ph, image1_ph, model='net-lin', net='alex') elif mode == 'perceptual_tf_vgg': distance_t = lpips_tf.lpips(image0_ph, image1_ph, model='net-lin', net='vgg') else: raise sess = tf.Session() if mode == 'l2_with_gradient': import demo import tensorflow as tf output = tf.placeholder(tf.float32, shape=[1, None, None, None]) gradient = demo.image_gradients(output) sess = tf.Session() files1 = os.listdir(dir1) files2 = os.listdir(dir2) img_files1 = sorted([ file for file in files1 if file.endswith('.png') or file.endswith('.jpg') ]) img_files2 = sorted([ file for file in files2 if file.endswith('.png') or file.endswith('.jpg') ]) if '--prefix' in sys.argv: prefix_idx = sys.argv.index('--prefix') prefix = sys.argv[prefix_idx + 1] img_files2 = [file for file in img_files2 if file.startswith(prefix)] if mask is not None: mask_files = [] for dir in sorted(mask): files3 = os.listdir(dir) add_mask_files = sorted([ os.path.join(dir, file) for file in files3 if file.startswith('mask') ]) if len(add_mask_files) == 0: files_to_dilate = sorted([ file for file in files3 if file.startswith('g_intermediates') ]) for file in files_to_dilate: mask_arr = numpy.load(os.path.join(dir, file)) mask_arr = numpy.squeeze(mask_arr) mask_arr = mask_arr >= thre dilated_mask = dilation(mask_arr, disk(10)) dilated_filename = 'mask_' + file numpy.save(os.path.join(dir, dilated_filename), dilated_mask.astype('f')) add_mask_files.append(os.path.join(dir, dilated_filename)) mask_files += add_mask_files print(mask_files) assert len(mask_files) == len(img_files1) skip_first_n = 0 if '--skip_first_n' in sys.argv: try: skip_first_n = int(sys.argv[sys.argv.index('--skip_first_n') + 1]) except: skip_first_n = 0 skip_last_n = 0 if '--skip_last_n' in sys.argv: try: skip_last_n = int(sys.argv[sys.argv.index('--skip_last_n') + 1]) except: skip_last_n = 0 img_files2 = img_files2[skip_first_n:] if skip_last_n > 0: img_files2 = img_files2[:-skip_last_n] assert len(img_files1) == len(img_files2) # locate GT gradient directory if mode == 'l2_with_gradient': head, tail = os.path.split(dir2) gradient_gt_dir = os.path.join(head, tail[:-3] + 'grad') if not os.path.exists(gradient_gt_dir): printf("dir not found,", gradient_gt_dir) raise gradient_gt_files = os.listdir(gradient_gt_dir) gradient_gt_files = sorted( [file for file in gradient_gt_files if file.endswith('.npy')]) assert len(img_files1) == len(gradient_gt_files) vals = numpy.empty(len(img_files1)) #if mode == 'perceptual': # global model for ind in range(len(img_files1)): if mode == 'ssim' or mode == 'l2' or mode == 'l2_with_gradient': img1 = skimage.img_as_float( skimage.io.imread(os.path.join(dir1, img_files1[ind]))) img2 = skimage.img_as_float( skimage.io.imread(os.path.join(dir2, img_files2[ind]))) if mode == 'ssim': #vals[ind] = skimage.measure.compare_ssim(img1, img2, datarange=img2.max()-img2.min(), multichannel=True) metric_val = skimage.measure.compare_ssim( img1, img2, datarange=img2.max() - img2.min(), multichannel=True) else: #vals[ind] = numpy.mean((img1 - img2) ** 2) * 255.0 * 255.0 metric_val = ((img1 - img2)**2) * 255.0 * 255.0 if mode == 'l2_with_gradient': metric_val = numpy.mean(metric_val, axis=2) gradient_gt = numpy.load( os.path.join(gradient_gt_dir, gradient_gt_files[ind])) dx, dy = sess.run(gradient, feed_dict={ output: numpy.expand_dims(img1[..., ::-1], axis=0) }) #is_edge = skimage.feature.canny(skimage.color.rgb2gray(img1)) dx_ground = gradient_gt[:, :, :, 1:4] dy_ground = gradient_gt[:, :, :, 4:] edge_ground = gradient_gt[:, :, :, 0] gradient_loss_term = numpy.mean( (dx - dx_ground)**2.0 + (dy - dy_ground)**2.0, axis=3) metric_val += numpy.squeeze( 0.2 * 255.0 * 255.0 * gradient_loss_term * edge_ground * edge_ground.size / numpy.sum(edge_ground)) #if mode == 'l2' and mask is not None: # img_diff = (img1 - img2) ** 2.0 # mask_img = numpy.load(mask_files[ind]) # img_diff *= numpy.expand_dims(mask_img, axis=2) # vals[ind] = (numpy.sum(img_diff) / numpy.sum(mask_img * 3)) * 255.0 * 255.0 elif mode == 'perceptual': img1 = util.im2tensor( util.load_image(os.path.join(dir1, img_files1[ind]))) img2 = util.im2tensor( util.load_image(os.path.join(dir2, img_files2[ind]))) #vals[ind] = numpy.mean(model.forward(img1, img2)[0]) metric_val = numpy.expand_dims(model.forward(img1, img2), axis=2) elif mode.startswith('perceptual_tf'): img1 = np.expand_dims(skimage.img_as_float( skimage.io.imread(os.path.join(dir1, img_files1[ind]))), axis=0) img2 = np.expand_dims(skimage.img_as_float( skimage.io.imread(os.path.join(dir2, img_files2[ind]))), axis=0) metric_val = sess.run(distance_t, feed_dict={ image0_ph: img1, image1_ph: img2 }) else: raise if mask is not None: assert mode in ['l2', 'perceptual'] mask_img = numpy.load(mask_files[ind]) metric_val *= numpy.expand_dims(mask_img, axis=2) vals[ind] = numpy.sum(metric_val) / (numpy.sum(mask_img) * metric_val.shape[2]) else: vals[ind] = numpy.mean(metric_val) mode = mode + ('_mask' if mask is not None else '') filename_all = mode + '_all.txt' filename_breakdown = mode + '_breakdown.txt' filename_single = mode + '.txt' numpy.savetxt(os.path.join(dir1, filename_all), vals, fmt="%f, ") target = open(os.path.join(dir1, filename_single), 'w') target.write("%f" % numpy.mean(vals)) target.close() if len(img_files1) == 30: target = open(os.path.join(dir1, filename_breakdown), 'w') target.write("%f, %f, %f" % (numpy.mean( vals[:5]), numpy.mean(vals[5:10]), numpy.mean(vals[10:]))) target.close() if mode in ['l2_with_gradient', 'perceptual_tf']: sess.close() return vals
def validate(val_dirs: ValidationDirs, images_iterator: ImagesIterator, flags: OutputFlags): """ Saves in val_dirs.log_dir/val/dataset_name/measures.csv: - `img_name,bpp,psnr,ms-ssim forall img_name` """ print(_VALIDATION_INFO_STR) validated_checkpoints = val_dirs.get_validated_checkpoints( ) # :: [10000, 18000, ..., 256000], ie, [int] all_ckpts = Saver.all_ckpts_with_iterations(val_dirs.ckpt_dir) if len(all_ckpts) == 0: print('No checkpoints found in {}'.format(val_dirs.ckpt_dir)) return # if ckpt_step is -1, then all_ckpt[:-1:flags.ckpt_step] === [] because of how strides work ckpt_to_check = all_ckpts[:-1:flags.ckpt_step] + [ all_ckpts[-1] ] # every ckpt_step-th checkpoint plus the last one if flags.ckpt_step == -1: assert len(ckpt_to_check) == 1 print('Validating {}/{} checkpoints (--ckpt_step {})...'.format( len(ckpt_to_check), len(all_ckpts), flags.ckpt_step)) missing_checkpoints = [(ckpt_itr, ckpt_path) for ckpt_itr, ckpt_path in ckpt_to_check if ckpt_itr not in validated_checkpoints] if len(missing_checkpoints) == 0: print('All checkpoints validated, stopping...') return # --- # create networks autoencoder_config_path, probclass_config_path = logdir_helpers.config_paths_from_log_dir( val_dirs.log_dir, base_dirs=[constants.CONFIG_BASE_AE, constants.CONFIG_BASE_PC]) ae_config, ae_config_rel_path = config_parser.parse( autoencoder_config_path) pc_config, pc_config_rel_path = config_parser.parse(probclass_config_path) ae_cls = autoencoder.get_network_cls(ae_config) pc_cls = probclass.get_network_cls(pc_config) # Instantiate autoencoder and probability classifier ae = ae_cls(ae_config) pc = pc_cls(pc_config, num_centers=ae_config.num_centers) x_val_ph = tf.placeholder(tf.uint8, (3, None, None), name='x_val_ph') x_val_uint8 = tf.expand_dims(x_val_ph, 0, name='batch') x_val = tf.to_float(x_val_uint8, name='x_val') x_val_normalized = tf.div( tf.subtract(x_val, tf.reduce_min(x_val)), tf.subtract(tf.reduce_max(x_val), tf.reduce_min(x_val))) enc_out_val = ae.encode(x_val, is_training=False) x_out_val = ae.decode(enc_out_val.qhard, is_training=False) x_out_val_normalized = tf.div( tf.subtract(x_out_val, tf.reduce_min(x_out_val)), tf.subtract(tf.reduce_max(x_out_val), tf.reduce_min(x_out_val))) bc_val = pc.bitcost(enc_out_val.qbar, enc_out_val.symbols, is_training=False, pad_value=pc.auto_pad_value(ae)) bpp_val = bits.bitcost_to_bpp(bc_val, x_val) x_out_val_uint8 = tf.cast(x_out_val, tf.uint8, name='x_out_val_uint8') # Using numpy implementation due to dynamic shapes msssim_val = ms_ssim_np.tf_msssim_np(x_val_uint8, x_out_val_uint8, data_format='NCHW') psnr_val = psnr_np(x_val_uint8, x_out_val_uint8) restorer = Saver(val_dirs.ckpt_dir, var_list=Saver.get_var_list_of_ckpt_dir( val_dirs.ckpt_dir)) # create fetch_dict fetch_dict = { 'bpp': bpp_val, 'ms-ssim': msssim_val, 'psnr': psnr_val, } if flags.real_bpp: fetch_dict['sym'] = enc_out_val.symbols # NCHW if flags.save_ours: fetch_dict['img_out'] = x_out_val_uint8 # --- fw = tf.summary.FileWriter(val_dirs.out_dir, graph=tf.get_default_graph()) def full_summary_tag(summary_name): return '/'.join(['val', images_iterator.dataset_name, summary_name]) # Distance try: codec_distance_ms_ssim = CodecDistance(images_iterator.dataset_name, codec='bpg', metric='ms-ssim') codec_distance_psnr = CodecDistance(images_iterator.dataset_name, codec='bpg', metric='psnr') except CodecDistanceReadException as e: # no codec distance values stored for the current setup print('*** Distance to BPG not available for {}:\n{}'.format( images_iterator.dataset_name, str(e))) codec_distance_ms_ssim = None codec_distance_psnr = None # Note that for each checkpoint, the structure of the network will be the same. Thus the pad depending image # loading can be cached. lpips_ph1 = tf.placeholder(tf.float32) lpips_ph2 = tf.placeholder(tf.float32) distance_t = lpips_tf.lpips(lpips_ph1, lpips_ph2, model='net-lin', net='alex') distance_t = tf.Print(distance_t, [distance_t]) # create session with tf_helpers.create_session() as sess: if flags.real_bpp: pred = probclass.PredictionNetwork(pc, pc_config, ae.get_centers_variable(), sess) checker = probclass.ProbclassNetworkTesting(pc, ae, sess) bpp_fetcher = bpp_helpers.BppFetcher(pred, checker) fetcher = sess.make_callable(fetch_dict, feed_list=[x_val_ph]) last_ckpt_itr = missing_checkpoints[-1][0] for ckpt_itr, ckpt_path in missing_checkpoints: if not ckpt_still_exists(ckpt_path): # May happen if job is still training print('Checkpoint disappeared: {}'.format(ckpt_path)) continue print(_CKPT_ITR_INFO_STR.format(ckpt_itr)) restorer.restore_ckpt(sess, ckpt_path) values_aggregator = ValuesAggregator('bpp', 'ms-ssim', 'psnr') # truncates the previous measures.csv file! This way, only the last valid checkpoint is saved. measures_writer = MeasuresWriter(val_dirs.out_dir) # ---------------------------------------- # iterate over images # images are padded to work with current auto encoder for img_i, (img_name, img_content) in enumerate( images_iterator.iter_imgs( pad=ae.get_subsampling_factor())): otp = fetcher(img_content) measures_writer.append(img_name, otp) img_content_normalized = ( img_content - np.min(img_content)) / (np.max(img_content) - np.min(img_content)) img_content_normalized = np.transpose( np.expand_dims(img_content_normalized, axis=0), [0, 2, 3, 1]) out_img_content_normalized = (otp['img_out'] - np.min( otp['img_out'])) / (np.max(otp['img_out']) - np.min(otp['img_out'])) #print(out_img_content_normalized.shape) out_img_content_normalized = np.transpose( out_img_content_normalized, [0, 2, 3, 1]) sess.run(distance_t, feed_dict={ lpips_ph1: img_content_normalized, lpips_ph2: out_img_content_normalized }) if flags.real_bpp: # Calculate bpp_real, bpp_theory = bpp_fetcher.get_bpp( otp['sym'], bpp_helpers.num_pixels_in_image(img_content)) # Logging bpp_loss = otp['bpp'] diff_percent_tr = (bpp_theory / bpp_real) * 100 diff_percent_lt = (bpp_loss / bpp_theory) * 100 print('BPP: Real {:.5f}\n' ' Theoretical: {:.5f} [{:5.1f}% of real]\n' ' Loss: {:.5f} [{:5.1f}% of real]'.format( bpp_real, bpp_theory, diff_percent_tr, bpp_loss, diff_percent_lt)) assert abs( bpp_theory - bpp_loss ) < 1e-3, 'Expected bpp_theory to match loss! Got {} and {}'.format( bpp_theory, bpp_loss) if flags.save_ours and ckpt_itr == last_ckpt_itr: save_img(img_name, otp['img_out'], val_dirs) values_aggregator.update(otp) print('{: 10d} {img_name} | Mean: {avgs}'.format( img_i, img_name=img_name, avgs=values_aggregator.averages_str()), end=('\r' if not flags.real_bpp else '\n'), flush=True) measures_writer.close() print() # add newline avgs = values_aggregator.averages() avg_bpp, avg_ms_ssim, avg_psnr = avgs['bpp'], avgs[ 'ms-ssim'], avgs['psnr'] tf_helpers.log_values( fw, [(full_summary_tag('avg_bpp'), avg_bpp), (full_summary_tag('avg_ms_ssim'), avg_ms_ssim), (full_summary_tag('avg_psnr'), avg_psnr)], iteration=ckpt_itr) if codec_distance_ms_ssim and codec_distance_psnr: try: d_ms_ssim = codec_distance_ms_ssim.distance( avg_bpp, avg_ms_ssim) d_pnsr = codec_distance_psnr.distance(avg_bpp, avg_psnr) print('Distance to BPG: {:.3f} ms-ssim // {:.3f} psnr'. format(d_ms_ssim, d_pnsr)) tf_helpers.log_values( fw, [(full_summary_tag('distance_BPG_MS-SSIM'), d_ms_ssim), (full_summary_tag('distance_BPG_PSNR'), d_pnsr)], iteration=ckpt_itr) except ValueError as e: # out of range errors from distance calls print(e) val_dirs.add_validated_checkpoint(ckpt_itr) print('Validation completed {}'.format(val_dirs))
def main(): args, save_dir, load_dir = check_args(parse_arguments()) ### open session config = tf.ConfigProto() config.gpu_options.allow_growth = True with tf.Session(config=config) as sess: sess.as_default() ### load pre-trained model network_path = sorted(glob(os.path.join(load_dir, 'network-*.pkl')))[-1] with open(network_path, 'rb') as file: print('Loading networks from "%s"...' % network_path) G, D, Gs = pickle.load(file) Gs.print_layers() D.print_layers() ### define variables global BATCH_SIZE BATCH_SIZE = args.batch_size z_dim = G.input_shape[-1] resolution = G.output_shape[-1] x = tf.placeholder(tf.float32, shape=(BATCH_SIZE, resolution, resolution, 3)) labels = tf.zeros([BATCH_SIZE, 0], tf.float32) ### initialization init_val_ph = None init_val = {'pos': None, 'neg': None} if args.initialize_type == 'zero': z = tf.Variable(tf.zeros([BATCH_SIZE, z_dim], tf.float32), name='latent_z') elif args.initialize_type == 'random': np.random.seed(RANDOM_SEED) init_val_np = np.random.normal(size=(z_dim, )) init_val_np = init_val_np / np.sqrt( np.mean(np.square(init_val_np)) + 1e-8) init = np.tile(init_val_np, (BATCH_SIZE, 1)).astype(np.float32) z = tf.Variable(init, name='latent_z') elif args.initialize_type == 'nn': init_val = {} init_val['pos'] = np.load(os.path.join(args.nn_dir, 'pos_z.npy'))[:, 0, :] init_val['neg'] = np.load(os.path.join(args.nn_dir, 'neg_z.npy'))[:, 0, :] init_val_ph = tf.placeholder(dtype=tf.float32, name='init_ph', shape=(BATCH_SIZE, z_dim)) z = tf.Variable(init_val_ph, name='latent_z') else: raise NotImplementedError ### get the reconstruction (x_hat) with tf.variable_scope(Gs.scope, reuse=True): assert tf.get_variable_scope().name == Gs.scope with tf.control_dependencies( None): # ignore surrounding control_dependencies inputs = [z, labels] x_hat = Gs._build_func(*inputs, is_template_graph=True, **Gs.static_kwargs) x_hat = tf.transpose(x_hat, perm=[0, 2, 3, 1]) x_hat = tf.clip_by_value(x_hat, -1., 1.) ### loss if args.distance == 'l2': print('Use distance: l2') loss_l2 = tf.reduce_mean(tf.square(x_hat - x), axis=[1, 2, 3]) vec_loss = loss_l2 vec_losses = {'l2': loss_l2} elif args.distance == 'l2-lpips': print('Use distance: lpips + l2') loss_l2 = tf.reduce_mean(tf.square(x_hat - x), axis=[1, 2, 3]) loss_lpips = lpips_tf.lpips(x_hat, x, normalize=False, model='net-lin', net='vgg', version='0.1') vec_losses = {'l2': loss_l2, 'lpips': loss_lpips} vec_loss = loss_l2 + LAMBDA2 * loss_lpips else: raise NotImplementedError ## regularizer norm = tf.reduce_sum(tf.square(z), axis=1) norm_penalty = (norm - z_dim)**2 if args.if_norm_reg: loss = tf.reduce_mean( vec_loss) + LAMBDA3 * tf.reduce_mean(norm_penalty) vec_losses['norm'] = norm_penalty else: loss = tf.reduce_mean(vec_loss) ### set up optimizer opt = tf.contrib.opt.ScipyOptimizerInterface( loss, var_list=[z], method='L-BFGS-B', options={'maxfun': args.maxfunc}) ### load query images pos_data_paths = get_filepaths_from_dir(args.pos_data_dir, ext='png')[:args.data_num] pos_query_imgs = np.array( [read_image(f, resolution) for f in pos_data_paths]) neg_data_paths = get_filepaths_from_dir(args.neg_data_dir, ext='png')[:args.data_num] neg_query_imgs = np.array( [read_image(f, resolution) for f in neg_data_paths]) ### run the optimization on query images query_loss, query_z, query_xhat = optimize_z( sess, z, x, x_hat, init_val_ph, init_val['pos'], pos_query_imgs, check_folder(os.path.join(save_dir, 'pos_results')), opt, vec_loss, vec_losses) save_files(save_dir, ['pos_loss'], [query_loss]) query_loss, query_z, query_xhat = optimize_z( sess, z, x, x_hat, init_val_ph, init_val['neg'], neg_query_imgs, check_folder(os.path.join(save_dir, 'neg_results')), opt, vec_loss, vec_losses) save_files(save_dir, ['neg_loss'], [query_loss])
image2_ph = tf.compat.v1.placeholder(tf.float32) image3_ph = tf.compat.v1.placeholder(tf.float32) image4_ph = tf.compat.v1.placeholder(tf.float32) image5_ph = tf.compat.v1.placeholder(tf.float32) image6_ph = tf.compat.v1.placeholder(tf.float32) image7_ph = tf.compat.v1.placeholder(tf.float32) image8_ph = tf.compat.v1.placeholder(tf.float32) image9_ph = tf.compat.v1.placeholder(tf.float32) image10_ph = tf.compat.v1.placeholder(tf.float32) image11_ph = tf.compat.v1.placeholder(tf.float32) image12_ph = tf.compat.v1.placeholder(tf.float32) image13_ph = tf.compat.v1.placeholder(tf.float32) image14_ph = tf.compat.v1.placeholder(tf.float32) image15_ph = tf.compat.v1.placeholder(tf.float32) distance_t1 = lpips_tf.lpips(image0_ph, image1_ph, model='net-lin', net='alex') distance_t2 = lpips_tf.lpips(image0_ph, image2_ph, model='net-lin', net='alex') distance_t3 = lpips_tf.lpips(image0_ph, image3_ph, model='net-lin', net='alex') distance_t4 = lpips_tf.lpips(image0_ph, image4_ph, model='net-lin', net='alex') distance_t5 = lpips_tf.lpips(image0_ph, image5_ph, model='net-lin', net='alex') distance_t6 = lpips_tf.lpips(image0_ph, image6_ph, model='net-lin', net='alex') distance_t7 = lpips_tf.lpips(image0_ph, image7_ph, model='net-lin', net='alex') distance_t8 = lpips_tf.lpips(image0_ph, image8_ph, model='net-lin', net='alex') distance_t9 = lpips_tf.lpips(image0_ph, image9_ph, model='net-lin', net='alex') distance_t10 = lpips_tf.lpips(image0_ph, image10_ph, model='net-lin', net='alex') distance_t11 = lpips_tf.lpips(image0_ph, image11_ph, model='net-lin',
def main(): args, save_dir, load_dir = check_args(parse_arguments()) config_path = os.path.join(load_dir, 'params.pkl') if os.path.exists(config_path): config = pickle.load(open(config_path, 'rb')) OUTPUT_SIZE = config['OUTPUT_SIZE'] GAN_TYPE = config['Architecture'] Z_DIM = config['Z_DIM'] else: OUTPUT_SIZE = 64 GAN_TYPE = 'good' Z_DIM = 128 ### set up the generator and the discriminator Generator, Discriminator = GeneratorAndDiscriminator(GAN_TYPE) ### open session config = tf.ConfigProto() config.gpu_options.allow_growth = True with tf.Session(config=config) as sess: ### define variables global BATCH_SIZE BATCH_SIZE = args.batch_size x = tf.placeholder(tf.float32, shape=(None, OUTPUT_SIZE, OUTPUT_SIZE, 3), name='x') ### initialization init_val_ph = None init_val = {'pos': None, 'neg': None} if args.initialize_type == 'zero': z = tf.Variable(tf.zeros([BATCH_SIZE, Z_DIM], tf.float32), name='latent_z') elif args.initialize_type == 'random': np.random.seed(RANDOM_SEED) init_val_np = np.random.normal(size=(Z_DIM, )) init = np.tile(init_val_np, (BATCH_SIZE, 1)).astype(np.float32) z = tf.Variable(init, name='latent_z') elif args.initialize_type == 'nn': init_val['pos'] = np.load(os.path.join(args.nn_dir, 'pos_z.npy'))[:, 0, :] init_val['neg'] = np.load(os.path.join(args.nn_dir, 'neg_z.npy'))[:, 0, :] init_val_ph = tf.placeholder(dtype=tf.float32, name='init_ph', shape=(BATCH_SIZE, Z_DIM)) z = tf.Variable(init_val_ph, name='latent_z') else: raise NotImplementedError ### get the reconstruction (x_hat) x_hat = Generator(BATCH_SIZE, noise=z, is_training=False, z_dim=Z_DIM) x_hat = tf.reshape(x_hat, [-1, 3, OUTPUT_SIZE, OUTPUT_SIZE]) x_hat = tf.transpose(x_hat, perm=[0, 2, 3, 1]) ### load model vars = [v for v in tf.global_variables() if 'latent_z' not in v.name] saver = tf.train.Saver(vars) sess.run(tf.initialize_variables(vars)) if_load, counter = load_model_from_checkpoint(load_dir, saver, sess) assert if_load is True ### loss if args.distance == 'l2': print('use distance: l2') loss_l2 = tf.reduce_mean(tf.square(x_hat - x), axis=[1, 2, 3]) vec_loss = loss_l2 vec_losses = {'l2': loss_l2} elif args.distance == 'l2-lpips': print('use distance: lpips + l2') loss_l2 = tf.reduce_mean(tf.square(x_hat - x), axis=[1, 2, 3]) loss_lpips = lpips_tf.lpips(x_hat, x, normalize=False, model='net-lin', net='vgg', version='0.1') vec_losses = {'l2': loss_l2, 'lpips': loss_lpips} vec_loss = loss_l2 + LAMBDA2 * loss_lpips else: raise NotImplementedError ## regularizer norm = tf.reduce_sum(tf.square(z), axis=1) norm_penalty = (norm - Z_DIM)**2 if args.if_norm_reg: loss = tf.reduce_mean( vec_loss) + LAMBDA3 * tf.reduce_mean(norm_penalty) vec_losses['norm'] = norm_penalty else: loss = tf.reduce_mean(vec_loss) ### set up optimizer opt = tf.contrib.opt.ScipyOptimizerInterface( loss, var_list=[z], method='Powell', options={'maxiter': args.maxiter}) ### load query images pos_data_paths = get_filepaths_from_dir(args.pos_data_dir, ext='png')[:args.data_num] pos_query_imgs = np.array( [read_image(f, OUTPUT_SIZE) for f in pos_data_paths]) neg_data_paths = get_filepaths_from_dir(args.neg_data_dir, ext='png')[:args.data_num] neg_query_imgs = np.array( [read_image(f, OUTPUT_SIZE) for f in neg_data_paths]) ### run the optimization on query images query_loss, query_z, query_xhat = optimize_z( sess, z, x, x_hat, init_val_ph, init_val['pos'], pos_query_imgs, check_folder(os.path.join(save_dir, 'pos_results')), opt, vec_loss, vec_losses) save_files(save_dir, ['pos_loss'], [query_loss]) query_loss, query_z, query_xhat = optimize_z( sess, z, x, x_hat, init_val_ph, init_val['neg'], neg_query_imgs, check_folder(os.path.join(save_dir, 'neg_results')), opt, vec_loss, vec_losses) save_files(save_dir, ['neg_loss'], [query_loss])
# get test IDs test_fns = glob.glob(gt_dir + '/%d*.ARW' % d_id) test_ids = [int(os.path.basename(test_fn)[0:5]) for test_fn in test_fns] misaligned = [10034, 10045, 10172] sess = tf.Session() in_image = tf.placeholder(tf.float32, [None, None, None, None, 4]) in_image_low = tf.placeholder(tf.float32, [None, None, None, None, 4]) gt_image = tf.placeholder(tf.float32, [None, None, None, 3]) coarse_outs = burst_nets.coarse_net(in_image_low) out_image = burst_nets.fine_net(in_image, coarse_outs) distance_t = lpips_tf.lpips(gt_image, out_image, model='net-lin', net='alex') t_vars = tf.trainable_variables() saver = tf.train.Saver(t_vars) sess.run(tf.global_variables_initializer()) ckpt = tf.train.get_checkpoint_state(checkpoint_dir) if ckpt: print('loaded ' + ckpt.model_checkpoint_path) saver.restore(sess, ckpt.model_checkpoint_path) ssim_list = [] psnr_list = [] lpips_list = [] time_list = [] ratio_list = []