Added IK functionality to retargeting. If rig contains an IK constraint, in addition...
[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):
37     bonemap = {}
38     for bone in perf_arm.bones:
39         bonemap[bone.name] = bone.map
40     #root is the root of the enduser
41     root = "root"
42     # creation of a reverse map
43     # multiple keys get mapped to list values
44     bonemapr = {}
45     for key, value in bonemap.items():
46         if not value in bonemapr:
47             if isinstance(bonemap[key], tuple):
48                 for key_x in bonemap[key]:
49                     bonemapr[key_x] = [key]
50             else:
51                 bonemapr[bonemap[key]] = [key]
52         else:
53             bonemapr[bonemap[key]].append(key)
54     return bonemap, bonemapr, root
55 # list of empties created to keep track of "original"
56 # position data
57 # in final product, these locations can be stored as custom props
58 # these help with constraining, etc.
59
60 #creation of intermediate armature
61 # the intermediate armature has the hiearchy of the end user,
62 # does not have rotation inheritence
63 # and bone roll is identical to the performer
64 # its purpose is to copy over the rotations
65 # easily while concentrating on the hierarchy changes
66
67
68 def createIntermediate(performer_obj, enduser_obj, bonemap, bonemapr, root, s_frame, e_frame, scene):
69     #creates and keyframes an empty with its location
70     #the original position of the tail bone
71     #useful for storing the important data in the original motion
72     #i.e. using this empty to IK the chain to that pos / DEBUG
73     def locOfOriginal(inter_bone, perf_bone):
74         if not perf_bone.name + "Org" in bpy.data.objects:
75             bpy.ops.object.add()
76             empty = bpy.context.active_object
77             empty.name = perf_bone.name + "Org"
78             empty.empty_draw_size = 0.01
79         empty = bpy.data.objects[perf_bone.name + "Org"]
80         offset = perf_bone.vector
81         if inter_bone.length == 0 or perf_bone.length == 0:
82             scaling = 1
83         else:
84             scaling = perf_bone.length / inter_bone.length
85         offset /= scaling
86         empty.location = inter_bone.head + offset
87         empty.keyframe_insert("location")
88
89     #Simple 1to1 retarget of a bone
90     def singleBoneRetarget(inter_bone, perf_bone):
91             perf_world_rotation = perf_bone.matrix * performer_obj.matrix_world
92             inter_world_base_rotation = inter_bone.bone.matrix_local * inter_obj.matrix_world
93             inter_world_base_inv = Matrix(inter_world_base_rotation)
94             inter_world_base_inv.invert()
95             return (inter_world_base_inv.to_3x3() * perf_world_rotation.to_3x3()).to_4x4()
96
97     #uses 1to1 and interpolation/averaging to match many to 1 retarget
98     def manyPerfToSingleInterRetarget(inter_bone, performer_bones_s):
99         retarget_matrices = [singleBoneRetarget(inter_bone, perf_bone) for perf_bone in performer_bones_s]
100         lerp_matrix = Matrix()
101         for i in range(len(retarget_matrices) - 1):
102             first_mat = retarget_matrices[i]
103             next_mat = retarget_matrices[i + 1]
104             lerp_matrix = first_mat.lerp(next_mat, 0.5)
105         return lerp_matrix
106
107     #determines the type of hierachy change needed and calls the
108     #right function
109     def retargetPerfToInter(inter_bone):
110         if inter_bone.name in bonemapr:
111             perf_bone_name = bonemapr[inter_bone.name]
112             #is it a 1 to many?
113             if isinstance(bonemap[perf_bone_name[0]], tuple):
114                 perf_bone = performer_bones[perf_bone_name[0]]
115                 if inter_bone.name == bonemap[perf_bone_name[0]][0]:
116                     locOfOriginal(inter_bone, perf_bone)
117             else:
118                 # then its either a many to 1 or 1 to 1
119
120                 if len(perf_bone_name) > 1:
121                     performer_bones_s = [performer_bones[name] for name in perf_bone_name]
122                     #we need to map several performance bone to a single
123                     for perf_bone in performer_bones_s:
124                         locOfOriginal(inter_bone, perf_bone)
125                     inter_bone.matrix_basis = manyPerfToSingleInterRetarget(inter_bone, performer_bones_s)
126                 else:
127                     perf_bone = performer_bones[perf_bone_name[0]]
128                     locOfOriginal(inter_bone, perf_bone)
129                     inter_bone.matrix_basis = singleBoneRetarget(inter_bone, perf_bone)
130
131         inter_bone.keyframe_insert("rotation_quaternion")
132         for child in inter_bone.children:
133             retargetPerfToInter(child)
134
135     #creates the intermediate armature object
136     inter_obj = enduser_obj.copy()
137     inter_obj.data = inter_obj.data.copy()  # duplicate data
138     bpy.context.scene.objects.link(inter_obj)
139     inter_obj.name = "intermediate"
140     bpy.context.scene.objects.active = inter_obj
141     bpy.ops.object.mode_set(mode='EDIT')
142     #resets roll
143     bpy.ops.armature.calculate_roll(type='Z')
144     bpy.ops.object.mode_set(mode="OBJECT")
145     inter_obj.data.name = "inter_arm"
146     inter_arm = inter_obj.data
147     performer_bones = performer_obj.pose.bones
148     inter_bones = inter_obj.pose.bones
149     #clears inheritance
150     for inter_bone in inter_bones:
151         inter_bone.bone.use_inherit_rotation = False
152
153     for t in range(s_frame, e_frame):
154         scene.frame_set(t)
155         inter_bone = inter_bones[root]
156         retargetPerfToInter(inter_bone)
157
158     return inter_obj, inter_arm
159
160 # this procedure copies the rotations over from the intermediate
161 # armature to the end user one.
162 # As the hierarchies are 1 to 1, this is a simple matter of
163 # copying the rotation, while keeping in mind bone roll, parenting, etc.
164 # TODO: Control Bones: If a certain bone is constrained in a way
165 #       that its rotation is determined by another (a control bone)
166 #       We should determine the right pos of the control bone.
167 #       Scale: ? Should work but needs testing.
168
169
170 def retargetEnduser(inter_obj, enduser_obj, root, s_frame, e_frame, scene):
171     inter_bones = inter_obj.pose.bones
172     end_bones = enduser_obj.pose.bones
173
174     def bakeTransform(end_bone):
175         src_bone = inter_bones[end_bone.name]
176         trg_bone = end_bone
177         bake_matrix = src_bone.matrix
178         rest_matrix = trg_bone.bone.matrix_local
179
180         if trg_bone.parent and trg_bone.bone.use_inherit_rotation:
181             parent_mat = src_bone.parent.matrix
182             parent_rest = trg_bone.parent.bone.matrix_local
183             parent_rest_inv = parent_rest.copy()
184             parent_rest_inv.invert()
185             parent_mat_inv = parent_mat.copy()
186             parent_mat_inv.invert()
187             bake_matrix = parent_mat_inv * bake_matrix
188             rest_matrix = parent_rest_inv * rest_matrix
189
190         rest_matrix_inv = rest_matrix.copy()
191         rest_matrix_inv.invert()
192         bake_matrix = rest_matrix_inv * bake_matrix
193         trg_bone.matrix_basis = bake_matrix
194         end_bone.keyframe_insert("rotation_quaternion")
195
196         for bone in end_bone.children:
197             bakeTransform(bone)
198
199     for t in range(s_frame, e_frame):
200         scene.frame_set(t)
201         end_bone = end_bones[root]
202         bakeTransform(end_bone)
203
204 #recieves the performer feet bones as a variable
205 # by "feet" I mean those bones that have plants
206 # (they don't move, despite root moving) somewhere in the animation.
207
208
209 def copyTranslation(performer_obj, enduser_obj, perfFeet, bonemap, bonemapr, root, s_frame, e_frame, scene):
210     endFeet = [bonemap[perfBone] for perfBone in perfFeet]
211     perfRoot = bonemapr[root][0]
212     locDictKeys = perfFeet + endFeet + [perfRoot]
213     perf_bones = performer_obj.pose.bones
214     end_bones = enduser_obj.pose.bones
215
216     def tailLoc(bone):
217         return bone.center + (bone.vector / 2)
218
219     #Step 1 - we create a dict that contains these keys:
220     #(Performer) Hips, Feet
221     #(End user) Feet
222     # where the values are their world position on each (1,120) frame
223
224     locDict = {}
225     for key in locDictKeys:
226         locDict[key] = []
227
228     for t in range(scene.frame_start, scene.frame_end):
229         scene.frame_set(t)
230         for bone in perfFeet:
231             locDict[bone].append(tailLoc(perf_bones[bone]))
232         locDict[perfRoot].append(tailLoc(perf_bones[perfRoot]))
233         for bone in endFeet:
234             locDict[bone].append(tailLoc(end_bones[bone]))
235
236     # now we take our locDict and analyze it.
237     # we need to derive all chains
238
239     locDeriv = {}
240     for key in locDictKeys:
241         locDeriv[key] = []
242
243     for key in locDict.keys():
244         graph = locDict[key]
245         for t in range(len(graph) - 1):
246             x = graph[t]
247             xh = graph[t + 1]
248             locDeriv[key].append(xh - x)
249
250     # now find the plant frames, where perfFeet don't move much
251
252     linearAvg = []
253
254     for key in perfFeet:
255         for i in range(len(locDeriv[key]) - 1):
256             v = locDeriv[key][i]
257             hipV = locDeriv[perfRoot][i]
258             endV = locDeriv[bonemap[key]][i]
259             if (v.length < 0.1):
260                 #this is a plant frame.
261                 #lets see what the original hip delta is, and the corresponding
262                 #end bone's delta
263                 if endV.length != 0:
264                     linearAvg.append(hipV.length / endV.length)
265     if linearAvg:
266         avg = sum(linearAvg) / len(linearAvg)
267         print("retargeted root motion should be " + str(1 / avg) + " of original")
268
269         bpy.ops.object.add()
270         stride_bone = bpy.context.active_object
271         stride_bone.name = "stride_bone"
272         for t in range(s_frame, e_frame):
273             scene.frame_set(t)
274             newTranslation = (tailLoc(perf_bones[perfRoot]) / avg)
275             stride_bone.location = newTranslation
276             stride_bone.keyframe_insert("location")
277
278
279 def IKRetarget(bonemap, bonemapr, performer_obj, enduser_obj, s_frame, e_frame, scene):
280     end_bones = enduser_obj.pose.bones
281     for pose_bone in end_bones:
282         if "IK" in [constraint.type for constraint in pose_bone.constraints]:
283             # set constraint target to corresponding empty if targetless,
284             # if not, keyframe current target to corresponding empty
285             perf_bone = bonemapr[pose_bone.name]
286             if isinstance(perf_bone, list):
287                 perf_bone = bonemapr[pose_bone.name][-1]
288             end_empty = bpy.data.objects[perf_bone + "Org"]
289             ik_constraint = [constraint for constraint in pose_bone.constraints if constraint.type == "IK"][0]
290             if not ik_constraint.target:
291                 ik_constraint.target = end_empty
292             else:
293                 #Bone target
294                 target_is_bone = False
295                 if ik_constraint.subtarget:
296                     target = ik_constraint.target.pose.bones[ik_constraint.subtarget]
297                     target.bone.use_local_location = False
298                     target_is_bone = True
299                 else:
300                     target = ik_constraint.target
301                 for t in range(s_frame, e_frame):
302                     scene.frame_set(t)
303                     if target_is_bone:
304                         final_loc = end_empty.location - target.bone.matrix_local.to_translation()
305                     else:
306                         final_loc = end_empty.location
307                     target.location = final_loc
308                     target.keyframe_insert("location")
309             ik_constraint.mute = False
310
311
312 def turnOffIK(enduser_obj):
313     end_bones = enduser_obj.pose.bones
314     for pose_bone in end_bones:
315         if pose_bone.is_in_ik_chain:
316             pass
317             # TODO:
318             # set stiffness according to place on chain
319             # and values from analysis that is stored in the bone
320             #pose_bone.ik_stiffness_x = 0.5
321             #pose_bone.ik_stiffness_y = 0.5
322             #pose_bone.ik_stiffness_z = 0.5
323         if "IK" in [constraint.type for constraint in pose_bone.constraints]:
324             ik_constraint = [constraint for constraint in pose_bone.constraints if constraint.type == "IK"][0]
325             ik_constraint.mute = True
326
327
328 def totalRetarget():
329     print("retargeting...")
330     enduser_obj = bpy.context.active_object
331     performer_obj = [obj for obj in bpy.context.selected_objects if obj != enduser_obj]
332     if enduser_obj is None or len(performer_obj) != 1:
333         print("Need active and selected armatures")
334     else:
335         performer_obj = performer_obj[0]
336     perf_arm = performer_obj.data
337     end_arm = enduser_obj.data
338     scene = bpy.context.scene
339     s_frame = scene.frame_start
340     e_frame = scene.frame_end
341     bonemap, bonemapr, root = createDictionary(perf_arm)
342     turnOffIK(enduser_obj)
343     inter_obj, inter_arm = createIntermediate(performer_obj, enduser_obj, bonemap, bonemapr, root, s_frame, e_frame, scene)
344     retargetEnduser(inter_obj, enduser_obj, root, s_frame, e_frame, scene)
345     copyTranslation(performer_obj, enduser_obj, ["RightFoot", "LeftFoot"], bonemap, bonemapr, root, s_frame, e_frame, scene)
346     IKRetarget(bonemap, bonemapr, performer_obj, enduser_obj, s_frame, e_frame, scene)
347     bpy.ops.object.mode_set(mode='OBJECT')
348     bpy.ops.object.select_name(name=inter_obj.name, extend=False)
349     bpy.ops.object.delete()
350
351 if __name__ == "__main__":
352     totalRetarget()