| """SE(3) and quaternion utilities. |
| |
| Conventions |
| ----------- |
| * Quaternions are **scalar-first** ``[w, x, y, z]`` to match gsplat's |
| rasterization API (so gaussian rotations flow straight to the rasterizer). |
| * Camera coordinates are **OpenCV** style: x-right, y-down, z-forward; depth is |
| the camera-space z. A ``cam2world`` (a.k.a. extrinsic / pose) matrix maps a |
| point in camera coordinates to world coordinates; ``world2cam`` is its inverse |
| and is what the rasterizer consumes as ``viewmats``. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import torch |
|
|
|
|
| def normalize_quat(quat: torch.Tensor, eps: float = 1e-8) -> torch.Tensor: |
| """Normalize a ``[..., 4]`` quaternion to unit norm.""" |
| return quat / quat.norm(dim=-1, keepdim=True).clamp_min(eps) |
|
|
|
|
| def quat_to_rotmat(quat: torch.Tensor) -> torch.Tensor: |
| """Convert scalar-first quaternions ``[..., 4]`` to rotations ``[..., 3, 3]``.""" |
| quat = normalize_quat(quat) |
| w, x, y, z = quat.unbind(-1) |
| |
| tx, ty, tz = 2 * x, 2 * y, 2 * z |
| twx, twy, twz = tx * w, ty * w, tz * w |
| txx, txy, txz = tx * x, ty * x, tz * x |
| tyy, tyz, tzz = ty * y, tz * y, tz * z |
| R = torch.stack( |
| [ |
| 1 - (tyy + tzz), txy - twz, txz + twy, |
| txy + twz, 1 - (txx + tzz), tyz - twx, |
| txz - twy, tyz + twx, 1 - (txx + tyy), |
| ], |
| dim=-1, |
| ) |
| return R.reshape(quat.shape[:-1] + (3, 3)) |
|
|
|
|
| def rotmat_to_quat(R: torch.Tensor) -> torch.Tensor: |
| """Convert rotations ``[..., 3, 3]`` to scalar-first quaternions ``[..., 4]``.""" |
| m = R.reshape(R.shape[:-2] + (9,)) |
| m00, m01, m02, m10, m11, m12, m20, m21, m22 = m.unbind(-1) |
| trace = m00 + m11 + m22 |
| |
| q0 = torch.stack([1 + trace, m21 - m12, m02 - m20, m10 - m01], dim=-1) |
| q1 = torch.stack([m21 - m12, 1 + m00 - m11 - m22, m01 + m10, m02 + m20], dim=-1) |
| q2 = torch.stack([m02 - m20, m01 + m10, 1 - m00 + m11 - m22, m12 + m21], dim=-1) |
| q3 = torch.stack([m10 - m01, m02 + m20, m12 + m21, 1 - m00 - m11 + m22], dim=-1) |
| cond = torch.stack([trace, m00, m11, m22], dim=-1) |
| idx = cond.argmax(dim=-1, keepdim=True) |
| stacked = torch.stack([q0, q1, q2, q3], dim=-2) |
| quat = torch.gather(stacked, -2, idx.unsqueeze(-1).expand(idx.shape + (4,))).squeeze(-2) |
| return normalize_quat(quat) |
|
|
|
|
| def quat_multiply(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: |
| """Hamilton product of two scalar-first quaternions.""" |
| aw, ax, ay, az = a.unbind(-1) |
| bw, bx, by, bz = b.unbind(-1) |
| return torch.stack( |
| [ |
| aw * bw - ax * bx - ay * by - az * bz, |
| aw * bx + ax * bw + ay * bz - az * by, |
| aw * by - ax * bz + ay * bw + az * bx, |
| aw * bz + ax * by - ay * bx + az * bw, |
| ], |
| dim=-1, |
| ) |
|
|
|
|
| def se3_inverse(T: torch.Tensor) -> torch.Tensor: |
| """Invert a ``[..., 4, 4]`` rigid transform.""" |
| R = T[..., :3, :3] |
| t = T[..., :3, 3:] |
| Rt = R.transpose(-1, -2) |
| out = torch.zeros_like(T) |
| out[..., :3, :3] = Rt |
| out[..., :3, 3:] = -Rt @ t |
| out[..., 3, 3] = 1.0 |
| return out |
|
|
|
|
| def apply_se3(T: torch.Tensor, points: torch.Tensor) -> torch.Tensor: |
| """Apply ``[..., 4, 4]`` transform to ``[..., N, 3]`` points (broadcasting on batch).""" |
| R = T[..., :3, :3] |
| t = T[..., :3, 3] |
| return torch.einsum("...ij,...nj->...ni", R, points) + t.unsqueeze(-2) |
|
|
