# python3.7 """Manipulates synthesized or real images with existing boundary. Support StyleGAN2 and StyleGAN3. """ import os.path import argparse import numpy as np from tqdm import tqdm import torch from models import build_model from utils.visualizers.html_visualizer import HtmlVisualizer from utils.image_utils import save_image from utils.parsing_utils import parse_index from utils.image_utils import postprocess_image from utils.custom_utils import to_numpy, linear_interpolate from utils.custom_utils import make_transform def parse_args(): """Parses arguments.""" parser = argparse.ArgumentParser() group = parser.add_argument_group('General options.') group.add_argument('weight_path', type=str, help='Weight path to the pre-trained model.') group.add_argument('boundary_path', type=str, help='Path to the attribute vectors.') group.add_argument('--save_dir', type=str, default=None, help='Directory to save the results. If not specified, ' 'the results will be saved to ' '`work_dirs/{TASK_SPECIFIC}/` by default.') group.add_argument('--job', type=str, default='manipulations', help='Name for the job. (default: manipulations)') group.add_argument('--seed', type=int, default=4, help='Seed for sampling. (default: 4)') group.add_argument('--nums', type=int, default=10, help='Number of samples to synthesized. (default: 10)') group.add_argument('--img_size', type=int, default=1024, help='Size of the synthesized images. (default: 1024)') group.add_argument('--vis_size', type=int, default=256, help='Size of the visualize images. (default: 256)') group.add_argument('--w_dim', type=int, default=512, help='Dimension of the latent w. (default: 512)') group.add_argument('--batch_size', type=int, default=4, help='Batch size. (default: 4)') group.add_argument('--save_jpg', action='store_true', default=False, help='Whether to save raw image. (default: False)') group.add_argument('-d', '--data_name', type=str, default='ffhq', help='Name of the datasets. (default: ffhq)') group.add_argument('--latent_path', type=str, default='', help='Path to the given latent codes. (default: None)') group.add_argument('--trunc_psi', type=float, default=0.7, help='Psi factor used for truncation. (default: 0.7)') group.add_argument('--trunc_layers', type=int, default=8, help='Number of layers to perform truncation.' ' (default: 8)') group.add_argument('--name', type=str, default='resefa', help='Name of help save the results.') group = parser.add_argument_group('StyleGAN2') group.add_argument('--stylegan2', action='store_true', help='Whether or not using StyleGAN2. (default: False)') group.add_argument('--scale_stylegan2', type=float, default=1.0, help='Scale for the number of channel fro stylegan2.') group.add_argument('--randomize_noise', type=str, default='const', help='Noise type when editing. (const or random)') group = parser.add_argument_group('StyleGAN3') group.add_argument('--stylegan3', action='store_true', help='Whether or not using StyleGAN3. (default: False)') group.add_argument('--cfg', type=str, default='T', help='Config of the stylegan3 (T/R)') group.add_argument('--scale_stylegan3r', type=float, default=2.0, help='Scale for the number of channel for stylegan3 R.') group.add_argument('--scale_stylegan3t', type=float, default=1.0, help='Scale for the number of channel for stylegan3 T.') group.add_argument('--tx', type=float, default=0, help='Translate X-coordinate. (default: 0.0)') group.add_argument('--ty', type=float, default=0, help='Translate Y-coordinate. (default: 0.0)') group.add_argument('--rotate', type=float, default=0, help='Rotation angle in degrees. (default: 0)') group = parser.add_argument_group('Manipulation') group.add_argument('--mani_layers', type=str, default='4,5,6,7', help='The layers will be manipulated.' '(default: 4,5,6,7). For the eyebrow and lipstick,' 'using [8-11] layers instead.') group.add_argument('--step', type=int, default=7, help='Number of manipulation steps. (default: 7)') group.add_argument('--start', type=int, default=0, help='The start index of the manipulation directions.') group.add_argument('--end', type=int, default=1, help='The end index of the manipulation directions.') group.add_argument('--start_distance', type=float, default=-10.0, help='Start distance for manipulation. (default: -10.0)') group.add_argument('--end_distance', type=float, default=10.0, help='End distance for manipulation. (default: 10.0)') return parser.parse_args() def main(): """Main function.""" args = parse_args() # Parse model configuration. assert (args.stylegan2 and not args.stylegan3) or \ (not args.stylegan2 and args.stylegan3) checkpoint_path = args.weight_path boundary_path = args.boundary_path assert os.path.exists(checkpoint_path) assert os.path.exists(boundary_path) boundary_name = os.path.splitext(os.path.basename(boundary_path))[0] job_disc = '' if args.stylegan2: config = dict(model_type='StyleGAN2Generator', resolution=args.img_size, w_dim=args.w_dim, fmaps_base=int(args.scale_stylegan2 * (32 << 10)), fmaps_max=512,) job_disc += 'stylegan2' else: if args.stylegan3 and args.cfg == 'R': config = dict(model_type='StyleGAN3Generator', resolution=args.img_size, w_dim=args.w_dim, fmaps_base=int(args.scale_stylegan3r * (32 << 10)), fmaps_max=1024, use_radial_filter=True,) job_disc += 'stylegan3r' elif args.stylegan3 and args.cfg == 'T': config = dict(model_type='StyleGAN3Generator', resolution=args.img_size, w_dim=args.w_dim, fmaps_base=int(args.scale_stylegan3t * (32 << 10)), fmaps_max=512, use_radial_filter=False, kernel_size=3,) job_disc += 'stylegan3t' else: raise TypeError(f'StyleGAN3 config type error, need `R/T`,' f' but got {args.cfg} instead.') # Get work directory and job name. save_dir = args.save_dir or f'work_dirs/{args.job}/{args.data_name}' os.makedirs(save_dir, exist_ok=True) job_name = f'seed_{args.seed}_num_{args.nums}_{job_disc}_{boundary_name}' os.makedirs(f'{save_dir}/{job_name}', exist_ok=True) print('Building generator...') generator = build_model(**config) print(f'Loading checkpoint from `{checkpoint_path}` ...') checkpoint = torch.load(checkpoint_path, map_location='cpu')['models'] if 'generator_smooth' in checkpoint: generator.load_state_dict(checkpoint['generator_smooth']) else: generator.load_state_dict(checkpoint['generator']) generator = generator.eval().cuda() print('Finish loading checkpoint.') if args.stylegan3 and hasattr(generator.synthesis, 'early_layer'): m = make_transform(args.tx, args.ty, args.rotate) m = np.linalg.inv(m) generator.synthesis.early_layer.transform.copy_(torch.from_numpy(m)) np.random.seed(args.seed) torch.manual_seed(args.seed) if os.path.exists(args.latent_path): print(f'Load latent codes from {args.latent_path}') latent_zs = np.load(args.latent_path) latent_zs = latent_zs[:args.nums] else: print('Sampling latent code randomly') latent_zs = np.random.randn(args.nums, generator.z_dim) latent_zs = torch.from_numpy(latent_zs.astype(np.float32)) latent_zs = latent_zs.cuda() num_images = latent_zs.shape[0] wp = [] for idx in range(0, num_images, args.batch_size): latent_z = latent_zs[idx:idx+args.batch_size] latent_w_ = generator.mapping(latent_z, None)['wp'] wp.append(latent_w_) wp = torch.cat(wp, dim=0) trunc_psi = args.trunc_psi trunc_layers = args.trunc_layers if trunc_psi < 1.0 and trunc_layers > 0: w_avg = generator.w_avg w_avg = w_avg.reshape(1, -1, generator.w_dim)[:, :trunc_layers] wp[:, :trunc_layers] = w_avg.lerp(wp[:, :trunc_layers], trunc_psi) print(f'Shape of the latent ws: {wp.shape}') image_list = [] for i in range(num_images): image_list.append(f'{i:06d}') print('Loading boundary.') directions = np.load(boundary_path) layer_index = parse_index(args.mani_layers) if not layer_index: layer_index = list(range(generator.num_layers - 1)) print(f'Manipulating on layers `{layer_index}`.') vis_size = None if args.vis_size == 0 else args.vis_size delta_num = args.end - args.start visualizer = HtmlVisualizer(num_rows=num_images * delta_num, num_cols=args.step + 2, image_size=vis_size) visualizer.set_headers( ['Name', 'Origin'] + [f'Step {i:02d}' for i in range(1, args.step + 1)] ) # Manipulate images. print('Start manipulation.') for row in tqdm(range(num_images)): latent_w = wp[row:row+1] images_ori = generator.synthesis(latent_w)['image'] images_ori = postprocess_image(to_numpy(images_ori)) if args.save_jpg: save_image(f'{save_dir}/{job_name}/{row:06d}_orin.jpg', images_ori[0]) for num_direc in range(args.start, args.end): html_row = num_direc - args.start direction = directions[num_direc:num_direc+1] direction = np.tile(direction, [1, generator.num_layers, 1]) visualizer.set_cell(row * delta_num + html_row, 0, text=f'{image_list[row]}_{num_direc:03d}') visualizer.set_cell(row * delta_num + html_row, 1, image=images_ori[0]) mani_codes = linear_interpolate(latent_code=to_numpy(latent_w), boundary=direction, layer_index=layer_index, start_distance=args.start_distance, end_distance=args.end_distance, steps=args.step) mani_codes = torch.from_numpy(mani_codes.astype(np.float32)).cuda() for idx in range(0, mani_codes.shape[0], args.batch_size): codes_ = mani_codes[idx:idx+args.batch_size] images_ = generator.synthesis(codes_)['image'] images_ = postprocess_image(to_numpy(images_)) for i in range(images_.shape[0]): visualizer.set_cell(row * delta_num + html_row, idx+i+2, image=images_[i]) if args.save_jpg: save_image(f'{save_dir}/{job_name}/{row:06d}_ind_' f'{num_direc:06d}_mani_{idx+i:06d}.jpg', images_[i]) # Save results. np.save(f'{save_dir}/{job_name}/latent_codes.npy', to_numpy(wp)) visualizer.save(f'{save_dir}/{job_name}_{args.name}.html') if __name__ == '__main__': main()