first commit

.gitignore

@@ -0,0 +1,2 @@
+output.txt
+

ik_deepmimic_test.py

@@ -0,0 +1,361 @@
+import bpy
+import json
+from mathutils import Vector
+context = bpy.context
+scene = context.scene
+
+def generate_DM_skeleton():
+    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"
+    
+    # Rename armature
+    armature = obj.data
+    armature.name = "DeepMimicSkeleton"
+    # 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)))
+    
+    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():
+    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']
+
+    # Root
+    root_con = skel.pose.bones['root'].constraints.new('COPY_TRANSFORMS')
+    root_con.target = bpy.data.objects['rotation_prob']
+    root_con.subtarget = "Hip"
+
+    # Spine
+    spine_con = skel.pose.bones['spine'].constraints.new('COPY_ROTATION')
+    spine_con.target = bpy.data.objects['rotation_prob']
+    spine_con.subtarget = "Hip"
+    
+    # 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 = {
+        # "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"
+    }
+
+    for bone in bonedict.keys():
+        damped_track(arm, target_arm, bone, bonedict[bone])
+
+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 = {
+        "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"
+    }
+
+    for bone in bonedict.keys():
+        inverse_kinematic_constraint(arm, target_arm, bone, bonedict[bone][0], bonedict[bone][1])
+
+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='POSE')
+
+    frames = []
+    loopText = "none"
+
+    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()
+        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)
+
+DMS_name = 'DMS'
+generate_DM_skeleton()
+constrain_DM_skeleton()
+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)
+
+# 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)