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/HMP/holden/Animation.py

@@ -0,0 +1,675 @@
+import operator
+
+import numpy as np
+import numpy.core.umath as ut               # (for numpy 1.25)
+# import numpy.core.umath_tests as ut
+
+import AnimationStructure
+from Quaternions import Quaternions
+
+
+class Animation:
+    """
+    Animation is a numpy-like wrapper for animation data
+
+    Animation data consists of several arrays consisting
+    of F frames and J joints.
+
+    The animation is specified by
+
+        rotations : (F, J) Quaternions | Joint Rotations
+        positions : (F, J, 3) ndarray  | Joint Positions
+
+    The base pose is specified by
+
+        orients   : (J) Quaternions    | Joint Orientations
+        offsets   : (J, 3) ndarray     | Joint Offsets
+
+    And the skeletal structure is specified by
+
+        parents   : (J) ndarray        | Joint Parents
+    """
+
+    def __init__(self, rotations, positions, orients, offsets, parents):
+
+        self.rotations = rotations
+        self.positions = positions
+        self.orients = orients
+        self.offsets = offsets
+        self.parents = parents
+
+    def __op__(self, op, other):
+        return Animation(
+            op(self.rotations, other.rotations),
+            op(self.positions, other.positions),
+            op(self.orients, other.orients),
+            op(self.offsets, other.offsets),
+            op(self.parents, other.parents))
+
+    def __iop__(self, op, other):
+        self.rotations = op(self.roations, other.rotations)
+        self.positions = op(self.roations, other.positions)
+        self.orients = op(self.orients, other.orients)
+        self.offsets = op(self.offsets, other.offsets)
+        self.parents = op(self.parents, other.parents)
+        return self
+
+    def __sop__(self, op):
+        return Animation(
+            op(self.rotations),
+            op(self.positions),
+            op(self.orients),
+            op(self.offsets),
+            op(self.parents))
+
+    def __add__(self, other): return self.__op__(operator.add, other)
+    def __sub__(self, other): return self.__op__(operator.sub, other)
+    def __mul__(self, other): return self.__op__(operator.mul, other)
+    def __div__(self, other): return self.__op__(operator.div, other)
+
+    def __abs__(self): return self.__sop__(operator.abs)
+    def __neg__(self): return self.__sop__(operator.neg)
+
+    def __iadd__(self, other): return self.__iop__(operator.iadd, other)
+    def __isub__(self, other): return self.__iop__(operator.isub, other)
+    def __imul__(self, other): return self.__iop__(operator.imul, other)
+    def __idiv__(self, other): return self.__iop__(operator.idiv, other)
+
+    def __len__(self): return len(self.rotations)
+
+    def __getitem__(self, k):
+        if isinstance(k, tuple):
+            return Animation(
+                self.rotations[k],
+                self.positions[k],
+                self.orients[k[1:]],
+                self.offsets[k[1:]],
+                self.parents[k[1:]])
+        else:
+            return Animation(
+                self.rotations[k],
+                self.positions[k],
+                self.orients,
+                self.offsets,
+                self.parents)
+
+    def __setitem__(self, k, v):
+        if isinstance(k, tuple):
+            self.rotations.__setitem__(k, v.rotations)
+            self.positions.__setitem__(k, v.positions)
+            self.orients.__setitem__(k[1:], v.orients)
+            self.offsets.__setitem__(k[1:], v.offsets)
+            self.parents.__setitem__(k[1:], v.parents)
+        else:
+            self.rotations.__setitem__(k, v.rotations)
+            self.positions.__setitem__(k, v.positions)
+            self.orients.__setitem__(k, v.orients)
+            self.offsets.__setitem__(k, v.offsets)
+            self.parents.__setitem__(k, v.parents)
+
+    @property
+    def shape(self): return (self.rotations.shape[0], self.rotations.shape[1])
+
+    def copy(self): return Animation(
+        self.rotations.copy(), self.positions.copy(),
+        self.orients.copy(), self.offsets.copy(),
+        self.parents.copy())
+
+    def repeat(self, *args, **kw):
+        return Animation(
+            self.rotations.repeat(*args, **kw),
+            self.positions.repeat(*args, **kw),
+            self.orients, self.offsets, self.parents)
+
+    def ravel(self):
+        return np.hstack([
+            self.rotations.log().ravel(),
+            self.positions.ravel(),
+            self.orients.log().ravel(),
+            self.offsets.ravel()])
+
+    @classmethod
+    def unravel(clas, anim, shape, parents):
+        nf, nj = shape
+        rotations = anim[nf*nj*0:nf*nj*3]
+        positions = anim[nf*nj*3:nf*nj*6]
+        orients = anim[nf*nj*6+nj*0:nf*nj*6+nj*3]
+        offsets = anim[nf*nj*6+nj*3:nf*nj*6+nj*6]
+        return cls(
+            Quaternions.exp(rotations), positions,
+            Quaternions.exp(orients), offsets,
+            parents.copy())
+
+
+""" Maya Interaction """
+
+
+def load_to_maya(anim, names=None, radius=0.5):
+    """
+    Load Animation Object into Maya as Joint Skeleton
+    loads each frame as a new keyfame in maya.
+
+    If the animation is too slow or too fast perhaps
+    the framerate needs adjusting before being loaded
+    such that it matches the maya scene framerate.
+
+
+    Parameters
+    ----------
+
+    anim : Animation
+        Animation to load into Scene
+
+    names : [str]
+        Optional list of Joint names for Skeleton
+
+    Returns
+    -------
+
+    List of Maya Joint Nodes loaded into scene
+    """
+
+    import pymel.core as pm
+
+    joints = []
+    frames = range(1, len(anim)+1)
+
+    if names is None:
+        names = ["joint_" + str(i) for i in range(len(anim.parents))]
+
+    for i, offset, orient, parent, name in zip(range(len(anim.offsets)), anim.offsets, anim.orients, anim.parents, names):
+
+        if parent < 0:
+            pm.select(d=True)
+        else:
+            pm.select(joints[parent])
+
+        joint = pm.joint(n=name, p=offset, relative=True, radius=radius)
+        joint.setOrientation([orient[1], orient[2], orient[3], orient[0]])
+
+        curvex = pm.nodetypes.AnimCurveTA(n=name + "_rotateX")
+        curvey = pm.nodetypes.AnimCurveTA(n=name + "_rotateY")
+        curvez = pm.nodetypes.AnimCurveTA(n=name + "_rotateZ")
+
+        jrotations = (-Quaternions(orient[np.newaxis]) * anim.rotations[:, i]).euler()
+        curvex.addKeys(frames, jrotations[:, 0])
+        curvey.addKeys(frames, jrotations[:, 1])
+        curvez.addKeys(frames, jrotations[:, 2])
+
+        pm.connectAttr(curvex.output, joint.rotateX)
+        pm.connectAttr(curvey.output, joint.rotateY)
+        pm.connectAttr(curvez.output, joint.rotateZ)
+
+        offsetx = pm.nodetypes.AnimCurveTU(n=name + "_translateX")
+        offsety = pm.nodetypes.AnimCurveTU(n=name + "_translateY")
+        offsetz = pm.nodetypes.AnimCurveTU(n=name + "_translateZ")
+
+        offsetx.addKeys(frames, anim.positions[:, i, 0])
+        offsety.addKeys(frames, anim.positions[:, i, 1])
+        offsetz.addKeys(frames, anim.positions[:, i, 2])
+
+        pm.connectAttr(offsetx.output, joint.translateX)
+        pm.connectAttr(offsety.output, joint.translateY)
+        pm.connectAttr(offsetz.output, joint.translateZ)
+
+        joints.append(joint)
+
+    return joints
+
+
+def load_from_maya(root, start, end):
+    """
+    Load Animation Object from Maya Joint Skeleton    
+
+    Parameters
+    ----------
+
+    root : PyNode
+        Root Joint of Maya Skeleton
+
+    start, end : int, int
+        Start and End frame index of Maya Animation
+
+    Returns
+    -------
+
+    animation : Animation
+        Loaded animation from maya
+
+    names : [str]
+        Joint names from maya   
+    """
+
+    import pymel.core as pm
+
+    original_time = pm.currentTime(q=True)
+    pm.currentTime(start)
+
+    """ Build Structure """
+
+    names, parents = AnimationStructure.load_from_maya(root)
+    descendants = AnimationStructure.descendants_list(parents)
+    orients = Quaternions.id(len(names))
+    offsets = np.array([pm.xform(j, q=True, translation=True) for j in names])
+
+    for j, name in enumerate(names):
+        scale = pm.xform(pm.PyNode(name), q=True, scale=True, relative=True)
+        if len(descendants[j]) == 0:
+            continue
+        offsets[descendants[j]] *= scale
+
+    """ Load Animation """
+
+    eulers = np.zeros((end-start, len(names), 3))
+    positions = np.zeros((end-start, len(names), 3))
+    rotations = Quaternions.id((end-start, len(names)))
+
+    for i in range(end-start):
+
+        pm.currentTime(start+i+1, u=True)
+
+        scales = {}
+
+        for j, name, parent in zip(range(len(names)), names, parents):
+
+            node = pm.PyNode(name)
+
+            if i == 0 and pm.hasAttr(node, 'jointOrient'):
+                ort = node.getOrientation()
+                orients[j] = Quaternions(np.array([ort[3], ort[0], ort[1], ort[2]]))
+
+            if pm.hasAttr(node, 'rotate'):
+                eulers[i, j] = np.radians(pm.xform(node, q=True, rotation=True))
+            if pm.hasAttr(node, 'translate'):
+                positions[i, j] = pm.xform(node, q=True, translation=True)
+            if pm.hasAttr(node, 'scale'):
+                scales[j] = pm.xform(node, q=True, scale=True, relative=True)
+
+        for j in scales:
+            if len(descendants[j]) == 0:
+                continue
+            positions[i, descendants[j]] *= scales[j]
+
+        positions[i, 0] = pm.xform(root, q=True, translation=True, worldSpace=True)
+
+    rotations = orients[np.newaxis] * Quaternions.from_euler(eulers, order='xyz', world=True)
+
+    """ Done """
+
+    pm.currentTime(original_time)
+
+    return Animation(rotations, positions, orients, offsets, parents), names
+
+# local transformation matrices
+
+
+def transforms_local(anim):
+    """
+    Computes Animation Local Transforms
+
+    As well as a number of other uses this can
+    be used to compute global joint transforms,
+    which in turn can be used to compete global
+    joint positions
+
+    Parameters
+    ----------
+
+    anim : Animation
+        Input animation
+
+    Returns
+    -------
+
+    transforms : (F, J, 4, 4) ndarray
+
+        For each frame F, joint local
+        transforms for each joint J
+    """
+
+    transforms = anim.rotations.transforms()
+    transforms = np.concatenate([transforms, np.zeros(transforms.shape[:2] + (3, 1))], axis=-1)
+    transforms = np.concatenate([transforms, np.zeros(transforms.shape[:2] + (1, 4))], axis=-2)
+    # the last column is filled with the joint positions!
+    transforms[:, :, 0:3, 3] = anim.positions
+    transforms[:, :, 3:4, 3] = 1.0
+    return transforms
+
+
+def transforms_multiply(t0s, t1s):
+    """
+    Transforms Multiply
+
+    Multiplies two arrays of animation transforms
+
+    Parameters
+    ----------
+
+    t0s, t1s : (F, J, 4, 4) ndarray
+        Two arrays of transforms
+        for each frame F and each
+        joint J
+
+    Returns
+    -------
+
+    transforms : (F, J, 4, 4) ndarray
+        Array of transforms for each
+        frame F and joint J multiplied
+        together
+    """
+
+    return ut.matrix_multiply(t0s, t1s)
+
+
+def transforms_inv(ts):
+    fts = ts.reshape(-1, 4, 4)
+    fts = np.array(list(map(lambda x: np.linalg.inv(x), fts)))
+    return fts.reshape(ts.shape)
+
+
+def transforms_blank(anim):
+    """
+    Blank Transforms
+
+    Parameters
+    ----------
+
+    anim : Animation
+        Input animation
+
+    Returns
+    -------
+
+    transforms : (F, J, 4, 4) ndarray
+        Array of identity transforms for 
+        each frame F and joint J
+    """
+
+    ts = np.zeros(anim.shape + (4, 4))
+    ts[:, :, 0, 0] = 1.0
+    ts[:, :, 1, 1] = 1.0
+    ts[:, :, 2, 2] = 1.0
+    ts[:, :, 3, 3] = 1.0
+    return ts
+
+# global transformation matrices
+
+
+def transforms_global(anim):
+    """
+    Global Animation Transforms
+
+    This relies on joint ordering
+    being incremental. That means a joint
+    J1 must not be a ancestor of J0 if
+    J0 appears before J1 in the joint
+    ordering.
+
+    Parameters
+    ----------
+
+    anim : Animation
+        Input animation
+
+    Returns
+    ------
+
+    transforms : (F, J, 4, 4) ndarray
+        Array of global transforms for 
+        each frame F and joint J
+    """
+
+    joints = np.arange(anim.shape[1])
+    parents = np.arange(anim.shape[1])
+    locals = transforms_local(anim)
+    globals = transforms_blank(anim)
+
+    globals[:, 0] = locals[:, 0]
+
+    for i in range(1, anim.shape[1]):
+        globals[:, i] = transforms_multiply(globals[:, anim.parents[i]], locals[:, i])
+
+    return globals
+
+# !!! useful!
+
+
+def positions_global(anim):
+    """
+    Global Joint Positions
+
+    Given an animation compute the global joint
+    positions at at every frame
+
+    Parameters
+    ----------
+
+    anim : Animation
+        Input animation
+
+    Returns
+    -------
+
+    positions : (F, J, 3) ndarray
+        Positions for every frame F 
+        and joint position J
+    """
+
+    # get the last column -- corresponding to the coordinates
+    positions = transforms_global(anim)[:, :, :, 3]
+    return positions[:, :, :3] / positions[:, :, 3, np.newaxis]
+
+
+""" Rotations """
+
+
+def rotations_global(anim):
+    """
+    Global Animation Rotations
+
+    This relies on joint ordering
+    being incremental. That means a joint
+    J1 must not be a ancestor of J0 if
+    J0 appears before J1 in the joint
+    ordering.
+
+    Parameters
+    ----------
+
+    anim : Animation
+        Input animation
+
+    Returns
+    -------
+
+    points : (F, J) Quaternions
+        global rotations for every frame F 
+        and joint J
+    """
+
+    joints = np.arange(anim.shape[1])
+    parents = np.arange(anim.shape[1])
+    locals = anim.rotations
+    globals = Quaternions.id(anim.shape)
+
+    globals[:, 0] = locals[:, 0]
+
+    for i in range(1, anim.shape[1]):
+        globals[:, i] = globals[:, anim.parents[i]] * locals[:, i]
+
+    return globals
+
+
+def rotations_parents_global(anim):
+    rotations = rotations_global(anim)
+    rotations = rotations[:, anim.parents]
+    rotations[:, 0] = Quaternions.id(len(anim))
+    return rotations
+
+
+def rotations_load_to_maya(rotations, positions, names=None):
+    """
+    Load Rotations into Maya
+
+    Loads a Quaternions array into the scene
+    via the representation of axis
+
+    Parameters
+    ----------
+
+    rotations : (F, J) Quaternions 
+        array of rotations to load
+        into the scene where
+            F = number of frames
+            J = number of joints
+
+    positions : (F, J, 3) ndarray 
+        array of positions to load
+        rotation axis at where:
+            F = number of frames
+            J = number of joints
+
+    names : [str]
+        List of joint names
+
+    Returns
+    -------
+
+    maxies : Group
+        Grouped Maya Node of all Axis nodes
+    """
+
+    import pymel.core as pm
+
+    if names is None:
+        names = ["joint_" + str(i) for i in range(rotations.shape[1])]
+
+    maxis = []
+    frames = range(1, len(positions)+1)
+    for i, name in enumerate(names):
+
+        name = name + "_axis"
+        axis = pm.group(
+            pm.curve(p=[(0, 0, 0), (1, 0, 0)], d=1, n=name+'_axis_x'),
+            pm.curve(p=[(0, 0, 0), (0, 1, 0)], d=1, n=name+'_axis_y'),
+            pm.curve(p=[(0, 0, 0), (0, 0, 1)], d=1, n=name+'_axis_z'),
+            n=name)
+
+        axis.rotatePivot.set((0, 0, 0))
+        axis.scalePivot.set((0, 0, 0))
+        axis.childAtIndex(0).overrideEnabled.set(1)
+        axis.childAtIndex(0).overrideColor.set(13)
+        axis.childAtIndex(1).overrideEnabled.set(1)
+        axis.childAtIndex(1).overrideColor.set(14)
+        axis.childAtIndex(2).overrideEnabled.set(1)
+        axis.childAtIndex(2).overrideColor.set(15)
+
+        curvex = pm.nodetypes.AnimCurveTA(n=name + "_rotateX")
+        curvey = pm.nodetypes.AnimCurveTA(n=name + "_rotateY")
+        curvez = pm.nodetypes.AnimCurveTA(n=name + "_rotateZ")
+
+        arotations = rotations[:, i].euler()
+        curvex.addKeys(frames, arotations[:, 0])
+        curvey.addKeys(frames, arotations[:, 1])
+        curvez.addKeys(frames, arotations[:, 2])
+
+        pm.connectAttr(curvex.output, axis.rotateX)
+        pm.connectAttr(curvey.output, axis.rotateY)
+        pm.connectAttr(curvez.output, axis.rotateZ)
+
+        offsetx = pm.nodetypes.AnimCurveTU(n=name + "_translateX")
+        offsety = pm.nodetypes.AnimCurveTU(n=name + "_translateY")
+        offsetz = pm.nodetypes.AnimCurveTU(n=name + "_translateZ")
+
+        offsetx.addKeys(frames, positions[:, i, 0])
+        offsety.addKeys(frames, positions[:, i, 1])
+        offsetz.addKeys(frames, positions[:, i, 2])
+
+        pm.connectAttr(offsetx.output, axis.translateX)
+        pm.connectAttr(offsety.output, axis.translateY)
+        pm.connectAttr(offsetz.output, axis.translateZ)
+
+        maxis.append(axis)
+
+    return pm.group(*maxis, n='RotationAnimation')
+
+
+""" Offsets & Orients """
+
+
+def orients_global(anim):
+
+    joints = np.arange(anim.shape[1])
+    parents = np.arange(anim.shape[1])
+    locals = anim.orients
+    globals = Quaternions.id(anim.shape[1])
+
+    globals[:, 0] = locals[:, 0]
+
+    for i in range(1, anim.shape[1]):
+        globals[:, i] = globals[:, anim.parents[i]] * locals[:, i]
+
+    return globals
+
+
+def offsets_transforms_local(anim):
+
+    transforms = anim.orients[np.newaxis].transforms()
+    transforms = np.concatenate([transforms, np.zeros(transforms.shape[:2] + (3, 1))], axis=-1)
+    transforms = np.concatenate([transforms, np.zeros(transforms.shape[:2] + (1, 4))], axis=-2)
+    transforms[:, :, 0:3, 3] = anim.offsets[np.newaxis]
+    transforms[:, :, 3:4, 3] = 1.0
+    return transforms
+
+
+def offsets_transforms_global(anim):
+
+    joints = np.arange(anim.shape[1])
+    parents = np.arange(anim.shape[1])
+    locals = offsets_transforms_local(anim)
+    globals = transforms_blank(anim)
+
+    globals[:, 0] = locals[:, 0]
+
+    for i in range(1, anim.shape[1]):
+        globals[:, i] = transforms_multiply(globals[:, anim.parents[i]], locals[:, i])
+
+    return globals
+
+
+def offsets_global(anim):
+    offsets = offsets_transforms_global(anim)[:, :, :, 3]
+    return offsets[0, :, :3] / offsets[0, :, 3, np.newaxis]
+
+
+""" Lengths """
+
+
+def offset_lengths(anim):
+    return np.sum(anim.offsets[1:]**2.0, axis=1)**0.5
+
+
+def position_lengths(anim):
+    return np.sum(anim.positions[:, 1:]**2.0, axis=2)**0.5
+
+
+""" Skinning """
+
+
+def skin(anim, rest, weights, mesh, maxjoints=4):
+
+    full_transforms = transforms_multiply(
+        transforms_global(anim),
+        transforms_inv(transforms_global(rest[0:1])))
+
+    weightids = np.argsort(-weights, axis=1)[:, :maxjoints]
+    weightvls = np.array(list(map(lambda w, i: w[i], weights, weightids)))
+    weightvls = weightvls / weightvls.sum(axis=1)[..., np.newaxis]
+
+    verts = np.hstack([mesh, np.ones((len(mesh), 1))])
+    verts = verts[np.newaxis, :, np.newaxis, :, np.newaxis]
+    verts = transforms_multiply(full_transforms[:, weightids], verts)
+    verts = (verts[:, :, :, :3] / verts[:, :, :, 3:4])[:, :, :, :, 0]
+
+    return np.sum(weightvls[np.newaxis, :, :, np.newaxis] * verts, axis=2)