| |
| """Script that synthesizes images with pre-trained models. |
| |
| Support StyleGAN2 and StyleGAN3. |
| """ |
|
|
| import os |
| import argparse |
| from tqdm import tqdm |
| import numpy as np |
|
|
| import torch |
| from models import build_model |
| from utils.visualizers.html_visualizer import HtmlVisualizer |
| from utils.image_utils import save_image, resize_image |
| from utils.image_utils import postprocess_image |
| from utils.custom_utils import to_numpy |
|
|
|
|
| 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('--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='synthesize', |
| help='Name for the job. (default: synthesize)') |
| group.add_argument('--seed', type=int, default=4, |
| help='Seed for sampling. (default: 4)') |
| group.add_argument('--nums', type=int, default=100, |
| help='Number of samples to synthesized. (default: 100)') |
| 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 = 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 synthesizing. (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)') |
| return parser.parse_args() |
|
|
|
|
| def main(): |
| """Main function.""" |
| args = parse_args() |
| |
| assert (args.stylegan2 and not args.stylegan3) or \ |
| (not args.stylegan2 and args.stylegan3) |
| 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.') |
|
|
| |
| 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}' |
| os.makedirs(f'{save_dir}/{job_name}', exist_ok=True) |
|
|
| |
| print('Building generator...') |
| generator = build_model(**config) |
| synthesis_kwargs = dict(trunc_psi=args.trunc_psi, |
| trunc_layers=args.trunc_layers,) |
| |
| checkpoint_path = args.weight_path |
| 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.') |
|
|
| np.random.seed(args.seed) |
| torch.manual_seed(args.seed) |
| if os.path.exists(args.latent_path): |
| latent_zs = np.load(args.latent_path) |
| latent_zs = latent_zs[:args.nums] |
| else: |
| latent_zs = np.random.randn(args.nums, generator.z_dim) |
| num_images = latent_zs.shape[0] |
| latent_zs = torch.from_numpy(latent_zs.astype(np.float32)) |
| html = HtmlVisualizer(grid_size=num_images) |
| print(f'Synthesizing {num_images} images ...') |
| latent_ws = [] |
| for batch_idx in tqdm(range(0, num_images, args.batch_size)): |
| latent_z = latent_zs[batch_idx:batch_idx + args.batch_size] |
| latent_z = latent_z.cuda() |
| with torch.no_grad(): |
| g_outputs = generator(latent_z, **synthesis_kwargs) |
| g_image = to_numpy(g_outputs['image']) |
| images = postprocess_image(g_image) |
| for idx in range(images.shape[0]): |
| sub_idx = batch_idx + idx |
| img = images[idx] |
| row_idx, col_idx = divmod(sub_idx, html.num_cols) |
| image = resize_image(img, (args.vis_size, args.vis_size)) |
| html.set_cell(row_idx, col_idx, image=image, |
| text=f'Sample {sub_idx:06d}') |
| if args.save_jpg: |
| save_path = f'{save_dir}/{job_name}/{sub_idx:06d}.jpg' |
| save_image(save_path, img) |
| latent_ws.append(to_numpy(g_outputs['wp'])) |
| latent_ws = np.concatenate(latent_ws, axis=0) |
| print(f'shape of the latent code: {latent_ws.shape}') |
| np.save(f'{save_dir}/{job_name}/latent_codes.npy', latent_ws) |
| html.save(f'{save_dir}/{job_name}.html') |
| print(f'Finish synthesizing {num_images} samples.') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
