Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,854 Bytes
476e0f0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
#!/usr/bin/env python
# coding=utf-8
import torch
def rotation_matrix(a, b):
"""Compute the rotation matrix that rotates vector a to vector b.
Args:
a: The vector to rotate.
b: The vector to rotate to.
Returns:
The rotation matrix.
"""
a = a / torch.linalg.norm(a)
b = b / torch.linalg.norm(b)
v = torch.cross(a, b)
c = torch.dot(a, b)
# If vectors are exactly opposite, we add a little noise to one of them
if c < -1 + 1e-8:
eps = (torch.rand(3) - 0.5) * 0.01
return rotation_matrix(a + eps, b)
s = torch.linalg.norm(v)
skew_sym_mat = torch.Tensor(
[
[0, -v[2], v[1]],
[v[2], 0, -v[0]],
[-v[1], v[0], 0],
]
)
return torch.eye(3) + skew_sym_mat + skew_sym_mat @ skew_sym_mat * ((1 - c) / (s**2 + 1e-8))
def auto_orient_and_center_poses(
poses, method="up", center_poses=True
):
"""Orients and centers the poses. We provide two methods for orientation: pca and up.
pca: Orient the poses so that the principal component of the points is aligned with the axes.
This method works well when all of the cameras are in the same plane.
up: Orient the poses so that the average up vector is aligned with the z axis.
This method works well when images are not at arbitrary angles.
Args:
poses: The poses to orient.
method: The method to use for orientation.
center_poses: If True, the poses are centered around the origin.
Returns:
The oriented poses.
"""
translation = poses[..., :3, 3]
mean_translation = torch.mean(translation, dim=0)
translation_diff = translation - mean_translation
if center_poses:
translation = mean_translation
else:
translation = torch.zeros_like(mean_translation)
if method == "pca":
_, eigvec = torch.linalg.eigh(translation_diff.T @ translation_diff)
eigvec = torch.flip(eigvec, dims=(-1,))
if torch.linalg.det(eigvec) < 0:
eigvec[:, 2] = -eigvec[:, 2]
transform = torch.cat([eigvec, eigvec @ -translation[..., None]], dim=-1)
oriented_poses = transform @ poses
if oriented_poses.mean(axis=0)[2, 1] < 0:
oriented_poses[:, 1:3] = -1 * oriented_poses[:, 1:3]
elif method == "up":
up = torch.mean(poses[:, :3, 1], dim=0)
up = up / torch.linalg.norm(up)
rotation = rotation_matrix(up, torch.Tensor([0, 0, 1]))
transform = torch.cat([rotation, rotation @ -translation[..., None]], dim=-1)
oriented_poses = transform @ poses
elif method == "none":
transform = torch.eye(4)
transform[:3, 3] = -translation
transform = transform[:3, :]
oriented_poses = transform @ poses
return oriented_poses, transform
|