Class Test_Generate_Cam_Rays

tests/plib/test_utils.py:118–172

class Test_Generate_Cam_Rays(unittest.TestCase):

    def test_1(self):
        m = 1
        width_px, height_px = 100, 50
        f = 40.
        cam_poses = rigid_motion.generate_random_camera_poses(
            n=m,
            max_angle=180.,
            min_r=0.5,
            max_r=1.5,
        )
        cam_poses = utils.to_tensor(cam_poses, dtype=torch.float)
        cam_poses = torch.stack(cam_poses, dim=0)  # (m, 4, 4)

        intrinsics = torch.tensor([
            [f, 0., width_px * 0.5],
            [0., f, height_px * 0.5],
            [0, 0, 1],
        ]).expand(m ,3, 3)  # (m, 3, 3)

        ray_origins_w, ray_directions_w = utils.generate_camera_rays(
            cam_poses=cam_poses,
            intrinsics=intrinsics,
            width_px=width_px,
            height_px=height_px,
            subsample=1,
        )

        # o3d ray
        H_w2c = torch.linalg.inv(cam_poses)  # (m, 4, 4)
        o3d_ray_origins = []
        o3d_ray_directions = []
        for i in range(m):
            rays = o3d.t.geometry.RaycastingScene.create_rays_pinhole(
                intrinsic_matrix=intrinsics[i].numpy(),
                extrinsic_matrix=H_w2c[i].numpy(),
                width_px=width_px,
                height_px=height_px,
            )  # (target_height_px, target_width_px, 6)  [ox, oy, oz, dx, dy, dz]  origin is the pinhole
            rays = torch.from_numpy(rays.numpy())  # (h, w, 6)
            ray_o = rays[:, :, :3]
            ray_d = rays[:, :, 3:]

            o3d_ray_origins.append(ray_o)
            o3d_ray_directions.append(ray_d)
        o3d_ray_origins = torch.stack(o3d_ray_origins, dim=0)  # (m, h, w, 3)
        o3d_ray_directions = torch.stack(o3d_ray_directions, dim=0)  # (m, h, w, 3)
        o3d_ray_directions = o3d_ray_directions / torch.linalg.vector_norm(o3d_ray_directions, dim=-1, keepdims=True)


        assert torch.allclose(o3d_ray_origins, ray_origins_w)
        assert torch.allclose(o3d_ray_directions, ray_directions_w, rtol=1e-4, atol=1e-4)

        diff = torch.linalg.vector_norm(o3d_ray_directions - ray_directions_w, dim=-1)