BvhToDM_Blender/BvhToDM.py

import bpy
import json
from mathutils import Vector
from typing import Union, List
context = bpy.context
scene = context.scene

# TODO: Based on a dictionary with corresponding bone names (supplied by the user): generate the DeepMimicSkeleton with the exact bone lengths as the example BVH file!
# This should lead to almost perfectly accurate IK solutions!

target = "rotation_prob"
skel_name = "DeepMimicSkeleton"
# Couple every DM_bone to BVH bone
# Structure: (Target_bone_name: str, inherit_bone_length: bool, inherit_bone_pos: bool, use_connect: bool, bone_direction_vec: list[float])

class BoneConstructionData:
    def __init__(self, target_name: str,
                inherit_bone_length: bool,
                inherit_bone_pos: bool,
                use_connect: bool,
                bone_direction_vec: List[float] = [0.0, 0.0, 1.0],
                use_IK: bool = False,
                use_DT: bool = False,
                IK_chain_length: int = 0):
        self.target_name = target_name
        self.inherit_bone_length = inherit_bone_length
        self.inherit_bone_pos = inherit_bone_pos
        self.use_connect = use_connect
        self.bone_direction_vec = bone_direction_vec
        self.use_IK = use_IK
        self.use_DT = use_DT
        self.IK_chain_length = IK_chain_length

bonedict = {
    "root": BoneConstructionData("Hip", False, False, False, [0.0, 0.0, 1.0]),
    "spine": BoneConstructionData("Hip", False, False, True, [0.0, 0.0, 1.0]),
    "chest": BoneConstructionData("ShoulderCenter", False, False, True, [0.0, 0.0, 1.0], use_DT=True),
    "head": BoneConstructionData("Head", False, False, True, [0.0, 0.0, 1.0], use_DT=True),
    "right_upper_arm": BoneConstructionData("ShoulderRight", True, True, False, [0.0, 0.0, -1.0], use_IK=True, IK_chain_length=1),
    "right_lower_arm": BoneConstructionData("ElbowRight", True, False, True, [0.0, 0.0, -1.0], use_IK=True, IK_chain_length=2),
    "right_hand": BoneConstructionData("WristRight", True, False, True, [0.0, 0.0, -1.0], use_DT=True),
    "left_upper_arm": BoneConstructionData("ShoulderLeft", True, True, False, [0.0, 0.0, -1.0], use_IK=True, IK_chain_length=1),
    "left_lower_arm": BoneConstructionData("ElbowLeft", True, False, True, [0.0, 0.0, -1.0], use_IK=True, IK_chain_length=2),
    "left_hand": BoneConstructionData("WristLeft", True, False, True, [0.0, 0.0, -1.0], use_DT=True),
    "right_thigh": BoneConstructionData("HipRight", True, True, False, [0.0, 0.0, -1.0], use_IK=True, IK_chain_length=1),
    "right_shin": BoneConstructionData("KneeRight", True, False, True, [0.0, 0.0, -1.0], use_IK=True, IK_chain_length=2),
    "right_foot": BoneConstructionData("AnkleRight", True, False, True, [0.0, -1.0, 0.0], use_DT=True),
    "left_thigh": BoneConstructionData("HipLeft", True, True, False, [0.0, 0.0, -1.0], use_IK=True, IK_chain_length=1),
    "left_shin": BoneConstructionData("KneeLeft", True, False, True, [0.0, 0.0, -1.0], use_IK=True, IK_chain_length=2),
    "left_foot": BoneConstructionData("AnkleLeft", True, False, True, [0.0, -1.0, 0.0], use_DT=True),
}

child_dict = {
    "root": ["spine", "right_thigh", "left_thigh"],
    "spine": ["chest"],
    "chest": ["head", "right_upper_arm", "left_upper_arm"],
    "head": [],
    "right_upper_arm": ["right_lower_arm"],
    "right_lower_arm": ["right_hand"],
    "right_hand": [],
    "left_upper_arm": ["left_lower_arm"],
    "left_lower_arm": ["left_hand"],
    "left_hand": [],
    "right_thigh": ["right_shin"],
    "right_shin": ["right_foot"],
    "right_foot": [],
    "left_thigh": ["left_shin"],
    "left_shin": ["left_foot"],
    "left_foot": [],
}

