Animation.py

##############################
#
# based on http://theorangeduck.com/page/deep-learning-framework-character-motion-synthesis-and-editing
#
##############################

import operator

import numpy as np
import numpy.core.umath_tests as ut

try:
    from . import AnimationStructure
    from .Quaternions import Quaternions
except:
    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):

            try:
                # since joints are re-indexed, we need to take spacial care on parents indices
                reindexed_parents = reindex(self.parents, k[1])
            except:  # we will reach the 'except' if k[1] is Nones, i.e., if the sub-selection is over frames rather then over joints
                reindexed_parents = self.parents

            return Animation(
                self.rotations[k],
                self.positions[k],
                self.orients[k[1:]],
                self.offsets[k[1:]],
                reindexed_parents)
        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())

    # def reorder(self, new_order):
    #     self.rotations = self.rotations[:, new_order]
    #     self.positions = self.positions[:, new_order]
    #     self.offsets = self.offsets[new_order]
    #     if self.orients.shape[0] > 0:
    #         self.orients = self.orients[:, new_order]
    #
    #     # reorder parents
    #     sorted_order_inversed = {num: i for i, num in enumerate(new_order)}
    #     sorted_order_inversed[-1] = -1
    #     self.parents = np.array([sorted_order_inversed[self.parents[i]] for i in new_order])

    # def sort(self, names):
    #     children = AnimationStructure.children_list(self.parents)
    #     sorted_order = np.zeros(self.parents.shape, dtype=np.int)
    #     root_idx = np.where(self.parents == -1)[0][0]
    #     sorted_order[0] = root_idx
    #
    #     def get_sorted_order(sorted_order, parent_out_idx, parent_in_idx, children):
    #         out_idx = parent_out_idx  # return same index in case there are no children
    #         for child in children[parent_in_idx]:
    #             out_idx = out_idx + 1
    #             sorted_order[out_idx] = child
    #             sorted_order, out_idx = get_sorted_order(sorted_order, out_idx, child, children)
    #         return sorted_order, out_idx
    #
    #     sorted_order, _ = get_sorted_order(sorted_order, 0, root_idx, children)
    #
    #     anim = self.copy()
    #     anim.reorder(sorted_order)
    #     if names is not None:
    #         names = names[sorted_order]
    #
    #     return anim, names


""" 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


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


def transforms_global(anim):
    """
    Global Animation Transforms

    This relies on joint ordering
    being incremental. That means a joint
    J1 must not be an 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


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
    """

    positions = transforms_global(anim)[:, :, :, 3]
    return positions[:, :, :3] / positions[:, :, 3,
                                 np.newaxis]  # transforms_global(anim)[:,:,3,3] is considered here as a normalizer, but in practice it is 1


""" 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]


def offsets_from_positions(positions, parents):
    is_one_frame = False
    if positions.ndim == 2:
        positions = positions[np.newaxis]
        is_one_frame = True
    offsets = positions.copy()
    root_idx = np.where(parents == -1)[0][0]
    idx = np.delete(np.arange(positions.shape[1]), root_idx)
    offsets[:, idx] = positions[:, idx] - positions[:, parents[idx]]
    if is_one_frame:
        offsets = offsets[0]
    return offsets


def animation_from_offsets(offsets, parents, shape=None):
    sorted_order = AnimationStructure.get_sorted_order(parents)
    offsets = offsets[sorted_order]

    # reorder parents
    parents = reindex(parents, sorted_order)

    if shape is None:
        shape = (1, offsets.shape[0])
    orients = Quaternions.id(0)
    anim_positions = np.repeat(offsets[np.newaxis], shape[0], axis=0)
    rotations = Quaternions.id((shape[0], shape[1]))

    anim = Animation(rotations, anim_positions, orients, offsets, parents)
    return anim, sorted_order, parents


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


def reindex(orig_idx, sub_idx):
    # key: new index, value: old index
    order_inversed = {num: i for i, num in enumerate(sub_idx)}
    order_inversed[-1] = -1
    new_idx = np.array([order_inversed[orig_idx[i]] for i in sub_idx])
    return new_idx