Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

third_party/Dyn-HaMR/dyn-hamr/geometry/__init__.py

@@ -0,0 +1,5 @@
+from . import camera
+from . import mesh
+from . import pcl
+from . import plane
+from . import rotation

third_party/Dyn-HaMR/dyn-hamr/geometry/camera.py

@@ -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)

third_party/Dyn-HaMR/dyn-hamr/geometry/mesh.py

@@ -0,0 +1,145 @@
+import os
+import numpy as np
+import trimesh
+
+
+def get_mesh_bb(mesh):
+    """
+    :param mesh - trimesh mesh object
+    returns bb_min (3), bb_max (3)
+    """
+    bb_min = mesh.vertices.max(axis=0)
+    bb_max = mesh.vertices.min(axis=0)
+    return bb_min, bb_max
+
+
+def get_scene_bb(meshes):
+    """
+    :param mesh_seqs - (potentially nested) list of trimesh objects
+    returns bb_min (3), bb_max (3)
+    """
+    if isinstance(meshes, trimesh.Trimesh):
+        return get_mesh_bb(meshes)
+
+    bb_mins, bb_maxs = zip(*[get_scene_bb(mesh) for mesh in meshes])
+    bb_mins = np.stack(bb_mins, axis=0)
+    bb_maxs = np.stack(bb_maxs, axis=0)
+    return bb_mins.min(axis=0), bb_maxs.max(axis=0)
+
+
+def make_batch_mesh(verts, faces, colors):
+    """
+    convenience function to make batch of meshes
+    meshs have same faces in batch, verts have same color in mesh
+    :param verts (B, V, 3)
+    :param faces (F, 3)
+    :param colors (B, 3)
+    """
+    B, V, _ = verts.shape
+    return [make_mesh(verts[b], faces, colors[b, None].expand(V, -1)) for b in range(B)]
+
+
+def make_mesh(verts, faces, colors=None, yup=True):
+    """
+    create a trimesh object for the faces and vertices
+    :param verts (V, 3) tensor
+    :param faces (F, 3) tensor
+    :param colors (optional) (V, 3) tensor
+    :param yup (optional bool) whether or not to save with Y up
+    """
+    verts = verts.detach().cpu().numpy()
+    faces = faces.detach().cpu().numpy()
+    if yup:
+        verts = np.array([1, -1, -1])[None, :] * verts
+    if colors is None:
+        colors = np.ones_like(verts) * 0.5
+    else:
+        colors = colors.detach().cpu().numpy()
+    return trimesh.Trimesh(
+        vertices=verts, faces=faces, vertex_colors=colors, process=False
+    )
+
+
+def save_mesh_scenes(out_dir, scenes):
+    """
+    :param scenes, list of scenes (list of meshes)
+    """
+    assert isinstance(scenes, list)
+    assert isinstance(scenes[0], list)
+    B = len(scenes[0])
+    if B == 1:
+        save_meshes_to_obj(out_dir, [x[0] for x in scenes])
+    else:
+        save_scenes_to_glb(out_dir, scenes)
+
+
+def save_scenes_to_glb(out_dir, scenes):
+    """
+    Saves a list of scenes (list of meshes) each to glb files
+    """
+    os.makedirs(out_dir, exist_ok=True)
+    for t, meshes in enumerate(scenes):
+        save_meshes_to_glb(f"{out_dir}/scene_{t:03d}.glb", meshes)
+
+
+def save_meshes_to_glb(path, meshes, names=None):
+    """
+    put trimesh meshes in a scene and export to glb
+    """
+    if names is not None:
+        assert len(meshes) == len(names)
+
+    scene = trimesh.Scene()
+    for i, mesh in enumerate(meshes):
+        name = f"mesh_{i:03d}" if names is None else names[i]
+        scene.add_geometry(mesh, node_name=name)
+
+    with open(path, "wb") as f:
+        f.write(trimesh.exchange.gltf.export_glb(scene, include_normals=True))
+
+
+def save_meshes_to_obj(out_dir, meshes, names=None):
+    if names is not None:
+        assert len(meshes) == len(names)
+
+    os.makedirs(out_dir, exist_ok=True)
+    for i, mesh in enumerate(meshes):
+        name = f"mesh_{i:03d}" if names is None else names[i]
+        path = os.path.join(out_dir, f"{name}.obj")
+        with open(path, "w") as f:
+            mesh.export(f, file_type="obj", include_color=False, include_normals=True)
+
+def vertices_to_trimesh(vertices, faces, mesh_base_color=(1.0, 1.0, 0.9), 
+                        rot_axis=[1,0,0], rot_angle=0, is_right=1):
+    # material = pyrender.MetallicRoughnessMaterial(
+    #     metallicFactor=0.0,
+    #     alphaMode='OPAQUE',
+    #     baseColorFactor=(*mesh_base_color, 1.0))
+    faces_new = np.array([[92, 38, 234],
+                            [234, 38, 239],
+                            [38, 122, 239],
+                            [239, 122, 279],
+                            [122, 118, 279],
+                            [279, 118, 215],
+                            [118, 117, 215],
+                            [215, 117, 214],
+                            [117, 119, 214],
+                            [214, 119, 121],
+                            [119, 120, 121],
+                            [121, 120, 78],
+                            [120, 108, 78],
+                            [78, 108, 79]])
+    faces = np.concatenate([faces, faces_new], axis=0)
+    vertex_colors = np.array([(*mesh_base_color, 1.0)] * vertices.shape[0])
+    # print(vertices.shape, self.faces.shape)
+    mesh = trimesh.Trimesh(vertices.copy(), faces.copy(), vertex_colors=vertex_colors)
+    # mesh = trimesh.Trimesh(vertices.copy(), self.faces.copy())
+    
+    rot = trimesh.transformations.rotation_matrix(
+            np.radians(rot_angle), rot_axis)
+    mesh.apply_transform(rot)
+
+    rot = trimesh.transformations.rotation_matrix(
+        np.radians(180), [1, 0, 0])
+    mesh.apply_transform(rot)
+    return mesh

