MIRPAMO/bpy/retarget.py

"""MIRPAMO stage 2: retarget motion clips onto a rigged character.

Replicates Mixamo's animation retargeter locally:
  - Rotational data transfer: for every mapped bone, the source bone's
    world-space rotation *delta from rest* is applied to the target bone's
    rest orientation. Pure quaternion transfer — no positional copying —
    so limb proportions never stretch.
  - Scale compensation: root (Hips) translation is scaled by the ratio of
    target hip height to source hip height, so a short character takes
    short steps from the same clip.

Bone mapping: exact name match first, then normalized fuzzy match
("mixamorig:LeftArm" == "LeftArm" == "left_arm"), plus optional explicit
map JSON {"source_bone": "target_bone", "skip_me": null}.

Usage:
  blender -b --python bpy/retarget.py -- --rig rigged.glb \
      --clip walk.bvh [--clip run.fbx ...] --out out.glb \
      [--bonemap bonemaps/cmu_mb_to_mixamo.json] [--inplace] [--fps 30]
"""
import argparse
import json
import os
import re
import sys

import bpy
from mathutils import Matrix, Vector

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from common import (clean_scene, die, export_glb, find_armature, find_meshes,
                    import_any, stage_args)


def norm_name(n):
    n = n.lower()
    n = re.sub(r"mixamo[_:]?rig", "", n)
    n = re.sub(r"[^a-z0-9]", "", n)
    return n


def build_bone_map(src_arm, tgt_arm, explicit):
    """Map source bone name -> target bone name."""
    tgt_names = {b.name for b in tgt_arm.data.bones}
    tgt_by_norm = {norm_name(b.name): b.name for b in tgt_arm.data.bones}
    mapping = {}
    for sb in src_arm.data.bones:
        if explicit and sb.name in explicit:
            t = explicit[sb.name]
            if t and t in tgt_names:
                mapping[sb.name] = t
            continue  # explicit null = deliberately skipped
        if sb.name in tgt_names:
            mapping[sb.name] = sb.name
        elif norm_name(sb.name) in tgt_by_norm:
            mapping[sb.name] = tgt_by_norm[norm_name(sb.name)]
    return mapping


def hip_bone(arm, mapping_values=None):
    candidates = mapping_values or [b.name for b in arm.data.bones]
    for b in candidates:
        if norm_name(b) in ("hips", "hip", "pelvis", "root"):
            return b
    # fall back to a root bone (no parent) that has children
    for b in arm.data.bones:
        if b.parent is None and b.children:
            return b.name
    return None


def world_rest(arm_obj, bone_name):
    return arm_obj.matrix_world @ arm_obj.data.bones[bone_name].matrix_local


