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

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-09-15 14:13:49 +08:00
commit 7e4ef6f98b
7002 changed files with 1452249 additions and 0 deletions
@@ -0,0 +1,152 @@
+"""Based on Daniel Holden code from:
+   A Deep Learning Framework for Character Motion Synthesis and Editing
+   (http://www.ipab.inf.ed.ac.uk/cgvu/motionsynthesis.pdf)
+"""
+
+import os
+
+import numpy as np
+import torch
+import torch.nn as nn
+from rotations import euler_angles_to_matrix, quaternion_to_matrix, rotation_6d_to_matrix
+
+
+class ForwardKinematicsLayer(nn.Module):
+    """ Forward Kinematics Layer Class """
+
+    def __init__(self, args=None, parents=None, positions=None, device=None):
+        super().__init__()
+        self.b_idxs = None
+        if device is None:
+            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        else:
+            self.device = device
+        if parents is None and positions is None:
+            # Load SMPL skeleton (their joint order is different from the one we use for bvh export)
+            # smpl_fname = os.path.join(args.smpl.smpl_body_model, args.data.gender, 'model.npz')
+            smpl_fname = os.path.join(args.smpl.smpl_body_model, "MANO_RIGHT.pkl")
+            if not os.path.exists(smpl_fname):
+                raise FileNotFoundError(f'SMPL file {os.path.abspath(smpl_fname)} does not exist. Please check the path., pwd: {os.getcwd()}')
+            smpl_data = np.load(smpl_fname, encoding='latin1', allow_pickle=True)
+            self.parents = torch.from_numpy(smpl_data['kintree_table'][0].astype(np.int32)).to(self.device)
+            self.parents = self.parents.long()
+            self.positions = torch.from_numpy(smpl_data['J'].astype(np.float32)).to(self.device)
+            self.positions[1:] -= self.positions[self.parents[1:]]
+        else:
+            self.parents = torch.from_numpy(parents).to(self.device)
+            self.parents = self.parents.long()
+            self.positions = torch.from_numpy(positions).to(self.device)
+            self.positions = self.positions.float()
+        self.positions[0] = 0
+
+    def rotate(self, t0s, t1s):
+        return torch.matmul(t0s, t1s)
+
+    def identity_rotation(self, rotations):
+        diagonal = torch.diag(torch.tensor([1.0, 1.0, 1.0, 1.0])).to(self.device)
+        diagonal = torch.reshape(
+            diagonal, torch.Size([1] * len(rotations.shape[:2]) + [4, 4]))
+        ts = diagonal.repeat(rotations.shape[:2] + torch.Size([1, 1]))
+        return ts
+
+    def make_fast_rotation_matrices(self, positions, rotations):
+        if len(rotations.shape) == 4 and rotations.shape[-2:] == torch.Size([3, 3]):
+            rot_matrices = rotations
+        elif rotations.shape[-1] == 3:
+            rot_matrices = euler_angles_to_matrix(rotations, convention='XYZ')
+        elif rotations.shape[-1] == 4:
+            rot_matrices = quaternion_to_matrix(rotations)
+        elif rotations.shape[-1] == 6:
+            rot_matrices = rotation_6d_to_matrix(rotations)
+        else:
+            raise NotImplementedError(f'Unimplemented rotation representation in FK layer, shape of {rotations.shape}')
+
+        rot_matrices = torch.cat([rot_matrices, positions[..., None]], dim=-1)
+        zeros = torch.zeros(rot_matrices.shape[:-2] + torch.Size([1, 3])).to(self.device)
+        ones = torch.ones(rot_matrices.shape[:-2] + torch.Size([1, 1])).to(self.device)
+        zerosones = torch.cat([zeros, ones], dim=-1)
+        rot_matrices = torch.cat([rot_matrices, zerosones], dim=-2)
+        return rot_matrices
+
+    def rotate_global(self, parents, positions, rotations):
+        locals = self.make_fast_rotation_matrices(positions, rotations)
+        globals = self.identity_rotation(rotations)
+
+        globals = torch.cat([locals[:, 0:1], globals[:, 1:]], dim=1)
+        b_size = positions.shape[0]
+        if self.b_idxs is None:
+            self.b_idxs = torch.LongTensor(np.arange(b_size)).to(self.device)
+        elif self.b_idxs.shape[-1] != b_size:
+            self.b_idxs = torch.LongTensor(np.arange(b_size)).to(self.device)
+
+        for i in range(1, positions.shape[1]):
+            globals[:, i] = self.rotate(
+                globals[self.b_idxs, parents[i]], locals[:, i])
+
+        return globals
+
+    def get_tpose_joints(self, offsets, parents):
+        num_joints = len(parents)
+        joints = [offsets[:, 0]]
+        for j in range(1, len(parents)):
+            joints.append(joints[parents[j]] + offsets[:, j])
+
+        return torch.stack(joints, dim=1)
+
+    def canonical_to_local(self, canonical_xform, global_orient=None):
+        """
+        Args:
+            canonical_xform: (B, J, 3, 3)
+            global_orient: (B, 3, 3)
+
+        Returns:
+            local_xform: (B, J, 3, 3)
+        """
+        local_xform = torch.zeros_like(canonical_xform)
+
+        if global_orient is None:
+            global_xform = canonical_xform
+        else:
+            global_xform = torch.matmul(global_orient.unsqueeze(1), canonical_xform)
+        for i in range(global_xform.shape[1]):
+            if i == 0:
+                local_xform[:, i] = global_xform[:, i]
+            else:
+                local_xform[:, i] = torch.bmm(torch.linalg.inv(global_xform[:, self.parents[i]]), global_xform[:, i])
+
+        return local_xform
+
+    def global_to_local(self, global_xform):
+        """
+        Args:
+            global_xform: (B, J, 3, 3)
+
+        Returns:
+            local_xform: (B, J, 3, 3)
+        """
+        local_xform = torch.zeros_like(global_xform)
+
+        for i in range(global_xform.shape[1]):
+            if i == 0:
+                local_xform[:, i] = global_xform[:, i]
+            else:
+                local_xform[:, i] = torch.bmm(torch.linalg.inv(global_xform[:, self.parents[i]]), global_xform[:, i])
+
+        return local_xform
+
+    def forward(self, rotations, positions=None):
+        """
+        Args:
+            rotations (B, J, D)
+
+        Returns:
+            The global position of each joint after FK (B, J, 3)
+        """
+        # Get the full transform with rotations for skinning
+        b_size = rotations.shape[0]
+        if positions is None:
+            positions = self.positions.repeat(b_size, 1, 1)
+        transforms = self.rotate_global(self.parents, positions, rotations)
+        coordinates = transforms[:, :, :3, 3] / transforms[:, :, 3:, 3]
+
+        return coordinates, transforms