Geometry

def convert(
    rotation: Float[torch.Tensor, "B D"],
    translation: Float[torch.Tensor, "B 3"],
    parameterization: Parameterization,
    convention: str = None,
    degrees: bool = True,
    isocenter: Float[torch.Tensor, "3"] | None = None,
) -> Float[torch.Tensor, "B 4 4"]:
    """Convert a rotation parameterization + camera center into a batch of camera
    poses in SE(3). The `translation` is interpreted as the camera center in world
    coordinates, i.e., the resulting matrix stores `t = R @ translation`.

    Args:
        rotation: Rotation parameters, shape depends on parameterization:

            - `EULER`:                (B, 3) Euler angles
            - `QUATERNION`:           (B, 4) unit quaternion (XYZW)
            - `QUATERNION_ADJUGATE`:  (B, 10) upper-tri of 4x4 symmetric matrix
            - `ROTATION_9D`:          (B, 9) flattened 3x3 matrix (projected via SVD)
            - `SE3_LOG`:              (B, 3) rotation part of se(3) logarithm
        translation: Camera center in world coordinates, shape (B, 3).
            For `SE3_LOG` this is the log-translation (coupled via the V-matrix).
        parameterization: Which rotation representation to use.
        convention: Required for EULER only. 3-letter string from {X, Y, Z}, e.g., "XYZ".
        degrees: If `True` and parameterization is EULER, interpret angles in degrees.
        isocenter: If provided, the rotation is applied about this point
            instead of the world origin. The resulting translation is adjusted
            so that the camera orbits around `isocenter`.

    Returns:
        Batched SE(3) matrices of shape (B, 4, 4).
    """
    parameterization = Parameterization(parameterization)

    if parameterization == Parameterization.SE3_LOG:
        return _se3_exp_map(rotation, translation)

    R = rotation_to_matrix(rotation, parameterization, convention, degrees)
    t = torch.einsum("bij, bj -> bi", R, translation)
    if isocenter is not None:
        t = t + isocenter
    return make_se3(R, t)

def transform_point(
    xform: Float[torch.Tensor, "B *_ N+1 N+1"],
    v: Float[torch.Tensor, "B *_ N"],
) -> Float[torch.Tensor, "B *_ N"]:
    """Apply homogeneous transformation to points.

    Args:
        xform: Transformation matrix of shape (B, ..., N+1, N+1).
        v: Points of shape (B, ..., N).

    Returns:
        Transformed points of shape (B, ..., N).
    """
    return _dehomo(torch.einsum("b...ij,b...j->b...i", xform, _homo(v)))