def find_parent(bone_name, child_dict) -> Union[str, None]:
    for bone in child_dict.keys():
        if bone_name in child_dict[bone]:
            return bone
    return None

def find_child(bone_name, child_dict) -> Union[str, None]:
    result = child_dict[bone_name]
    if len(result) == 1:
        return result[0]
    else:
        return None

def tail_relative(bone, rel_vec: List[float]):
    bone.tail = [bone.head[x] + rel_vec[x] for x in range(3)]

def generate_humanoid_skeleton(target, name, bonedict, childdict):
    # First fill up dictionary with edit_bone positions
    # (Since it seems to be difficult to read the edit_bones of an armature while editing another one...)
    target_bone_head_dict = {}
    bpy.data.objects[target].select_set(True)
    if bpy.ops.object.mode_set.poll():
        bpy.ops.object.mode_set(mode='EDIT')
        for edit_bone in bpy.context.object.data.edit_bones:
            target_bone_head_dict[edit_bone.name] = [x for x in edit_bone.head]
        bpy.ops.object.mode_set(mode='OBJECT')
    
    # Add the armature
    bpy.ops.object.armature_add(enter_editmode=True)
    obj = bpy.data.objects["Armature"]
    obj.name = name
    
    # Rename armature
    armature = obj.data
    armature.name = name
    # Delete first auto-generated bone
    bpy.ops.armature.select_all(action='SELECT')
    bpy.ops.armature.delete()
    
    for bone_name in bonedict.keys():
        construction_data: BoneConstructionData = bonedict[bone_name]
        if bone_name == "root":
            root = armature.edit_bones.new("root")
            root.tail = (0.0, 0.0, 0.001)
        else:
            bone = armature.edit_bones.new(bone_name)
            parent = find_parent(bone_name, childdict)
            if parent is not None:
                bone.parent = armature.edit_bones[parent]
            bone.use_connect = construction_data.use_connect
            if construction_data.inherit_bone_pos:
                bone.head = target_bone_head_dict[construction_data.target_name]
            if construction_data.inherit_bone_length:
                length = bpy.data.objects[target].data.bones[construction_data.target_name].length
                tail_relative(bone, [length * x for x in construction_data.bone_direction_vec])
            else:
                # TODO: find out some way to guess the correct bone lengths if there are no corresponding bones in the target armature...
                if bone_name in ["spine", "chest", "head"]:
                    tail_relative(bone, [0.0, 0.0, 0.22])
                pass

    
