facebook/StyleNeRF
34
1# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved2 3# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.4#5# NVIDIA CORPORATION and its licensors retain all intellectual property6# and proprietary rights in and to this software, related documentation7# and any modifications thereto. Any use, reproduction, disclosure or8# distribution of this software and related documentation without an express9# license agreement from NVIDIA CORPORATION is strictly prohibited.10 11"""Generate images using pretrained network pickle."""12 13import os14import re15import time16import glob17from typing import List, Optional18 19import click20import dnnlib21import numpy as np22import PIL.Image23import torch24import imageio25import legacy26from renderer import Renderer27 28#----------------------------------------------------------------------------29 30def num_range(s: str) -> List[int]:31 '''Accept either a comma separated list of numbers 'a,b,c' or a range 'a-c' and return as a list of ints.'''32 33 range_re = re.compile(r'^(\d+)-(\d+)$')34 m = range_re.match(s)35 if m:36 return list(range(int(m.group(1)), int(m.group(2))+1))37 vals = s.split(',')38 return [int(x) for x in vals]39 40#----------------------------------------------------------------------------41os.environ['PYOPENGL_PLATFORM'] = 'egl'42 43@click.command()44@click.pass_context45@click.option('--network', 'network_pkl', help='Network pickle filename', required=True)46@click.option('--seeds', type=num_range, help='List of random seeds')47@click.option('--trunc', 'truncation_psi', type=float, help='Truncation psi', default=1, show_default=True)48@click.option('--class', 'class_idx', type=int, help='Class label (unconditional if not specified)')49@click.option('--noise-mode', help='Noise mode', type=click.Choice(['const', 'random', 'none']), default='const', show_default=True)50@click.option('--projected-w', help='Projection result file', type=str, metavar='FILE')51@click.option('--outdir', help='Where to save the output images', type=str, required=True, metavar='DIR')52@click.option('--render-program', default=None, show_default=True)53@click.option('--render-option', default=None, type=str, help="e.g. up_256, camera, depth")54@click.option('--n_steps', default=8, type=int, help="number of steps for each seed")55@click.option('--no-video', default=False)56@click.option('--relative_range_u_scale', default=1.0, type=float, help="relative scale on top of the original range u")57def generate_images(58 ctx: click.Context,59 network_pkl: str,60 seeds: Optional[List[int]],61 truncation_psi: float,62 noise_mode: str,63 outdir: str,64 class_idx: Optional[int],65 projected_w: Optional[str],66 render_program=None,67 render_option=None,68 n_steps=8,69 no_video=False,70 relative_range_u_scale=1.071):72 73 74 device = torch.device('cuda')75 if os.path.isdir(network_pkl):76 network_pkl = sorted(glob.glob(network_pkl + '/*.pkl'))[-1]77 print('Loading networks from "%s"...' % network_pkl)78 79 with dnnlib.util.open_url(network_pkl) as f:80 network = legacy.load_network_pkl(f)81 G = network['G_ema'].to(device) # type: ignore82 D = network['D'].to(device)83 # from fairseq import pdb;pdb.set_trace()84 os.makedirs(outdir, exist_ok=True)85 86 # Labels.87 label = torch.zeros([1, G.c_dim], device=device)88 if G.c_dim != 0:89 if class_idx is None:90 ctx.fail('Must specify class label with --class when using a conditional network')91 label[:, class_idx] = 192 else:93 if class_idx is not None:94 print ('warn: --class=lbl ignored when running on an unconditional network')95 96 # avoid persistent classes... 97 from training.networks import Generator98 # from training.stylenerf import Discriminator99 from torch_utils import misc100 with torch.no_grad():101 G2 = Generator(*G.init_args, **G.init_kwargs).to(device)102 misc.copy_params_and_buffers(G, G2, require_all=False)103 # D2 = Discriminator(*D.init_args, **D.init_kwargs).to(device)104 # misc.copy_params_and_buffers(D, D2, require_all=False)105 G2 = Renderer(G2, D, program=render_program)106 107 # Generate images.108 all_imgs = []109 110 def stack_imgs(imgs):111 img = torch.stack(imgs, dim=2)112 return img.reshape(img.size(0) * img.size(1), img.size(2) * img.size(3), 3)113 114 def proc_img(img): 115 return (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).cpu()116 117 if projected_w is not None:118 ws = np.load(projected_w)119 ws = torch.tensor(ws, device=device) # pylint: disable=not-callable120 img = G2(styles=ws, truncation_psi=truncation_psi, noise_mode=noise_mode, render_option=render_option)121 assert isinstance(img, List)122 imgs = [proc_img(i) for i in img]123 all_imgs += [imgs]124 125 else:126 for seed_idx, seed in enumerate(seeds):127 print('Generating image for seed %d (%d/%d) ...' % (seed, seed_idx, len(seeds)))128 G2.set_random_seed(seed)129 z = torch.from_numpy(np.random.RandomState(seed).randn(2, G.z_dim)).to(device)130 relative_range_u = [0.5 - 0.5 * relative_range_u_scale, 0.5 + 0.5 * relative_range_u_scale]131 outputs = G2(132 z=z,133 c=label,134 truncation_psi=truncation_psi,135 noise_mode=noise_mode,136 render_option=render_option,137 n_steps=n_steps,138 relative_range_u=relative_range_u,139 return_cameras=True)140 if isinstance(outputs, tuple):141 img, cameras = outputs142 else:143 img = outputs144 145 if isinstance(img, List):146 imgs = [proc_img(i) for i in img]147 if not no_video:148 all_imgs += [imgs]149 150 curr_out_dir = os.path.join(outdir, 'seed_{:0>6d}'.format(seed))151 os.makedirs(curr_out_dir, exist_ok=True)152 153 if (render_option is not None) and ("gen_ibrnet_metadata" in render_option):154 intrinsics = []155 poses = []156 _, H, W, _ = imgs[0].shape157 for i, camera in enumerate(cameras):158 intri, pose, _, _ = camera159 focal = (H - 1) * 0.5 / intri[0, 0, 0].item()160 intri = np.diag([focal, focal, 1.0, 1.0]).astype(np.float32)161 intri[0, 2], intri[1, 2] = (W - 1) * 0.5, (H - 1) * 0.5162 163 pose = pose.squeeze().detach().cpu().numpy() @ np.diag([1, -1, -1, 1]).astype(np.float32)164 intrinsics.append(intri)165 poses.append(pose)166 167 intrinsics = np.stack(intrinsics, axis=0)168 poses = np.stack(poses, axis=0)169 170 np.savez(os.path.join(curr_out_dir, 'cameras.npz'), intrinsics=intrinsics, poses=poses)171 with open(os.path.join(curr_out_dir, 'meta.conf'), 'w') as f:172 f.write('depth_range = {}\ntest_hold_out = {}\nheight = {}\nwidth = {}'.173 format(G2.generator.synthesis.depth_range, 2, H, W))174 175 img_dir = os.path.join(curr_out_dir, 'images_raw')176 os.makedirs(img_dir, exist_ok=True)177 for step, img in enumerate(imgs):178 PIL.Image.fromarray(img[0].detach().cpu().numpy(), 'RGB').save(f'{img_dir}/{step:03d}.png')179 180 else:181 img = proc_img(img)[0]182 PIL.Image.fromarray(img.numpy(), 'RGB').save(f'{outdir}/seed_{seed:0>6d}.png')183 184 if len(all_imgs) > 0 and (not no_video):185 # write to video186 timestamp = time.strftime('%Y%m%d.%H%M%S',time.localtime(time.time()))187 seeds = ','.join([str(s) for s in seeds]) if seeds is not None else 'projected'188 network_pkl = network_pkl.split('/')[-1].split('.')[0]189 all_imgs = [stack_imgs([a[k] for a in all_imgs]).numpy() for k in range(len(all_imgs[0]))]190 imageio.mimwrite(f'{outdir}/{network_pkl}_{timestamp}_{seeds}.mp4', all_imgs, fps=30, quality=8)191 outdir = f'{outdir}/{network_pkl}_{timestamp}_{seeds}'192 os.makedirs(outdir, exist_ok=True)193 for step, img in enumerate(all_imgs):194 PIL.Image.fromarray(img, 'RGB').save(f'{outdir}/{step:04d}.png')195 196 197#----------------------------------------------------------------------------198 199if __name__ == "__main__":200 generate_images() # pylint: disable=no-value-for-parameter201 202#----------------------------------------------------------------------------203