third_party/Dyn-HaMR/dyn-hamr/geometry/pcl.py

@@ -0,0 +1,60 @@
+import numpy as np
+import torch
+
+
+def read_pcl_tensor(path):
+    pcl_np = read_pcl(path)
+    return torch.from_numpy(pcl_np)
+
+
+def align_pcl(Y, X, weight=None, fixed_scale=False):
+    """align similarity transform to align X with Y using umeyama method
+    X' = s * R * X + t is aligned with Y
+    :param Y (*, N, 3) first trajectory
+    :param X (*, N, 3) second trajectory
+    :param weight (*, N, 1) optional weight of valid correspondences
+    :returns s (*, 1), R (*, 3, 3), t (*, 3)
+    """
+    *dims, N, _ = Y.shape
+    N = torch.ones(*dims, 1, 1) * N
+
+    if weight is not None:
+        Y = Y * weight
+        X = X * weight
+        N = weight.sum(dim=-2, keepdim=True)  # (*, 1, 1)
+
+    # subtract mean
+    my = Y.sum(dim=-2) / N[..., 0]  # (*, 3)
+    mx = X.sum(dim=-2) / N[..., 0]
+    y0 = Y - my[..., None, :]  # (*, N, 3)
+    x0 = X - mx[..., None, :]
+
+    if weight is not None:
+        y0 = y0 * weight
+        x0 = x0 * weight
+
+    # correlation
+    C = torch.matmul(y0.transpose(-1, -2), x0) / N  # (*, 3, 3)
+    U, D, Vh = torch.linalg.svd(C)  # (*, 3, 3), (*, 3), (*, 3, 3)
+
+    S = torch.eye(3).reshape(*(1,) * (len(dims)), 3, 3).repeat(*dims, 1, 1)
+    neg = torch.det(U) * torch.det(Vh.transpose(-1, -2)) < 0
+    S[neg, 2, 2] = -1
+
+    R = torch.matmul(U, torch.matmul(S, Vh))  # (*, 3, 3)
+
+    D = torch.diag_embed(D)  # (*, 3, 3)
+    if fixed_scale:
+        s = torch.ones(*dims, 1, device=Y.device, dtype=torch.float32)
+    else:
+        var = torch.sum(torch.square(x0), dim=(-1, -2), keepdim=True) / N  # (*, 1, 1)
+        s = (
+            torch.diagonal(torch.matmul(D, S), dim1=-2, dim2=-1).sum(
+                dim=-1, keepdim=True
+            )
+            / var[..., 0]
+        )  # (*, 1)
+
+    t = my - s * torch.matmul(R, mx[..., None])[..., 0]  # (*, 3)
+
+    return s, R, t