def constrain_DM_skeleton(bonedict):
    if bpy.ops.object.mode_set.poll():
        bpy.ops.object.mode_set(mode='POSE')
    # Note that pose_bones are under: bpy.data.objects[...].pose.bones
    skel = bpy.data.objects[skel_name]

    # Root
    root_con = skel.pose.bones['root'].constraints.new('COPY_TRANSFORMS')
    root_con.target = bpy.data.objects[target]
    root_con.subtarget = bonedict['root'].target_name

    # Spine
    spine_con = skel.pose.bones['spine'].constraints.new('COPY_ROTATION')
    spine_con.target = bpy.data.objects[target]
    spine_con.subtarget = bonedict['spine'].target_name
    
    # Fixing rotation in elbow and knee joints:
    # 1) limit rotation constraint
    # 2) y- and z-rotations constrain from 0 to 0 degrees. "Fixed"
    # 3) Elbows: limit rotation in "Local space" from -180 to 0 degrees.
    #    Knees: from 0 to 180 degrees

    # Knees
    rightshin = skel.pose.bones['right_shin']
    rightshin.ik_max_x = 3.1415
    rightshin.ik_min_x = 0
    rightshin.use_ik_limit_x = True
    rightshin.lock_ik_y = True
    rightshin.lock_ik_z = True
    rightshin.rotation_mode = 'AXIS_ANGLE'

    leftshin = skel.pose.bones['left_shin']
    leftshin.ik_max_x = 3.1415
    leftshin.ik_min_x = 0
    leftshin.use_ik_limit_x = True
    leftshin.lock_ik_y = True
    leftshin.lock_ik_z = True
    leftshin.rotation_mode = 'AXIS_ANGLE'

    # Elbows
    right_lower_arm = skel.pose.bones['right_lower_arm']
    right_lower_arm.ik_max_x = 0
    right_lower_arm.ik_min_x = -3.1415
    right_lower_arm.use_ik_limit_x = True
    right_lower_arm.lock_ik_y = True
    right_lower_arm.lock_ik_z = True
    right_lower_arm.rotation_mode = 'AXIS_ANGLE'

    left_lower_arm = skel.pose.bones['left_lower_arm']
    left_lower_arm.ik_max_x = 0
    left_lower_arm.ik_min_x = -3.1415
    left_lower_arm.use_ik_limit_x = True
    left_lower_arm.lock_ik_y = True
    left_lower_arm.lock_ik_z = True
    left_lower_arm.rotation_mode = 'AXIS_ANGLE'

    # IK settings
    skel.pose.ik_solver = "ITASC"
    skel.pose.ik_param.mode = "SIMULATION"
    
def scale_skeleton():
    # Scale so that the distance from ankle to shoulders is the same

    # Set both objects to edit mode
    if bpy.ops.object.mode_set.poll():
        bpy.ops.object.mode_set(mode='EDIT')

    shoulder_DM = "right_upper_arm"
    shoulder_mocap = "ShoulderRight"
    ankle_DM = "right_foot"
    ankle_mocap = "AnkleRight"

    DMS = bpy.data.armatures['DeepMimicSkeleton']
    mocap_skel = bpy.data.armatures['rotation_prob']

    # Only take into account z-component
    mocap_height = mocap_skel.bones[shoulder_mocap].head_local[2] - mocap_skel.bones[ankle_mocap].head_local[2]
    DM_height = DMS.bones[shoulder_DM].head_local[2] - DMS.bones[ankle_DM].head_local[2]
    
    scale_factor = mocap_height/DM_height

    # Select all bones
    bpy.data.objects['DMS'].select_set(True)
    bpy.ops.armature.select_all(action='SELECT')

    # Scale
    bpy.ops.transform.resize(value=(scale_factor, scale_factor, scale_factor))

def damped_track(arm, target_arm, bone: str, target_bone: str):
    constraint = arm.pose.bones[bone].constraints.new('DAMPED_TRACK')
    constraint.target = target_arm
    constraint.subtarget = target_bone
    constraint.head_tail = 1.0

def follow_mocap_damped_tracks(arm, target_arm, bonedict):
    for bone in bonedict.keys():
        bone_data = bonedict[bone]
        if bone_data.use_DT:
            damped_track(arm, target_arm, bone, bone_data.target_name)

def inverse_kinematic_constraint(arm, target_arm, bone: str, target_bone: str, chain_count: int):
    constraint = arm.pose.bones[bone].constraints.new('IK')
    constraint.target = target_arm
    constraint.subtarget = target_bone
    constraint.use_tail = True
    constraint.use_location = True
    constraint.use_stretch = True
    constraint.iterations = 500
    constraint.chain_count = chain_count
    
def follow_mocap_IK(arm, target_arm, bonedict, child_dict):
    for bone in bonedict.keys():
        bone_data = bonedict[bone]
        if bone_data.use_IK:
            inverse_kinematic_constraint(arm, target_arm, bone, bonedict[find_child(bone, child_dict)].target_name, bone_data.IK_chain_length)

