|
import torch |
|
import torch.nn as nn |
|
|
|
import numpy as np |
|
|
|
from additional_modules.eg3d.networks import TriPlaneGenerator |
|
from additional_modules.eg3d.camera_utils import LookAtPoseSampler, IntrinsicsSampler |
|
|
|
|
|
NETWORK_PKL = 'experiments/pretrained_models/eg3d_ffhq_rebalance.pth' |
|
|
|
|
|
class EG3DSampler(nn.Module): |
|
def __init__(self): |
|
super().__init__() |
|
|
|
eg3d_data = torch.load(NETWORK_PKL, map_location='cpu') |
|
G = TriPlaneGenerator(**eg3d_data['init_kwargs']) |
|
G.init_kwargs = eg3d_data['init_kwargs'] |
|
G.load_state_dict(eg3d_data['state_dict']) |
|
G.neural_rendering_resolution = eg3d_data['neural_rendering_resolution'] |
|
G.rendering_kwargs = eg3d_data['rendering_kwargs'] |
|
G.rendering_kwargs['ray_start'] = 2.0 |
|
G.rendering_kwargs['ray_end'] = 3.5 |
|
G.rendering_kwargs['depth_resolution'] = 52 |
|
G.rendering_kwargs['depth_resolution_importance'] = 60 |
|
self.G = G |
|
|
|
self.pose_sampler = LookAtPoseSampler() |
|
self.intrinsics_sampler = IntrinsicsSampler() |
|
|
|
self.register_buffer('lookat_position', torch.tensor([0, 0, 0])) |
|
|
|
def render(self, z, yaw, pitch): |
|
device = self.lookat_position.device |
|
lookat_position = self.lookat_position.unsqueeze(0) |
|
|
|
cam2world_pose = self.pose_sampler.sample( |
|
yaw, pitch, 2.7, |
|
lookat_position, |
|
yaw, pitch, 0.0, |
|
batch_size=1, device=device |
|
) |
|
intrinsics = self.intrinsics_sampler.sample( |
|
18.837, 0.5, |
|
0.0, 0.0, |
|
batch_size=1, device=device |
|
) |
|
|
|
radius = torch.linalg.vector_norm(cam2world_pose[:, :3, 3], dim=1, keepdim=True) |
|
conditioning_cam2world_pose = self.pose_sampler.sample( |
|
np.pi/2, np.pi/2, radius, |
|
lookat_position, |
|
np.pi/2, np.pi/2, 0, |
|
batch_size=1, device=device |
|
) |
|
|
|
camera_params = torch.cat([cam2world_pose.reshape(-1, 16), intrinsics.reshape(-1, 9)], 1) |
|
conditioning_params = torch.cat([conditioning_cam2world_pose.reshape(-1, 16), intrinsics.reshape(-1, 9)], 1) |
|
|
|
ws = self.G.mapping(z, conditioning_params, truncation_psi=0.7, truncation_cutoff=14) |
|
out = self.G.synthesis(ws, camera_params) |
|
|
|
return out['image'] |
|
|
|
@torch.no_grad() |
|
def forward(self, num_views, batch_size, z=None): |
|
device = self.lookat_position.device |
|
lookat_position = self.lookat_position.unsqueeze(0).repeat(batch_size, 1) |
|
|
|
if z is None: |
|
z = torch.randn((batch_size, self.G.z_dim), device=device) |
|
assert z.shape[0] == batch_size |
|
|
|
all_out = [] |
|
for view_idx in range(num_views): |
|
cam2world_pose = self.pose_sampler.sample( |
|
0.71, 1.11, 2.7, |
|
lookat_position, |
|
2.42, 2.02, 0.1, |
|
batch_size=batch_size, device=device |
|
) |
|
intrinsics = self.intrinsics_sampler.sample( |
|
18.837, 0.5, |
|
1.5, 0.02, |
|
batch_size=batch_size, device=device |
|
) |
|
|
|
radius = torch.linalg.vector_norm(cam2world_pose[:, :3, 3], dim=1, keepdim=True) |
|
conditioning_cam2world_pose = self.pose_sampler.sample( |
|
np.pi/2, np.pi/2, radius, |
|
lookat_position, |
|
np.pi/2, np.pi/2, 0, |
|
batch_size=batch_size, device=device |
|
) |
|
|
|
camera_params = torch.cat([cam2world_pose.reshape(-1, 16), intrinsics.reshape(-1, 9)], 1) |
|
conditioning_params = torch.cat([conditioning_cam2world_pose.reshape(-1, 16), intrinsics.reshape(-1, 9)], 1) |
|
|
|
ws = self.G.mapping(z, conditioning_params, truncation_psi=0.7, truncation_cutoff=14) |
|
out = self.G.synthesis(ws, camera_params) |
|
out['cam2world'] = cam2world_pose |
|
out['intrinsics'] = intrinsics |
|
|
|
all_out.append(out) |
|
|
|
return all_out |
|
|
