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}'))