def retarget_clip(src_arm, tgt_arm, mapping, clip_name, inplace, frame_range):
    scene = bpy.context.scene
    f0, f1 = frame_range
    print(f"[mirpamo] retargeting '{clip_name}' frames {f0}..{f1} "
          f"({len(mapping)} bones mapped)")

    tgt_arm.animation_data_create()
    action = bpy.data.actions.new(clip_name)
    action.use_fake_user = True
    tgt_arm.animation_data.action = action

    # rest matrices (world space)
    src_rest = {s: world_rest(src_arm, s) for s in mapping}
    tgt_rest = {t: world_rest(tgt_arm, t) for t in mapping.values()}

    # scale compensation: hip heights at rest
    s_hip = hip_bone(src_arm, list(mapping.keys()))
    t_hip = mapping.get(s_hip) if s_hip else None
    scale = 1.0
    if s_hip and t_hip:
        sh = (src_arm.matrix_world @
              src_arm.data.bones[s_hip].head_local).z
        th = (tgt_arm.matrix_world @
              tgt_arm.data.bones[t_hip].head_local).z
        if sh > 1e-6:
            scale = th / sh
    print(f"[mirpamo] hip scale compensation: {scale:.4f}")

    # target bones in parent-before-child order
    ordered = []
    def walk(bone):
        if bone.name in tgt_rest:
            ordered.append(bone.name)
        for c in bone.children:
            walk(c)
    for b in tgt_arm.data.bones:
        if b.parent is None:
            walk(b)
    tgt_to_src = {t: s for s, t in mapping.items()}

    mw_tgt_inv = tgt_arm.matrix_world.inverted()
    depsgraph = bpy.context.evaluated_depsgraph_get()

    for f in range(f0, f1 + 1):
        scene.frame_set(f)
        src_ev = src_arm.evaluated_get(depsgraph)
        src_world = {s: src_ev.matrix_world @ src_ev.pose.bones[s].matrix
                     for s in mapping}

        # desired world matrix per target bone, computed top-down
        desired = {}
        for t in ordered:
            s = tgt_to_src[t]
            delta = (src_world[s] @ src_rest[s].inverted())
            m = delta @ tgt_rest[t]
            desired[t] = m

        for t in ordered:
            pb = tgt_arm.pose.bones[t]
            bone = tgt_arm.data.bones[t]
            m_des = mw_tgt_inv @ desired[t]  # into armature space
            if bone.parent:
                m_parent_des = mw_tgt_inv @ desired.get(
                    bone.parent.name,
                    tgt_arm.matrix_world @ tgt_arm.pose.bones[bone.parent.name].matrix)
                basis = (bone.matrix_local.inverted() @
                         bone.parent.matrix_local @
                         m_parent_des.inverted() @ m_des)
            else:
                basis = bone.matrix_local.inverted() @ m_des

            loc, rot, _ = basis.decompose()
            pb.rotation_mode = "QUATERNION"
            pb.rotation_quaternion = rot
            pb.keyframe_insert("rotation_quaternion", frame=f)

            if t == t_hip:
                # scale-compensated root translation
                src_hip_world = src_world[tgt_to_src[t]].to_translation()
                rest_hip_world = tgt_rest[t].to_translation()
                target_world = src_hip_world * scale
                if inplace:
                    target_world.x = rest_hip_world.x
                    target_world.y = rest_hip_world.y
                m = desired[t].copy()
                m.translation = target_world
                m_des2 = mw_tgt_inv @ m
                basis2 = bone.matrix_local.inverted() @ m_des2 \
                    if not bone.parent else basis
                pb.location = basis2.to_translation()
                pb.keyframe_insert("location", frame=f)

    return action


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--rig", required=True, help="rigged character (glb/fbx)")
    ap.add_argument("--clip", action="append", required=True,
                    help="motion clip (bvh/fbx/glb); repeatable")
    ap.add_argument("--out", required=True)
    ap.add_argument("--bonemap", default=None)
    ap.add_argument("--inplace", action="store_true",
                    help="strip horizontal root motion")
    ap.add_argument("--fps", type=int, default=None)
    args = ap.parse_args(stage_args())

    explicit = None
    if args.bonemap:
        with open(args.bonemap) as f:
            explicit = json.load(f)

    clean_scene()
    rig_objs = import_any(args.rig)
    tgt_arm = find_armature(rig_objs)
    if not tgt_arm:
        die("rig file contains no armature — run autorig first")
    tgt_meshes = find_meshes(rig_objs)
    if args.fps:
        bpy.context.scene.render.fps = args.fps

    for clip_path in args.clip:
        clip_objs = import_any(clip_path)
        src_arm = find_armature(clip_objs)
        if not src_arm:
            die(f"{clip_path}: no armature/motion found")

        mapping = build_bone_map(src_arm, tgt_arm, explicit)
        if len(mapping) < 3:
            die(f"{clip_path}: only {len(mapping)} bones mapped — "
                "provide a --bonemap")

        ad = src_arm.animation_data
        if not ad or not ad.action:
            die(f"{clip_path}: armature has no animation")
        src_action = ad.action
        fr = src_action.frame_range
        frame_range = (int(fr[0]), int(fr[1]))

        clip_name = os.path.splitext(os.path.basename(clip_path))[0]
        action = retarget_clip(src_arm, tgt_arm, mapping, clip_name,
                               args.inplace, frame_range)

        # drop the source skeleton + its action so only retargeted
        # animations reach the export
        for o in clip_objs:
            bpy.data.objects.remove(o, do_unlink=True)
        bpy.data.actions.remove(src_action)
        action.name = clip_name  # reclaim the name if it collided

    export_glb(args.out, objects=[tgt_arm] + tgt_meshes)
    print("[mirpamo] retarget OK")


main()