Function pinhole_projection

plib/utils.py:1199–1239

Compute the image coordinates of the 3D points in the world. Args: xyz_w: (*b_shape, *n_shape, 3) the points in world coordinate intrinsics: (*b_shape, *m_shape, 3, 3) the camera intrinsics H_c2w: (*b_shape, *m_shape, 4, 4) homege

(
        xyz_w: torch.Tensor,
        intrinsics: torch.Tensor,
        H_c2w: torch.Tensor,
        dim_b: int = 0,
)

def pinhole_projection(
        xyz_w: torch.Tensor,
        intrinsics: torch.Tensor,
        H_c2w: torch.Tensor,
        dim_b: int = 0,
) -> torch.Tensor:
    """
    Compute the image coordinates of the 3D points in the world.

    Args:
        xyz_w:
            (*b_shape, *n_shape, 3) the points in world coordinate
        intrinsics:
            (*b_shape, *m_shape, 3, 3) the camera intrinsics
        H_c2w:
            (*b_shape, *m_shape, 4, 4) homegeneous matrix (camera -> world)
        dim_b:
            length of b_shape
    Returns:
        (*b_shape, *m_shape, *n_shape, 2) (col, row) on the images, can be outside of the image boundary
    """
    *bn_shape, _ = xyz_w.shape
    *bm_shape, _, _ = intrinsics.shape
    n_shape = bn_shape[dim_b:]
    m_shape = bm_shape[dim_b:]
    assert bn_shape[:dim_b] == bm_shape[:dim_b]
    b_shape = bn_shape[:dim_b]

    H_w2c = rigid_motion.inv_homogeneous_tensors(H_c2w)  # (*b, *m, 4, 4)
    H_w2c = H_w2c.reshape(*b_shape, *m_shape, *([1] * len(n_shape)), 4, 4)  # (*b, *m, *n, 4, 4)
    intrinsics = intrinsics.reshape(*b_shape, *m_shape, *([1] * len(n_shape)), 3, 3)  # (*b, *m, *n, 3, 3)
    xyz_w = torch.cat(
        [
            xyz_w,  # (*b, *n, 3)
            torch.ones(*b_shape, *n_shape, 1, device=xyz_w.device, dtype=xyz_w.dtype),
        ], dim=-1)  # (*b, *n, 4)
    xyz_w = xyz_w.reshape(*b_shape, *([1] * len(m_shape)), *n_shape, 4, 1)  # (*b, *m, *n, 4, 1)
    xyz_c = H_w2c @ xyz_w  # (*b, *m, *n, 4, 1)
    uvw_c = intrinsics @ xyz_c[..., :3, :]  # (*b, *m, *n, 3, 1)
    uv_c = uvw_c[..., :2, 0] / uvw_c[..., 2:3, 0]  # (*b, *m, *n, 2)
    return uv_c


def find_corresponding_uv(