From f516ac8f66ceb634988fa01b1c2187533ec1329f Mon Sep 17 00:00:00 2001 From: Bart Moyaers Date: Mon, 10 Feb 2020 16:17:08 +0100 Subject: [PATCH] generalize script --- BvhToDM_blender.py | 292 +++++++++++++++++++++------------------------ 1 file changed, 138 insertions(+), 154 deletions(-) diff --git a/BvhToDM_blender.py b/BvhToDM_blender.py index d213cc3..68164b6 100644 --- a/BvhToDM_blender.py +++ b/BvhToDM_blender.py @@ -1,142 +1,153 @@ import bpy import json from mathutils import Vector +from typing import Union, List context = bpy.context scene = context.scene -def generate_DM_skeleton(): +# 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] + + # Add the armature if bpy.ops.object.mode_set.poll(): bpy.ops.object.mode_set(mode='OBJECT') - # TODO: by disabling the "inherit_rotation" property of some bones, the "matrix_base" does not take into account the rotation of the parent. - # This makes sense from an implementation standpoint, but not for our goal, namely: - # We want to disable inherit_rotation in order to get the IK-solver going to correct solutions (hips should rotate). - # But we want the calculated quaternion to be in LOCAL joint space, so we want it to be in relation to the parent's rotation. - # Proposed solution: re-enable the inherit-rotation option and use another IK solver instead. (itasc) bpy.ops.object.armature_add(enter_editmode=True) obj = bpy.data.objects["Armature"] - # obj = bpy.context.object - obj.name = "DMS" + obj.name = name # Rename armature armature = obj.data - armature.name = "DeepMimicSkeleton" + armature.name = name # Delete first auto-generated bone bpy.ops.armature.select_all(action='SELECT') bpy.ops.armature.delete() - # Start adding bones - root = armature.edit_bones.new("root") - root.tail = (0.0, 0.0, 0.001) - root.translate(Vector((0.0, 0.0, 0.0))) - # Spine, chest and head - spine = armature.edit_bones.new("spine") - spine.parent = root - spine.use_connect = True - spine.tail = (0.0, 0.0, 0.236151) - spine.translate(Vector((0.0, 0.0, 0.0))) + 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 + - chest = armature.edit_bones.new("chest") - chest.parent = spine - chest.use_connect = True - chest.tail = (chest.parent.tail[0], chest.parent.tail[1], chest.parent.tail[2] + 0.223894) - - head = armature.edit_bones.new("head") - head.parent = chest - head.use_connect = True - head.tail = (head.parent.tail[0], head.parent.tail[1], head.parent.tail[2] + 0.24) - - # Right arm - right_upper_arm = armature.edit_bones.new("right_upper_arm") - right_upper_arm.head = (chest.tail[0]-0.15, chest.tail[1], chest.tail[2]) - right_upper_arm.tail = (chest.tail[0]-0.15, chest.tail[1], chest.tail[2] - 0.28) - right_upper_arm.parent = chest - # right_upper_arm.use_inherit_rotation = False - - right_lower_arm = armature.edit_bones.new("right_lower_arm") - right_lower_arm.parent = right_upper_arm - right_lower_arm.use_connect = True - right_lower_arm.tail = (right_lower_arm.parent.tail[0], - right_lower_arm.parent.tail[1], - right_lower_arm.parent.tail[2] - 0.24) - - right_hand = armature.edit_bones.new("right_hand") - right_hand.parent = right_lower_arm - right_hand.use_connect = True - right_hand.tail = (right_hand.parent.tail[0], - right_hand.parent.tail[1], - right_hand.parent.tail[2] - 0.1) - - # Left arm - left_upper_arm = armature.edit_bones.new("left_upper_arm") - left_upper_arm.head = (chest.tail[0]+0.15, chest.tail[1], chest.tail[2]) - left_upper_arm.tail = (chest.tail[0]+0.15, chest.tail[1], chest.tail[2] - 0.28) - left_upper_arm.parent = chest - # left_upper_arm.use_inherit_rotation = False - - left_lower_arm = armature.edit_bones.new("left_lower_arm") - left_lower_arm.parent = left_upper_arm - left_lower_arm.use_connect = True - left_lower_arm.tail = (left_lower_arm.parent.tail[0], - left_lower_arm.parent.tail[1], - left_lower_arm.parent.tail[2] - 0.24) - - left_hand = armature.edit_bones.new("left_hand") - left_hand.parent = left_lower_arm - left_hand.use_connect = True - left_hand.tail = (left_hand.parent.tail[0], - left_hand.parent.tail[1], - left_hand.parent.tail[2] - 0.1) - - # Right leg - right_thigh = armature.edit_bones.new("right_thigh") - right_thigh.parent = root - right_thigh.head = (-0.084887, 0.0, 0.0) - right_thigh.tail = (-0.084887, 0.0, -0.421546) - # right_thigh.use_inherit_rotation = False - - right_shin = armature.edit_bones.new("right_shin") - right_shin.parent = right_thigh - right_shin.use_connect = True - right_shin.tail = (right_shin.parent.tail[0], right_shin.parent.tail[1], right_shin.parent.tail[2]-0.39) - - right_foot = armature.edit_bones.new("right_foot") - right_foot.parent = right_shin - right_foot.use_connect = True - right_foot.tail = (right_foot.parent.tail[0], right_foot.parent.tail[1]-0.15, right_foot.parent.tail[2]) - - # Left leg - left_thigh = armature.edit_bones.new("left_thigh") - left_thigh.parent = root - left_thigh.head = (0.084887, 0.0, 0.0) - left_thigh.tail = (0.084887, 0.0, -0.421546) - # left_thigh.use_inherit_rotation = False - - left_shin = armature.edit_bones.new("left_shin") - left_shin.parent = left_thigh - left_shin.use_connect = True - left_shin.tail = (left_shin.parent.tail[0], left_shin.parent.tail[1], left_shin.parent.tail[2]-0.39) - - left_foot = armature.edit_bones.new("left_foot") - left_foot.parent = left_shin - left_foot.use_connect = True - left_foot.tail = (left_foot.parent.tail[0], left_foot.parent.tail[1]-0.15, left_foot.parent.tail[2]) - -def constrain_DM_skeleton(): +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['DMS'] + skel = bpy.data.objects[skel_name] # Root root_con = skel.pose.bones['root'].constraints.new('COPY_TRANSFORMS') - root_con.target = bpy.data.objects['rotation_prob'] - root_con.subtarget = "Hip" + 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['rotation_prob'] - spine_con.subtarget = "Hip" + spine_con.target = bpy.data.objects[skel_name] + spine_con.subtarget = bonedict['spine'].target_name # Fixing rotation in elbow and knee joints: # 1) limit rotation constraint @@ -216,26 +227,11 @@ def damped_track(arm, target_arm, bone: str, target_bone: str): constraint.subtarget = target_bone constraint.head_tail = 1.0 -def follow_mocap_damped_tracks(arm, target_arm): - bonedict = { - # "right_upper_arm": "ShoulderRight", - # "right_lower_arm": "ElbowRight", - # "left_upper_arm": "ShoulderLeft", - # "left_lower_arm": "ElbowLeft", - # "right_thigh": "HipRight", - # "right_shin": "KneeRight", - # "left_thigh": "HipLeft", - # "left_shin": "KneeLeft", - "right_foot": "AnkleRight", - "left_foot": "AnkleLeft", - "right_hand": "WristRight", - "left_hand": "WristLeft", - "chest": "ShoulderCenter", - "head": "Head" - } - +def follow_mocap_damped_tracks(arm, target_arm, bonedict): for bone in bonedict.keys(): - damped_track(arm, target_arm, bone, bonedict[bone]) + 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') @@ -247,22 +243,11 @@ def inverse_kinematic_constraint(arm, target_arm, bone: str, target_bone: str, c constraint.iterations = 500 constraint.chain_count = chain_count -def follow_mocap_IK(arm, target_arm): - bonedict = { - "right_upper_arm": ("ElbowRight", 1), - "right_lower_arm": ("WristRight", 2), - "left_upper_arm": ("ElbowLeft", 1), - "left_lower_arm": ("WristLeft", 2), - "right_thigh": ("KneeRight", 1), - "right_shin": ("AnkleRight", 2), - "left_thigh": ("KneeLeft", 1), - "left_shin": ("AnkleLeft", 2), - # "chest": "ShoulderCenter", - # "head": "Head" - } - +def follow_mocap_IK(arm, target_arm, bonedict, child_dict): for bone in bonedict.keys(): - inverse_kinematic_constraint(arm, target_arm, bone, bonedict[bone][0], bonedict[bone][1]) + 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! @@ -344,18 +329,17 @@ def generate_frames(arm): return json.dumps(outputDict, indent=4) -DMS_name = 'DMS' -generate_DM_skeleton() -constrain_DM_skeleton() -scale_skeleton() +generate_humanoid_skeleton(target, skel_name, bonedict, child_dict) +constrain_DM_skeleton(bonedict) +# scale_skeleton() -DMS = bpy.data.objects[DMS_name] -target_arm = bpy.data.objects['rotation_prob'] -follow_mocap_damped_tracks(DMS, target_arm) -follow_mocap_IK(DMS, target_arm) +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) +# print(generated_text) with open("C:\\UntrackedGit\\BvhToDM_Blender"+ "\\output.txt", "w") as output: output.write(generated_text) \ No newline at end of file