third_party/Dyn-HaMR/dyn-hamr/geometry/plane.py

@@ -0,0 +1,101 @@
+import torch
+
+
+def fit_plane(points):
+    """
+    :param points (*, N, 3)
+    returns (*, 3) plane parameters (returns in normal * offset format)
+    """
+    *dims, N, D = points.shape
+    mean = points.mean(dim=-2, keepdim=True)
+    # (*, N, D), (*, D), (*, D, D)
+    U, S, Vh = torch.linalg.svd(points - mean)
+    normal = Vh[..., -1, :]  # (*, D)
+    offset = torch.einsum("...ij,...j->...i", points, normal)  # (*, N)
+    offset = offset.mean(dim=-1, keepdim=True)
+    return torch.cat([normal, offset], dim=-1)
+
+
+def get_plane_transform(up, ground_plane=None, xyz_orig=None):
+    """
+    get R, t rigid transform from plane and desired origin
+    :param up (3,) up vector of coordinate frame
+    :param ground_plane (4) (a, b, c, d) where a,b,c is the normal
+    :param xyz_orig (3) desired origin
+    """
+    R = torch.eye(3)
+    t = torch.zeros(3)
+    if ground_plane is None:
+        return R, t
+
+    # compute transform between world up vector and passed in floor
+    ground_plane = torch.as_tensor(ground_plane)
+    ground_plane = torch.sign(ground_plane[3]) * ground_plane
+
+    normal = ground_plane[:3]
+    normal = normal / torch.linalg.norm(normal)
+    v = torch.linalg.cross(up, normal)
+    ang_sin = torch.linalg.norm(v)
+    ang_cos = up.dot(normal)
+    skew_v = torch.as_tensor([[0.0, -v[2], v[1]], [v[2], 0.0, -v[0]], [-v[1], v[0], 0.0]])
+    R = torch.eye(3) + skew_v + (skew_v @ skew_v) * ((1.0 - ang_cos) / (ang_sin**2))
+
+    # project origin onto plane
+    if xyz_orig is None:
+        xyz_orig = torch.zeros(3)
+    t, _ = compute_plane_intersection(xyz_orig, -normal, ground_plane)
+
+    return R, t
+
+
+def parse_floor_plane(floor_plane):
+    """
+    Takes floor plane in the optimization form (Bx3 with a,b,c * d) and parses into
+    (a,b,c,d) from with (a,b,c) normal facing "up in the camera frame and d the offset.
+    """
+    floor_offset = torch.norm(floor_plane, dim=-1, keepdim=True)
+    floor_normal = floor_plane / floor_offset
+
+    # in camera system -y is up, so floor plane normal y component should never be positive
+    #       (assuming the camera is not sideways or upside down)
+    neg_mask = floor_normal[..., 1:2] > 0.0
+    floor_normal = torch.where(
+        neg_mask.expand_as(floor_normal), -floor_normal, floor_normal
+    )
+    floor_offset = torch.where(neg_mask, -floor_offset, floor_offset)
+    floor_plane_4d = torch.cat([floor_normal, floor_offset], dim=-1)
+
+    return floor_plane_4d
+
+
+def compute_plane_intersection(point, direction, plane):
+    """
+    Given a ray defined by a point in space and a direction,
+    compute the intersection point with the given plane.
+    Detect intersection in either direction or -direction.
+    Note, ray may not actually intersect with the plane.
+
+    Returns the intersection point and s where
+    point + s * direction = intersection_point. if s < 0 it means
+    -direction intersects.
+
+    - point : B x 3
+    - direction : B x 3
+    - plane : B x 4 (a, b, c, d) where (a, b, c) is the normal and (d) the offset.
+    """
+    dims = point.shape[:-1]
+    plane_normal = plane[..., :3]
+    plane_off = plane[..., 3]
+    s = (plane_off - bdot(plane_normal, point)) / (bdot(plane_normal, direction) + 1e-4)
+    itsct_pt = point + s.reshape((-1, 1)) * direction
+    return itsct_pt, s
+
+
+def bdot(A1, A2, keepdim=False):
+    """
+    Batched dot product.
+    - A1 : B x D
+    - A2 : B x D.
+    Returns B.
+    """
+    return (A1 * A2).sum(dim=-1, keepdim=keepdim)

