Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
+351
@@ -0,0 +1,351 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
|
||||
def perspective_projection(
|
||||
points, focal_length, camera_center, rotation=None, translation=None
|
||||
):
|
||||
"""
|
||||
Adapted from https://github.com/mkocabas/VIBE/blob/master/lib/models/spin.py
|
||||
This function computes the perspective projection of a set of points.
|
||||
Input:
|
||||
points (bs, N, 3): 3D points
|
||||
focal_length (bs, 2): Focal length
|
||||
camera_center (bs, 2): Camera center
|
||||
rotation (bs, 3, 3): OPTIONAL Camera rotation
|
||||
translation (bs, 3): OPTIONAL Camera translation
|
||||
"""
|
||||
batch_size = points.shape[0]
|
||||
K = torch.zeros([batch_size, 3, 3], device=points.device)
|
||||
K[:, 0, 0] = focal_length[:, 0]
|
||||
K[:, 1, 1] = focal_length[:, 1]
|
||||
K[:, 2, 2] = 1.0
|
||||
K[:, :-1, -1] = camera_center
|
||||
|
||||
if rotation is not None and translation is not None:
|
||||
# Transform points
|
||||
points = torch.einsum("bij,bkj->bki", rotation, points)
|
||||
points = points + translation.unsqueeze(1)
|
||||
|
||||
# Apply perspective distortion
|
||||
projected_points = points / points[..., 2:3]
|
||||
|
||||
# Apply camera intrinsics
|
||||
projected_points = torch.einsum("bij,bkj->bki", K, projected_points)
|
||||
|
||||
return projected_points[:, :, :-1]
|
||||
|
||||
|
||||
def reproject(points3d, cam_R, cam_t, cam_f, cam_center):
|
||||
"""
|
||||
reproject points3d into the scene cameras
|
||||
:param points3d (B, T, N, 3)
|
||||
:param cam_R (B, T, 3, 3)
|
||||
:param cam_t (B, T, 3)
|
||||
:param cam_f (T, 2)
|
||||
:param cam_center (T, 2)
|
||||
"""
|
||||
# print(points3d.shape, cam_R.shape, cam_t.shape, cam_f.shape, cam_center.shape)
|
||||
# print(cam_f[None, :, None].shape, cam_center[None, :, None].shape)
|
||||
# raise ValueError
|
||||
B, T, N, _ = points3d.shape
|
||||
points3d = torch.einsum("btij,btnj->btni", cam_R, points3d)
|
||||
points3d = points3d + cam_t[..., None, :] # (B, T, N, 3)
|
||||
points2d = points3d[..., :2] / points3d[..., 2:3]
|
||||
points2d = cam_f[None, :, None] * points2d + cam_center[None, :, None]
|
||||
return points2d
|
||||
|
||||
|
||||
def focal2fov(focal, R):
|
||||
"""
|
||||
:param focal, focal length
|
||||
:param R, either W / 2 or H / 2
|
||||
"""
|
||||
return 2 * np.arctan(R / focal)
|
||||
|
||||
|
||||
def fov2focal(fov, R):
|
||||
"""
|
||||
:param fov, field of view in radians
|
||||
:param R, either W / 2 or H / 2
|
||||
"""
|
||||
return R / np.tan(fov / 2)
|
||||
|
||||
|
||||
def compute_lookat_box(bb_min, bb_max, intrins):
|
||||
"""
|
||||
The center and distance to a scene with bb_min, bb_max
|
||||
to place a camera with given intrinsics
|
||||
:param bb_min (3,)
|
||||
:param bb_max (3,)
|
||||
:param intrinsics, (fx, fy, cx, cy) of camera
|
||||
:param view_angle (optional) viewing angle in radians (elevation)
|
||||
"""
|
||||
fx, fy, cx, cy = intrins
|
||||
bb_min, bb_max = torch.tensor(bb_min), torch.tensor(bb_max)
|
||||
center = 0.5 * (bb_min + bb_max)
|
||||
size = torch.linalg.norm(bb_max - bb_min)
|
||||
cam_dist = np.sqrt(fx**2 + fy**2) / np.sqrt(cx**2 + cy**2)
|
||||
cam_dist = 0.75 * size * cam_dist
|
||||
return center, cam_dist
|
||||
|
||||
|
||||
def lookat_origin(cam_dist, view_angle=-np.pi / 6):
|
||||
"""
|
||||
:param cam_dist (float)
|
||||
:param view_angle (float)
|
||||
"""
|
||||
cam_dist = np.abs(cam_dist)
|
||||
view_angle = np.abs(view_angle)
|
||||
pos = cam_dist * torch.tensor([0, np.sin(view_angle), np.cos(view_angle)])
|
||||
rot = rotx(view_angle)
|
||||
return rot, pos
|
||||
|
||||
|
||||
def lookat_matrix(source_pos, target_pos, up):
|
||||
"""
|
||||
IMPORTANT: USES RIGHT UP BACK XYZ CONVENTION
|
||||
:param source_pos (*, 3)
|
||||
:param target_pos (*, 3)
|
||||
:param up (3,)
|
||||
"""
|
||||
*dims, _ = source_pos.shape
|
||||
up = up.reshape(*(1,) * len(dims), 3)
|
||||
up = up / torch.linalg.norm(up, dim=-1, keepdim=True)
|
||||
back = normalize(target_pos - source_pos)
|
||||
right = normalize(torch.linalg.cross(up, back))
|
||||
up = normalize(torch.linalg.cross(back, right))
|
||||
R = torch.stack([right, up, back], dim=-1)
|
||||
return make_4x4_pose(R, source_pos)
|
||||
|
||||
|
||||
def normalize(x):
|
||||
return x / torch.linalg.norm(x, dim=-1, keepdim=True)
|
||||
|
||||
|
||||
def invert_camera(R, t):
|
||||
"""
|
||||
:param R (*, 3, 3)
|
||||
:param t (*, 3)
|
||||
returns Ri (*, 3, 3), ti (*, 3)
|
||||
"""
|
||||
R, t = torch.tensor(R), torch.tensor(t)
|
||||
Ri = R.transpose(-1, -2)
|
||||
ti = -torch.einsum("...ij,...j->...i", Ri, t)
|
||||
return Ri, ti
|
||||
|
||||
|
||||
def compose_cameras(R1, t1, R2, t2):
|
||||
"""
|
||||
composes [R1, t1] and [R2, t2]
|
||||
:param R1 (*, 3, 3)
|
||||
:param t1 (*, 3)
|
||||
:param R2 (*, 3, 3)
|
||||
:param t2 (*, 3)
|
||||
"""
|
||||
R = torch.einsum("...ij,...jk->...ik", R1, R2)
|
||||
t = t1 + torch.einsum("...ij,...j->...i", R1, t2)
|
||||
return R, t
|
||||
|
||||
|
||||
def matmul_nd(A, x):
|
||||
"""
|
||||
multiply batch matrix A to batch nd tensors
|
||||
:param A (B, m, n)
|
||||
:param x (B, *dims, m)
|
||||
"""
|
||||
B, m, n = A.shape
|
||||
assert len(A) == len(x)
|
||||
assert x.shape[-1] == m
|
||||
B, *dims, _ = x.shape
|
||||
return torch.matmul(A.reshape(B, *(1,) * len(dims), m, n), x[..., None])[..., 0]
|
||||
|
||||
|
||||
def view_matrix(z, up, pos):
|
||||
"""
|
||||
:param z (*, 3) up (*, 3) pos (*, 3)
|
||||
returns (*, 4, 4)
|
||||
"""
|
||||
*dims, _ = z.shape
|
||||
x = normalize(torch.linalg.cross(up, z))
|
||||
y = normalize(torch.linalg.cross(z, x))
|
||||
bottom = (
|
||||
torch.tensor([0, 0, 0, 1], dtype=torch.float32)
|
||||
.reshape(*(1,) * len(dims), 1, 4)
|
||||
.expand(*dims, 1, 4)
|
||||
)
|
||||
|
||||
return torch.cat([torch.stack([x, y, z, pos], dim=-1), bottom], dim=-2)
|
||||
|
||||
|
||||
def average_pose(poses):
|
||||
"""
|
||||
:param poses (N, 4, 4)
|
||||
returns average pose (4, 4)
|
||||
"""
|
||||
center = poses[:, :3, 3].mean(0)
|
||||
up = normalize(poses[:, :3, 1].sum(0))
|
||||
z = normalize(poses[:, :3, 2].sum(0))
|
||||
return view_matrix(z, up, center)
|
||||
|
||||
|
||||
def project_so3(M, eps=1e-4):
|
||||
"""
|
||||
:param M (N, *, 3, 3)
|
||||
"""
|
||||
N, *dims, _, _ = M.shape
|
||||
M = M * (1 + torch.rand(N, *dims, 1, 3, device=M.device))
|
||||
U, D, Vt = torch.linalg.svd(M) # (N, *, 3, 3), (N, *, 3), (N, *, 3, 3)
|
||||
detuvt = torch.linalg.det(torch.matmul(U, Vt)) # (N, *)
|
||||
S = torch.cat(
|
||||
[torch.ones(N, *dims, 2, device=M.device), detuvt[..., None]], dim=-1
|
||||
) # (N, *, 3)
|
||||
return torch.matmul(U, torch.matmul(torch.diag_embed(S), Vt))
|
||||
|
||||
|
||||
def make_translation(t):
|
||||
return make_4x4_pose(torch.eye(3), t)
|
||||
|
||||
|
||||
def make_rotation(rx=0, ry=0, rz=0, order="xyz"):
|
||||
Rx = rotx(rx)
|
||||
Ry = roty(ry)
|
||||
Rz = rotz(rz)
|
||||
if order == "xyz":
|
||||
R = Rz @ Ry @ Rx
|
||||
elif order == "xzy":
|
||||
R = Ry @ Rz @ Rx
|
||||
elif order == "yxz":
|
||||
R = Rz @ Rx @ Ry
|
||||
elif order == "yzx":
|
||||
R = Rx @ Rz @ Ry
|
||||
elif order == "zyx":
|
||||
R = Rx @ Ry @ Rz
|
||||
elif order == "zxy":
|
||||
R = Ry @ Rx @ Rz
|
||||
return make_4x4_pose(R, torch.zeros(3))
|
||||
|
||||
|
||||
def make_4x4_pose(R, t):
|
||||
"""
|
||||
:param R (*, 3, 3)
|
||||
:param t (*, 3)
|
||||
return (*, 4, 4)
|
||||
"""
|
||||
dims = R.shape[:-2]
|
||||
pose_3x4 = torch.cat([R, t.view(*dims, 3, 1)], dim=-1)
|
||||
bottom = (
|
||||
torch.tensor([0, 0, 0, 1], device=R.device)
|
||||
.reshape(*(1,) * len(dims), 1, 4)
|
||||
.expand(*dims, 1, 4)
|
||||
)
|
||||
return torch.cat([pose_3x4, bottom], dim=-2)
|
||||
|
||||
|
||||
def normalize(x):
|
||||
return x / torch.sqrt(torch.sum(x**2, dim=-1, keepdim=True))
|
||||
|
||||
|
||||
def rotx(theta):
|
||||
return torch.tensor(
|
||||
[
|
||||
[1, 0, 0],
|
||||
[0, np.cos(theta), -np.sin(theta)],
|
||||
[0, np.sin(theta), np.cos(theta)],
|
||||
],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
|
||||
def roty(theta):
|
||||
return torch.tensor(
|
||||
[
|
||||
[np.cos(theta), 0, np.sin(theta)],
|
||||
[0, 1, 0],
|
||||
[-np.sin(theta), 0, np.cos(theta)],
|
||||
],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
|
||||
def rotz(theta):
|
||||
return torch.tensor(
|
||||
[
|
||||
[np.cos(theta), -np.sin(theta), 0],
|
||||
[np.sin(theta), np.cos(theta), 0],
|
||||
[0, 0, 1],
|
||||
],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
|
||||
def relative_pose_c2w(Rwc1, Rwc2, twc1, twc2):
|
||||
"""
|
||||
compute relative pose from cam 1 to cam 2 given c2w pose matrices
|
||||
:param Rwc1, Rwc2 (N, 3, 3) cam1, cam2 to world rotations
|
||||
:param twc1, twc2 (N, 3) cam1, cam2 to world translations
|
||||
returns R21 (N, 3, 3) t21 (N, 3)
|
||||
"""
|
||||
twc1 = twc1.view(-1, 3, 1)
|
||||
twc2 = twc2.view(-1, 3, 1)
|
||||
Rc2w = Rwc2.transpose(-1, -2) # world to c2
|
||||
tc2w = -torch.matmul(Rc2w, twc2)
|
||||
Rc2c1 = torch.matmul(Rc2w, Rwc1)
|
||||
tc2c1 = tc2w + torch.matmul(Rc2w, twc1)
|
||||
return Rc2c1, tc2c1[..., 0]
|
||||
|
||||
|
||||
def relative_pose_w2c(Rc1w, Rc2w, tc1w, tc2w):
|
||||
"""
|
||||
compute relative pose from cam 1 to cam 2 given w2c camera matrices
|
||||
:param Rc1w, Rc2w (N, 3, 3) world to cam1, cam2 rotations
|
||||
:param tc1w, tc2w (N, 3) world to cam1, cam2 translations
|
||||
"""
|
||||
tc1w = tc1w.view(-1, 3, 1)
|
||||
tc2w = tc2w.view(-1, 3, 1)
|
||||
# we keep the world to cam transforms
|
||||
Rwc1 = Rc1w.transpose(-1, -2) # c1 to world
|
||||
twc1 = -torch.matmul(Rwc1, tc1w)
|
||||
Rc2c1 = torch.matmul(Rc2w, Rwc1) # c1 to c2
|
||||
tc2c1 = tc2w + torch.matmul(Rc2w, twc1)
|
||||
return Rc2c1, tc2c1[..., 0]
|
||||
|
||||
|
||||
def project(xyz_c, center, focal, eps=1e-5):
|
||||
"""
|
||||
:param xyz_c (*, 3) 3d point in camera coordinates
|
||||
:param focal (1)
|
||||
:param center (*, 2)
|
||||
return (*, 2)
|
||||
"""
|
||||
return focal * xyz_c[..., :2] / (xyz_c[..., 2:3] + eps) + center # (N, *, 2)
|
||||
|
||||
|
||||
def convert_yup(xyz):
|
||||
"""
|
||||
converts points in x right y down z forward to x right y up z back
|
||||
:param xyz (*, 3)
|
||||
"""
|
||||
x, y, z = torch.split(xyz[..., :3], 1, dim=-1)
|
||||
return torch.cat([x, -y, -z], dim=-1)
|
||||
|
||||
|
||||
def inv_project(uv, z, center, focal, yup=True):
|
||||
"""
|
||||
:param uv (*, 2)
|
||||
:param z (*, 1)
|
||||
:param center (*, 2)
|
||||
:param focal (1)
|
||||
:returns (*, 3)
|
||||
"""
|
||||
uv = uv - center
|
||||
if yup:
|
||||
return z * torch.cat(
|
||||
[uv[..., :1] / focal, -uv[..., 1:2] / focal, -torch.ones_like(uv[..., :1])],
|
||||
dim=-1,
|
||||
) # (N, *, 3)
|
||||
|
||||
return z * torch.cat(
|
||||
[uv / focal, torch.ones_like(uv[..., :1])],
|
||||
dim=-1,
|
||||
) # (N, *, 3)
|
||||
Reference in New Issue
Block a user