Changed name of Mocap constraints to mocap fixes, for user clarity.
[blender.git] / release / scripts / modules / retarget.py
index c7a482659ef6a90a079426ce4977fe23754b729e..bec7b8aaa3e532bee3460bb265d64c89549a41ea 100644 (file)
@@ -1,49 +1,65 @@
+# ##### BEGIN GPL LICENSE BLOCK #####
+#
+#  This program is free software; you can redistribute it and/or
+#  modify it under the terms of the GNU General Public License
+#  as published by the Free Software Foundation; either version 2
+#  of the License, or (at your option) any later version.
+#
+#  This program is distributed in the hope that it will be useful,
+#  but WITHOUT ANY WARRANTY; without even the implied warranty of
+#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#  GNU General Public License for more details.
+#
+#  You should have received a copy of the GNU General Public License
+#  along with this program; if not, write to the Free Software Foundation,
+#  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+#
+# ##### END GPL LICENSE BLOCK #####
+
+# <pep8 compliant>
+
 import bpy
 from mathutils import *
 from math import radians, acos
-performer_obj = bpy.data.objects["performer"]
-enduser_obj = bpy.data.objects["enduser"]
-scene = bpy.context.scene
-
-# dictionary of mapping
-bonemap = { "LeftFoot": ("DEF_Foot.L","DEF_Toes.L"),
-            "LeftUpLeg": "DEF_Thigh.L",
-            "Hips": "DEF_Hip",
-            "LowerBack": "DEF_Spine",
-            "Spine": "DEF_Torso",
-            "Neck": "DEF_Neck",
-            "Neck1": "DEF_Neck",
-            "Head": "DEF_Head",
-            "LeftShoulder": "DEF_Shoulder.L",
-            "LeftArm": "DEF_Forearm.L",
-            "LeftForeArm": "DEF_Arm.L",
-            "LeftHand": "DEF_Hand.L",
-            "RightShoulder": "DEF_Shoulder.R",
-            "RightArm": "DEF_Forearm.R",
-            "RightForeArm": "DEF_Arm.R",
-            "RightHand": "DEF_Hand.R",
-            "RightFoot": ("DEF_Foot.R","DEF_Toes.R"),
-            "RightUpLeg": "DEF_Thigh.R",
-            "RightLeg": "DEF_Shin.R",
-            "LeftLeg": "DEF_Shin.L"}
-# creation of a reverse map
-# multiple keys get mapped to list values
-bonemapr = {}
-for key in bonemap.keys():
-    if not bonemap[key] in bonemapr:
-        if type(bonemap[key])==type((0,0)):
-            for key_x in bonemap[key]:
-                bonemapr[key_x] = [key]
-        else:
-            bonemapr[bonemap[key]] = [key]
+import cProfile
+
+
+def hasIKConstraint(pose_bone):
+    #utility function / predicate, returns True if given bone has IK constraint
+    ik = [constraint for constraint in pose_bone.constraints if constraint.type == "IK"]
+    if ik:
+        return ik[0]
     else:
-        bonemapr[bonemap[key]].append(key)
-        
-# list of empties created to keep track of "original"
-# position data
-# in final product, these locations can be stored as custom props
+        return False
+
+
+def createDictionary(perf_arm, end_arm):
+    # clear any old data
+    for end_bone in end_arm.bones:
+        for mapping in end_bone.reverseMap:
+            end_bone.reverseMap.remove(0)
+
+    for perf_bone in perf_arm.bones:
+        #find its match and add perf_bone to the match's mapping
+        if perf_bone.map:
+            end_bone = end_arm.bones[perf_bone.map]
+            newMap = end_bone.reverseMap.add()
+            newMap.name = perf_bone.name
+            end_bone.foot = perf_bone.foot
+
+    #root is the root of the enduser
+    root = end_arm.bones[0].name
+    feetBones = [bone.name for bone in perf_arm.bones if bone.foot]
+    return feetBones, root
 
