Make the triplane available when call synthesis(). Useful for buildin…

…g loss function of live3D portrait.
NVlabs · Sep 28, 2023 · a33fe63 · a33fe63
1 parent 7cf1fd1
commit a33fe63
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/eg3d/gen_samples.py b/eg3d/gen_samples.py
@@ -156,7 +156,7 @@ def generate_images(
     cam2world_pose = LookAtPoseSampler.sample(3.14/2, 3.14/2, torch.tensor([0, 0, 0.2], device=device), radius=2.7, device=device)
     intrinsics = FOV_to_intrinsics(fov_deg, device=device)
 
-    # Generate images.
+    # Generate images of the different people with three different views
     for seed_idx, seed in enumerate(seeds):
         print('Generating image for seed %d (%d/%d) ...' % (seed, seed_idx, len(seeds)))
         z = torch.from_numpy(np.random.RandomState(seed).randn(1, G.z_dim)).to(device)

diff --git a/eg3d/training/triplane.py b/eg3d/training/triplane.py
@@ -86,7 +86,7 @@ def synthesis(self, ws, c, neural_rendering_resolution=None, update_emas=False,
         rgb_image = feature_image[:, :3]
         sr_image = self.superresolution(rgb_image, feature_image, ws, noise_mode=self.rendering_kwargs['superresolution_noise_mode'], **{k:synthesis_kwargs[k] for k in synthesis_kwargs.keys() if k != 'noise_mode'})
 
-        return {'image': sr_image, 'image_raw': rgb_image, 'image_depth': depth_image}
+        return {'image': sr_image, 'image_raw': rgb_image, 'image_depth': depth_image, 'triplane': planes}
 
     def sample(self, coordinates, directions, z, c, truncation_psi=1, truncation_cutoff=None, update_emas=False, **synthesis_kwargs):
         # Compute RGB features, density for arbitrary 3D coordinates. Mostly used for extracting shapes.