Function test_manual_pulsar_rendering

tests/pulsar_points_rendering_tests.py:113–139

(pcd=pcd, feat=feat, radius=radius)

def test_manual_pulsar_rendering(pcd=pcd, feat=feat, radius=radius):
    if isinstance(radius, torch.Tensor): radius = radius[..., 0]  # remove last dim to make it homogenous for float and tensor
    batch = add_batch(to_cuda(camera.to_batch()))
    K, R, T = batch.K, batch.R, batch.T
    C = -R.mT @ T

    acc = torch.ones_like(feat[..., :1])
    depth: torch.Tensor = (pcd - C.mT).norm(dim=-1, keepdim=True)  # world space depth
    feat = torch.cat([feat, acc, depth], dim=-1)  # B, N, 3 + C

    pcd = pcd @ R.mT + T.mT  # apply w2c conversion

    T = torch.zeros_like(T)
    R = torch.zeros_like(R)
    R[..., torch.arange(3), torch.arange(3)] = 1.0  # identity
    cam_params = get_pulsar_camera_params(R, T[..., 0], K, torch.as_tensor([HEIGHT, WIDTH], device=pcd.device)[None], batch.meta.n.item())

    # NOTE: DOUBLE CHECK MAX DEPTH USAGE
    feat = pulsar(pcd, feat, radius, cam_params, 5e-5, batch.meta.f.item(), batch.meta.n.item(), torch.zeros(5, device=pcd.device)).flip(1)

    rgb_map, acc_map, dpt_map = feat[..., :3], feat[..., 3:4], feat[..., 4:5]
    dpt_map = dpt_map + (1 - acc_map) * depth.max()
    rgb_map, acc_map, dpt_map = torch.cat([rgb_map, acc_map], dim=-1), torch.cat([acc_map, acc_map, acc_map, acc_map], dim=-1), torch.cat([dpt_map, dpt_map, dpt_map, acc_map], dim=-1)

    save_image('test_manual_pulsar_rendering_rgb.png', rgb_map[0].detach().cpu().numpy())
    save_image('test_manual_pulsar_rendering_dpt.png', dpt_map[0].detach().cpu().numpy())
    save_image('test_manual_pulsar_rendering_acc.png', acc_map[0].detach().cpu().numpy())


def test_pulsar_rendering(pcd=pcd, feat=feat, radius=radius):