-constraints = []
+
+def loadMapping(perf_arm, end_arm):
+
+    for end_bone in end_arm.bones:
+        #find its match and add perf_bone to the match's mapping
+        if end_bone.reverseMap:
+            for perf_bone in end_bone.reverseMap:
+                perf_arm.bones[perf_bone.name].map = end_bone.name
 
 #creation of intermediate armature
 # the intermediate armature has the hiearchy of the end user,
@@ -51,86 +67,360 @@ constraints = []
 # and bone roll is identical to the performer
 # its purpose is to copy over the rotations
 # easily while concentrating on the hierarchy changes
-def createIntermediate():
-    
+
+
+def createIntermediate(performer_obj, enduser_obj, root, s_frame, e_frame, scene):
     #creates and keyframes an empty with its location
     #the original position of the tail bone
     #useful for storing the important data in the original motion
-    #i.e. using this empty to IK the chain to that pos.
-    def locOfOriginal(inter_bone,perf_bone):
-        if not perf_bone.name+"Org" in bpy.data.objects:
-            bpy.ops.object.add()
-            empty = bpy.context.active_object
-            empty.name = perf_bone.name+"Org"
-        empty = bpy.data.objects[perf_bone.name+"Org"]
-        offset = perf_bone.vector
-        scaling = perf_bone.length / inter_bone.length
-        offset/=scaling
-        empty.location = inter_bone.head + offset
-        empty.keyframe_insert("location")
-    
+    #i.e. using this empty to IK the chain to that pos / DEBUG
+
     #Simple 1to1 retarget of a bone
-    def singleBoneRetarget(inter_bone,perf_bone):
-            perf_world_rotation = perf_bone.matrix * performer_obj.matrix_world         
+    def singleBoneRetarget(inter_bone, perf_bone):
+            perf_world_rotation = perf_bone.matrix * performer_obj.matrix_world
             inter_world_base_rotation = inter_bone.bone.matrix_local * inter_obj.matrix_world
             inter_world_base_inv = Matrix(inter_world_base_rotation)
             inter_world_base_inv.invert()
             return (inter_world_base_inv.to_3x3() * perf_world_rotation.to_3x3()).to_4x4()
-        
-    #uses 1to1 and interpolation/averaging to match many to 1 retarget    
-    def manyPerfToSingleInterRetarget(inter_bone,performer_bones_s):
-        retarget_matrices = [singleBoneRetarget(inter_bone,perf_bone) for perf_bone in performer_bones_s]
+
+    #uses 1to1 and interpolation/averaging to match many to 1 retarget
+    def manyPerfToSingleInterRetarget(inter_bone, performer_bones_s):
+        retarget_matrices = [singleBoneRetarget(inter_bone, perf_bone) for perf_bone in performer_bones_s]
         lerp_matrix = Matrix()
-        for i in range(len(retarget_matrices)-1):
+        for i in range(len(retarget_matrices) - 1):
             first_mat = retarget_matrices[i]
-            next_mat = retarget_matrices[i+1]
-            lerp_matrix = first_mat.lerp(next_mat,0.5)
+            next_mat = retarget_matrices[i + 1]
+            lerp_matrix = first_mat.lerp(next_mat, 0.5)
         return lerp_matrix
-    
-    #determines the type of hierachy change needed and calls the 
-    #right function        
+
+    #determines the type of hierachy change needed and calls the
+    #right function
     def retargetPerfToInter(inter_bone):
-        if inter_bone.name in bonemapr.keys():
-            perf_bone_name = bonemapr[inter_bone.name]
-            #is it a 1 to many?
-            if type(bonemap[perf_bone_name[0]])==type((0,0)):
-                perf_bone = performer_bones[perf_bone_name[0]]
-                if inter_bone.name == bonemap[perf_bone_name[0]][0]:
-                    locOfOriginal(inter_bone,perf_bone)
-            else:
+        if inter_bone.bone.reverseMap:
+            perf_bone_name = inter_bone.bone.reverseMap
+                # 1 to many not supported yet
                 # then its either a many to 1 or 1 to 1