third_party/Dyn-HaMR/dyn-hamr/geometry/rotation.py

@@ -0,0 +1,293 @@
+import torch
+import numpy as np
+from torch.nn import functional as F
+
+
+def batch_rodrigues(rot_vecs, epsilon=1e-8, dtype=torch.float32):
+    """
+    Taken from https://github.com/mkocabas/VIBE/blob/master/lib/utils/geometry.py
+    Calculates the rotation matrices for a batch of rotation vectors
+    - param rot_vecs: torch.tensor (N, 3) array of N axis-angle vectors
+    - returns R: torch.tensor (N, 3, 3) rotation matrices
+    """
+    batch_size = rot_vecs.shape[0]
+    device = rot_vecs.device
+
+    angle = torch.norm(rot_vecs + 1e-8, dim=1, keepdim=True)
+    rot_dir = rot_vecs / angle
+
+    cos = torch.unsqueeze(torch.cos(angle), dim=1)
+    sin = torch.unsqueeze(torch.sin(angle), dim=1)
+
+    # Bx1 arrays
+    rx, ry, rz = torch.split(rot_dir, 1, dim=1)
+    K = torch.zeros((batch_size, 3, 3), dtype=dtype, device=device)
+
+    zeros = torch.zeros((batch_size, 1), dtype=dtype, device=device)
+    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1).view(
+        (batch_size, 3, 3)
+    )
+
+    ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
+    rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K)
+    return rot_mat
+
+
+def quaternion_mul(q0, q1):
+    """
+    EXPECTS WXYZ
+    :param q0 (*, 4)
+    :param q1 (*, 4)
+    """
+    r0, r1 = q0[..., :1], q1[..., :1]
+    v0, v1 = q0[..., 1:], q1[..., 1:]
+    r = r0 * r1 - (v0 * v1).sum(dim=-1, keepdim=True)
+    v = r0 * v1 + r1 * v0 + torch.linalg.cross(v0, v1)
+    return torch.cat([r, v], dim=-1)
+
+
+def quaternion_inverse(q, eps=1e-8):
+    """
+    EXPECTS WXYZ
+    :param q (*, 4)
+    """
+    conj = torch.cat([q[..., :1], -q[..., 1:]], dim=-1)
+    mag = torch.square(q).sum(dim=-1, keepdim=True) + eps
+    return conj / mag
+
+
+def quaternion_slerp(t, q0, q1, eps=1e-8):
+    """
+    :param t (*, 1)  must be between 0 and 1
+    :param q0 (*, 4)
+    :param q1 (*, 4)
+    """
+    dims = q0.shape[:-1]
+    t = t.view(*dims, 1)
+
+    q0 = F.normalize(q0, p=2, dim=-1)
+    q1 = F.normalize(q1, p=2, dim=-1)
+    dot = (q0 * q1).sum(dim=-1, keepdim=True)
+
+    # make sure we give the shortest rotation path (< 180d)
+    neg = dot < 0
+    q1 = torch.where(neg, -q1, q1)
+    dot = torch.where(neg, -dot, dot)
+    angle = torch.acos(dot)
+
+    # if angle is too small, just do linear interpolation
+    collin = torch.abs(dot) > 1 - eps
+    fac = 1 / torch.sin(angle)
+    w0 = torch.where(collin, 1 - t, torch.sin((1 - t) * angle) * fac)
+    w1 = torch.where(collin, t, torch.sin(t * angle) * fac)
+    slerp = q0 * w0 + q1 * w1
+    return slerp
+
+
+def rotation_matrix_to_angle_axis(rotation_matrix):
+    """
+    This function is borrowed from https://github.com/kornia/kornia
+
+    Convert rotation matrix to Rodrigues vector
+    """
+    quaternion = rotation_matrix_to_quaternion(rotation_matrix)
+    aa = quaternion_to_angle_axis(quaternion)
+    aa[torch.isnan(aa)] = 0.0
+    return aa
+
+
+def quaternion_to_angle_axis(quaternion):
+    """
+    This function is borrowed from https://github.com/kornia/kornia
+
+    Convert quaternion vector to angle axis of rotation.
+    Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h
+
+    :param quaternion (*, 4) expects WXYZ
+    :returns angle_axis (*, 3)
+    """
+    # unpack input and compute conversion
+    q1 = quaternion[..., 1]
+    q2 = quaternion[..., 2]
+    q3 = quaternion[..., 3]
+    sin_squared_theta = q1 * q1 + q2 * q2 + q3 * q3
+
+    sin_theta = torch.sqrt(sin_squared_theta)
+    cos_theta = quaternion[..., 0]
+    two_theta = 2.0 * torch.where(
+        cos_theta < 0.0,
+        torch.atan2(-sin_theta, -cos_theta),
+        torch.atan2(sin_theta, cos_theta),
+    )
+
+    k_pos = two_theta / sin_theta
+    k_neg = 2.0 * torch.ones_like(sin_theta)
+    k = torch.where(sin_squared_theta > 0.0, k_pos, k_neg)
+
+    angle_axis = torch.zeros_like(quaternion)[..., :3]
+    angle_axis[..., 0] += q1 * k
+    angle_axis[..., 1] += q2 * k
+    angle_axis[..., 2] += q3 * k
+    return angle_axis
+
+
+def angle_axis_to_rotation_matrix(angle_axis):
+    """
+    :param angle_axis (*, 3)
+    return (*, 3, 3)
+    """
+    quat = angle_axis_to_quaternion(angle_axis)
+    return quaternion_to_rotation_matrix(quat)
+
+
+def quaternion_to_rotation_matrix(quaternion):
+    """
+    Convert a quaternion to a rotation matrix.
+    Taken from https://github.com/kornia/kornia, based on
+    https://github.com/matthew-brett/transforms3d/blob/8965c48401d9e8e66b6a8c37c65f2fc200a076fa/transforms3d/quaternions.py#L101
+    https://github.com/tensorflow/graphics/blob/master/tensorflow_graphics/geometry/transformation/rotation_matrix_3d.py#L247
+    :param quaternion (N, 4) expects WXYZ order
+    returns rotation matrix (N, 3, 3)
+    """
+    # normalize the input quaternion
+    quaternion_norm = F.normalize(quaternion, p=2, dim=-1, eps=1e-12)
+    *dims, _ = quaternion_norm.shape
+
+    # unpack the normalized quaternion components
+    w, x, y, z = torch.chunk(quaternion_norm, chunks=4, dim=-1)
+
+    # compute the actual conversion
+    tx = 2.0 * x
+    ty = 2.0 * y
+    tz = 2.0 * z
+    twx = tx * w
+    twy = ty * w
+    twz = tz * w
+    txx = tx * x
+    txy = ty * x
+    txz = tz * x
+    tyy = ty * y
+    tyz = tz * y
+    tzz = tz * z
+    one = torch.tensor(1.0)
+
+    matrix = torch.stack(
+        (
+            one - (tyy + tzz),
+            txy - twz,
+            txz + twy,
+            txy + twz,
+            one - (txx + tzz),
+            tyz - twx,
+            txz - twy,
+            tyz + twx,
+            one - (txx + tyy),
+        ),
+        dim=-1,
+    ).view(*dims, 3, 3)
+    return matrix
+
+
+def angle_axis_to_quaternion(angle_axis):
+    """
+    This function is borrowed from https://github.com/kornia/kornia
+    Convert angle axis to quaternion in WXYZ order
+    :param angle_axis (*, 3)
+    :returns quaternion (*, 4) WXYZ order
+    """
+    theta_sq = torch.sum(angle_axis**2, dim=-1, keepdim=True)  # (*, 1)
+    # need to handle the zero rotation case
+    valid = theta_sq > 0
+    theta = torch.sqrt(theta_sq)
+    half_theta = 0.5 * theta
+    ones = torch.ones_like(half_theta)
+    # fill zero with the limit of sin ax / x -> a
+    k = torch.where(valid, torch.sin(half_theta) / theta, 0.5 * ones)
+    w = torch.where(valid, torch.cos(half_theta), ones)
+    quat = torch.cat([w, k * angle_axis], dim=-1)
+    return quat
+
+
+def rotation_matrix_to_quaternion(rotation_matrix, eps=1e-6):
+    """
+    This function is borrowed from https://github.com/kornia/kornia
+    Convert rotation matrix to 4d quaternion vector
+    This algorithm is based on algorithm described in
+    https://github.com/KieranWynn/pyquaternion/blob/master/pyquaternion/quaternion.py#L201
+
+    :param rotation_matrix (N, 3, 3)
+    """
+    *dims, m, n = rotation_matrix.shape
+    rmat_t = torch.transpose(rotation_matrix.reshape(-1, m, n), -1, -2)
+
+    mask_d2 = rmat_t[:, 2, 2] < eps
+
+    mask_d0_d1 = rmat_t[:, 0, 0] > rmat_t[:, 1, 1]
+    mask_d0_nd1 = rmat_t[:, 0, 0] < -rmat_t[:, 1, 1]
+
+    t0 = 1 + rmat_t[:, 0, 0] - rmat_t[:, 1, 1] - rmat_t[:, 2, 2]
+    q0 = torch.stack(
+        [
+            rmat_t[:, 1, 2] - rmat_t[:, 2, 1],
+            t0,
+            rmat_t[:, 0, 1] + rmat_t[:, 1, 0],
+            rmat_t[:, 2, 0] + rmat_t[:, 0, 2],
+        ],
+        -1,
+    )
+    t0_rep = t0.repeat(4, 1).t()
+
+    t1 = 1 - rmat_t[:, 0, 0] + rmat_t[:, 1, 1] - rmat_t[:, 2, 2]
+    q1 = torch.stack(
+        [
+            rmat_t[:, 2, 0] - rmat_t[:, 0, 2],
+            rmat_t[:, 0, 1] + rmat_t[:, 1, 0],
+            t1,
+            rmat_t[:, 1, 2] + rmat_t[:, 2, 1],
+        ],
+        -1,
+    )
+    t1_rep = t1.repeat(4, 1).t()
+
+    t2 = 1 - rmat_t[:, 0, 0] - rmat_t[:, 1, 1] + rmat_t[:, 2, 2]
+    q2 = torch.stack(
+        [
+            rmat_t[:, 0, 1] - rmat_t[:, 1, 0],
+            rmat_t[:, 2, 0] + rmat_t[:, 0, 2],
+            rmat_t[:, 1, 2] + rmat_t[:, 2, 1],
+            t2,
+        ],
+        -1,
+    )
+    t2_rep = t2.repeat(4, 1).t()
+
+    t3 = 1 + rmat_t[:, 0, 0] + rmat_t[:, 1, 1] + rmat_t[:, 2, 2]
+    q3 = torch.stack(
+        [
+            t3,
+            rmat_t[:, 1, 2] - rmat_t[:, 2, 1],
+            rmat_t[:, 2, 0] - rmat_t[:, 0, 2],
+            rmat_t[:, 0, 1] - rmat_t[:, 1, 0],
+        ],
+        -1,
+    )
+    t3_rep = t3.repeat(4, 1).t()
+
+    mask_c0 = mask_d2 * mask_d0_d1
+    mask_c1 = mask_d2 * ~mask_d0_d1
+    mask_c2 = ~mask_d2 * mask_d0_nd1
+    mask_c3 = ~mask_d2 * ~mask_d0_nd1
+    mask_c0 = mask_c0.view(-1, 1).type_as(q0)
+    mask_c1 = mask_c1.view(-1, 1).type_as(q1)
+    mask_c2 = mask_c2.view(-1, 1).type_as(q2)
+    mask_c3 = mask_c3.view(-1, 1).type_as(q3)
+
+    q = q0 * mask_c0 + q1 * mask_c1 + q2 * mask_c2 + q3 * mask_c3
+    q /= torch.sqrt(
+        t0_rep * mask_c0
+        + t1_rep * mask_c1
+        + t2_rep * mask_c2  # noqa
+        + t3_rep * mask_c3
+    )  # noqa
+    q *= 0.5
+    return q.reshape(*dims, 4)