Merging trunk up to r38193.
[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 #TODO: Only selected bones get retargeted.
26 #      Selected Bones/chains get original pos empties,
27 #      if ppl want IK instead of FK
28 #      Some "magic" numbers - frame start and end,
29 #       eulers of all orders instead of just quats keyframed
30
31 # dictionary of mapping
32 # this is currently manuall input'ed, but willW
33 # be created from a more comfortable UI in the future
34
35
36 def createDictionary(perf_arm, end_arm):
37     bonemap = {}
38     #Bonemap: performer to enduser
39     for bone in perf_arm.bones:
40         bonemap[bone.name] = bone.map
41
42     # creation of a reverse map
43     # multiple keys get mapped to list values
44     #Bonemapr: enduser to performer
45     bonemapr = {}
46     for key, value in bonemap.items():
47         if not value in bonemapr:
48             if isinstance(bonemap[key], tuple):
49                 for key_x in bonemap[key]:
50                     bonemapr[key_x] = [key]
51             else:
52                 bonemapr[bonemap[key]] = [key]
53         else:
54             bonemapr[bonemap[key]].append(key)
55     #root is the root of the enduser
56     root = end_arm.bones[0].name
57     feetBones = [bone.name for bone in perf_arm.bones if bone.foot]
58     return bonemap, bonemapr, feetBones, root
59 # list of empties created to keep track of "original"
60 # position data
61 # in final product, these locations can be stored as custom props
62 # these help with constraining, etc.
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, bonemap, bonemapr, 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.name in bonemapr:
100             perf_bone_name = bonemapr[inter_bone.name]
101             #is it a 1 to many?
102             if isinstance(bonemap[perf_bone_name[0]], tuple):
103                 pass
104                 # 1 to many not supported yet
105             else:
106                 # then its either a many to 1 or 1 to 1
107
108                 if len(perf_bone_name) > 1:
109                     performer_bones_s = [performer_bones[name] for name in perf_bone_name]
110                     #we need to map several performance bone to a single
111                     inter_bone.matrix_basis = manyPerfToSingleInterRetarget(inter_bone, performer_bones_s)
112                 else:
113                     perf_bone = performer_bones[perf_bone_name[0]]
114                     inter_bone.matrix_basis = singleBoneRetarget(inter_bone, perf_bone)
115
116         inter_bone.keyframe_insert("rotation_quaternion")
117         for child in inter_bone.children:
118             retargetPerfToInter(child)
119
120     #creates the intermediate armature object
121     inter_obj = enduser_obj.copy()
122     inter_obj.data = inter_obj.data.copy()  # duplicate data
123     bpy.context.scene.objects.link(inter_obj)
124     inter_obj.name = "intermediate"
125     bpy.context.scene.objects.active = inter_obj
126     bpy.ops.object.mode_set(mode='EDIT')
127     #resets roll
128     bpy.ops.armature.calculate_roll(type='Z')
129     bpy.ops.object.mode_set(mode="OBJECT")
130     inter_obj.data.name = "inter_arm"
131     inter_arm = inter_obj.data
132     performer_bones = performer_obj.pose.bones
133     inter_bones = inter_obj.pose.bones
134     #clears inheritance
135     for inter_bone in inter_bones:
136         inter_bone.bone.use_inherit_rotation = False
137
138     for t in range(s_frame, e_frame):
139         scene.frame_set(t)
140         inter_bone = inter_bones[root]
141         retargetPerfToInter(inter_bone)
142
143     return inter_obj, inter_arm
144
145 # this procedure copies the rotations over from the intermediate
146 # armature to the end user one.
147 # As the hierarchies are 1 to 1, this is a simple matter of
148 # copying the rotation, while keeping in mind bone roll, parenting, etc.
149 # TODO: Control Bones: If a certain bone is constrained in a way
150 #       that its rotation is determined by another (a control bone)
151 #       We should determine the right pos of the control bone.
152 #       Scale: ? Should work but needs testing.
153
154
155 def retargetEnduser(inter_obj, enduser_obj, root, s_frame, e_frame, scene):
156     inter_bones = inter_obj.pose.bones
157     end_bones = enduser_obj.pose.bones
158
159     def bakeTransform(end_bone):
160         src_bone = inter_bones[end_bone.name]
161         trg_bone = end_bone
162         bake_matrix = src_bone.matrix
163         rest_matrix = trg_bone.bone.matrix_local
164
165         if trg_bone.parent and trg_bone.bone.use_inherit_rotation:
166             parent_mat = src_bone.parent.matrix
167             parent_rest = trg_bone.parent.bone.matrix_local
168             parent_rest_inv = parent_rest.copy()
169             parent_rest_inv.invert()
170             parent_mat_inv = parent_mat.copy()
171             parent_mat_inv.invert()
172             bake_matrix = parent_mat_inv * bake_matrix
173             rest_matrix = parent_rest_inv * rest_matrix
174
175         rest_matrix_inv = rest_matrix.copy()
176         rest_matrix_inv.invert()
177         bake_matrix = rest_matrix_inv * bake_matrix
178         trg_bone.matrix_basis = bake_matrix
179         end_bone.keyframe_insert("rotation_quaternion")
180
181         for bone in end_bone.children:
182             bakeTransform(bone)
183
184     for t in range(s_frame, e_frame):
185         scene.frame_set(t)
186         end_bone = end_bones[root]
187         end_bone.location = Vector((0, 0, 0))
188         end_bone.keyframe_insert("location")
189         bakeTransform(end_bone)
190
191 #recieves the performer feet bones as a variable
192 # by "feet" I mean those bones that have plants
193 # (they don't move, despite root moving) somewhere in the animation.
194
195
196 def copyTranslation(performer_obj, enduser_obj, perfFeet, bonemap, bonemapr, root, s_frame, e_frame, scene, enduser_obj_mat):
197
198     perf_bones = performer_obj.pose.bones
199     end_bones = enduser_obj.pose.bones
200
201     perfRoot = bonemapr[root][0]
202     endFeet = [bonemap[perfBone] for perfBone in perfFeet]
203     locDictKeys = perfFeet + endFeet + [perfRoot]
204
205     def tailLoc(bone):
206         return bone.center + (bone.vector / 2)
207
208     #Step 1 - we create a dict that contains these keys:
209     #(Performer) Hips, Feet
210     #(End user) Feet
211     # where the values are their world position on each (1,120) frame
212
213     locDict = {}
214     for key in locDictKeys:
215         locDict[key] = []
216
217     for t in range(scene.frame_start, scene.frame_end):
218         scene.frame_set(t)
219         for bone in perfFeet:
220             locDict[bone].append(tailLoc(perf_bones[bone]))
221         locDict[perfRoot].append(tailLoc(perf_bones[perfRoot]))
222         for bone in endFeet:
223             locDict[bone].append(tailLoc(end_bones[bone]))
224
225     # now we take our locDict and analyze it.
226     # we need to derive all chains
227
228     locDeriv = {}
229     for key in locDictKeys:
230         locDeriv[key] = []
231
232     for key in locDict.keys():
233         graph = locDict[key]
234         for t in range(len(graph) - 1):
235             x = graph[t]
236             xh = graph[t + 1]
237             locDeriv[key].append(xh - x)
238
239     # now find the plant frames, where perfFeet don't move much
240
241     linearAvg = []
242
243     for key in perfFeet:
244         for i in range(len(locDeriv[key]) - 1):
245             v = locDeriv[key][i]
246             hipV = locDeriv[perfRoot][i]
247             endV = locDeriv[bonemap[key]][i]
248             if (v.length < 0.1):
249                 #this is a plant frame.
250                 #lets see what the original hip delta is, and the corresponding
251                 #end bone's delta
252                 if endV.length != 0:
253                     linearAvg.append(hipV.length / endV.length)
254
255     bpy.ops.object.add()
256     stride_bone = bpy.context.active_object
257     stride_bone.name = "stride_bone"
258
259     if linearAvg:
260         avg = sum(linearAvg) / len(linearAvg)
261         scene.frame_set(s_frame)
262         initialPos = (tailLoc(perf_bones[perfRoot]) / avg)
263         for t in range(s_frame, e_frame):
264             scene.frame_set(t)
265             newTranslation = (tailLoc(perf_bones[perfRoot]) / avg)
266             stride_bone.location = (newTranslation - initialPos) * enduser_obj_mat
267             stride_bone.keyframe_insert("location")
268     return stride_bone
269
270
271 def IKRetarget(bonemap, bonemapr, performer_obj, enduser_obj, s_frame, e_frame, scene):
272     end_bones = enduser_obj.pose.bones
273     for pose_bone in end_bones:
274         if "IK" in [constraint.type for constraint in pose_bone.constraints]:
275             target_is_bone = False
276             # set constraint target to corresponding empty if targetless,
277             # if not, keyframe current target to corresponding empty
278             perf_bone = bonemapr[pose_bone.name]
279             if isinstance(perf_bone, list):
280                 perf_bone = bonemapr[pose_bone.name][-1]
281             orgLocTrg = originalLocationTarget(pose_bone)
282             ik_constraint = [constraint for constraint in pose_bone.constraints if constraint.type == "IK"][0]
283             if not ik_constraint.target:
284                 ik_constraint.target = orgLocTrg
285                 target = orgLocTrg
286
287             # There is a target now
288             if ik_constraint.subtarget:
289                 target = ik_constraint.target.pose.bones[ik_constraint.subtarget]
290                 target.bone.use_local_location = False
291                 target_is_bone = True
292             else:
293                 target = ik_constraint.target
294
295             for t in range(s_frame, e_frame):
296                 scene.frame_set(t)
297                 if target_is_bone:
298                     final_loc = pose_bone.tail - target.bone.matrix_local.to_translation()
299                 else:
300                     final_loc = pose_bone.tail
301                 target.location = final_loc
302                 target.keyframe_insert("location")
303             ik_constraint.mute = False
304
305
306 def turnOffIK(enduser_obj):
307     end_bones = enduser_obj.pose.bones
308     for pose_bone in end_bones:
309         if pose_bone.is_in_ik_chain:
310             pass
311             # TODO:
312             # set stiffness according to place on chain
313             # and values from analysis that is stored in the bone
314             #pose_bone.ik_stiffness_x = 0.5
315             #pose_bone.ik_stiffness_y = 0.5
316             #pose_bone.ik_stiffness_z = 0.5
317         if "IK" in [constraint.type for constraint in pose_bone.constraints]:
318             ik_constraint = [constraint for constraint in pose_bone.constraints if constraint.type == "IK"][0]
319             ik_constraint.mute = True
320
321
322 def cleanAndStoreObjMat(performer_obj, enduser_obj):
323     perf_obj_mat = performer_obj.matrix_world.copy()
324     enduser_obj_mat = enduser_obj.matrix_world.copy()
325     zero_mat = Matrix()  # Matrix(((0,0,0,0),(0,0,0,0),(0,0,0,0),(0,0,0,0)))
326     performer_obj.matrix_world = zero_mat
327     enduser_obj.matrix_world = zero_mat
328     return perf_obj_mat, enduser_obj_mat
329
330
331 def restoreObjMat(performer_obj, enduser_obj, perf_obj_mat, enduser_obj_mat, stride_bone):
332     pose_bones = enduser_obj.pose.bones
333     for pose_bone in pose_bones:
334         if pose_bone.name + "Org" in bpy.data.objects:
335             empty = bpy.data.objects[pose_bone.name + "Org"]
336             empty.parent = stride_bone
337     performer_obj.matrix_world = perf_obj_mat
338     enduser_obj.matrix_world = enduser_obj_mat
339     enduser_obj.parent = stride_bone
340
341
342 def originalLocationTarget(end_bone):
343     if not end_bone.name + "Org" in bpy.data.objects:
344         bpy.ops.object.add()
345         empty = bpy.context.active_object
346         empty.name = end_bone.name + "Org"
347         empty.empty_draw_size = 0.1
348         #empty.parent = enduser_obj
349     empty = bpy.data.objects[end_bone.name + "Org"]
350     return empty
351
352
353 def totalRetarget():
354     print("retargeting...")
355     enduser_obj = bpy.context.active_object
356     performer_obj = [obj for obj in bpy.context.selected_objects if obj != enduser_obj]
357     if enduser_obj is None or len(performer_obj) != 1:
358         print("Need active and selected armatures")
359     else:
360         performer_obj = performer_obj[0]
361     perf_arm = performer_obj.data
362     end_arm = enduser_obj.data
363     scene = bpy.context.scene
364     s_frame = scene.frame_start
365     e_frame = scene.frame_end
366     bonemap, bonemapr, feetBones, root = createDictionary(perf_arm, end_arm)
367     perf_obj_mat, enduser_obj_mat = cleanAndStoreObjMat(performer_obj, enduser_obj)
368     turnOffIK(enduser_obj)
369     inter_obj, inter_arm = createIntermediate(performer_obj, enduser_obj, bonemap, bonemapr, root, s_frame, e_frame, scene)
370     retargetEnduser(inter_obj, enduser_obj, root, s_frame, e_frame, scene)
371     stride_bone = copyTranslation(performer_obj, enduser_obj, feetBones, bonemap, bonemapr, root, s_frame, e_frame, scene, enduser_obj_mat)
372     IKRetarget(bonemap, bonemapr, performer_obj, enduser_obj, s_frame, e_frame, scene)
373     restoreObjMat(performer_obj, enduser_obj, perf_obj_mat, enduser_obj_mat, stride_bone)
374     bpy.ops.object.mode_set(mode='OBJECT')
375     bpy.ops.object.select_name(name=inter_obj.name, extend=False)
376     bpy.ops.object.delete()
377
378 if __name__ == "__main__":
379     totalRetarget()