-                
-                if len(perf_bone_name) > 1:
-                    performer_bones_s = [performer_bones[name] for name in perf_bone_name]
-                    #we need to map several performance bone to a single
-                    inter_bone.matrix_basis = manyPerfToSingleInterRetarget(inter_bone,performer_bones_s)
-                else:
-                    perf_bone = performer_bones[perf_bone_name[0]]
-                    inter_bone.matrix_basis = singleBoneRetarget(inter_bone,perf_bone)
-                    
+            if len(perf_bone_name) > 1:
+                performer_bones_s = [performer_bones[map.name] for map in perf_bone_name]
+                #we need to map several performance bone to a single
+                inter_bone.matrix_basis = manyPerfToSingleInterRetarget(inter_bone, performer_bones_s)
+            else:
+                perf_bone = performer_bones[perf_bone_name[0].name]
+                inter_bone.matrix_basis = singleBoneRetarget(inter_bone, perf_bone)
         inter_bone.keyframe_insert("rotation_quaternion")
-        for child in inter_bone.children:
-            retargetPerfToInter(child)
-            
-    #creates the intermediate armature object        
-    bpy.ops.object.select_name(name="enduser",extend=False)
-    bpy.ops.object.duplicate(linked=False)
-    bpy.context.active_object.name = "intermediate"
-    inter_obj = bpy.context.active_object
+
+    #creates the intermediate armature object
+    inter_obj = enduser_obj.copy()
+    inter_obj.data = inter_obj.data.copy()  # duplicate data
+    bpy.context.scene.objects.link(inter_obj)
+    inter_obj.name = "intermediate"
+    bpy.context.scene.objects.active = inter_obj
     bpy.ops.object.mode_set(mode='EDIT')
-    #resets roll 
+    #add some temporary connecting bones in case end user bones are not connected to their parents
+    print("creating temp bones")
+    for bone in inter_obj.data.edit_bones:
+        if not bone.use_connect and bone.parent:
+            if inter_obj.data.bones[bone.parent.name].reverseMap or inter_obj.data.bones[bone.name].reverseMap:
+                newBone = inter_obj.data.edit_bones.new("Temp")
+                newBone.head = bone.parent.tail
+                newBone.tail = bone.head
+                newBone.parent = bone.parent
+                bone.parent = newBone
+                bone.use_connect = True
+                newBone.use_connect = True
+    #resets roll
+    print("retargeting to intermediate")
     bpy.ops.armature.calculate_roll(type='Z')
     bpy.ops.object.mode_set(mode="OBJECT")
+    inter_obj.data.name = "inter_arm"
+    inter_arm = inter_obj.data
     performer_bones = performer_obj.pose.bones
-    inter_bones =  inter_obj.pose.bones
-    
+    inter_bones = inter_obj.pose.bones
     #clears inheritance
     for inter_bone in inter_bones:
