Advanced Retargeting option: If the end user armature is complex, on the level of...
[blender.git] / release / scripts / modules / retarget.py
1 # ##### BEGIN GPL LICENSE BLOCK #####
2 #
3 #  This program is free software; you can redistribute it and/or
4 #  modify it under the terms of the GNU General Public License
5 #  as published by the Free Software Foundation; either version 2
6 #  of the License, or (at your option) any later version.
7 #
8 #  This program is distributed in the hope that it will be useful,
9 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
10 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11 #  GNU General Public License for more details.
12 #
13 #  You should have received a copy of the GNU General Public License
14 #  along with this program; if not, write to the Free Software Foundation,
15 #  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
16 #
17 # ##### END GPL LICENSE BLOCK #####
18
19 # <pep8 compliant>
20
21 import bpy
22 from mathutils import *
23 from math import radians, acos
24 from bl_operators import nla
25 import cProfile
26
27
28 def hasIKConstraint(pose_bone):
29     #utility function / predicate, returns True if given bone has IK constraint
30     ik = [constraint for constraint in pose_bone.constraints if constraint.type == "IK"]
31     if ik:
32         return ik[0]
33     else:
34         return False
35
36
37 def createDictionary(perf_arm, end_arm):
38     # clear any old data
39     for end_bone in end_arm.bones:
40         for mapping in end_bone.reverseMap:
41             end_bone.reverseMap.remove(0)
42
43     for perf_bone in perf_arm.bones:
44         #find its match and add perf_bone to the match's mapping
45         if perf_bone.map:
46             end_bone = end_arm.bones[perf_bone.map]
47             newMap = end_bone.reverseMap.add()
48             newMap.name = perf_bone.name
49             end_bone.foot = perf_bone.foot
50
51     #root is the root of the enduser
52     root = end_arm.bones[0].name
53     feetBones = [bone.name for bone in perf_arm.bones if bone.foot]
54     return feetBones, root
55
56 def loadMapping(perf_arm, end_arm):
57
58     for end_bone in end_arm.bones:
59         #find its match and add perf_bone to the match's mapping
60         if end_bone.reverseMap:
61             for perf_bone in end_bone.reverseMap:
62                 perf_arm.bones[perf_bone.name].map = end_bone.name
63
64 #creation of intermediate armature
65 # the intermediate armature has the hiearchy of the end user,
66 # does not have rotation inheritence
67 # and bone roll is identical to the performer
68 # its purpose is to copy over the rotations
69 # easily while concentrating on the hierarchy changes
70
71
72 def createIntermediate(performer_obj, enduser_obj, root, s_frame, e_frame, scene):
73     #creates and keyframes an empty with its location
74     #the original position of the tail bone
75     #useful for storing the important data in the original motion
76     #i.e. using this empty to IK the chain to that pos / DEBUG
77
78     #Simple 1to1 retarget of a bone
79     def singleBoneRetarget(inter_bone, perf_bone):
80             perf_world_rotation = perf_bone.matrix * performer_obj.matrix_world
81             inter_world_base_rotation = inter_bone.bone.matrix_local * inter_obj.matrix_world
82             inter_world_base_inv = Matrix(inter_world_base_rotation)
83             inter_world_base_inv.invert()
84             return (inter_world_base_inv.to_3x3() * perf_world_rotation.to_3x3()).to_4x4()
85
86     #uses 1to1 and interpolation/averaging to match many to 1 retarget
87     def manyPerfToSingleInterRetarget(inter_bone, performer_bones_s):
88         retarget_matrices = [singleBoneRetarget(inter_bone, perf_bone) for perf_bone in performer_bones_s]
89         lerp_matrix = Matrix()
90         for i in range(len(retarget_matrices) - 1):
91             first_mat = retarget_matrices[i]
92             next_mat = retarget_matrices[i + 1]
93             lerp_matrix = first_mat.lerp(next_mat, 0.5)
94         return lerp_matrix
95
96     #determines the type of hierachy change needed and calls the
97     #right function
98     def retargetPerfToInter(inter_bone):
99         if inter_bone.bone.reverseMap:
100             perf_bone_name = inter_bone.bone.reverseMap
101                 # 1 to many not supported yet
102                 # then its either a many to 1 or 1 to 1
103             if len(perf_bone_name) > 1:
104                 performer_bones_s = [performer_bones[map.name] for map in perf_bone_name]
105                 #we need to map several performance bone to a single
106                 inter_bone.matrix_basis = manyPerfToSingleInterRetarget(inter_bone, performer_bones_s)
107             else:
108                 perf_bone = performer_bones[perf_bone_name[0].name]
109                 inter_bone.matrix_basis = singleBoneRetarget(inter_bone, perf_bone)
110         inter_bone.keyframe_insert("rotation_quaternion")
111
112     #creates the intermediate armature object
113     inter_obj = enduser_obj.copy()
114     inter_obj.data = inter_obj.data.copy()  # duplicate data
115     bpy.context.scene.objects.link(inter_obj)
116     inter_obj.name = "intermediate"
117     bpy.context.scene.objects.active = inter_obj
118     bpy.ops.object.mode_set(mode='EDIT')
119     #add some temporary connecting bones in case end user bones are not connected to their parents
120     print("creating temp bones")
121     for bone in inter_obj.data.edit_bones:
122         if not bone.use_connect and bone.parent:
123             if inter_obj.data.bones[bone.parent.name].reverseMap or inter_obj.data.bones[bone.name].reverseMap:
124                 newBone = inter_obj.data.edit_bones.new("Temp")
125                 newBone.head = bone.parent.tail
126                 newBone.tail = bone.head
127                 newBone.parent = bone.parent
128                 bone.parent = newBone
129                 bone.use_connect = True
130                 newBone.use_connect = True
131     #resets roll
132     print("retargeting to intermediate")
133     bpy.ops.armature.calculate_roll(type='Z')
134     bpy.ops.object.mode_set(mode="OBJECT")
135     inter_obj.data.name = "inter_arm"
136     inter_arm = inter_obj.data
137     performer_bones = performer_obj.pose.bones
138     inter_bones = inter_obj.pose.bones
139     #clears inheritance
140     for inter_bone in inter_bones:
141         if inter_bone.bone.reverseMap:
142             inter_bone.bone.use_inherit_rotation = False
143         else:
144             inter_bone.bone.use_inherit_rotation = True
145
146     for t in range(s_frame, e_frame):
147         if (t - s_frame) % 10 == 0:
148             print("First pass: retargeting frame {0}/{1}".format(t, e_frame - s_frame))
149         scene.frame_set(t)
150         for bone in inter_bones:
151             retargetPerfToInter(bone)
152
153     return inter_obj
154
155 # this procedure copies the rotations over from the intermediate
156 # armature to the end user one.
157 # As the hierarchies are 1 to 1, this is a simple matter of
158 # copying the rotation, while keeping in mind bone roll, parenting, etc.
159 # TODO: Control Bones: If a certain bone is constrained in a way
160 #       that its rotation is determined by another (a control bone)
161 #       We should determine the right pos of the control bone.
162 #       Scale: ? Should work but needs testing.
163
164
165 def retargetEnduser(inter_obj, enduser_obj, root, s_frame, e_frame, scene):
166     inter_bones = inter_obj.pose.bones
167     end_bones = enduser_obj.pose.bones
168
169     def bakeTransform(end_bone):
170         src_bone = inter_bones[end_bone.name]
171         trg_bone = end_bone
172         bake_matrix = src_bone.matrix
173         rest_matrix = trg_bone.bone.matrix_local
174
175         if trg_bone.parent and trg_bone.bone.use_inherit_rotation:
176             srcParent = src_bone.parent
177             if "Temp" in srcParent.name:
178                 srcParent = srcParent.parent
179             parent_mat = srcParent.matrix
180             parent_rest = trg_bone.parent.bone.matrix_local
181             parent_rest_inv = parent_rest.copy()
182             parent_rest_inv.invert()
183             parent_mat_inv = parent_mat.copy()
184             parent_mat_inv.invert()
185             bake_matrix = parent_mat_inv * bake_matrix
186             rest_matrix = parent_rest_inv * rest_matrix
187
188         rest_matrix_inv = rest_matrix.copy()
189         rest_matrix_inv.invert()
190         bake_matrix = rest_matrix_inv * bake_matrix
191         end_bone.matrix_basis = bake_matrix
192         rot_mode = end_bone.rotation_mode
193         if rot_mode == "QUATERNION":
194             end_bone.keyframe_insert("rotation_quaternion")
195         elif rot_mode == "AXIS_ANGLE":
196             end_bone.keyframe_insert("rotation_axis_angle")
197         else:
198             end_bone.keyframe_insert("rotation_euler")
199         if not end_bone.bone.use_connect:
200             end_bone.keyframe_insert("location")
201
202         for bone in end_bone.children:
203             bakeTransform(bone)
204
205     for t in range(s_frame, e_frame):
206         if (t - s_frame) % 10 == 0:
207             print("Second pass: retargeting frame {0}/{1}".format(t, e_frame - s_frame))
208         scene.frame_set(t)
209         end_bone = end_bones[root]
210         end_bone.location = Vector((0, 0, 0))
211         end_bone.keyframe_insert("location")
212         bakeTransform(end_bone)
213
214 #recieves the performer feet bones as a variable
215 # by "feet" I mean those bones that have plants
216 # (they don't move, despite root moving) somewhere in the animation.
217
218
219 def copyTranslation(performer_obj, enduser_obj, perfFeet, root, s_frame, e_frame, scene, enduser_obj_mat):
220
221     perf_bones = performer_obj.pose.bones
222     end_bones = enduser_obj.pose.bones
223
224     perfRoot = perf_bones[0].name
225     endFeet = [perf_bones[perfBone].bone.map for perfBone in perfFeet]
226     locDictKeys = perfFeet + endFeet + [perfRoot]
227
228     def tailLoc(bone):
229         return bone.center + (bone.vector / 2)
230
231     #Step 1 - we create a dict that contains these keys:
232     #(Performer) Hips, Feet
233     #(End user) Feet
234     # where the values are their world position on each frame in range (s,e)
235
236     locDict = {}
237     for key in locDictKeys:
238         locDict[key] = []
239
240     for t in range(scene.frame_start, scene.frame_end):
241         scene.frame_set(t)
242         for bone in perfFeet:
243             locDict[bone].append(tailLoc(perf_bones[bone]))
244         locDict[perfRoot].append(tailLoc(perf_bones[perfRoot]))
245         for bone in endFeet:
246             locDict[bone].append(tailLoc(end_bones[bone]))
247
248     # now we take our locDict and analyze it.
249     # we need to derive all chains
250
251     locDeriv = {}
252     for key in locDictKeys:
253         locDeriv[key] = []
254
255     for key in locDict.keys():
256         graph = locDict[key]
257         locDeriv[key] = [graph[t + 1] - graph[t] for t in range(len(graph) - 1)]
258
259     # now find the plant frames, where perfFeet don't move much
260
261     linearAvg = []
262
263     for key in perfFeet:
264         for i in range(len(locDeriv[key]) - 1):
265             v = locDeriv[key][i]
266             hipV = locDeriv[perfRoot][i]
267             endV = locDeriv[perf_bones[key].bone.map][i]
268             print(v.length,)
269             if (v.length < 0.1):
270                 #this is a plant frame.
271                 #lets see what the original hip delta is, and the corresponding
272                 #end bone's delta
273                 if endV.length != 0:
274                     linearAvg.append(hipV.length / endV.length)
275
276     action_name = performer_obj.animation_data.action.name
277     #is there a stride_bone?
278     if "stride_bone" in bpy.data.objects:
279         stride_action = bpy.data.actions.new("Stride Bone " + action_name)
280         stride_action.use_fake_user = True
281         stride_bone = enduser_obj.parent
282         stride_bone.animation_data.action = stride_action
283     else:
284         bpy.ops.object.add()
285         stride_bone = bpy.context.active_object
286         stride_bone.name = "stride_bone"
287     print(stride_bone)
288     stride_bone.location = Vector((0, 0, 0))
289     print(linearAvg)
290     if linearAvg:
291         #determine the average change in scale needed
292         avg = sum(linearAvg) / len(linearAvg)
293         scene.frame_set(s_frame)
294         initialPos = (tailLoc(perf_bones[perfRoot]) / avg) + stride_bone.location
295         for t in range(s_frame, e_frame):
296             scene.frame_set(t)
297             #calculate the new position, by dividing by the found ratio between performer and enduser
298             newTranslation = (tailLoc(perf_bones[perfRoot]) / avg)
299             stride_bone.location = enduser_obj_mat * (newTranslation - initialPos)
300             stride_bone.keyframe_insert("location")
301     else:
302         stride_bone.keyframe_insert("location")
303     stride_bone.animation_data.action.name = ("Stride Bone " + action_name)
304
305     return stride_bone
306
307
308 def IKRetarget(performer_obj, enduser_obj, s_frame, e_frame, scene):
309     end_bones = enduser_obj.pose.bones
310     for pose_bone in end_bones:
311         ik_constraint = hasIKConstraint(pose_bone)
312         if ik_constraint:
313             target_is_bone = False
314             # set constraint target to corresponding empty if targetless,
315             # if not, keyframe current target to corresponding empty
316             perf_bone = pose_bone.bone.reverseMap[-1].name
317             orgLocTrg = originalLocationTarget(pose_bone, enduser_obj)
318             if not ik_constraint.target:
319                 ik_constraint.target = orgLocTrg
320                 target = orgLocTrg
321
322             # There is a target now
323             if ik_constraint.subtarget:
324                 target = ik_constraint.target.pose.bones[ik_constraint.subtarget]
325                 target.bone.use_local_location = False
326                 target_is_bone = True
327             else:
328                 target = ik_constraint.target
329
330             # bake the correct locations for the ik target bones
331             for t in range(s_frame, e_frame):
332                 scene.frame_set(t)
333                 if target_is_bone:
334                     final_loc = pose_bone.tail - target.bone.matrix_local.to_translation()
335                 else:
336                     final_loc = pose_bone.tail
337                 target.location = final_loc
338                 target.keyframe_insert("location")
339             ik_constraint.mute = False
340     scene.frame_set(s_frame)
341
342
343 def turnOffIK(enduser_obj):
344     end_bones = enduser_obj.pose.bones
345     for pose_bone in end_bones:
346         if pose_bone.is_in_ik_chain:
347             pass
348             # TODO:
349             # set stiffness according to place on chain
350             # and values from analysis that is stored in the bone
351             #pose_bone.ik_stiffness_x = 0.5
352             #pose_bone.ik_stiffness_y = 0.5
353             #pose_bone.ik_stiffness_z = 0.5
354         ik_constraint = hasIKConstraint(pose_bone)
355         if ik_constraint:
356             ik_constraint.mute = True
357
358
359 #copy the object matrixes and clear them (to be reinserted later)
360 def cleanAndStoreObjMat(performer_obj, enduser_obj):
361     perf_obj_mat = performer_obj.matrix_world.copy()
362     enduser_obj_mat = enduser_obj.matrix_world.copy()
363     zero_mat = Matrix()
364     performer_obj.matrix_world = zero_mat
365     enduser_obj.matrix_world = zero_mat
366     return perf_obj_mat, enduser_obj_mat
367
368
369 #restore the object matrixes after parenting the auto generated IK empties
370 def restoreObjMat(performer_obj, enduser_obj, perf_obj_mat, enduser_obj_mat, stride_bone):
371     pose_bones = enduser_obj.pose.bones
372     for pose_bone in pose_bones:
373         if pose_bone.name + "Org" in bpy.data.objects:
374             empty = bpy.data.objects[pose_bone.name + "Org"]
375             empty.parent = stride_bone
376     performer_obj.matrix_world = perf_obj_mat
377     enduser_obj.parent = stride_bone
378     enduser_obj.matrix_world = enduser_obj_mat
379
380
381 #create (or return if exists) the related IK empty to the bone
382 def originalLocationTarget(end_bone, enduser_obj):
383     if not end_bone.name + "Org" in bpy.data.objects:
384         bpy.ops.object.add()
385         empty = bpy.context.active_object
386         empty.name = end_bone.name + "Org"
387         empty.empty_draw_size = 0.1
388         empty.parent = enduser_obj
389     empty = bpy.data.objects[end_bone.name + "Org"]
390     return empty
391
392
393 #create the specified NLA setup for base animation, constraints and tweak layer.
394 def NLASystemInitialize(enduser_arm, context):#enduser_obj, name):
395     enduser_obj = context.active_object
396     NLATracks = enduser_arm.mocapNLATracks[enduser_obj.data.active_mocap]
397     name = NLATracks.name
398     anim_data = enduser_obj.animation_data
399     s_frame = 0
400     print(name)
401     if ("Base " + name) in bpy.data.actions:
402         mocapAction = bpy.data.actions[("Base " + name)]
403     else:
404         print("That retargeted anim has no base action")
405     anim_data.use_nla = True
406     for track in anim_data.nla_tracks:
407         anim_data.nla_tracks.remove(track)
408     mocapTrack = anim_data.nla_tracks.new()
409     mocapTrack.name = "Base " + name
410     NLATracks.base_track = mocapTrack.name
411     mocapStrip = mocapTrack.strips.new("Base " + name, s_frame, mocapAction)
412     constraintTrack = anim_data.nla_tracks.new()
413     constraintTrack.name = "Auto fixes " + name
414     NLATracks.auto_fix_track = constraintTrack.name
415     if ("Auto fixes " + name) in bpy.data.actions:
416         constraintAction = bpy.data.actions[("Auto fixes " + name)]
417     else:
418         constraintAction = bpy.data.actions.new("Auto fixes " + name)
419         constraintAction.use_fake_user = True
420     constraintStrip = constraintTrack.strips.new("Auto fixes " + name, s_frame, constraintAction)
421     constraintStrip.extrapolation = "NOTHING"
422     userTrack = anim_data.nla_tracks.new()
423     userTrack.name = "Manual fixes " + name
424     NLATracks.manual_fix_track = userTrack.name
425     if ("Manual fixes " + name) in bpy.data.actions:
426         userAction = bpy.data.actions[("Manual fixes " + name)]
427     else:
428         userAction = bpy.data.actions.new("Manual fixes " + name)
429         userAction.use_fake_user = True
430     userStrip = userTrack.strips.new("Manual fixes " + name, s_frame, userAction)
431     userStrip.extrapolation = "HOLD"
432     #userStrip.blend_type = "MULITPLY" - doesn't work due to work, will be activated soon
433     anim_data.nla_tracks.active = constraintTrack
434     #anim_data.action = constraintAction
435     anim_data.action_extrapolation = "NOTHING"
436     #set the stride_bone's action
437     if "stride_bone" in bpy.data.objects:
438         stride_bone = bpy.data.objects["stride_bone"]
439         if NLATracks.stride_action:
440             stride_bone.animation_data.action = bpy.data.actions[NLATracks.stride_action]
441         else:
442             NLATracks.stride_action = stride_bone.animation_data.action.name
443             stride_bone.animation_data.action.use_fake_user = True
444     anim_data.action = None
445
446
447 def preAdvancedRetargeting(performer_obj, enduser_obj):
448     createDictionary(performer_obj.data, enduser_obj.data)
449     bones = enduser_obj.pose.bones
450     map_bones = [bone for bone in bones if bone.bone.reverseMap]
451     for bone in map_bones:
452         perf_bone = bone.bone.reverseMap[0].name
453         addLocalRot = False;
454         if bone.bone.use_connect or not bone.constraints:
455             locks = bone.lock_location
456             if not (locks[0] or locks[1] or locks[2]):  
457                 cons = bone.constraints.new('COPY_LOCATION')
458                 cons.name = "retargetTemp"
459                 cons.use_x = not locks[0]
460                 cons.use_y = not locks[1]
461                 cons.use_z = not locks[2]
462                 cons.target = performer_obj
463                 cons.subtarget = perf_bone
464                 addLocalRot = True
465
466        
467         cons2 = bone.constraints.new('COPY_ROTATION')
468         cons2.name = "retargetTemp"
469         locks = bone.lock_rotation
470         cons2.use_x = not locks[0]
471         cons2.use_y = not locks[1]
472         cons2.use_z = not locks[2]
473         cons2.target = performer_obj
474         cons2.subtarget = perf_bone
475
476         if addLocalRot:
477             for constraint in bone.constraints:
478                 if constraint.type == 'COPY_ROTATION':
479                     constraint.target_space = 'LOCAL'
480                     constraint.owner_space = 'LOCAL_WITH_PARENT'
481
482
483 def prepareForBake(enduser_obj):
484     bones = enduser_obj.pose.bones
485     for bone in bones:
486         bone.bone.select = False
487     map_bones = [bone for bone in bones if bone.bone.reverseMap]
488     for bone in map_bones:
489         for cons in bone.constraints:
490             if "retargetTemp" in cons.name:
491                 bone.bone.select = True
492
493 def cleanTempConstraints(enduser_obj):
494     bones = enduser_obj.pose.bones
495     map_bones = [bone for bone in bones if bone.bone.reverseMap]
496     for bone in map_bones:
497         for cons in bone.constraints:
498             if "retargetTemp" in cons.name:
499                 bone.constraints.remove(cons)
500
501 #Main function that runs the retargeting sequence.
502 #If advanced == True, we assume constraint's were already created
503 def totalRetarget(performer_obj, enduser_obj, scene, s_frame, e_frame):
504     perf_arm = performer_obj.data
505     end_arm = enduser_obj.data
506     advanced = end_arm.advancedRetarget
507     
508     try:
509         enduser_obj.animation_data.action = bpy.data.actions.new("temp")
510         enduser_obj.animation_data.action.use_fake_user = True
511     except:
512         print("no need to create new action")
513     
514     print("creating Dictionary")
515     feetBones, root = createDictionary(perf_arm, end_arm)
516     print("cleaning stuff up")
517     perf_obj_mat, enduser_obj_mat = cleanAndStoreObjMat(performer_obj, enduser_obj)
518     if not advanced:
519         turnOffIK(enduser_obj)
520         print("Creating intermediate armature (for first pass)")
521         inter_obj = createIntermediate(performer_obj, enduser_obj, root, s_frame, e_frame, scene)
522         print("First pass: retargeting from intermediate to end user")
523         retargetEnduser(inter_obj, enduser_obj, root, s_frame, e_frame, scene)
524     else:
525         prepareForBake(enduser_obj)
526         print("Retargeting pose (Advanced Retarget)")
527         nla.bake(s_frame, e_frame, action=enduser_obj.animation_data.action, only_selected=True, do_pose=True, do_object=False)
528     name = performer_obj.animation_data.action.name
529     enduser_obj.animation_data.action.name = "Base " + name
530     print("Second pass: retargeting root translation and clean up")
531     stride_bone = copyTranslation(performer_obj, enduser_obj, feetBones, root, s_frame, e_frame, scene, enduser_obj_mat)
532     if not advanced:
533         IKRetarget(performer_obj, enduser_obj, s_frame, e_frame, scene)
534     restoreObjMat(performer_obj, enduser_obj, perf_obj_mat, enduser_obj_mat, stride_bone)
535     bpy.ops.object.mode_set(mode='OBJECT')
536     if not advanced:
537         bpy.ops.object.select_name(name=inter_obj.name, extend=False)
538         bpy.ops.object.delete()
539     else:
540         cleanTempConstraints(enduser_obj)
541     bpy.ops.object.select_name(name=enduser_obj.name, extend=False)
542
543     if not name in [tracks.name for tracks in end_arm.mocapNLATracks]:
544         NLATracks = end_arm.mocapNLATracks.add()
545         NLATracks.name = name
546     else:
547         NLATracks = end_arm.mocapNLATracks[name]
548     end_arm.active_mocap = name
549     print("retargeting done!")
550
551 if __name__ == "__main__":
552     totalRetarget()