Spaces:
Running
Running
| import math | |
| import trimesh | |
| import numpy as np | |
| import random | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import raymarching | |
| from .utils import custom_meshgrid, get_audio_features, euler_angles_to_matrix, convert_poses | |
| def sample_pdf(bins, weights, n_samples, det=False): | |
| # This implementation is from NeRF | |
| # bins: [B, T], old_z_vals | |
| # weights: [B, T - 1], bin weights. | |
| # return: [B, n_samples], new_z_vals | |
| # Get pdf | |
| weights = weights + 1e-5 # prevent nans | |
| pdf = weights / torch.sum(weights, -1, keepdim=True) | |
| cdf = torch.cumsum(pdf, -1) | |
| cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], -1) | |
| # Take uniform samples | |
| if det: | |
| u = torch.linspace(0. + 0.5 / n_samples, 1. - 0.5 / n_samples, steps=n_samples).to(weights.device) | |
| u = u.expand(list(cdf.shape[:-1]) + [n_samples]) | |
| else: | |
| u = torch.rand(list(cdf.shape[:-1]) + [n_samples]).to(weights.device) | |
| # Invert CDF | |
| u = u.contiguous() | |
| inds = torch.searchsorted(cdf, u, right=True) | |
| below = torch.max(torch.zeros_like(inds - 1), inds - 1) | |
| above = torch.min((cdf.shape[-1] - 1) * torch.ones_like(inds), inds) | |
| inds_g = torch.stack([below, above], -1) # (B, n_samples, 2) | |
| matched_shape = [inds_g.shape[0], inds_g.shape[1], cdf.shape[-1]] | |
| cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), 2, inds_g) | |
| bins_g = torch.gather(bins.unsqueeze(1).expand(matched_shape), 2, inds_g) | |
| denom = (cdf_g[..., 1] - cdf_g[..., 0]) | |
| denom = torch.where(denom < 1e-5, torch.ones_like(denom), denom) | |
| t = (u - cdf_g[..., 0]) / denom | |
| samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0]) | |
| return samples | |
| def plot_pointcloud(pc, color=None): | |
| # pc: [N, 3] | |
| # color: [N, 3/4] | |
| print('[visualize points]', pc.shape, pc.dtype, pc.min(0), pc.max(0)) | |
| pc = trimesh.PointCloud(pc, color) | |
| # axis | |
| axes = trimesh.creation.axis(axis_length=4) | |
| # sphere | |
| sphere = trimesh.creation.icosphere(radius=1) | |
| trimesh.Scene([pc, axes, sphere]).show() | |
| class NeRFRenderer(nn.Module): | |
| def __init__(self, opt): | |
| super().__init__() | |
| self.opt = opt | |
| self.bound = opt.bound | |
| self.cascade = 1 + math.ceil(math.log2(opt.bound)) | |
| self.grid_size = 128 | |
| self.density_scale = 1 | |
| self.min_near = opt.min_near | |
| self.density_thresh = opt.density_thresh | |
| self.density_thresh_torso = opt.density_thresh_torso | |
| self.exp_eye = opt.exp_eye | |
| self.test_train = opt.test_train | |
| self.smooth_lips = opt.smooth_lips | |
| self.torso = opt.torso | |
| self.cuda_ray = opt.cuda_ray | |
| # prepare aabb with a 6D tensor (xmin, ymin, zmin, xmax, ymax, zmax) | |
| # NOTE: aabb (can be rectangular) is only used to generate points, we still rely on bound (always cubic) to calculate density grid and hashing. | |
| aabb_train = torch.FloatTensor([-opt.bound, -opt.bound/2, -opt.bound, opt.bound, opt.bound/2, opt.bound]) | |
| aabb_infer = aabb_train.clone() | |
| self.register_buffer('aabb_train', aabb_train) | |
| self.register_buffer('aabb_infer', aabb_infer) | |
| # individual codes | |
| self.individual_num = opt.ind_num | |
| self.individual_dim = opt.ind_dim | |
| if self.individual_dim > 0: | |
| self.individual_codes = nn.Parameter(torch.randn(self.individual_num, self.individual_dim) * 0.1) | |
| if self.torso: | |
| self.individual_dim_torso = opt.ind_dim_torso | |
| if self.individual_dim_torso > 0: | |
| self.individual_codes_torso = nn.Parameter(torch.randn(self.individual_num, self.individual_dim_torso) * 0.1) | |
| # optimize camera pose | |
| self.train_camera = self.opt.train_camera | |
| if self.train_camera: | |
| self.camera_dR = nn.Parameter(torch.zeros(self.individual_num, 3)) # euler angle | |
| self.camera_dT = nn.Parameter(torch.zeros(self.individual_num, 3)) # xyz offset | |
| # extra state for cuda raymarching | |
| # 3D head density grid | |
| density_grid = torch.zeros([self.cascade, self.grid_size ** 3]) # [CAS, H * H * H] | |
| density_bitfield = torch.zeros(self.cascade * self.grid_size ** 3 // 8, dtype=torch.uint8) # [CAS * H * H * H // 8] | |
| self.register_buffer('density_grid', density_grid) | |
| self.register_buffer('density_bitfield', density_bitfield) | |
| self.mean_density = 0 | |
| self.iter_density = 0 | |
| # 2D torso density grid | |
| if self.torso: | |
| density_grid_torso = torch.zeros([self.grid_size ** 2]) # [H * H] | |
| self.register_buffer('density_grid_torso', density_grid_torso) | |
| self.mean_density_torso = 0 | |
| # step counter | |
| step_counter = torch.zeros(16, 2, dtype=torch.int32) # 16 is hardcoded for averaging... | |
| self.register_buffer('step_counter', step_counter) | |
| self.mean_count = 0 | |
| self.local_step = 0 | |
| # decay for enc_a | |
| if self.smooth_lips: | |
| self.enc_a = None | |
| def forward(self, x, d): | |
| raise NotImplementedError() | |
| # separated density and color query (can accelerate non-cuda-ray mode.) | |
| def density(self, x): | |
| raise NotImplementedError() | |
| def color(self, x, d, mask=None, **kwargs): | |
| raise NotImplementedError() | |
| def reset_extra_state(self): | |
| if not self.cuda_ray: | |
| return | |
| # density grid | |
| self.density_grid.zero_() | |
| self.mean_density = 0 | |
| self.iter_density = 0 | |
| # step counter | |
| self.step_counter.zero_() | |
| self.mean_count = 0 | |
| self.local_step = 0 | |
| def run_cuda(self, rays_o, rays_d, auds, bg_coords, poses, eye=None, index=0, dt_gamma=0, bg_color=None, perturb=False, force_all_rays=False, max_steps=1024, T_thresh=1e-4, **kwargs): | |
| # rays_o, rays_d: [B, N, 3], assumes B == 1 | |
| # auds: [B, 16] | |
| # index: [B] | |
| # return: image: [B, N, 3], depth: [B, N] | |
| prefix = rays_o.shape[:-1] | |
| rays_o = rays_o.contiguous().view(-1, 3) | |
| rays_d = rays_d.contiguous().view(-1, 3) | |
| bg_coords = bg_coords.contiguous().view(-1, 2) | |
| # only add camera offset at training! | |
| if self.train_camera and (self.training or self.test_train): | |
| dT = self.camera_dT[index] # [1, 3] | |
| dR = euler_angles_to_matrix(self.camera_dR[index] / 180 * np.pi + 1e-8).squeeze(0) # [1, 3] --> [3, 3] | |
| rays_o = rays_o + dT | |
| rays_d = rays_d @ dR | |
| N = rays_o.shape[0] # N = B * N, in fact | |
| device = rays_o.device | |
| results = {} | |
| # pre-calculate near far | |
| nears, fars = raymarching.near_far_from_aabb(rays_o, rays_d, self.aabb_train if self.training else self.aabb_infer, self.min_near) | |
| nears = nears.detach() | |
| fars = fars.detach() | |
| # encode audio | |
| enc_a = self.encode_audio(auds) # [1, 64] | |
| if enc_a is not None and self.smooth_lips: | |
| if self.enc_a is not None: | |
| _lambda = 0.35 | |
| enc_a = _lambda * self.enc_a + (1 - _lambda) * enc_a | |
| self.enc_a = enc_a | |
| if self.individual_dim > 0: | |
| if self.training: | |
| ind_code = self.individual_codes[index] | |
| # use a fixed ind code for the unknown test data. | |
| else: | |
| ind_code = self.individual_codes[0] | |
| else: | |
| ind_code = None | |
| if self.training: | |
| # setup counter | |
| counter = self.step_counter[self.local_step % 16] | |
| counter.zero_() # set to 0 | |
| self.local_step += 1 | |
| xyzs, dirs, deltas, rays = raymarching.march_rays_train(rays_o, rays_d, self.bound, self.density_bitfield, self.cascade, self.grid_size, nears, fars, counter, self.mean_count, perturb, 128, force_all_rays, dt_gamma, max_steps) | |
| sigmas, rgbs, amb_aud, amb_eye, uncertainty = self(xyzs, dirs, enc_a, ind_code, eye) | |
| sigmas = self.density_scale * sigmas | |
| #print(f'valid RGB query ratio: {mask.sum().item() / mask.shape[0]} (total = {mask.sum().item()})') | |
| # weights_sum, ambient_sum, uncertainty_sum, depth, image = raymarching.composite_rays_train_uncertainty(sigmas, rgbs, ambient.abs().sum(-1), uncertainty, deltas, rays) | |
| weights_sum, amb_aud_sum, amb_eye_sum, uncertainty_sum, depth, image = raymarching.composite_rays_train_triplane(sigmas, rgbs, amb_aud.abs().sum(-1), amb_eye.abs().sum(-1), uncertainty, deltas, rays) | |
| # for training only | |
| results['weights_sum'] = weights_sum | |
| results['ambient_aud'] = amb_aud_sum | |
| results['ambient_eye'] = amb_eye_sum | |
| results['uncertainty'] = uncertainty_sum | |
| results['rays'] = xyzs, dirs, enc_a, ind_code, eye | |
| else: | |
| dtype = torch.float32 | |
| weights_sum = torch.zeros(N, dtype=dtype, device=device) | |
| depth = torch.zeros(N, dtype=dtype, device=device) | |
| image = torch.zeros(N, 3, dtype=dtype, device=device) | |
| amb_aud_sum = torch.zeros(N, dtype=dtype, device=device) | |
| amb_eye_sum = torch.zeros(N, dtype=dtype, device=device) | |
| uncertainty_sum = torch.zeros(N, dtype=dtype, device=device) | |
| n_alive = N | |
| rays_alive = torch.arange(n_alive, dtype=torch.int32, device=device) # [N] | |
| rays_t = nears.clone() # [N] | |
| step = 0 | |
| while step < max_steps: | |
| # count alive rays | |
| n_alive = rays_alive.shape[0] | |
| # exit loop | |
| if n_alive <= 0: | |
| break | |
| # decide compact_steps | |
| n_step = max(min(N // n_alive, 8), 1) | |
| xyzs, dirs, deltas = raymarching.march_rays(n_alive, n_step, rays_alive, rays_t, rays_o, rays_d, self.bound, self.density_bitfield, self.cascade, self.grid_size, nears, fars, 128, perturb if step == 0 else False, dt_gamma, max_steps) | |
| sigmas, rgbs, ambients_aud, ambients_eye, uncertainties = self(xyzs, dirs, enc_a, ind_code, eye) | |
| sigmas = self.density_scale * sigmas | |
| # raymarching.composite_rays_uncertainty(n_alive, n_step, rays_alive, rays_t, sigmas, rgbs, deltas, ambients, uncertainties, weights_sum, depth, image, ambient_sum, uncertainty_sum, T_thresh) | |
| raymarching.composite_rays_triplane(n_alive, n_step, rays_alive, rays_t, sigmas, rgbs, deltas, ambients_aud, ambients_eye, uncertainties, weights_sum, depth, image, amb_aud_sum, amb_eye_sum, uncertainty_sum, T_thresh) | |
| rays_alive = rays_alive[rays_alive >= 0] | |
| # print(f'step = {step}, n_step = {n_step}, n_alive = {n_alive}, xyzs: {xyzs.shape}') | |
| step += n_step | |
| torso_results = self.run_torso(rays_o, bg_coords, poses, index, bg_color) | |
| bg_color = torso_results['bg_color'] | |
| image = image + (1 - weights_sum).unsqueeze(-1) * bg_color | |
| image = image.view(*prefix, 3) | |
| image = image.clamp(0, 1) | |
| depth = torch.clamp(depth - nears, min=0) / (fars - nears) | |
| depth = depth.view(*prefix) | |
| amb_aud_sum = amb_aud_sum.view(*prefix) | |
| amb_eye_sum = amb_eye_sum.view(*prefix) | |
| results['depth'] = depth | |
| results['image'] = image # head_image if train, else com_image | |
| results['ambient_aud'] = amb_aud_sum | |
| results['ambient_eye'] = amb_eye_sum | |
| results['uncertainty'] = uncertainty_sum | |
| return results | |
| def run_torso(self, rays_o, bg_coords, poses, index=0, bg_color=None, **kwargs): | |
| # rays_o, rays_d: [B, N, 3], assumes B == 1 | |
| # auds: [B, 16] | |
| # index: [B] | |
| # return: image: [B, N, 3], depth: [B, N] | |
| rays_o = rays_o.contiguous().view(-1, 3) | |
| bg_coords = bg_coords.contiguous().view(-1, 2) | |
| N = rays_o.shape[0] # N = B * N, in fact | |
| device = rays_o.device | |
| results = {} | |
| # background | |
| if bg_color is None: | |
| bg_color = 1 | |
| # first mix torso with background | |
| if self.torso: | |
| # torso ind code | |
| if self.individual_dim_torso > 0: | |
| if self.training: | |
| ind_code_torso = self.individual_codes_torso[index] | |
| # use a fixed ind code for the unknown test data. | |
| else: | |
| ind_code_torso = self.individual_codes_torso[0] | |
| else: | |
| ind_code_torso = None | |
| # 2D density grid for acceleration... | |
| density_thresh_torso = min(self.density_thresh_torso, self.mean_density_torso) | |
| occupancy = F.grid_sample(self.density_grid_torso.view(1, 1, self.grid_size, self.grid_size), bg_coords.view(1, -1, 1, 2), align_corners=True).view(-1) | |
| mask = occupancy > density_thresh_torso | |
| # masked query of torso | |
| torso_alpha = torch.zeros([N, 1], device=device) | |
| torso_color = torch.zeros([N, 3], device=device) | |
| if mask.any(): | |
| torso_alpha_mask, torso_color_mask, deform = self.forward_torso(bg_coords[mask], poses, ind_code_torso) | |
| torso_alpha[mask] = torso_alpha_mask.float() | |
| torso_color[mask] = torso_color_mask.float() | |
| results['deform'] = deform | |
| # first mix torso with background | |
| bg_color = torso_color * torso_alpha + bg_color * (1 - torso_alpha) | |
| results['torso_alpha'] = torso_alpha | |
| results['torso_color'] = bg_color | |
| # print(torso_alpha.shape, torso_alpha.max().item(), torso_alpha.min().item()) | |
| results['bg_color'] = bg_color | |
| return results | |
| def mark_untrained_grid(self, poses, intrinsic, S=64): | |
| # poses: [B, 4, 4] | |
| # intrinsic: [3, 3] | |
| if not self.cuda_ray: | |
| return | |
| if isinstance(poses, np.ndarray): | |
| poses = torch.from_numpy(poses) | |
| B = poses.shape[0] | |
| fx, fy, cx, cy = intrinsic | |
| X = torch.arange(self.grid_size, dtype=torch.int32, device=self.density_bitfield.device).split(S) | |
| Y = torch.arange(self.grid_size, dtype=torch.int32, device=self.density_bitfield.device).split(S) | |
| Z = torch.arange(self.grid_size, dtype=torch.int32, device=self.density_bitfield.device).split(S) | |
| count = torch.zeros_like(self.density_grid) | |
| poses = poses.to(count.device) | |
| # 5-level loop, forgive me... | |
| for xs in X: | |
| for ys in Y: | |
| for zs in Z: | |
| # construct points | |
| xx, yy, zz = custom_meshgrid(xs, ys, zs) | |
| coords = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1) # [N, 3], in [0, 128) | |
| indices = raymarching.morton3D(coords).long() # [N] | |
| world_xyzs = (2 * coords.float() / (self.grid_size - 1) - 1).unsqueeze(0) # [1, N, 3] in [-1, 1] | |
| # cascading | |
| for cas in range(self.cascade): | |
| bound = min(2 ** cas, self.bound) | |
| half_grid_size = bound / self.grid_size | |
| # scale to current cascade's resolution | |
| cas_world_xyzs = world_xyzs * (bound - half_grid_size) | |
| # split batch to avoid OOM | |
| head = 0 | |
| while head < B: | |
| tail = min(head + S, B) | |
| # world2cam transform (poses is c2w, so we need to transpose it. Another transpose is needed for batched matmul, so the final form is without transpose.) | |
| cam_xyzs = cas_world_xyzs - poses[head:tail, :3, 3].unsqueeze(1) | |
| cam_xyzs = cam_xyzs @ poses[head:tail, :3, :3] # [S, N, 3] | |
| # query if point is covered by any camera | |
| mask_z = cam_xyzs[:, :, 2] > 0 # [S, N] | |
| mask_x = torch.abs(cam_xyzs[:, :, 0]) < cx / fx * cam_xyzs[:, :, 2] + half_grid_size * 2 | |
| mask_y = torch.abs(cam_xyzs[:, :, 1]) < cy / fy * cam_xyzs[:, :, 2] + half_grid_size * 2 | |
| mask = (mask_z & mask_x & mask_y).sum(0).reshape(-1) # [N] | |
| # update count | |
| count[cas, indices] += mask | |
| head += S | |
| # mark untrained grid as -1 | |
| self.density_grid[count == 0] = -1 | |
| #print(f'[mark untrained grid] {(count == 0).sum()} from {resolution ** 3 * self.cascade}') | |
| def update_extra_state(self, decay=0.95, S=128): | |
| # call before each epoch to update extra states. | |
| if not self.cuda_ray: | |
| return | |
| # use random auds (different expressions should have similar density grid...) | |
| rand_idx = random.randint(0, self.aud_features.shape[0] - 1) | |
| auds = get_audio_features(self.aud_features, self.att, rand_idx).to(self.density_bitfield.device) | |
| # encode audio | |
| enc_a = self.encode_audio(auds) | |
| ### update density grid | |
| if not self.torso: # forbid updating head if is training torso... | |
| tmp_grid = torch.zeros_like(self.density_grid) | |
| # use a random eye area based on training dataset's statistics... | |
| if self.exp_eye: | |
| eye = self.eye_area[[rand_idx]].to(self.density_bitfield.device) # [1, 1] | |
| else: | |
| eye = None | |
| # full update | |
| X = torch.arange(self.grid_size, dtype=torch.int32, device=self.density_bitfield.device).split(S) | |
| Y = torch.arange(self.grid_size, dtype=torch.int32, device=self.density_bitfield.device).split(S) | |
| Z = torch.arange(self.grid_size, dtype=torch.int32, device=self.density_bitfield.device).split(S) | |
| for xs in X: | |
| for ys in Y: | |
| for zs in Z: | |
| # construct points | |
| xx, yy, zz = custom_meshgrid(xs, ys, zs) | |
| coords = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1) # [N, 3], in [0, 128) | |
| indices = raymarching.morton3D(coords).long() # [N] | |
| xyzs = 2 * coords.float() / (self.grid_size - 1) - 1 # [N, 3] in [-1, 1] | |
| # cascading | |
| for cas in range(self.cascade): | |
| bound = min(2 ** cas, self.bound) | |
| half_grid_size = bound / self.grid_size | |
| # scale to current cascade's resolution | |
| cas_xyzs = xyzs * (bound - half_grid_size) | |
| # add noise in [-hgs, hgs] | |
| cas_xyzs += (torch.rand_like(cas_xyzs) * 2 - 1) * half_grid_size | |
| # query density | |
| sigmas = self.density(cas_xyzs, enc_a, eye)['sigma'].reshape(-1).detach().to(tmp_grid.dtype) | |
| sigmas *= self.density_scale | |
| # assign | |
| tmp_grid[cas, indices] = sigmas | |
| # dilate the density_grid (less aggressive culling) | |
| tmp_grid = raymarching.morton3D_dilation(tmp_grid) | |
| # ema update | |
| valid_mask = (self.density_grid >= 0) & (tmp_grid >= 0) | |
| self.density_grid[valid_mask] = torch.maximum(self.density_grid[valid_mask] * decay, tmp_grid[valid_mask]) | |
| self.mean_density = torch.mean(self.density_grid.clamp(min=0)).item() # -1 non-training regions are viewed as 0 density. | |
| self.iter_density += 1 | |
| # convert to bitfield | |
| density_thresh = min(self.mean_density, self.density_thresh) | |
| self.density_bitfield = raymarching.packbits(self.density_grid, density_thresh, self.density_bitfield) | |
| ### update torso density grid | |
| if self.torso: | |
| tmp_grid_torso = torch.zeros_like(self.density_grid_torso) | |
| # random pose, random ind_code | |
| rand_idx = random.randint(0, self.poses.shape[0] - 1) | |
| # pose = convert_poses(self.poses[[rand_idx]]).to(self.density_bitfield.device) | |
| pose = self.poses[[rand_idx]].to(self.density_bitfield.device) | |
| if self.opt.ind_dim_torso > 0: | |
| ind_code = self.individual_codes_torso[[rand_idx]] | |
| else: | |
| ind_code = None | |
| X = torch.arange(self.grid_size, dtype=torch.int32, device=self.density_bitfield.device).split(S) | |
| Y = torch.arange(self.grid_size, dtype=torch.int32, device=self.density_bitfield.device).split(S) | |
| half_grid_size = 1 / self.grid_size | |
| for xs in X: | |
| for ys in Y: | |
| xx, yy = custom_meshgrid(xs, ys) | |
| coords = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1)], dim=-1) # [N, 2], in [0, 128) | |
| indices = (coords[:, 1] * self.grid_size + coords[:, 0]).long() # NOTE: xy transposed! | |
| xys = 2 * coords.float() / (self.grid_size - 1) - 1 # [N, 2] in [-1, 1] | |
| xys = xys * (1 - half_grid_size) | |
| # add noise in [-hgs, hgs] | |
| xys += (torch.rand_like(xys) * 2 - 1) * half_grid_size | |
| # query density | |
| alphas, _, _ = self.forward_torso(xys, pose, ind_code) # [N, 1] | |
| # assign | |
| tmp_grid_torso[indices] = alphas.squeeze(1).float() | |
| # dilate | |
| tmp_grid_torso = tmp_grid_torso.view(1, 1, self.grid_size, self.grid_size) | |
| # tmp_grid_torso = F.max_pool2d(tmp_grid_torso, kernel_size=3, stride=1, padding=1) | |
| tmp_grid_torso = F.max_pool2d(tmp_grid_torso, kernel_size=5, stride=1, padding=2) | |
| tmp_grid_torso = tmp_grid_torso.view(-1) | |
| self.density_grid_torso = torch.maximum(self.density_grid_torso * decay, tmp_grid_torso) | |
| self.mean_density_torso = torch.mean(self.density_grid_torso).item() | |
| # density_thresh_torso = min(self.density_thresh_torso, self.mean_density_torso) | |
| # print(f'[density grid torso] min={self.density_grid_torso.min().item():.4f}, max={self.density_grid_torso.max().item():.4f}, mean={self.mean_density_torso:.4f}, occ_rate={(self.density_grid_torso > density_thresh_torso).sum() / (128**2):.3f}') | |
| ### update step counter | |
| total_step = min(16, self.local_step) | |
| if total_step > 0: | |
| self.mean_count = int(self.step_counter[:total_step, 0].sum().item() / total_step) | |
| self.local_step = 0 | |
| #print(f'[density grid] min={self.density_grid.min().item():.4f}, max={self.density_grid.max().item():.4f}, mean={self.mean_density:.4f}, occ_rate={(self.density_grid > 0.01).sum() / (128**3 * self.cascade):.3f} | [step counter] mean={self.mean_count}') | |
| def get_audio_grid(self, S=128): | |
| # call before each epoch to update extra states. | |
| if not self.cuda_ray: | |
| return | |
| # use random auds (different expressions should have similar density grid...) | |
| rand_idx = random.randint(0, self.aud_features.shape[0] - 1) | |
| auds = get_audio_features(self.aud_features, self.att, rand_idx).to(self.density_bitfield.device) | |
| # encode audio | |
| enc_a = self.encode_audio(auds) | |
| tmp_grid = torch.zeros_like(self.density_grid) | |
| # use a random eye area based on training dataset's statistics... | |
| if self.exp_eye: | |
| eye = self.eye_area[[rand_idx]].to(self.density_bitfield.device) # [1, 1] | |
| else: | |
| eye = None | |
| # full update | |
| X = torch.arange(self.grid_size, dtype=torch.int32, device=self.density_bitfield.device).split(S) | |
| Y = torch.arange(self.grid_size, dtype=torch.int32, device=self.density_bitfield.device).split(S) | |
| Z = torch.arange(self.grid_size, dtype=torch.int32, device=self.density_bitfield.device).split(S) | |
| for xs in X: | |
| for ys in Y: | |
| for zs in Z: | |
| # construct points | |
| xx, yy, zz = custom_meshgrid(xs, ys, zs) | |
| coords = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1) # [N, 3], in [0, 128) | |
| indices = raymarching.morton3D(coords).long() # [N] | |
| xyzs = 2 * coords.float() / (self.grid_size - 1) - 1 # [N, 3] in [-1, 1] | |
| # cascading | |
| for cas in range(self.cascade): | |
| bound = min(2 ** cas, self.bound) | |
| half_grid_size = bound / self.grid_size | |
| # scale to current cascade's resolution | |
| cas_xyzs = xyzs * (bound - half_grid_size) | |
| # add noise in [-hgs, hgs] | |
| cas_xyzs += (torch.rand_like(cas_xyzs) * 2 - 1) * half_grid_size | |
| # query density | |
| aud_norms = self.density(cas_xyzs.to(tmp_grid.dtype), enc_a, eye)['ambient_aud'].reshape(-1).detach().to(tmp_grid.dtype) | |
| # assign | |
| tmp_grid[cas, indices] = aud_norms | |
| # dilate the density_grid (less aggressive culling) | |
| tmp_grid = raymarching.morton3D_dilation(tmp_grid) | |
| return tmp_grid | |
| # # ema update | |
| # valid_mask = (self.density_grid >= 0) & (tmp_grid >= 0) | |
| # self.density_grid[valid_mask] = torch.maximum(self.density_grid[valid_mask] * decay, tmp_grid[valid_mask]) | |
| def get_eye_grid(self, S=128): | |
| # call before each epoch to update extra states. | |
| if not self.cuda_ray: | |
| return | |
| # use random auds (different expressions should have similar density grid...) | |
| rand_idx = random.randint(0, self.aud_features.shape[0] - 1) | |
| auds = get_audio_features(self.aud_features, self.att, rand_idx).to(self.density_bitfield.device) | |
| # encode audio | |
| enc_a = self.encode_audio(auds) | |
| tmp_grid = torch.zeros_like(self.density_grid) | |
| # use a random eye area based on training dataset's statistics... | |
| if self.exp_eye: | |
| eye = self.eye_area[[rand_idx]].to(self.density_bitfield.device) # [1, 1] | |
| else: | |
| eye = None | |
| # full update | |
| X = torch.arange(self.grid_size, dtype=torch.int32, device=self.density_bitfield.device).split(S) | |
| Y = torch.arange(self.grid_size, dtype=torch.int32, device=self.density_bitfield.device).split(S) | |
| Z = torch.arange(self.grid_size, dtype=torch.int32, device=self.density_bitfield.device).split(S) | |
| for xs in X: | |
| for ys in Y: | |
| for zs in Z: | |
| # construct points | |
| xx, yy, zz = custom_meshgrid(xs, ys, zs) | |
| coords = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1) # [N, 3], in [0, 128) | |
| indices = raymarching.morton3D(coords).long() # [N] | |
| xyzs = 2 * coords.float() / (self.grid_size - 1) - 1 # [N, 3] in [-1, 1] | |
| # cascading | |
| for cas in range(self.cascade): | |
| bound = min(2 ** cas, self.bound) | |
| half_grid_size = bound / self.grid_size | |
| # scale to current cascade's resolution | |
| cas_xyzs = xyzs * (bound - half_grid_size) | |
| # add noise in [-hgs, hgs] | |
| cas_xyzs += (torch.rand_like(cas_xyzs) * 2 - 1) * half_grid_size | |
| # query density | |
| eye_norms = self.density(cas_xyzs.to(tmp_grid.dtype), enc_a, eye)['ambient_eye'].reshape(-1).detach().to(tmp_grid.dtype) | |
| # assign | |
| tmp_grid[cas, indices] = eye_norms | |
| # dilate the density_grid (less aggressive culling) | |
| tmp_grid = raymarching.morton3D_dilation(tmp_grid) | |
| return tmp_grid | |
| # # ema update | |
| # valid_mask = (self.density_grid >= 0) & (tmp_grid >= 0) | |
| # self.density_grid[valid_mask] = torch.maximum(self.density_grid[valid_mask] * decay, tmp_grid[valid_mask]) | |
| def render(self, rays_o, rays_d, auds, bg_coords, poses, staged=False, max_ray_batch=4096, **kwargs): | |
| # rays_o, rays_d: [B, N, 3], assumes B == 1 | |
| # auds: [B, 29, 16] | |
| # eye: [B, 1] | |
| # bg_coords: [1, N, 2] | |
| # return: pred_rgb: [B, N, 3] | |
| _run = self.run_cuda | |
| B, N = rays_o.shape[:2] | |
| device = rays_o.device | |
| # never stage when cuda_ray | |
| if staged and not self.cuda_ray: | |
| # not used | |
| raise NotImplementedError | |
| else: | |
| results = _run(rays_o, rays_d, auds, bg_coords, poses, **kwargs) | |
| return results | |
| def render_torso(self, rays_o, rays_d, auds, bg_coords, poses, staged=False, max_ray_batch=4096, **kwargs): | |
| # rays_o, rays_d: [B, N, 3], assumes B == 1 | |
| # auds: [B, 29, 16] | |
| # eye: [B, 1] | |
| # bg_coords: [1, N, 2] | |
| # return: pred_rgb: [B, N, 3] | |
| _run = self.run_torso | |
| B, N = rays_o.shape[:2] | |
| device = rays_o.device | |
| # never stage when cuda_ray | |
| if staged and not self.cuda_ray: | |
| # not used | |
| raise NotImplementedError | |
| else: | |
| results = _run(rays_o, bg_coords, poses, **kwargs) | |
| return results |