How to evaluate TensoRF on a dense RGBA grid?

[UPstartDeveloper]

Hello - I noticed there is a function getDenseAlpha() in tensorBase.py to output what the alpha values TensoRF would predict, in a dense 3D voxel grid.

I am wondering if it is a good idea to extend this function to also include outputting the RGB values for each of the cells in the grid (maybe call it getDenseRGBA())? I am thinking we could use such a grid in building a RT renderer (#7), in addition to having some kind of acceleration structure.

Do you perhaps have an idea on how this could work @apchenstu ?

[apchenstu]

Did you mean view-dependent RGB values? That would be a good idea I think, but I am struggling with other things recently:(

[UPstartDeveloper]

Hi again - just wanted to give an update and get feedback if possible:

I added this new function in tensorBase.py - alot of the code is the same as in the forward() method, so I'll only show a few lines and TODOs where it's a little unfinished:

def forward_based_on_ray_points(
        self,
        rays_pts,  # [n, 3] tensor
        white_bg=True,
        is_train=False,
        ndc_ray=False,
    ):
        """
        Returns only the RGB and alpha values for each point in a dense grid,
        based on only the XYZ locations of the grid cell centers.
        """
        if ndc_ray:
            # TODO - figure this out later
        else:
            # pretending like the each ray point is a ray origin
            xyz_sampled = rays_pts.reshape((
                rays_pts.shape[0], 1, rays_pts.shape[1]
            ))
            # we only take 1 sample (and we let viewsdirs = 0, so as not to move from the origin)
            z_vals = torch.ones((xyz_sampled.shape[0], 1)).to(xyz_sampled.device)
            mask_out_bbox = (
                (xyz_sampled < self.aabb[0]) | (xyz_sampled > self.aabb[1])
            ).any(dim=-1)
            ray_valid = ~mask_out_bbox
            dists = torch.cat(
                (z_vals[:, 1:] - z_vals[:, :-1], torch.zeros_like(z_vals[:, :1])),
                dim=-1,
            )
        ...
        
        # pass in None for the viewdirs - b/c none the render modules need it
        if app_mask.any():
            app_features = self.compute_appfeature(xyz_sampled[app_mask])
            valid_rgbs = self.renderModule( 
                xyz_sampled[app_mask], None, app_features
            )
            rgb[app_mask] = valid_rgbs

        ...
        # TODO: unsure about how to get depth maps, as it requires the viewdir?
        # with torch.no_grad():
        #     depth_map = torch.sum(weight * z_vals, -1)
        #     # depth_map = depth_map + (1.0 - acc_map) * rays_chunk[..., -1]
        #     depth_map = (
        #         depth_map + (1.0 - acc_map) * torch.ones_like(xyz_sampled)[..., -1]
        #     )  # just a guess

        return rgb_map, None, alpha

Thank you Anpei for confirming this would be at least possible!