def quat_bvh_to_DM(quat):
    # transform x -> -z and z -> -x, except for ankles and root!
    return [quat[0], -quat[3], -quat[2], -quat[1]]

def quat_bvh_to_DM_ankles(quat):
    # transform x -> -z and z -> x
    return [quat[0], quat[3], -quat[2], quat[1]]

def quat_bvh_to_DM_root(quat):
    return [quat[0], quat[3], quat[2], -quat[1]]

def pos_blender_DM(pos, scale):
    # y -> z, z -> -y?
    return [-scale*pos[1], scale*pos[2], -scale*pos[0]]

def generate_frame(arm):

    joints = [
        "seconds", "hip", "hip", "chest", "neck", "right hip", "right knee", "right ankle",
        "right shoulder", "right elbow", "left hip", "left knee", "left ankle", "left shoulder", 
        "left elbow"
    ]

    bone_names = [
        "chest", "head", "right_thigh", "right_shin", "right_foot",
        "right_upper_arm", "right_lower_arm", "left_thigh", "left_shin", "left_foot", "left_upper_arm", 
        "left_lower_arm"
    ]

    joints_1d = ["right_shin", "left_shin", "right_lower_arm", "left_lower_arm"]
    ankles = ["right_foot", "left_foot"]

    result = []
    result.append(1.0 / bpy.context.scene.render.fps)

    translation_scale = 1.2
    pos = pos_blender_DM(arm.pose.bones['root'].head, translation_scale)
    result.extend(pos)
    root_quat = quat_bvh_to_DM_root(arm.pose.bones['root'].rotation_quaternion)
    result.extend(root_quat)

    for bone_name in bone_names:
        if bone_name in joints_1d:
            # invert knee angles (no clue why)
            if bone_name in joints_1d[:2]:
                angle = -arm.pose.bones[bone_name].rotation_axis_angle[0]
            else:
                angle = arm.pose.bones[bone_name].rotation_axis_angle[0]
            result.append(angle)
        else:
            # TODO: these rotations are absolute rotations! But we need local rotations... In the local bone frame...
            if bone_name in ankles:
                quat_DM = quat_bvh_to_DM_ankles(arm.pose.bones[bone_name].rotation_quaternion)
            else:
                quat_DM = quat_bvh_to_DM(arm.pose.bones[bone_name].rotation_quaternion)
            result.extend([x for x in quat_DM])

    return result

def generate_frames(arm):
    if bpy.ops.object.mode_set.poll():
        bpy.ops.object.mode_set(mode='EDIT')
        # Select all bones
        bpy.data.objects[skel_name].select_set(True)
        bpy.ops.armature.select_all(action='SELECT')

    if bpy.ops.object.mode_set.poll():
        bpy.ops.object.mode_set(mode='POSE')

    frames = []
    loopText = "none"
    for frame in range(bpy.context.scene.frame_start, bpy.context.scene.frame_end + 1):
    # for frame in range(bpy.context.scene.frame_end + 1):
        bpy.context.scene.frame_set(frame)
        # Apply visual transform to pose
        bpy.ops.pose.visual_transform_apply()
        # bpy.context.view_layer.update()
        frames.append(generate_frame(arm))

    # Output in dictionary format for easy json dump
    outputDict = {
        "Loop": loopText,  # "none" or "wrap"
        "Frames": frames
    }

    return json.dumps(outputDict, indent=4)

generate_humanoid_skeleton(target, skel_name, bonedict, child_dict)
constrain_DM_skeleton(bonedict)
# scale_skeleton()

DMS = bpy.data.objects[skel_name]
target_arm = bpy.data.objects[target]
follow_mocap_damped_tracks(DMS, target_arm, bonedict)
follow_mocap_IK(DMS, target_arm, bonedict, child_dict)

# Save all quaternions in DeepMimic format, frame per frame!
generated_text = generate_frames(DMS)
# print(generated_text)
with open("C:\\UntrackedGit\\BvhToDM_Blender"+ "\\output.txt", "w") as output:
    output.write(generated_text)