Method render_debug_images

mapping/src/mapping.py:517–585

(self, current_frame, batch_size=200000)

    @torch.no_grad()
    def render_debug_images(self, current_frame, batch_size=200000):
        rgb = current_frame.rgb
        depth = current_frame.depth
        rotation = (current_frame.get_ref_pose().cuda() @ current_frame.get_d_pose().cuda())[:3, :3]
        ind = current_frame.stamp
        w, h = self.render_res
        final_outputs = dict()

        decoder = self.decoder.cuda()
        map_states = {}
        for k, v in self.map_states.items():
            map_states[k] = v.cuda()

        rays_d = current_frame.get_rays(w, h).cuda()
        rays_d = rays_d @ rotation.transpose(-1, -2)

        rays_o = (current_frame.get_ref_pose().cuda() @ current_frame.get_d_pose().cuda())[:3, 3]
        rays_o = rays_o.unsqueeze(0).expand_as(rays_d)

        rays_o = rays_o.reshape(1, -1, 3).contiguous()
        rays_d = rays_d.reshape(1, -1, 3)
        torch.cuda.empty_cache()

        batch_size = batch_size
        ray_mask_list = []
        color_list = []
        depth_list = []
        # To prevent memory overflow, batch_size can be given according to the video memory
        for batch_iter in range(0, rays_o.shape[1], batch_size):
            final_outputs = render_rays(
                rays_o[:, batch_iter:batch_iter + batch_size, :].clone(),
                rays_d[:, batch_iter:batch_iter + batch_size, :].clone(),
                map_states,
                decoder,
                self.step_size,
                self.voxel_size,
                self.sdf_truncation,
                self.max_voxel_hit,
                self.max_distance,
                chunk_size=500000000,
                return_raw=True,
                eval=True
            )
            if final_outputs["color"] == None:
                ray_mask_list.append(final_outputs["ray_mask"])
                continue
            ray_mask_list.append(final_outputs["ray_mask"])
            depth_list.append(final_outputs["depth"])
            color_list.append(final_outputs["color"])

        ray_mask_input = torch.cat(ray_mask_list, dim=1)

        if len(depth_list) == 0:
            return None, None, None

        depth_input = torch.cat(depth_list)
        color_input = torch.cat(color_list, dim=0)