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