def get_gan_model(model_name): """ :param model_name: Please refer `GAN_MODELS` :return: gan_model(nn.Module or nn.Sequential) """ gan = build_generator(model_name) if model_name.startswith('pggan'): gan_list = list(gan.net.children()) remove_index = PGGAN_Inter_Output_Layer_1024 if model_name == 'pggan_celebahq' else PGGAN_Inter_Output_Layer_256 for output_index in remove_index: gan_list.pop(output_index) return nn.Sequential(*gan_list) elif model_name.startswith('style'): return gan
if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--model-dir", default="expr/semantics_layerwise") parser.add_argument("--out-dir", default="results/layerplot/layer_loss") parser.add_argument("--place", default="paper", help="paper | appendix") parser.add_argument("--viz-model", default="LSE") parser.add_argument("--repeat", default=2, type=int) parser.add_argument("--gpu-id", default=0, type=int) parser.add_argument("--show-weight", default=0, type=int) args = parser.parse_args() G_names = "pggan_celebahq,stylegan_celebahq,stylegan2_ffhq,pggan_church,stylegan_church,stylegan2_church,pggan_bedroom,stylegan_bedroom,stylegan2_bedroom" # setup and constants data_dir = args.model_dir Gs = {G_name: build_generator(G_name).net for G_name in G_names.split(",")} is_face = "ffhq" in G_names or "celebahq" in G_names unet = FaceSegmenter() def P_from_name(G_name): if "ffhq" in G_name or "celebahq" in G_name: return unet return SceneSegmenter(model_name=G_name) Ps = {G_name: P_from_name(G_name) for G_name in G_names.split(",")} label_list = CELEBA_CATEGORY if is_face else [] n_class = len(label_list) N_repeat = args.repeat model_dirs = glob.glob(f"{data_dir}/*") model_files = [d for d in model_dirs if os.path.isdir(d)] model_files = [glob.glob(f"{f}/*.pth") for f in model_files]
def main(): """Main function.""" args = parse_args() work_dir = args.output_dir or f'{args.model_name}_rescore' logger_name = f'{args.model_name}_rescore_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 args.num <= 0: raise ValueError(f'Argument `num` should be specified as a positive ' f'number, but `{args.num}` received!') latent_codes = model.easy_sample(num=args.num, latent_space_type='z') latent_codes = model.easy_synthesize(latent_codes=latent_codes, latent_space_type='z', generate_style=False, generate_image=False) for key, val in latent_codes.items(): np.save(os.path.join(work_dir, f'{key}.npy'), val) logger.info(f'Initializing predictor.') predictor = build_predictor(args.predictor_name) boundaries = parse_boundary_list(args.boundary_list_path) logger.info(f'========================================') logger.info(f'Rescoring.') score_changing = [] for boundary_info, boundary_path in boundaries.items(): logger.info(f'----------------------------------------') boundary_name, space_type = boundary_info logger.info( f'Boundary `{boundary_name}` from {space_type.upper()} space.') prefix = f'{boundary_name}_{space_type}' attr_idx = predictor.attribute_name_to_idx[boundary_name] try: boundary_file = np.load(boundary_path, allow_pickle=True).item() boundary = boundary_file['boundary'] except ValueError: boundary = np.load(boundary_path) np.save(os.path.join(work_dir, f'{prefix}_boundary.npy'), boundary) if space_type == 'z': layerwise_manipulation = False is_code_layerwise = False is_boundary_layerwise = False num_layers = 0 strength = 1.0 else: layerwise_manipulation = True is_code_layerwise = True is_boundary_layerwise = (space_type == 'wp') num_layers = model.num_layers if args.layerwise_rescoring else 0 if space_type == 'w': strength = get_layerwise_manipulation_strength( model.num_layers, model.truncation_psi, model.truncation_layers) else: strength = 1.0 space_type = 'wp' codes = [] codes.append(latent_codes[space_type][:, np.newaxis]) for l in range(-1, num_layers): codes.append( manipulate(latent_codes[space_type], boundary, start_distance=2.0, end_distance=2.0, step=1, layerwise_manipulation=layerwise_manipulation, num_layers=model.num_layers, manipulate_layers=None if l < 0 else l, is_code_layerwise=is_code_layerwise, is_boundary_layerwise=is_boundary_layerwise, layerwise_manipulation_strength=strength)) codes = np.concatenate(codes, axis=1) scores = [] for i in tqdm(range(args.num), leave=False): images = model.easy_synthesize(latent_codes=codes[i], latent_space_type=space_type, generate_style=False, generate_image=True)['image'] scores.append( predictor.easy_predict(images)['attribute'][:, attr_idx]) scores = np.stack(scores, axis=0) np.save(os.path.join(work_dir, f'{prefix}_scores.npy'), scores) delta = scores[:, 1] - scores[:, 0] delta[delta < 0] = 0 score_changing.append((boundary_name, np.mean(delta))) if num_layers: layerwise_score_changing = [] for l in range(num_layers): delta = scores[:, l + 2] - scores[:, 0] delta[delta < 0] = 0 layerwise_score_changing.append( (f'Layer {l:02d}', np.mean(delta))) layerwise_score_changing.sort(key=lambda x: x[1], reverse=True) for layer_name, delta_score in layerwise_score_changing: logger.info(f' {layer_name}: {delta_score:7.4f}') logger.info(f'----------------------------------------') logger.info(f'Most relevant semantics:') score_changing.sort(key=lambda x: x[1], reverse=True) for boundary_name, delta_score in score_changing: logger.info(f' {boundary_name.ljust(15)}: {delta_score:7.4f}')
type=str, default='6-11', help='Indices of the layers to perform manipulation. ' 'Active ONLY when `layerwise_manipulation` is set ' 'as `True`. If not specified, all layers will be ' 'manipulated. More than one layers should be ' 'separated by `,`. (default: None)') args = parser.parse_args() work_dir = 'manipulation_results' os.makedirs(work_dir, exist_ok=True) prefix = f'{args.model_name}_{args.boundary_name}' logger = setup_logger(work_dir, '', 'logger') 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: 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(
def setup_generator(model_name): global generator setup_model(model_name) generator = build_generator(models[model_name]['name'])
# PGGAN Generator. if args.pggan or args.all: print('==== PGGAN Generator Test ====') if args.verbose: model_list = [] for model_name, model_setting in MODEL_POOL.items(): if model_setting['gan_type'] == 'pggan': model_list.append(model_name) else: model_list = ['pggan_celebahq', 'pggan_bedroom'] for model_name in model_list: logger = setup_logger( work_dir=RESULT_DIR, logfile_name=f'{model_name}_generator_test.log', logger_name=f'{model_name}_generator_logger') G = build_generator(model_name, logger=logger) G.batch_size = TEST_BATCH_SIZE z = G.easy_sample(args.test_num) x = G.easy_synthesize(z)['image'] visualizer = HtmlPageVisualizer(grid_size=args.test_num) for i in range(visualizer.num_rows): for j in range(visualizer.num_cols): visualizer.set_cell(i, j, image=x[i * visualizer.num_cols + j]) visualizer.save(f'{RESULT_DIR}/{model_name}_generator_test.html') del G print('Pass!') TEST_FLAG = True # StyleGAN Generator.
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 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')
else: mIoU, c_iou = read_results(res_file) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument('--eval-file', type=str, default='results/scs', help='The path to the experiment result file.') parser.add_argument('--out-dir', type=str, default='results/scs', help='The output directory.') parser.add_argument('--gpu-id', default='0', help='Which GPU(s) to use. (default: `0`)') args = parser.parse_args() set_cuda_devices(args.gpu_id) Gs = { "stylegan2_ffhq" : build_generator("stylegan2_ffhq").net, "stylegan2_bedroom" : build_generator("stylegan2_bedroom").net, "stylegan2_church" : build_generator("stylegan2_church").net } Ps = { "stylegan2_ffhq" : FaceSegmenter(), "stylegan2_bedroom" : SceneSegmenter(model_name="stylegan2_bedroom"), "stylegan2_church" : SceneSegmenter(model_name="stylegan2_church") } if ".pth" in args.eval_file: eval_single(Gs, Ps, args.eval_file) else: eval_files = glob.glob(f"{args.eval_file}/*.pth") eval_files.sort() for eval_file in eval_files:
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')
from predictors.scene_segmenter import SceneSegmenter print(f"=> Loading from {args.SE}") if "baseline" in args.SE: SE_name = args.SE pred = True else: SE = load_semantic_extractor(args.SE) SE.cuda().eval() SE_name = args.SE[args.SE.rfind("/") + 1 : args.SE.rfind(".pth")] pred = False print(SE_name) G_name = listkey_convert(args.SE, ["stylegan2_ffhq", "stylegan2_bedroom", "stylegan2_church"]) out_name = SE_name if pred else f"{G_name}_{SE_name}" if os.path.exists(f"{args.out_dir}/{out_name}.pth"): print(f"=> Skip {out_name}") exit(0) G = build_generator(G_name).net P = FaceSegmenter() if "ffhq" in G_name else SceneSegmenter(model_name=G_name) labels = read_labels(G_name, G, P) P_ = P if pred else SE if pred: out_name = SE_name z, wp = SCS.sseg_se(P_, G, labels, n_iter=args.n_iter, n_init=args.n_init, pred=pred, repeat=args.repeat, latent_strategy=args.latent_strategy) torch.save([z, wp], f"{args.out_dir}/{out_name}.pth")
def load_model(model_name, logger=None): model_load_state = st.text('Loading GAN model...') model = build_generator(model_name, logger=logger) model_load_state.empty() return model
def main(): """Main function.""" args = parse_args() file_name = args.output_file work_dir = args.output_dir or f'{args.model_name}_synthesis' logger_name = f'{args.model_name}_synthesis_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] if args.generate_prediction: logger.info(f'Initializing predictor.') predictor = build_predictor(args.predictor_name) if args.generate_html: viz_size = None if args.viz_size == 0 else args.viz_size visualizer = HtmlPageVisualizer(num_rows=args.html_row, num_cols=args.html_col, grid_size=total_num, viz_size=viz_size) logger.info(f'Generating {total_num} samples.') results = defaultdict(list) predictions = defaultdict(list) pbar = tqdm(total=total_num, leave=False) for inputs in model.get_batch_inputs(latent_codes): outputs = model.easy_synthesize( latent_codes=inputs, latent_space_type=args.latent_space_type, generate_style=args.generate_style, generate_image=not args.skip_image) for key, val in outputs.items(): if key == 'image': if args.generate_prediction: pred_outputs = predictor.easy_predict(val) for pred_key, pred_val in pred_outputs.items(): predictions[pred_key].append(pred_val) for image in val: if args.save_raw_synthesis: dest = os.path.join(work_dir, f'{pbar.n:06d}.jpg') if file_name != "": dest = os.path.join(work_dir, file_name) print('saving image to ', dest) save_image(dest, image) if args.generate_html: row_idx = pbar.n // visualizer.num_cols col_idx = pbar.n % visualizer.num_cols visualizer.set_cell(row_idx, col_idx, image=image) pbar.update(1) else: results[key].append(val) if 'image' not in outputs: pbar.update(inputs.shape[0]) pbar.close() logger.info(f'Saving results.') if args.generate_html: visualizer.save(os.path.join(work_dir, args.html_name)) for key, val in results.items(): np.save(os.path.join(work_dir, f'{key}.npy'), np.concatenate(val, axis=0)) if predictions: if args.predictor_name == 'scene': # Categories categories = np.concatenate(predictions['category'], axis=0) detailed_categories = { 'score': categories, 'name_to_idx': predictor.category_name_to_idx, 'idx_to_name': predictor.category_idx_to_name, } np.save(os.path.join(work_dir, 'category.npy'), detailed_categories) # Attributes attributes = np.concatenate(predictions['attribute'], axis=0) detailed_attributes = { 'score': attributes, 'name_to_idx': predictor.attribute_name_to_idx, 'idx_to_name': predictor.attribute_idx_to_name, } np.save(os.path.join(work_dir, 'attribute.npy'), detailed_attributes) else: for key, val in predictions.items(): np.save(os.path.join(work_dir, f'{key}.npy'), np.concatenate(val, axis=0))
def main(): """Main function.""" args = parse_args() set_cuda_devices(args.gpu_id) work_dir = args.output_dir os.system("mkdir " + work_dir) model = build_generator( args.model_name, truncation_psi=None if args.truncation < 0 else args.truncation) if os.path.isfile(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!') latent_codes = model.easy_sample( num=args.num, latent_space_type=args.latent_space_type) total_num = latent_codes.shape[0] if args.generate_prediction: predictor = build_predictor(args.predictor_name) if args.generate_html: visualizer = HtmlPageVisualizer(num_rows=args.html_row, num_cols=args.html_col, grid_size=total_num, viz_size=args.viz_size) results = defaultdict(list) predictions = defaultdict(list) pbar = tqdm(total=total_num, leave=False) for inputs in model.get_batch_inputs(latent_codes): outputs = model.easy_synthesize( latent_codes=inputs, latent_space_type=args.latent_space_type, generate_style=args.generate_style, generate_image=not args.skip_image) for key, val in outputs.items(): if key == 'image': if args.generate_prediction: pred_outputs = predictor.easy_predict(val) for pred_key, pred_val in pred_outputs.items(): predictions[pred_key].append(pred_val) for image in val: if args.save_raw_synthesis: save_image(os.path.join(work_dir, f'{pbar.n:06d}.jpg'), image) if args.generate_html: row_idx = pbar.n // visualizer.num_cols col_idx = pbar.n % visualizer.num_cols visualizer.set_cell(row_idx, col_idx, text=f'Sample {pbar.n:06d}', image=image) pbar.update(1) else: results[key].append(val) if 'image' not in outputs: pbar.update(inputs.shape[0]) pbar.close() if args.generate_html: visualizer.save(os.path.join(work_dir, args.html_name)) for key, val in results.items(): np.save(os.path.join(work_dir, f'{key}.npy'), np.concatenate(val, axis=0)) if predictions: print(len(predictions)) predictor.save(predictions, work_dir)
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}'))