def main(): """Main function.""" args = parse_args() os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id assert os.path.exists(args.image_list) image_list_name = os.path.splitext(os.path.basename(args.image_list))[0] output_dir = args.output_dir or f'results/inversion/{image_list_name}' logger = setup_logger(output_dir, 'inversion.log', 'inversion_logger') logger.info(f'Loading model.') inverter = StyleGANInverter( args.model_name, learning_rate=args.learning_rate, iteration=args.num_iterations, reconstruction_loss_weight=1.0, perceptual_loss_weight=args.loss_weight_feat, regularization_loss_weight=args.loss_weight_enc, logger=logger) image_size = inverter.G.resolution # Load image list. logger.info(f'Loading image list.') image_list = [] with open(args.image_list, 'r') as f: for line in f: image_list.append(line.strip()) # Initialize visualizer. save_interval = args.num_iterations // args.num_results headers = ['Name', 'Original Image', 'Encoder Output'] for step in range(1, args.num_iterations + 1): if step == args.num_iterations or step % save_interval == 0: headers.append(f'Step {step:06d}') viz_size = None if args.viz_size == 0 else args.viz_size visualizer = HtmlPageVisualizer( num_rows=len(image_list), num_cols=len(headers), viz_size=viz_size) visualizer.set_headers(headers) # Invert images. logger.info(f'Start inversion.') latent_codes = [] for img_idx in tqdm(range(len(image_list)), leave=False): image_path = image_list[img_idx] image_name = os.path.splitext(os.path.basename(image_path))[0] image = resize_image(load_image(image_path), (image_size, image_size)) code, viz_results = inverter.easy_invert(image, num_viz=args.num_results) latent_codes.append(code) save_image(f'{output_dir}/{image_name}_ori.png', image) save_image(f'{output_dir}/{image_name}_enc.png', viz_results[1]) save_image(f'{output_dir}/{image_name}_inv.png', viz_results[-1]) visualizer.set_cell(img_idx, 0, text=image_name) visualizer.set_cell(img_idx, 1, image=image) for viz_idx, viz_img in enumerate(viz_results[1:]): visualizer.set_cell(img_idx, viz_idx + 2, image=viz_img) # Save results. os.system(f'cp {args.image_list} {output_dir}/image_list.txt') np.save(f'{output_dir}/inverted_codes.npy', np.concatenate(latent_codes, axis=0)) visualizer.save(f'{output_dir}/inversion.html')
def Vis(bname, suffix, out, rownames=None, colnames=None): num_images = out.shape[0] step = out.shape[1] if colnames is None: colnames = [f'Step {i:02d}' for i in range(1, step + 1)] if rownames is None: rownames = [str(i) for i in range(num_images)] visualizer = HtmlPageVisualizer(num_rows=num_images, num_cols=step + 1, viz_size=256) visualizer.set_headers(['Name'] + colnames) for i in range(num_images): visualizer.set_cell(i, 0, text=rownames[i]) for i in range(num_images): for k in range(step): image = out[i, k, :, :, :] visualizer.set_cell(i, 1 + k, image=image) # Save results. visualizer.save(f'./html/' + bname + '_' + suffix + '.html')
def main(): """Main function.""" args = parse_args() os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id style_dir = args.style_dir style_dir_name = os.path.basename(style_dir.rstrip('/')) assert os.path.exists(style_dir) assert os.path.exists(f'{style_dir}/image_list.txt') assert os.path.exists(f'{style_dir}/inverted_codes.npy') content_dir = args.content_dir content_dir_name = os.path.basename(content_dir.rstrip('/')) assert os.path.exists(content_dir) assert os.path.exists(f'{content_dir}/image_list.txt') assert os.path.exists(f'{content_dir}/inverted_codes.npy') output_dir = args.output_dir or 'results/style_mixing' job_name = f'{style_dir_name}_STYLIZE_{content_dir_name}' logger = setup_logger(output_dir, f'{job_name}.log', f'{job_name}_logger') # Load model. logger.info(f'Loading generator.') tflib.init_tf({'rnd.np_random_seed': 1000}) with open(args.model_path, 'rb') as f: _, _, _, Gs = pickle.load(f) # Build graph. logger.info(f'Building graph.') sess = tf.get_default_session() num_layers, latent_dim = Gs.components.synthesis.input_shape[1:3] wp = tf.placeholder( tf.float32, [args.batch_size, num_layers, latent_dim], name='latent_code') x = Gs.components.synthesis.get_output_for(wp, randomize_noise=False) mix_layers = list(range(args.mix_layer_start_idx, num_layers)) # Load image and codes. logger.info(f'Loading images and corresponding inverted latent codes.') style_list = [] with open(f'{style_dir}/image_list.txt', 'r') as f: for line in f: name = os.path.splitext(os.path.basename(line.strip()))[0] assert os.path.exists(f'{style_dir}/{name}_ori.png') style_list.append(name) logger.info(f'Loading inverted latent codes.') style_codes = np.load(f'{style_dir}/inverted_codes.npy') assert style_codes.shape[0] == len(style_list) num_styles = style_codes.shape[0] content_list = [] with open(f'{content_dir}/image_list.txt', 'r') as f: for line in f: name = os.path.splitext(os.path.basename(line.strip()))[0] assert os.path.exists(f'{content_dir}/{name}_ori.png') content_list.append(name) logger.info(f'Loading inverted latent codes.') content_codes = np.load(f'{content_dir}/inverted_codes.npy') assert content_codes.shape[0] == len(content_list) num_contents = content_codes.shape[0] # Mix styles. logger.info(f'Start style mixing.') viz_size = None if args.viz_size == 0 else args.viz_size visualizer = HtmlPageVisualizer( num_rows=num_styles + 1, num_cols=num_contents + 1, viz_size=viz_size) visualizer.set_headers( ['Style'] + [f'Content {i:03d}' for i in range(num_contents)] ) for style_idx, style_name in enumerate(style_list): style_image = load_image(f'{style_dir}/{style_name}_ori.png') visualizer.set_cell(style_idx + 1, 0, image=style_image) for content_idx, content_name in enumerate(content_list): content_image = load_image(f'{content_dir}/{content_name}_ori.png') visualizer.set_cell(0, content_idx + 1, image=content_image) codes = mix_style(style_codes=style_codes, content_codes=content_codes, num_layers=num_layers, mix_layers=mix_layers) inputs = np.zeros((args.batch_size, num_layers, latent_dim), np.float32) for style_idx in tqdm(range(num_styles), leave=False): output_images = [] for idx in range(0, num_contents, args.batch_size): batch = codes[style_idx, idx:idx + args.batch_size] inputs[0:len(batch)] = batch images = sess.run(x, feed_dict={wp: inputs}) output_images.append(images[0:len(batch)]) output_images = adjust_pixel_range(np.concatenate(output_images, axis=0)) for content_idx, output_image in enumerate(output_images): visualizer.set_cell(style_idx + 1, content_idx + 1, image=output_image) # Save results. visualizer.save(f'{output_dir}/{job_name}.html')
def main(): """Main function.""" args = parse_args() os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id assert os.path.exists(args.image_list) image_list_name = os.path.splitext(os.path.basename(args.image_list))[0] output_dir = args.output_dir or f'results/inversion/{image_list_name}' logger = setup_logger(output_dir, 'inversion.log', 'inversion_logger') logger.info(f'Loading model.') tflib.init_tf({'rnd.np_random_seed': 1000}) with open(args.model_path, 'rb') as f: E, _, _, Gs = pickle.load(f) # Get input size. image_size = E.input_shape[2] assert image_size == E.input_shape[3] # Build graph. logger.info(f'Building graph.') sess = tf.get_default_session() input_shape = E.input_shape input_shape[0] = args.batch_size x = tf.placeholder(tf.float32, shape=input_shape, name='real_image') x_255 = (tf.transpose(x, [0, 2, 3, 1]) + 1) / 2 * 255 latent_shape = Gs.components.synthesis.input_shape latent_shape[0] = args.batch_size wp = tf.get_variable(shape=latent_shape, name='latent_code') x_rec = Gs.components.synthesis.get_output_for(wp, randomize_noise=False) x_rec_255 = (tf.transpose(x_rec, [0, 2, 3, 1]) + 1) / 2 * 255 if args.random_init: logger.info(f' Use random initialization for optimization.') wp_rnd = tf.random.normal(shape=latent_shape, name='latent_code_init') setter = tf.assign(wp, wp_rnd) else: logger.info( f' Use encoder output as the initialization for optimization.') w_enc = E.get_output_for(x, is_training=False) wp_enc = tf.reshape(w_enc, latent_shape) setter = tf.assign(wp, wp_enc) # Settings for optimization. logger.info(f'Setting configuration for optimization.') perceptual_model = PerceptualModel([image_size, image_size], False) x_feat = perceptual_model(x_255) x_rec_feat = perceptual_model(x_rec_255) loss_feat = tf.reduce_mean(tf.square(x_feat - x_rec_feat), axis=[1]) loss_pix = tf.reduce_mean(tf.square(x - x_rec), axis=[1, 2, 3]) if args.domain_regularizer: logger.info(f' Involve encoder for optimization.') w_enc_new = E.get_output_for(x_rec, is_training=False) wp_enc_new = tf.reshape(w_enc_new, latent_shape) loss_enc = tf.reduce_mean(tf.square(wp - wp_enc_new), axis=[1, 2]) else: logger.info(f' Do NOT involve encoder for optimization.') loss_enc = 0 loss = (loss_pix + args.loss_weight_feat * loss_feat + args.loss_weight_enc * loss_enc) optimizer = tf.train.AdamOptimizer(learning_rate=args.learning_rate) train_op = optimizer.minimize(loss, var_list=[wp]) tflib.init_uninitialized_vars() # Load image list. logger.info(f'Loading image list.') image_list = [] with open(args.image_list, 'r') as f: for line in f: image_list.append(line.strip()) # Invert images. logger.info(f'Start inversion.') save_interval = args.num_iterations // args.num_results headers = ['Name', 'Original Image', 'Encoder Output'] for step in range(1, args.num_iterations + 1): if step == args.num_iterations or step % save_interval == 0: headers.append(f'Step {step:06d}') viz_size = None if args.viz_size == 0 else args.viz_size visualizer = HtmlPageVisualizer(num_rows=len(image_list), num_cols=len(headers), viz_size=viz_size) visualizer.set_headers(headers) images = np.zeros(input_shape, np.uint8) names = ['' for _ in range(args.batch_size)] latent_codes_enc = [] latent_codes = [] for img_idx in tqdm(range(0, len(image_list), args.batch_size), leave=False): # Load inputs. batch = image_list[img_idx:img_idx + args.batch_size] for i, image_path in enumerate(batch): image = resize_image(load_image(image_path), (image_size, image_size)) images[i] = np.transpose(image, [2, 0, 1]) names[i] = os.path.splitext(os.path.basename(image_path))[0] inputs = images.astype(np.float32) / 255 * 2.0 - 1.0 # Run encoder. sess.run([setter], {x: inputs}) outputs = sess.run([wp, x_rec]) latent_codes_enc.append(outputs[0][0:len(batch)]) outputs[1] = adjust_pixel_range(outputs[1]) for i, _ in enumerate(batch): image = np.transpose(images[i], [1, 2, 0]) save_image(f'{output_dir}/{names[i]}_ori.png', image) save_image(f'{output_dir}/{names[i]}_enc.png', outputs[1][i]) visualizer.set_cell(i + img_idx, 0, text=names[i]) visualizer.set_cell(i + img_idx, 1, image=image) visualizer.set_cell(i + img_idx, 2, image=outputs[1][i]) # Optimize latent codes. col_idx = 3 for step in tqdm(range(1, args.num_iterations + 1), leave=False): sess.run(train_op, {x: inputs}) if step == args.num_iterations or step % save_interval == 0: outputs = sess.run([wp, x_rec]) outputs[1] = adjust_pixel_range(outputs[1]) for i, _ in enumerate(batch): if step == args.num_iterations: save_image(f'{output_dir}/{names[i]}_inv.png', outputs[1][i]) visualizer.set_cell(i + img_idx, col_idx, image=outputs[1][i]) col_idx += 1 latent_codes.append(outputs[0][0:len(batch)]) # Save results. os.system(f'cp {args.image_list} {output_dir}/image_list.txt') np.save(f'{output_dir}/encoded_codes.npy', np.concatenate(latent_codes_enc, axis=0)) np.save(f'{output_dir}/inverted_codes.npy', np.concatenate(latent_codes, axis=0)) visualizer.save(f'{output_dir}/inversion.html')
try: boundary_file = np.load(path, allow_pickle=True).item() boundary = boundary_file['boundary'] manipulate_layers = boundary_file['meta_data']['manipulate_layers'] except ValueError: boundary = np.load(path) manipulate_layers = args.manipulate_layers logger.info(f' Manipulating on layers `{manipulate_layers}`.') np.save(os.path.join(work_dir, f'{prefix}_boundary.npy'), boundary) step = args.step + int(args.step % 2 == 0) # Make sure it is an odd number. visualizer = HtmlPageVisualizer(num_rows=total_num, num_cols=step + 1) visualizer.set_headers([''] + [f'Step {i - step // 2}' for i in range(step // 2)] + ['Origin'] + [f'Step {i + 1}' for i in range(step // 2)]) for n in range(total_num): visualizer.set_cell(n, 0, text=f'Sample {n:05d}') strength = get_layerwise_manipulation_strength(model.num_layers, model.truncation_psi, model.truncation_layers) codes = manipulate(latent_codes=latent_codes['wp'], boundary=boundary, start_distance=args.start_distance, end_distance=args.end_distance, step=step, layerwise_manipulation=True, num_layers=model.num_layers, manipulate_layers=manipulate_layers,
def main(): """Main function.""" args = parse_args() os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id src_dir = args.src_dir src_dir_name = os.path.basename(src_dir.rstrip('/')) assert os.path.exists(src_dir) assert os.path.exists(f'{src_dir}/image_list.txt') assert os.path.exists(f'{src_dir}/inverted_codes.npy') dst_dir = args.dst_dir dst_dir_name = os.path.basename(dst_dir.rstrip('/')) assert os.path.exists(dst_dir) assert os.path.exists(f'{dst_dir}/image_list.txt') assert os.path.exists(f'{dst_dir}/inverted_codes.npy') output_dir = args.output_dir or 'results/interpolation' job_name = f'{src_dir_name}_TO_{dst_dir_name}' logger = setup_logger(output_dir, f'{job_name}.log', f'{job_name}_logger') # Load model. logger.info(f'Loading generator.') tflib.init_tf({'rnd.np_random_seed': 1000}) with open(args.model_path, 'rb') as f: _, _, _, Gs = pickle.load(f) # Build graph. logger.info(f'Building graph.') sess = tf.get_default_session() num_layers, latent_dim = Gs.components.synthesis.input_shape[1:3] wp = tf.placeholder(tf.float32, [args.batch_size, num_layers, latent_dim], name='latent_code') x = Gs.components.synthesis.get_output_for(wp, randomize_noise=False) # Load image and codes. logger.info(f'Loading images and corresponding inverted latent codes.') src_list = [] with open(f'{src_dir}/image_list.txt', 'r') as f: for line in f: name = os.path.splitext(os.path.basename(line.strip()))[0] assert os.path.exists(f'{src_dir}/{name}_ori.png') src_list.append(name) src_codes = np.load(f'{src_dir}/inverted_codes.npy') assert src_codes.shape[0] == len(src_list) num_src = src_codes.shape[0] dst_list = [] with open(f'{dst_dir}/image_list.txt', 'r') as f: for line in f: name = os.path.splitext(os.path.basename(line.strip()))[0] assert os.path.exists(f'{dst_dir}/{name}_ori.png') dst_list.append(name) dst_codes = np.load(f'{dst_dir}/inverted_codes.npy') assert dst_codes.shape[0] == len(dst_list) num_dst = dst_codes.shape[0] # Interpolate images. logger.info(f'Start interpolation.') step = args.step + 2 viz_size = None if args.viz_size == 0 else args.viz_size visualizer = HtmlPageVisualizer(num_rows=num_src * num_dst, num_cols=step + 2, viz_size=viz_size) visualizer.set_headers(['Source', 'Source Inversion'] + [f'Step {i:02d}' for i in range(1, step - 1)] + ['Target Inversion', 'Target']) inputs = np.zeros((args.batch_size, num_layers, latent_dim), np.float32) for src_idx in tqdm(range(num_src), leave=False): src_code = src_codes[src_idx:src_idx + 1] src_path = f'{src_dir}/{src_list[src_idx]}_ori.png' codes = interpolate(src_codes=np.repeat(src_code, num_dst, axis=0), dst_codes=dst_codes, step=step) for dst_idx in tqdm(range(num_dst), leave=False): dst_path = f'{dst_dir}/{dst_list[dst_idx]}_ori.png' output_images = [] for idx in range(0, step, args.batch_size): batch = codes[dst_idx, idx:idx + args.batch_size] inputs[0:len(batch)] = batch images = sess.run(x, feed_dict={wp: inputs}) output_images.append(images[0:len(batch)]) output_images = adjust_pixel_range( np.concatenate(output_images, axis=0)) row_idx = src_idx * num_dst + dst_idx visualizer.set_cell(row_idx, 0, image=load_image(src_path)) visualizer.set_cell(row_idx, step + 1, image=load_image(dst_path)) for s, output_image in enumerate(output_images): if s == 5 and row_idx == 2: save_image(f'./results/interpolation/005_int.png', output_image) visualizer.set_cell(row_idx, s + 1, image=output_image) # Save results. visualizer.save(f'{output_dir}/{job_name}.html')
def main(): """Main function.""" args = parse_args() os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id image_dir = args.image_dir image_dir_name = os.path.basename(image_dir.rstrip('/')) assert os.path.exists(image_dir) assert os.path.exists(f'{image_dir}/image_list.txt') assert os.path.exists(f'{image_dir}/inverted_codes.npy') boundary_path = args.boundary_path assert os.path.exists(boundary_path) boundary_name = os.path.splitext(os.path.basename(boundary_path))[0] output_dir = args.output_dir or 'results/manipulation' job_name = f'{boundary_name.upper()}_{image_dir_name}' logger = setup_logger(output_dir, f'{job_name}.log', f'{job_name}_logger') # Load model. logger.info(f'Loading generator.') generator = build_generator(args.model_name) # Load image, codes, and boundary. logger.info(f'Loading images and corresponding inverted latent codes.') image_list = [] with open(f'{image_dir}/image_list.txt', 'r') as f: for line in f: name = os.path.splitext(os.path.basename(line.strip()))[0] assert os.path.exists(f'{image_dir}/{name}_ori.png') assert os.path.exists(f'{image_dir}/{name}_inv.png') image_list.append(name) latent_codes = np.load(f'{image_dir}/inverted_codes.npy') assert latent_codes.shape[0] == len(image_list) num_images = latent_codes.shape[0] logger.info(f'Loading boundary.') boundary_file = np.load(boundary_path, allow_pickle=True)[()] if isinstance(boundary_file, dict): boundary = boundary_file['boundary'] manipulate_layers = boundary_file['meta_data']['manipulate_layers'] else: boundary = boundary_file manipulate_layers = args.manipulate_layers if manipulate_layers: logger.info(f' Manipulating on layers `{manipulate_layers}`.') else: logger.info(f' Manipulating on ALL layers.') # Manipulate images. logger.info(f'Start manipulation.') step = args.step viz_size = None if args.viz_size == 0 else args.viz_size visualizer = HtmlPageVisualizer(num_rows=num_images, num_cols=step + 3, viz_size=viz_size) visualizer.set_headers(['Name', 'Origin', 'Inverted'] + [f'Step {i:02d}' for i in range(1, step + 1)]) for img_idx, img_name in enumerate(image_list): ori_image = load_image(f'{image_dir}/{img_name}_ori.png') inv_image = load_image(f'{image_dir}/{img_name}_inv.png') visualizer.set_cell(img_idx, 0, text=img_name) visualizer.set_cell(img_idx, 1, image=ori_image) visualizer.set_cell(img_idx, 2, image=inv_image) codes = manipulate(latent_codes=latent_codes, boundary=boundary, start_distance=args.start_distance, end_distance=args.end_distance, step=step, layerwise_manipulation=True, num_layers=generator.num_layers, manipulate_layers=manipulate_layers, is_code_layerwise=True, is_boundary_layerwise=True) for img_idx in tqdm(range(num_images), leave=False): output_images = generator.easy_synthesize( codes[img_idx], latent_space_type='wp')['image'] for s, output_image in enumerate(output_images): visualizer.set_cell(img_idx, s + 3, image=output_image) # Save results. visualizer.save(f'{output_dir}/{job_name}.html')
def main(): """Main function.""" args = parse_args() os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id assert os.path.exists(args.target_list) target_list_name = os.path.splitext(os.path.basename(args.target_list))[0] assert os.path.exists(args.context_list) context_list_name = os.path.splitext(os.path.basename( args.context_list))[0] output_dir = args.output_dir or f'results/diffusion' job_name = f'{target_list_name}_TO_{context_list_name}' logger = setup_logger(output_dir, f'{job_name}.log', f'{job_name}_logger') logger.info(f'Loading model.') tflib.init_tf({'rnd.np_random_seed': 1000}) with open(args.model_path, 'rb') as f: E, _, _, Gs = pickle.load(f) # Get input size. image_size = E.input_shape[2] assert image_size == E.input_shape[3] crop_size = args.crop_size crop_x = args.center_x - crop_size // 2 crop_y = args.center_y - crop_size // 2 mask = np.zeros((1, image_size, image_size, 3), dtype=np.float32) mask[:, crop_y:crop_y + crop_size, crop_x:crop_x + crop_size, :] = 1.0 # Build graph. logger.info(f'Building graph.') sess = tf.get_default_session() input_shape = E.input_shape input_shape[0] = args.batch_size x = tf.placeholder(tf.float32, shape=input_shape, name='real_image') x_mask = (tf.transpose(x, [0, 2, 3, 1]) + 1) * mask - 1 x_mask_255 = (x_mask + 1) / 2 * 255 latent_shape = Gs.components.synthesis.input_shape latent_shape[0] = args.batch_size wp = tf.get_variable(shape=latent_shape, name='latent_code') x_rec = Gs.components.synthesis.get_output_for(wp, randomize_noise=False) x_rec_mask = (tf.transpose(x_rec, [0, 2, 3, 1]) + 1) * mask - 1 x_rec_mask_255 = (x_rec_mask + 1) / 2 * 255 w_enc = E.get_output_for(x, phase=False) wp_enc = tf.reshape(w_enc, latent_shape) setter = tf.assign(wp, wp_enc) # Settings for optimization. logger.info(f'Setting configuration for optimization.') perceptual_model = PerceptualModel([image_size, image_size], False) x_feat = perceptual_model(x_mask_255) x_rec_feat = perceptual_model(x_rec_mask_255) loss_feat = tf.reduce_mean(tf.square(x_feat - x_rec_feat), axis=[1]) loss_pix = tf.reduce_mean(tf.square(x_mask - x_rec_mask), axis=[1, 2, 3]) loss = loss_pix + args.loss_weight_feat * loss_feat optimizer = tf.train.AdamOptimizer(learning_rate=args.learning_rate) train_op = optimizer.minimize(loss, var_list=[wp]) tflib.init_uninitialized_vars() # Load image list. logger.info(f'Loading target images and context images.') target_list = [] with open(args.target_list, 'r') as f: for line in f: target_list.append(line.strip()) num_targets = len(target_list) context_list = [] with open(args.context_list, 'r') as f: for line in f: context_list.append(line.strip()) num_contexts = len(context_list) num_pairs = num_targets * num_contexts # Invert images. logger.info(f'Start diffusion.') save_interval = args.num_iterations // args.num_results headers = [ 'Target Image', 'Context Image', 'Stitched Image', 'Encoder Output' ] for step in range(1, args.num_iterations + 1): if step == args.num_iterations or step % save_interval == 0: headers.append(f'Step {step:06d}') viz_size = None if args.viz_size == 0 else args.viz_size visualizer = HtmlPageVisualizer(num_rows=num_pairs, num_cols=len(headers), viz_size=viz_size) visualizer.set_headers(headers) images = np.zeros(input_shape, np.uint8) latent_codes_enc = [] latent_codes = [] for target_idx in tqdm(range(num_targets), desc='Target ID', leave=False): # Load target. target_image = resize_image(load_image(target_list[target_idx]), (image_size, image_size)) visualizer.set_cell(target_idx * num_contexts, 0, image=target_image) for context_idx in tqdm(range(0, num_contexts, args.batch_size), desc='Context ID', leave=False): row_idx = target_idx * num_contexts + context_idx batch = context_list[context_idx:context_idx + args.batch_size] for i, context_image_path in enumerate(batch): context_image = resize_image(load_image(context_image_path), (image_size, image_size)) visualizer.set_cell(row_idx + i, 1, image=context_image) context_image[crop_y:crop_y + crop_size, crop_x:crop_x + crop_size] = (target_image[crop_y:crop_y + crop_size, crop_x:crop_x + crop_size]) visualizer.set_cell(row_idx + i, 2, image=context_image) images[i] = np.transpose(context_image, [2, 0, 1]) inputs = images.astype(np.float32) / 255 * 2.0 - 1.0 # Run encoder. sess.run([setter], {x: inputs}) outputs = sess.run([wp, x_rec]) latent_codes_enc.append(outputs[0][0:len(batch)]) outputs[1] = adjust_pixel_range(outputs[1]) for i, _ in enumerate(batch): visualizer.set_cell(row_idx + i, 3, image=outputs[1][i]) # Optimize latent codes. col_idx = 4 for step in tqdm(range(1, args.num_iterations + 1), leave=False): sess.run(train_op, {x: inputs}) if step == args.num_iterations or step % save_interval == 0: outputs = sess.run([wp, x_rec]) outputs[1] = adjust_pixel_range(outputs[1]) for i, _ in enumerate(batch): visualizer.set_cell(row_idx + i, col_idx, image=outputs[1][i]) col_idx += 1 latent_codes.append(outputs[0][0:len(batch)]) # Save results. code_shape = [num_targets, num_contexts] + list(latent_shape[1:]) np.save(f'{output_dir}/{job_name}_encoded_codes.npy', np.concatenate(latent_codes_enc, axis=0).reshape(code_shape)) np.save(f'{output_dir}/{job_name}_inverted_codes.npy', np.concatenate(latent_codes, axis=0).reshape(code_shape)) visualizer.save(f'{output_dir}/{job_name}.html')
def main(): """Main function.""" args = parse_args() os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id src_dir = args.src_dir src_dir_name = os.path.basename(src_dir.rstrip('/')) assert os.path.exists(src_dir) assert os.path.exists(f'{src_dir}/image_list.txt') assert os.path.exists(f'{src_dir}/inverted_codes.npy') dst_dir = args.dst_dir dst_dir_name = os.path.basename(dst_dir.rstrip('/')) assert os.path.exists(dst_dir) assert os.path.exists(f'{dst_dir}/image_list.txt') assert os.path.exists(f'{dst_dir}/inverted_codes.npy') output_dir = args.output_dir or 'results/interpolation' job_name = f'{src_dir_name}_TO_{dst_dir_name}' logger = setup_logger(output_dir, f'{job_name}.log', f'{job_name}_logger') # Load model. logger.info(f'Loading generator.') generator = build_generator(args.model_name) # Load image and codes. logger.info(f'Loading images and corresponding inverted latent codes.') src_list = [] with open(f'{src_dir}/image_list.txt', 'r') as f: for line in f: name = os.path.splitext(os.path.basename(line.strip()))[0] assert os.path.exists(f'{src_dir}/{name}_ori.png') src_list.append(name) src_codes = np.load(f'{src_dir}/inverted_codes.npy') assert src_codes.shape[0] == len(src_list) num_src = src_codes.shape[0] dst_list = [] with open(f'{dst_dir}/image_list.txt', 'r') as f: for line in f: name = os.path.splitext(os.path.basename(line.strip()))[0] assert os.path.exists(f'{dst_dir}/{name}_ori.png') dst_list.append(name) dst_codes = np.load(f'{dst_dir}/inverted_codes.npy') assert dst_codes.shape[0] == len(dst_list) num_dst = dst_codes.shape[0] # Interpolate images. logger.info(f'Start interpolation.') step = args.step + 2 viz_size = None if args.viz_size == 0 else args.viz_size visualizer = HtmlPageVisualizer(num_rows=num_src * num_dst, num_cols=step + 2, viz_size=viz_size) visualizer.set_headers(['Source', 'Source Inversion'] + [f'Step {i:02d}' for i in range(1, step - 1)] + ['Target Inversion', 'Target']) for src_idx in tqdm(range(num_src), leave=False): src_code = src_codes[src_idx:src_idx + 1] src_path = f'{src_dir}/{src_list[src_idx]}_ori.png' codes = interpolate(src_codes=np.repeat(src_code, num_dst, axis=0), dst_codes=dst_codes, step=step) for dst_idx in tqdm(range(num_dst), leave=False): dst_path = f'{dst_dir}/{dst_list[dst_idx]}_ori.png' output_images = generator.easy_synthesize( codes[dst_idx], latent_space_type='wp')['image'] row_idx = src_idx * num_dst + dst_idx visualizer.set_cell(row_idx, 0, image=load_image(src_path)) visualizer.set_cell(row_idx, step + 1, image=load_image(dst_path)) for s, output_image in enumerate(output_images): visualizer.set_cell(row_idx, s + 1, image=output_image) # Save results. visualizer.save(f'{output_dir}/{job_name}.html')
def main(): """Main function.""" args = parse_args() os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id style_dir = args.style_dir style_dir_name = os.path.basename(style_dir.rstrip('/')) assert os.path.exists(style_dir) assert os.path.exists(f'{style_dir}/image_list.txt') assert os.path.exists(f'{style_dir}/inverted_codes.npy') content_dir = args.content_dir content_dir_name = os.path.basename(content_dir.rstrip('/')) assert os.path.exists(content_dir) assert os.path.exists(f'{content_dir}/image_list.txt') assert os.path.exists(f'{content_dir}/inverted_codes.npy') output_dir = args.output_dir or 'results/style_mixing' job_name = f'{style_dir_name}_STYLIZE_{content_dir_name}' logger = setup_logger(output_dir, f'{job_name}.log', f'{job_name}_logger') # Load model. logger.info(f'Loading generator.') generator = build_generator(args.model_name) mix_layers = list(range(args.mix_layer_start_idx, generator.num_layers)) # Load image and codes. logger.info(f'Loading images and corresponding inverted latent codes.') style_list = [] with open(f'{style_dir}/image_list.txt', 'r') as f: for line in f: name = os.path.splitext(os.path.basename(line.strip()))[0] assert os.path.exists(f'{style_dir}/{name}_ori.png') style_list.append(name) logger.info(f'Loading inverted latent codes.') style_codes = np.load(f'{style_dir}/inverted_codes.npy') assert style_codes.shape[0] == len(style_list) num_styles = style_codes.shape[0] content_list = [] with open(f'{content_dir}/image_list.txt', 'r') as f: for line in f: name = os.path.splitext(os.path.basename(line.strip()))[0] assert os.path.exists(f'{content_dir}/{name}_ori.png') content_list.append(name) logger.info(f'Loading inverted latent codes.') content_codes = np.load(f'{content_dir}/inverted_codes.npy') assert content_codes.shape[0] == len(content_list) num_contents = content_codes.shape[0] # Mix styles. logger.info(f'Start style mixing.') viz_size = None if args.viz_size == 0 else args.viz_size visualizer = HtmlPageVisualizer(num_rows=num_styles + 1, num_cols=num_contents + 1, viz_size=viz_size) visualizer.set_headers(['Style'] + [f'Content {i:03d}' for i in range(num_contents)]) for style_idx, style_name in enumerate(style_list): style_image = load_image(f'{style_dir}/{style_name}_ori.png') visualizer.set_cell(style_idx + 1, 0, image=style_image) for content_idx, content_name in enumerate(content_list): content_image = load_image(f'{content_dir}/{content_name}_ori.png') visualizer.set_cell(0, content_idx + 1, image=content_image) codes = mix_style(style_codes=style_codes, content_codes=content_codes, num_layers=generator.num_layers, mix_layers=mix_layers) for style_idx in tqdm(range(num_styles), leave=False): output_images = generator.easy_synthesize( codes[style_idx], latent_space_type='wp')['image'] for content_idx, output_image in enumerate(output_images): visualizer.set_cell(style_idx + 1, content_idx + 1, image=output_image) # Save results. visualizer.save(f'{output_dir}/{job_name}.html')
def main(): """Main function.""" args = parse_args() os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id assert os.path.exists(args.target_list) target_list_name = os.path.splitext(os.path.basename(args.target_list))[0] assert os.path.exists(args.context_list) context_list_name = os.path.splitext(os.path.basename(args.context_list))[0] output_dir = args.output_dir or f'results/diffusion' job_name = f'{target_list_name}_TO_{context_list_name}' logger = setup_logger(output_dir, f'{job_name}.log', f'{job_name}_logger') logger.info(f'Loading model.') inverter = StyleGANInverter( args.model_name, learning_rate=args.learning_rate, iteration=args.num_iterations, reconstruction_loss_weight=1.0, perceptual_loss_weight=args.loss_weight_feat, regularization_loss_weight=0.0, logger=logger) image_size = inverter.G.resolution # Load image list. logger.info(f'Loading target images and context images.') target_list = [] with open(args.target_list, 'r') as f: for line in f: target_list.append(line.strip()) num_targets = len(target_list) context_list = [] with open(args.context_list, 'r') as f: for line in f: context_list.append(line.strip()) num_contexts = len(context_list) num_pairs = num_targets * num_contexts # Initialize visualizer. save_interval = args.num_iterations // args.num_results headers = ['Target Image', 'Context Image', 'Stitched Image', 'Encoder Output'] for step in range(1, args.num_iterations + 1): if step == args.num_iterations or step % save_interval == 0: headers.append(f'Step {step:06d}') viz_size = None if args.viz_size == 0 else args.viz_size visualizer = HtmlPageVisualizer( num_rows=num_pairs, num_cols=len(headers), viz_size=viz_size) visualizer.set_headers(headers) # Diffuse images. logger.info(f'Start diffusion.') latent_codes = [] for target_idx in tqdm(range(num_targets), desc='Target ID', leave=False): # Load target. target_image = resize_image(load_image(target_list[target_idx]), (image_size, image_size)) visualizer.set_cell(target_idx * num_contexts, 0, image=target_image) for context_idx in tqdm(range(num_contexts), desc='Context ID', leave=False): row_idx = target_idx * num_contexts + context_idx context_image = resize_image(load_image(context_list[context_idx]), (image_size, image_size)) visualizer.set_cell(row_idx, 1, image=context_image) code, viz_results = inverter.easy_diffuse(target=target_image, context=context_image, center_x=args.center_x, center_y=args.center_y, crop_x=args.crop_size, crop_y=args.crop_size, num_viz=args.num_results) for viz_idx, viz_img in enumerate(viz_results): visualizer.set_cell(row_idx, viz_idx + 2, image=viz_img) latent_codes.append(code) # Save results. os.system(f'cp {args.target_list} {output_dir}/target_list.txt') os.system(f'cp {args.context_list} {output_dir}/context_list.txt') np.save(f'{output_dir}/{job_name}_inverted_codes.npy', np.concatenate(latent_codes, axis=0)) visualizer.save(f'{output_dir}/{job_name}.html')
def main(): """Main function.""" args = parse_args() os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id image_dir = args.image_dir image_dir_name = os.path.basename(image_dir.rstrip('/')) assert os.path.exists(image_dir) assert os.path.exists(f'{image_dir}/image_list.txt') assert os.path.exists(f'{image_dir}/inverted_codes.npy') boundary_path = args.boundary_path assert os.path.exists(boundary_path) boundary_name = os.path.splitext(os.path.basename(boundary_path))[0] output_dir = args.output_dir or 'results/manipulation' job_name = f'{boundary_name.upper()}_{image_dir_name}' logger = setup_logger(output_dir, f'{job_name}.log', f'{job_name}_logger') # Load model. logger.info(f'Loading generator.') tflib.init_tf({'rnd.np_random_seed': 1000}) with open(args.model_path, 'rb') as f: _, _, _, Gs = pickle.load(f) # Build graph. logger.info(f'Building graph.') sess = tf.get_default_session() num_layers, latent_dim = Gs.components.synthesis.input_shape[1:3] wp = tf.placeholder(tf.float32, [args.batch_size, num_layers, latent_dim], name='latent_code') x = Gs.components.synthesis.get_output_for(wp, randomize_noise=False) # Load image, codes, and boundary. logger.info(f'Loading images and corresponding inverted latent codes.') image_list = [] with open(f'{image_dir}/image_list.txt', 'r') as f: for line in f: name = os.path.splitext(os.path.basename(line.strip()))[0] assert os.path.exists(f'{image_dir}/{name}_ori.png') assert os.path.exists(f'{image_dir}/{name}_inv.png') image_list.append(name) latent_codes = np.load(f'{image_dir}/inverted_codes.npy') assert latent_codes.shape[0] == len(image_list) num_images = latent_codes.shape[0] logger.info(f'Loading boundary.') boundary_file = np.load(boundary_path, allow_pickle=True)[()] if isinstance(boundary_file, dict): boundary = boundary_file['boundary'] manipulate_layers = boundary_file['meta_data']['manipulate_layers'] else: boundary = boundary_file manipulate_layers = args.manipulate_layers if manipulate_layers: logger.info(f' Manipulating on layers `{manipulate_layers}`.') else: logger.info(f' Manipulating on ALL layers.') # Manipulate images. logger.info(f'Start manipulation.') step = args.step viz_size = None if args.viz_size == 0 else args.viz_size visualizer = HtmlPageVisualizer(num_rows=num_images, num_cols=step + 3, viz_size=viz_size) visualizer.set_headers(['Name', 'Origin', 'Inverted'] + [f'Step {i:02d}' for i in range(1, step + 1)]) for img_idx, img_name in enumerate(image_list): ori_image = load_image(f'{image_dir}/{img_name}_ori.png') inv_image = load_image(f'{image_dir}/{img_name}_inv.png') visualizer.set_cell(img_idx, 0, text=img_name) visualizer.set_cell(img_idx, 1, image=ori_image) visualizer.set_cell(img_idx, 2, image=inv_image) codes = manipulate(latent_codes=latent_codes, boundary=boundary, start_distance=args.start_distance, end_distance=args.end_distance, step=step, layerwise_manipulation=True, num_layers=num_layers, manipulate_layers=manipulate_layers, is_code_layerwise=True, is_boundary_layerwise=True) inputs = np.zeros((args.batch_size, num_layers, latent_dim), np.float32) for img_idx in tqdm(range(num_images), leave=False): output_images = [] for idx in range(0, step, args.batch_size): batch = codes[img_idx, idx:idx + args.batch_size] inputs[0:len(batch)] = batch images = sess.run(x, feed_dict={wp: inputs}) output_images.append(images[0:len(batch)]) output_images = adjust_pixel_range( np.concatenate(output_images, axis=0)) for s, output_image in enumerate(output_images): visualizer.set_cell(img_idx, s + 3, image=output_image) # Save results. visualizer.save(f'{output_dir}/{job_name}.html')
def main(): """Main function.""" args = parse_args() os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id assert os.path.exists(args.image_list) image_list_name = os.path.splitext(os.path.basename(args.image_list))[0] output_dir = args.output_dir or f'results/ghfeat/{image_list_name}' logger = setup_logger(output_dir, 'extract_feature.log', 'inversion_logger') logger.info(f'Loading model.') tflib.init_tf({'rnd.np_random_seed': 1000}) with open(args.model_path, 'rb') as f: E, _, _, Gs = pickle.load(f) # Get input size. image_size = E.input_shape[2] assert image_size == E.input_shape[3] G_args = EasyDict(func_name='training.networks_stylegan.G_synthesis') G_style_mod = tflib.Network('G_StyleMod', resolution=image_size, label_size=0, **G_args) Gs_vars_pairs = { name: tflib.run(val) for name, val in Gs.components.synthesis.vars.items() } for g_name, g_val in G_style_mod.vars.items(): tflib.set_vars({g_val: Gs_vars_pairs[g_name]}) # Build graph. logger.info(f'Building graph.') sess = tf.get_default_session() input_shape = E.input_shape input_shape[0] = args.batch_size x = tf.placeholder(tf.float32, shape=input_shape, name='real_image') ghfeat = E.get_output_for(x, is_training=False) x_rec = G_style_mod.get_output_for(ghfeat, randomize_noise=False) # Load image list. logger.info(f'Loading image list.') image_list = [] with open(args.image_list, 'r') as f: for line in f: image_list.append(line.strip()) # Extract GH-Feat from images. logger.info(f'Start feature extraction.') headers = ['Name', 'Original Image', 'Encoder Output'] viz_size = None if args.viz_size == 0 else args.viz_size visualizer = HtmlPageVisualizer(num_rows=len(image_list), num_cols=len(headers), viz_size=viz_size) visualizer.set_headers(headers) images = np.zeros(input_shape, np.uint8) names = ['' for _ in range(args.batch_size)] features = [] for img_idx in tqdm(range(0, len(image_list), args.batch_size), leave=False): # Load inputs. batch = image_list[img_idx:img_idx + args.batch_size] for i, image_path in enumerate(batch): image = resize_image(load_image(image_path), (image_size, image_size)) images[i] = np.transpose(image, [2, 0, 1]) names[i] = os.path.splitext(os.path.basename(image_path))[0] inputs = images.astype(np.float32) / 255 * 2.0 - 1.0 # Run encoder. outputs = sess.run([ghfeat, x_rec], {x: inputs}) features.append(outputs[0][0:len(batch)]) outputs[1] = adjust_pixel_range(outputs[1]) for i, _ in enumerate(batch): image = np.transpose(images[i], [1, 2, 0]) save_image(f'{output_dir}/{names[i]}_ori.png', image) save_image(f'{output_dir}/{names[i]}_enc.png', outputs[1][i]) visualizer.set_cell(i + img_idx, 0, text=names[i]) visualizer.set_cell(i + img_idx, 1, image=image) visualizer.set_cell(i + img_idx, 2, image=outputs[1][i]) # Save results. os.system(f'cp {args.image_list} {output_dir}/image_list.txt') np.save(f'{output_dir}/ghfeat.npy', np.concatenate(features, axis=0)) visualizer.save(f'{output_dir}/reconstruction.html')
def main(): """Main function.""" args = parse_args() work_dir = args.output_dir or f'{args.model_name}_manipulation' logger_name = f'{args.model_name}_manipulation_logger' logger = setup_logger(work_dir, args.logfile_name, logger_name) logger.info(f'Initializing generator.') model = build_generator(args.model_name, logger=logger) logger.info(f'Preparing latent codes.') if os.path.isfile(args.latent_codes_path): logger.info(f' Load latent codes from `{args.latent_codes_path}`.') latent_codes = np.load(args.latent_codes_path) latent_codes = model.preprocess( latent_codes=latent_codes, latent_space_type=args.latent_space_type) else: if args.num <= 0: raise ValueError( f'Argument `num` should be specified as a positive ' f'number since the latent code path ' f'`{args.latent_codes_path}` does not exist!') logger.info(f' Sample latent codes randomly.') latent_codes = model.easy_sample( num=args.num, latent_space_type=args.latent_space_type) total_num = latent_codes.shape[0] latent_codes = model.easy_synthesize( latent_codes=latent_codes, latent_space_type=args.latent_space_type, generate_style=False, generate_image=False) for key, val in latent_codes.items(): np.save(os.path.join(work_dir, f'{key}.npy'), val) boundaries = parse_boundary_list(args.boundary_list_path) step = args.step + int(args.step % 2 == 0) # Make sure it is an odd number. for boundary_info, boundary_path in boundaries.items(): boundary_name, space_type = boundary_info logger.info( f'Boundary `{boundary_name}` from {space_type.upper()} space.') prefix = f'{boundary_name}_{space_type}' if args.generate_html: viz_size = None if args.viz_size == 0 else args.viz_size visualizer = HtmlPageVisualizer(num_rows=total_num, num_cols=step + 1, viz_size=viz_size) visualizer.set_headers( [''] + [f'Step {i - step // 2}' for i in range(step // 2)] + ['Origin'] + [f'Step {i + 1}' for i in range(step // 2)]) if args.generate_video: setup_images = model.easy_synthesize( latent_codes=latent_codes[args.latent_space_type], latent_space_type=args.latent_space_type)['image'] fusion_kwargs = { 'row': args.row, 'col': args.col, 'row_spacing': args.row_spacing, 'col_spacing': args.col_spacing, 'border_left': args.border_left, 'border_right': args.border_right, 'border_top': args.border_top, 'border_bottom': args.border_bottom, 'black_background': not args.white_background, 'image_size': None if args.viz_size == 0 else args.viz_size, } setup_image = fuse_images(setup_images, **fusion_kwargs) video_writer = VideoWriter(os.path.join( work_dir, f'{prefix}_{args.video_name}'), frame_height=setup_image.shape[0], frame_width=setup_image.shape[1], fps=args.fps) logger.info(f' Loading boundary.') try: boundary_file = np.load(boundary_path, allow_pickle=True).item() boundary = boundary_file['boundary'] manipulate_layers = boundary_file['meta_data']['manipulate_layers'] except ValueError: boundary = np.load(boundary_path) manipulate_layers = args.manipulate_layers logger.info(f' Manipulating on layers `{manipulate_layers}`.') np.save(os.path.join(work_dir, f'{prefix}_boundary.npy'), boundary) if args.layerwise_manipulation and space_type != 'z': layerwise_manipulation = True is_code_layerwise = True is_boundary_layerwise = (space_type == 'wp') if (not args.disable_manipulation_truncation ) and space_type == 'w': strength = get_layerwise_manipulation_strength( model.num_layers, model.truncation_psi, model.truncation_layers) else: strength = 1.0 space_type = 'wp' else: if args.layerwise_manipulation: logger.warning(f' Skip layer-wise manipulation for boundary ' f'`{boundary_name}` from Z space. Traditional ' f'manipulation is used instead.') layerwise_manipulation = False is_code_layerwise = False is_boundary_layerwise = False strength = 1.0 codes = manipulate(latent_codes=latent_codes[space_type], boundary=boundary, start_distance=args.start_distance, end_distance=args.end_distance, step=step, layerwise_manipulation=layerwise_manipulation, num_layers=model.num_layers, manipulate_layers=manipulate_layers, is_code_layerwise=is_code_layerwise, is_boundary_layerwise=is_boundary_layerwise, layerwise_manipulation_strength=strength) np.save( os.path.join(work_dir, f'{prefix}_manipulated_{space_type}.npy'), codes) logger.info(f' Start manipulating.') for s in tqdm(range(step), leave=False): images = model.easy_synthesize( latent_codes=codes[:, s], latent_space_type=space_type)['image'] if args.generate_video: video_writer.write(fuse_images(images, **fusion_kwargs)) for n, image in enumerate(images): if args.save_raw_synthesis: save_image( os.path.join(work_dir, f'{prefix}_{n:05d}_{s:03d}.jpg'), image) if args.generate_html: visualizer.set_cell(n, s + 1, image=image) if s == 0: visualizer.set_cell(n, 0, text=f'Sample {n:05d}') if args.generate_html: visualizer.save( os.path.join(work_dir, f'{prefix}_{args.html_name}'))