Method render

modules/radnerfs/radnerf_torso.py:86–199

(self, rays_o, rays_d, cond, bg_coords, poses, index=0, dt_gamma=0, bg_color=None, perturb=False, force_all_rays=False, max_steps=1024, T_thresh=1e-4, **kwargs)

        return alpha, color, dx

    def render(self, rays_o, rays_d, cond, bg_coords, poses, 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
        # cond: [B, 29, 16]
        # bg_coords: [1, N, 2]
        # return: pred_rgb: [B, N, 3]

        ### run head nerf with no_grad to get the renderred head
        with torch.no_grad():
            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)
            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
            cond_feat = self.cal_cond_feat(cond) # [1, 64]
            if self.individual_embedding_dim > 0:
                if self.training:
                    ind_code = self.individual_embeddings[index]
                # use a fixed ind code for the unknown test data.
                else:
                    ind_code = self.individual_embeddings[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)
                # xyzs, dirs, deltas, rays, points2rays = 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, ambient = self(xyzs, dirs, cond_feat, ind_code)
                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, depth, image = raymarching.composite_rays_train(sigmas, rgbs, ambient.abs().sum(-1), deltas, rays)
                # for training only
                results['weights_sum'] = weights_sum
                results['ambient'] = ambient_sum
            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)
                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