File size: 9,312 Bytes
112139c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | import sys
# sys.path.insert(0, os.path.abspath(os.path.join(__file__, '..')))
sys.path.append(".")
sys.path.append("..")
import math
import torch
from torch import nn
from configs.paths_config import model_paths
from PIL import Image
import numpy as np
from torch_utils import misc
from models.model import EG3DInvEncoder
from models.model import TriGenerator, TriplaneRenderer
from configs import paths_config
from models.eg3d.triplane import TriPlaneGenerator
from models.eg3d.camera_utils import FOV_to_intrinsics, LookAtPoseSampler
from models.eg3d.dual_discriminator import DualDiscriminator
class TriEncoder(nn.Module):
def __init__(self, mode="LT"):
super(TriEncoder, self).__init__()
self.device = "cuda"
self.mode = mode
self.set_encoder()
self.set_eg3d()
def set_eg3d(self):
rendering_options = {
'image_resolution': 128,
'disparity_space_sampling': False,
'clamp_mode': 'softplus',
'sr_antialias': True,
}
rendering_options.update({
'depth_resolution': 48, # number of uniform samples to take per ray.
'depth_resolution_importance': 48, # number of importance samples to take per ray.
'ray_start': 2.25, # near point along each ray to start taking samples.
'ray_end': 3.3, # far point along each ray to stop taking samples.
'box_warp': 1, # the side-length of the bounding box spanned by the tri-planes; box_warp=1 means [-0.5, -0.5, -0.5] -> [0.5, 0.5, 0.5].
'avg_camera_radius': 2.7, # used only in the visualizer to specify camera orbit radius.
'avg_camera_pivot': [0, 0, 0.2], # used only in the visualizer to control center of camera rotation.
})
img_resolution, img_channels = 512, 3
with open(paths_config.pretrained_model_path, 'rb') as f:
# with open(paths_config.model_paths["lpff_weighs"], 'rb') as f:
resume_data = torch.load(f, weights_only=False)
self.G = TriPlaneGenerator(*resume_data["G_init_args"], **resume_data["G_init_kwargs"]).eval().requires_grad_(False).to(self.device)
self.G.load_state_dict(resume_data['G'])
self.G.neural_rendering_resolution = resume_data["G_neural_rendering_resolution"]
self.G.rendering_kwargs = resume_data["G_renderinig_kwargs"]
self.D = DualDiscriminator(*resume_data["D_init_args"], **resume_data["D_init_kwargs"]).eval().requires_grad_(False).to(self.device)
self.D.load_state_dict(resume_data['D'])
# self.triplane_generator = TriGenerator(z_dim=self.G.z_dim, c_dim=self.G.c_dim, w_dim=self.G.w_dim, mapping_kwargs=mapping_kwargs, **synthesis_kwargs)
self.triplane_generator = TriGenerator(*resume_data["G_init_args"], **resume_data["G_init_kwargs"]).eval().requires_grad_(False).to(self.device)
misc.copy_params_and_buffers(self.G.backbone, self.triplane_generator.backbone, require_all=True)
self.triplane_renderer = TriplaneRenderer(img_resolution=img_resolution, img_channels=img_channels, rendering_kwargs=rendering_options).eval().requires_grad_(False).to(self.device)
misc.copy_params_and_buffers(self.G.decoder, self.triplane_renderer.decoder, require_all=True)
misc.copy_params_and_buffers(self.G.superresolution, self.triplane_renderer.superresolution, require_all=True)
# misc.copy_params_and_buffers(self.G, triplane_generator, require_all=False)
# misc.copy_params_and_buffers(self.G, triplane_renderer, require_all=False)
# self.triplane_renderer.requires_grad_(True)
self.triplane_renderer.neural_rendering_resolution = self.G.neural_rendering_resolution
self.triplane_generator.rendering_kwargs = self.G.rendering_kwargs
self.triplane_renderer.rendering_kwargs = self.G.rendering_kwargs
@staticmethod
def FOV_cxy_to_intrinsics(fov_deg, cx, cy, device='cuda'):
"""Converts FOV and image center to a 3x3 camera intrinsic matrix."""
focal_length = float(1 / (math.tan(fov_deg * 3.14159 / 360) * 1.414))
intrinsics = torch.tensor([[focal_length, 0, cx], [0, focal_length, cy], [0, 0, 1]], device=device)
return intrinsics
def eural_to_camera(self, batch_size, pitch, yaw, fov_deg=18.837, cx=0.5, cy=0.5):
# intrinsics = FOV_to_intrinsics(fov_deg, device=self.device).reshape(-1, 9).repeat(batch_size, 1)
intrinsics = self.FOV_cxy_to_intrinsics(fov_deg, cx, cy, device=self.device).reshape(-1, 9).repeat(batch_size, 1)
cam_pivot = torch.tensor(self.triplane_renderer.rendering_kwargs.get('avg_camera_pivot', [0, 0, 0.2]), device=self.device)
cam_radius = self.triplane_renderer.rendering_kwargs.get('avg_camera_radius', 2.7)
cam2world_pose = LookAtPoseSampler.sample(yaw, pitch, cam_pivot, radius=cam_radius, batch_size=batch_size, device=self.device)
camera_params = torch.cat([cam2world_pose.reshape(-1, 16), intrinsics], 1)
return camera_params
@torch.no_grad()
def sample_triplane(self, batch_size, pitch, yaw, fov_deg=18.837, cx=0.5, cy=0.5):
z = torch.randn((batch_size, self.G.z_dim)).to(self.device)
truncation_psi = 1
truncation_cutoff = 14
camera_params = self.eural_to_camera(batch_size, pitch, yaw, fov_deg, cx, cy)
ws = self.triplane_generator.mapping(z, camera_params, truncation_psi=truncation_psi, truncation_cutoff=truncation_cutoff)
triplanes = self.triplane_generator.synthesis(ws)
# img = self.triplane_renderer.synthesis(triplanes, camera_params)['image']
# gt = self.G.synthesis(ws, camera_params)['image']
gt = self.G.synthesis(ws, camera_params)
return triplanes, gt, camera_params, ws
@torch.no_grad()
def render_from_pretrain(self, batch_size, pitch, yaw, ws, fov_deg=18.837, cx=0.5, cy=0.5):
camera_params = self.eural_to_camera(batch_size, pitch, yaw, fov_deg, cx, cy)
gt = self.G.synthesis(ws, camera_params)
return gt, camera_params
@torch.no_grad()
def sample_from_synthesis(self, batch_size, pitch, yaw):
z = torch.randn((batch_size, self.G.z_dim)).to(self.device)
cam_pivot = torch.tensor(self.G.rendering_kwargs.get('avg_camera_pivot', [0, 0, 0.2]), device=self.device)
cam_radius = self.G.rendering_kwargs.get('avg_camera_radius', 2.7)
truncation_psi = 1
truncation_cutoff = 14
fov_deg = 18.837
intrinsics = FOV_to_intrinsics(fov_deg, device=self.device).reshape(-1, 9).repeat(batch_size, 1)
# cam2world_pose = LookAtPoseSampler.sample(np.pi/2 + angle_y, np.pi/2 + angle_p, cam_pivot, radius=cam_radius, device=device)
cam2world_pose = LookAtPoseSampler.sample(yaw, pitch, cam_pivot, radius=cam_radius, batch_size=batch_size, device=self.device)
conditioning_cam2world_pose = LookAtPoseSampler.sample(yaw, pitch, cam_pivot, radius=cam_radius, batch_size=batch_size, device=self.device)
camera_params = torch.cat([cam2world_pose.reshape(-1, 16), intrinsics], 1)
conditioning_params = torch.cat([conditioning_cam2world_pose.reshape(-1, 16), intrinsics], 1)
ws = self.G.mapping(z, conditioning_params, truncation_psi=truncation_psi, truncation_cutoff=truncation_cutoff)
img = self.G.synthesis(ws, camera_params)['image']
# visualize for debug
# Image.fromarray(((1 + img[0].clamp(-1,1)).cpu().numpy().transpose(1, 2, 0) / 2 * 255).astype(np.uint8)).save('test.png')
return img
def set_encoder(self):
self.encoder = EG3DInvEncoder(in_channels=5, encoder_name="resnet34", encoder_depth=3, mode=self.mode, use_bn=False)
def forward(self, x, c):
'''
Input x: [B, 5, 512, 512]
Output dict{'image', 'image_raw', ...}
'''
x = self.encoder(x) # [B, 96, 256, 256]
x = x.contiguous()
x = self.triplane_renderer(x, c)
return x
def debug():
B = 2
encoder = TriEncoder()
# a = torch.rand((B, 5, 512, 512))
# c = torch.rand((B, 25))
# b = encoder(a, c)
# print(b['image'].shape)
# rand_pitch = (torch.rand(1) - 0.5) * torch.pi
# rand_yaw = (torch.rand(1) - 0.5) * torch.pi
# print(encoder.sample_from_synthesis(B, rand_pitch.cuda(), rand_yaw.cuda()).shape)
for i in range(100):
# rand_pitch = (torch.rand(1) - 0.5) * torch.pi/2 + torch.pi/2
# rand_yaw = (torch.rand(1) - 0.5) * torch.pi/2 + torch.pi/2
rand_pitch = (torch.rand(1) - 0.5) * (26.0 / 180 * torch.pi) + torch.pi/2
rand_yaw = (torch.rand(1) - 0.5) * (39.0 / 180 * torch.pi) + torch.pi/2
rand_cx = (torch.rand(1) - 0.5) * 0.2 + 0.5
rand_cy = (torch.rand(1) - 0.5) * 0.2 + 0.5
rand_fov_deg = (torch.rand(1) - 0.5) * 4.8 + 18.837
# img = encoder.sample_from_synthesis(B, rand_pitch.cuda(), rand_yaw.cuda())
triplanes, gt, camera_params, ws = encoder.sample_triplane(B, rand_pitch.cuda(), rand_yaw.cuda(), rand_fov_deg.cuda(), rand_cx.cuda(), rand_cy.cuda())
Image.fromarray(((1 + gt['image'][0].clamp(-1,1)).cpu().numpy().transpose(1, 2, 0) / 2 * 255).astype(np.uint8)).save('test' + str(i) + '.png')
if __name__ == '__main__':
debug()
|