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