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