-        inter_bone.bone.use_inherit_rotation = False
-        
-    for t in range(1,150):
+        if inter_bone.bone.reverseMap:
+            inter_bone.bone.use_inherit_rotation = False
+        else:
+            inter_bone.bone.use_inherit_rotation = True
+
+    for t in range(s_frame, e_frame):
+        if (t - s_frame) % 10 == 0:
+            print("First pass: retargeting frame {0}/{1}".format(t, e_frame - s_frame))
+        scene.frame_set(t)
+        for bone in inter_bones:
+            retargetPerfToInter(bone)
+
+    return inter_obj
+
+# this procedure copies the rotations over from the intermediate
+# armature to the end user one.
+# As the hierarchies are 1 to 1, this is a simple matter of
+# copying the rotation, while keeping in mind bone roll, parenting, etc.
+# TODO: Control Bones: If a certain bone is constrained in a way
+#       that its rotation is determined by another (a control bone)
+#       We should determine the right pos of the control bone.
+#       Scale: ? Should work but needs testing.
+
+
+def retargetEnduser(inter_obj, enduser_obj, root, s_frame, e_frame, scene):
+    inter_bones = inter_obj.pose.bones
+    end_bones = enduser_obj.pose.bones
+
+    def bakeTransform(end_bone):
+        src_bone = inter_bones[end_bone.name]
+        trg_bone = end_bone
+        bake_matrix = src_bone.matrix
+        rest_matrix = trg_bone.bone.matrix_local
+
+        if trg_bone.parent and trg_bone.bone.use_inherit_rotation:
+            srcParent = src_bone.parent
+            if "Temp" in srcParent.name:
+                srcParent = srcParent.parent
+            parent_mat = srcParent.matrix
+            parent_rest = trg_bone.parent.bone.matrix_local
+            parent_rest_inv = parent_rest.copy()
+            parent_rest_inv.invert()
+            parent_mat_inv = parent_mat.copy()
+            parent_mat_inv.invert()
+            bake_matrix = parent_mat_inv * bake_matrix
+            rest_matrix = parent_rest_inv * rest_matrix
+
+        rest_matrix_inv = rest_matrix.copy()
+        rest_matrix_inv.invert()
+        bake_matrix = rest_matrix_inv * bake_matrix
+        end_bone.matrix_basis = bake_matrix
+        rot_mode = end_bone.rotation_mode
+        if rot_mode == "QUATERNION":
+            end_bone.keyframe_insert("rotation_quaternion")
+        elif rot_mode == "AXIS_ANGLE":
+            end_bone.keyframe_insert("rotation_axis_angle")
+        else:
+            end_bone.keyframe_insert("rotation_euler")
+        if not end_bone.bone.use_connect:
+            end_bone.keyframe_insert("location")
+
+        for bone in end_bone.children:
+            bakeTransform(bone)
+
+    for t in range(s_frame, e_frame):
+        if (t - s_frame) % 10 == 0:
+            print("Second pass: retargeting frame {0}/{1}".format(t, e_frame - s_frame))
+        scene.frame_set(t)
+        end_bone = end_bones[root]
+        end_bone.location = Vector((0, 0, 0))
+        end_bone.keyframe_insert("location")
+        bakeTransform(end_bone)
+
+#recieves the performer feet bones as a variable
+# by "feet" I mean those bones that have plants
+# (they don't move, despite root moving) somewhere in the animation.
+
+
+def copyTranslation(performer_obj, enduser_obj, perfFeet, root, s_frame, e_frame, scene, enduser_obj_mat):
+
+    perf_bones = performer_obj.pose.bones
+    end_bones = enduser_obj.pose.bones
+
+    perfRoot = perf_bones[0].name
+    endFeet = [perf_bones[perfBone].bone.map for perfBone in perfFeet]
+    locDictKeys = perfFeet + endFeet + [perfRoot]
+
+    def tailLoc(bone):
+        return bone.center + (bone.vector / 2)
+
+    #Step 1 - we create a dict that contains these keys:
+    #(Performer) Hips, Feet
+    #(End user) Feet
+    # where the values are their world position on each frame in range (s,e)
+
+    locDict = {}
+    for key in locDictKeys:
+        locDict[key] = []
+
+    for t in range(scene.frame_start, scene.frame_end):
         scene.frame_set(t)
