Method test_gui_with_editable_data

tasks/radnerfs/radnerf.py:338–385

(self, pose, intrinsics, W, H, cond_wins, index=0, bg_color=None, spp=1, downscale=1)

    ### GUI utils
    def test_gui_with_editable_data(self, pose, intrinsics, W, H, cond_wins, index=0, bg_color=None, spp=1, downscale=1):
    # def test_gui_with_edited_data(self, pose, intrinsics, W, H, cond_wins, index=0, bg_color=None, downscale=1):
        # render resolution (may need downscale to for better frame rate)
        rH = int(H * downscale)
        rW = int(W * downscale)
        intrinsics = intrinsics * downscale

        cond_wins = cond_wins.cuda()
        pose = torch.from_numpy(pose).unsqueeze(0).cuda()
        rays = get_rays(pose, intrinsics, rH, rW, -1)
        bg_coords = get_bg_coords(rH, rW, 'cuda')

        sample = {
            'rays_o': rays['rays_o'].cuda(),
            'rays_d': rays['rays_d'].cuda(),
            'H': rH,
            'W': rW,
            'cond_wins': cond_wins,
            'idx': [index], # support choosing index for individual codes
            'pose': convert_poses(pose),
            'bg_coords': bg_coords,
            'bg_img': bg_color.cuda()
        }

        self.model.eval()
        with torch.no_grad():
            with torch.cuda.amp.autocast(enabled=hparams['amp']):
                # here spp is used as perturb random seed!
                # face: do not perturb for the first spp, else lead to scatters.
                infer_outputs = self.run_model(sample, infer=True)
            preds = infer_outputs['rgb_map'].reshape([1,rH, rW, 3])
            preds_depth = infer_outputs['depth_map'].reshape([1, rH, rW])

        # interpolation to the original resolution
        if downscale != 1:
            # TODO: have to permute twice with torch...
            preds = F.interpolate(preds.permute(0, 3, 1, 2), size=(H, W), mode='bilinear').permute(0, 2, 3, 1).contiguous()
            preds_depth = F.interpolate(preds_depth.unsqueeze(1), size=(H, W), mode='nearest').squeeze(1)

        pred = preds[0].detach().cpu().numpy()
        pred_depth = preds_depth[0].detach().cpu().numpy()

        outputs = {
            'image': pred,
            'depth': pred_depth,
        }

        return outputs

    # [GUI] test with provided data
    def test_gui_with_data(self, sample, target_W, target_H):