-        inter_bone = inter_bones["DEF_Hip"]
-        retargetPerfToInter(inter_bone)
-         
-createIntermediate()   
\ No newline at end of file
+        for bone in perfFeet:
+            locDict[bone].append(tailLoc(perf_bones[bone]))
+        locDict[perfRoot].append(tailLoc(perf_bones[perfRoot]))
+        for bone in endFeet:
+            locDict[bone].append(tailLoc(end_bones[bone]))
+
+    # now we take our locDict and analyze it.
+    # we need to derive all chains
+
+    locDeriv = {}
+    for key in locDictKeys:
+        locDeriv[key] = []
+
+    for key in locDict.keys():
+        graph = locDict[key]
+        locDeriv[key] = [graph[t + 1] - graph[t] for t in range(len(graph) - 1)]
+
+    # now find the plant frames, where perfFeet don't move much
+
+    linearAvg = []
+
+    for key in perfFeet:
+        for i in range(len(locDeriv[key]) - 1):
+            v = locDeriv[key][i]
+            hipV = locDeriv[perfRoot][i]
+            endV = locDeriv[perf_bones[key].bone.map][i]
+            if (v.length < 0.1):
+                #this is a plant frame.
+                #lets see what the original hip delta is, and the corresponding
+                #end bone's delta
+                if endV.length != 0:
+                    linearAvg.append(hipV.length / endV.length)
+
+    bpy.ops.object.add()
+    stride_bone = bpy.context.active_object
+    stride_bone.name = "stride_bone"
+
+    if linearAvg:
+        #determine the average change in scale needed
+        avg = sum(linearAvg) / len(linearAvg)
+        scene.frame_set(s_frame)
+        initialPos = (tailLoc(perf_bones[perfRoot]) / avg)
+        for t in range(s_frame, e_frame):
+            scene.frame_set(t)
+            #calculate the new position, by dividing by the found ratio between performer and enduser
+            newTranslation = (tailLoc(perf_bones[perfRoot]) / avg)
+            stride_bone.location = (newTranslation - initialPos) * enduser_obj_mat
+            stride_bone.keyframe_insert("location")
+    return stride_bone
+
+
+def IKRetarget(performer_obj, enduser_obj, s_frame, e_frame, scene):
+    end_bones = enduser_obj.pose.bones
+    for pose_bone in end_bones:
+        ik_constraint = hasIKConstraint(pose_bone)
+        if ik_constraint:
+            target_is_bone = False
+            # set constraint target to corresponding empty if targetless,
+            # if not, keyframe current target to corresponding empty
+            perf_bone = pose_bone.bone.reverseMap[-1].name
+            orgLocTrg = originalLocationTarget(pose_bone)
+            if not ik_constraint.target:
+                ik_constraint.target = orgLocTrg
+                target = orgLocTrg
+
+            # There is a target now
+            if ik_constraint.subtarget:
+                target = ik_constraint.target.pose.bones[ik_constraint.subtarget]
+                target.bone.use_local_location = False
+                target_is_bone = True
+            else:
+                target = ik_constraint.target
+
+            # bake the correct locations for the ik target bones
+            for t in range(s_frame, e_frame):
+                scene.frame_set(t)
+                if target_is_bone:
+                    final_loc = pose_bone.tail - target.bone.matrix_local.to_translation()
+                else:
+                    final_loc = pose_bone.tail
+                target.location = final_loc
+                target.keyframe_insert("location")
+            ik_constraint.mute = False
+
+
+def turnOffIK(enduser_obj):
+    end_bones = enduser_obj.pose.bones
+    for pose_bone in end_bones:
+        if pose_bone.is_in_ik_chain:
+            pass
+            # TODO:
+            # set stiffness according to place on chain
+            # and values from analysis that is stored in the bone
+            #pose_bone.ik_stiffness_x = 0.5
+            #pose_bone.ik_stiffness_y = 0.5
+            #pose_bone.ik_stiffness_z = 0.5
+        ik_constraint = hasIKConstraint(pose_bone)
+        if ik_constraint:
+            ik_constraint.mute = True
+
+
+#copy the object matrixes and clear them (to be reinserted later)
+def cleanAndStoreObjMat(performer_obj, enduser_obj):
+    perf_obj_mat = performer_obj.matrix_world.copy()
+    enduser_obj_mat = enduser_obj.matrix_world.copy()
+    zero_mat = Matrix()
+    performer_obj.matrix_world = zero_mat
+    enduser_obj.matrix_world = zero_mat
+    return perf_obj_mat, enduser_obj_mat
+
+
+#restore the object matrixes after parenting the auto generated IK empties
+def restoreObjMat(performer_obj, enduser_obj, perf_obj_mat, enduser_obj_mat, stride_bone):
+    pose_bones = enduser_obj.pose.bones
+    for pose_bone in pose_bones:
+        if pose_bone.name + "Org" in bpy.data.objects:
+            empty = bpy.data.objects[pose_bone.name + "Org"]
+            empty.parent = stride_bone
+    performer_obj.matrix_world = perf_obj_mat
+    enduser_obj.matrix_world = enduser_obj_mat
+    enduser_obj.parent = stride_bone
+
+
+#create (or return if exists) the related IK empty to the bone
+def originalLocationTarget(end_bone):
+    if not end_bone.name + "Org" in bpy.data.objects:
+        bpy.ops.object.add()
+        empty = bpy.context.active_object
+        empty.name = end_bone.name + "Org"
+        empty.empty_draw_size = 0.1
+        #empty.parent = enduser_obj
+    empty = bpy.data.objects[end_bone.name + "Org"]
+    return empty
+
+
+#create the specified NLA setup for base animation, constraints and tweak layer.
+def NLASystemInitialize(enduser_obj, s_frame):
+    anim_data = enduser_obj.animation_data
+    mocapAction = anim_data.action
+    mocapAction.name = "Base Mocap"
+    anim_data.use_nla = True
+    mocapTrack = anim_data.nla_tracks.new()
+    mocapTrack.name = "Base Mocap Track"
+    mocapStrip = mocapTrack.strips.new("Base Mocap", s_frame, mocapAction)
+    constraintTrack = anim_data.nla_tracks.new()
+    constraintTrack.name = "Mocap fixes"
+    constraintAction = bpy.data.actions.new("Mocap fixes")
+    constraintStrip = constraintTrack.strips.new("Mocap fixes", s_frame, constraintAction)
+    constraintStrip.extrapolation = "NOTHING"
+    userTrack = anim_data.nla_tracks.new()
+    userTrack.name = "Mocap manual fix"
+    userAction = bpy.data.actions.new("Mocap manual fix")
+    userStrip = userTrack.strips.new("Mocap manual fix", s_frame, userAction)
+    userStrip.extrapolation = "HOLD"
+    #userStrip.blend_type = "MULITPLY" - doesn't work due to work, will be activated soon
+    anim_data.nla_tracks.active = constraintTrack
+    anim_data.action = constraintAction
+    anim_data.action_extrapolation = "NOTHING"
+
+
+#Main function that runs the retargeting sequence.
+def totalRetarget(performer_obj, enduser_obj, scene, s_frame, e_frame):
+    perf_arm = performer_obj.data
+    end_arm = enduser_obj.data
+    print("creating Dictionary")
+    feetBones, root = createDictionary(perf_arm, end_arm)
+    print("cleaning stuff up")
+    perf_obj_mat, enduser_obj_mat = cleanAndStoreObjMat(performer_obj, enduser_obj)
+    turnOffIK(enduser_obj)
+    print("Creating intermediate armature (for first pass)")
+    inter_obj = createIntermediate(performer_obj, enduser_obj, root, s_frame, e_frame, scene)
+    print("First pass: retargeting from intermediate to end user")
+    retargetEnduser(inter_obj, enduser_obj, root, s_frame, e_frame, scene)
+    print("Second pass: retargeting root translation and clean up")
+    stride_bone = copyTranslation(performer_obj, enduser_obj, feetBones, root, s_frame, e_frame, scene, enduser_obj_mat)
+    IKRetarget(performer_obj, enduser_obj, s_frame, e_frame, scene)
+    restoreObjMat(performer_obj, enduser_obj, perf_obj_mat, enduser_obj_mat, stride_bone)
+    bpy.ops.object.mode_set(mode='OBJECT')
+    bpy.ops.object.select_name(name=inter_obj.name, extend=False)
+    bpy.ops.object.delete()
+    NLASystemInitialize(enduser_obj, s_frame)
+    print("retargeting done!")
+
+if __name__ == "__main__":
+    totalRetarget()