Bugfixing and modification to ensure functionality when called from UI. Related to...
[blender.git] / release / scripts / modules / retarget.py
1 import bpy
2 from mathutils import *
3 from math import radians, acos
4
5 #TODO: Only selected bones get retargeted.
6 #      Selected Bones/chains get original pos empties, if ppl want IK instead of FK
7 #      Some "magic" numbers - frame start and end, eulers of all orders instead of just quats keyframed
8
9 # dictionary of mapping
10 # this is currently manuall input'ed, but will
11 # be created from a more comfortable UI in the future
12
13 def createDictionary(perf_arm):
14     bonemap = {}
15     #perf: enduser
16     for bone in perf_arm.bones:
17         bonemap[bone.name] = bone.map
18                 
19     #root is the root of the enduser
20     root = "root"
21     # creation of a reverse map
22     # multiple keys get mapped to list values
23     bonemapr = {}
24     for key in bonemap.keys():
25         if not bonemap[key] in bonemapr:
26             if type(bonemap[key])==type((0,0)):
27                 for key_x in bonemap[key]:
28                     bonemapr[key_x] = [key]
29             else:
30                 bonemapr[bonemap[key]] = [key]
31         else:
32             bonemapr[bonemap[key]].append(key)
33     
34     return bonemap, bonemapr, root       
35 # list of empties created to keep track of "original"
36 # position data
37 # in final product, these locations can be stored as custom props
38 # these help with constraining, etc.
39
40 #creation of intermediate armature
41 # the intermediate armature has the hiearchy of the end user,
42 # does not have rotation inheritence
43 # and bone roll is identical to the performer
44 # its purpose is to copy over the rotations
45 # easily while concentrating on the hierarchy changes
46 def createIntermediate(performer_obj,endu_name,bonemap, bonemapr,root,s_frame,e_frame,scene):
47     
48     #creates and keyframes an empty with its location
49     #the original position of the tail bone
50     #useful for storing the important data in the original motion
51     #i.e. using this empty to IK the chain to that pos / DEBUG
52     def locOfOriginal(inter_bone,perf_bone):
53         pass
54 #        if not perf_bone.name+"Org" in bpy.data.objects:
55 #            bpy.ops.object.add()
56 #            empty = bpy.context.active_object
57 #            empty.name = perf_bone.name+"Org"
58 #            empty.empty_draw_size = 0.01
59 #        empty = bpy.data.objects[perf_bone.name+"Org"]
60 #        offset = perf_bone.vector
61 #        if inter_bone.length == 0 or perf_bone.length == 0:
62 #            scaling = 1
63 #        else:
64 #            scaling = perf_bone.length / inter_bone.length
65 #        offset/=scaling
66 #        empty.location = inter_bone.head + offset
67 #        empty.keyframe_insert("location")
68     
69     #Simple 1to1 retarget of a bone
70     def singleBoneRetarget(inter_bone,perf_bone):
71             perf_world_rotation = perf_bone.matrix * performer_obj.matrix_world         
72             inter_world_base_rotation = inter_bone.bone.matrix_local * inter_obj.matrix_world
73             inter_world_base_inv = Matrix(inter_world_base_rotation)
74             inter_world_base_inv.invert()
75             return (inter_world_base_inv.to_3x3() * perf_world_rotation.to_3x3()).to_4x4()
76         
77     #uses 1to1 and interpolation/averaging to match many to 1 retarget    
78     def manyPerfToSingleInterRetarget(inter_bone,performer_bones_s):
79         retarget_matrices = [singleBoneRetarget(inter_bone,perf_bone) for perf_bone in performer_bones_s]
80         lerp_matrix = Matrix()
81         for i in range(len(retarget_matrices)-1):
82             first_mat = retarget_matrices[i]
83             next_mat = retarget_matrices[i+1]
84             lerp_matrix = first_mat.lerp(next_mat,0.5)
85         return lerp_matrix
86     
87     #determines the type of hierachy change needed and calls the 
88     #right function        
89     def retargetPerfToInter(inter_bone):
90         if inter_bone.name in bonemapr.keys():
91             perf_bone_name = bonemapr[inter_bone.name]
92             #is it a 1 to many?
93             if type(bonemap[perf_bone_name[0]])==type((0,0)):
94                 perf_bone = performer_bones[perf_bone_name[0]]
95                 if inter_bone.name == bonemap[perf_bone_name[0]][0]:
96                     locOfOriginal(inter_bone,perf_bone)
97             else:
98                 # then its either a many to 1 or 1 to 1
99                 
100                 if len(perf_bone_name) > 1:
101                     performer_bones_s = [performer_bones[name] for name in perf_bone_name]
102                     #we need to map several performance bone to a single
103                     for perf_bone in performer_bones_s:
104                         locOfOriginal(inter_bone,perf_bone)
105                     inter_bone.matrix_basis = manyPerfToSingleInterRetarget(inter_bone,performer_bones_s)
106                 else:
107                     perf_bone = performer_bones[perf_bone_name[0]]
108                     locOfOriginal(inter_bone,perf_bone)
109                     inter_bone.matrix_basis = singleBoneRetarget(inter_bone,perf_bone)
110                     
111         inter_bone.keyframe_insert("rotation_quaternion")
112         for child in inter_bone.children:
113             retargetPerfToInter(child)
114             
115     #creates the intermediate armature object        
116     bpy.ops.object.select_name(name=endu_name,extend=False)
117     bpy.ops.object.duplicate(linked=False)
118     bpy.context.active_object.name = "intermediate"
119     inter_obj = bpy.context.active_object
120     bpy.ops.object.mode_set(mode='EDIT')
121     #resets roll 
122     bpy.ops.armature.calculate_roll(type='Z')
123     bpy.ops.object.mode_set(mode="OBJECT")
124     inter_arm = bpy.data.armatures[endu_name+".001"]
125     inter_arm.name = "inter_arm"
126     performer_bones = performer_obj.pose.bones
127     inter_bones =  inter_obj.pose.bones
128     
129     #clears inheritance
130     for inter_bone in inter_bones:
131         inter_bone.bone.use_inherit_rotation = False
132         
133     for t in range(s_frame,e_frame):
134         scene.frame_set(t)
135         inter_bone = inter_bones[root]
136         retargetPerfToInter(inter_bone)
137         
138     return inter_obj,inter_arm
139
140 # this procedure copies the rotations over from the intermediate
141 # armature to the end user one.
142 # As the hierarchies are 1 to 1, this is a simple matter of 
143 # copying the rotation, while keeping in mind bone roll, parenting, etc.
144 # TODO: Control Bones: If a certain bone is constrained in a way
145 #       that its rotation is determined by another (a control bone)
146 #       We should determine the right pos of the control bone.
147 #       Scale: ? Should work but needs testing.          
148 def retargetEnduser(inter_obj, enduser_obj,root,s_frame,e_frame,scene):
149     inter_bones =  inter_obj.pose.bones
150     end_bones = enduser_obj.pose.bones
151     
152     def bakeTransform(end_bone):
153         src_bone = inter_bones[end_bone.name]
154         trg_bone = end_bone
155         bake_matrix = src_bone.matrix
156         rest_matrix = trg_bone.bone.matrix_local
157         
158         if trg_bone.parent and trg_bone.bone.use_inherit_rotation:
159             parent_mat = src_bone.parent.matrix
160             parent_rest = trg_bone.parent.bone.matrix_local
161             parent_rest_inv = parent_rest.copy()
162             parent_rest_inv.invert()
163             parent_mat_inv = parent_mat.copy()
164             parent_mat_inv.invert()
165             bake_matrix = parent_mat_inv * bake_matrix
166             rest_matrix = parent_rest_inv * rest_matrix
167             
168         rest_matrix_inv = rest_matrix.copy()
169         rest_matrix_inv.invert()
170         bake_matrix = rest_matrix_inv * bake_matrix
171         trg_bone.matrix_basis = bake_matrix
172         end_bone.keyframe_insert("rotation_quaternion")
173         
174         for bone in end_bone.children:
175             bakeTransform(bone)
176         
177     for t in range(s_frame,e_frame):
178         scene.frame_set(t)
179         end_bone = end_bones[root]
180         bakeTransform(end_bone)
181
182 #recieves the performer feet bones as a variable
183 # by "feet" I mean those bones that have plants
184 # (they don't move, despite root moving) somewhere in the animation.
185 def copyTranslation(performer_obj,enduser_obj,perfFeet,bonemap,bonemapr,root,s_frame,e_frame,scene):
186     endFeet = [bonemap[perfBone] for perfBone in perfFeet]
187     perfRoot = bonemapr[root][0]
188     locDictKeys = perfFeet+endFeet+[perfRoot]
189     perf_bones = performer_obj.pose.bones
190     end_bones = enduser_obj.pose.bones
191     
192     def tailLoc(bone):
193         return bone.center+(bone.vector/2)
194     
195     #Step 1 - we create a dict that contains these keys:
196     #(Performer) Hips, Feet
197     #(End user) Feet
198     # where the values are their world position on each (1,120) frame
199     
200     locDict = {}
201     for key in locDictKeys:
202         locDict[key] = []    
203     
204     for t in range(scene.frame_start,scene.frame_end):
205         scene.frame_set(t)
206         for bone in perfFeet:
207             locDict[bone].append(tailLoc(perf_bones[bone]))
208         locDict[perfRoot].append(tailLoc(perf_bones[perfRoot]))
209         for bone in endFeet:
210             locDict[bone].append(tailLoc(end_bones[bone]))
211             
212     # now we take our locDict and analyze it.
213     # we need to derive all chains 
214     
215     locDeriv = {}
216     for key in locDictKeys:
217         locDeriv[key] = []
218     
219     for key in locDict.keys():
220         graph = locDict[key]
221         for t in range(len(graph)-1):
222             x = graph[t]
223             xh = graph[t+1]
224             locDeriv[key].append(xh-x)
225             
226     # now find the plant frames, where perfFeet don't move much
227     
228     linearAvg = []
229     
230     for key in perfFeet:
231         for i in range(len(locDeriv[key])-1):
232             v = locDeriv[key][i]
233             hipV = locDeriv[perfRoot][i]
234             endV = locDeriv[bonemap[key]][i]
235             if (v.length<0.1):
236                 #this is a plant frame.
237                 #lets see what the original hip delta is, and the corresponding
238                 #end bone's delta
239                 if endV.length!=0:
240                     linearAvg.append(hipV.length/endV.length)
241     if linearAvg:
242         avg = sum(linearAvg)/len(linearAvg)
243         print("retargeted root motion should be "+ str(1/avg)+ " of original")
244         
245         bpy.ops.object.add()
246         stride_bone = bpy.context.active_object
247         stride_bone.name = "stride_bone"
248         bpy.ops.object.select_name(name=stride_bone.name,extend=False)
249         bpy.ops.object.select_name(name=enduser_obj.name,extend=True)
250         bpy.ops.object.mode_set(mode='POSE')
251         bpy.ops.pose.select_all(action='DESELECT')
252         root_bone = end_bones[root]
253         root_bone.bone.select = True
254         bpy.ops.pose.constraint_add_with_targets(type='CHILD_OF')
255         for t in range(s_frame,e_frame):
256             scene.frame_set(t)     
257             newTranslation = (tailLoc(perf_bones[perfRoot])/avg)
258             stride_bone.location = newTranslation
259             stride_bone.keyframe_insert("location")
260             
261
262
263 def totalRetarget():
264     perf_name = bpy.context.scene.performer
265     endu_name = bpy.context.scene.enduser
266     performer_obj = bpy.data.objects[perf_name]
267     enduser_obj = bpy.data.objects[endu_name]
268     end_arm = bpy.data.armatures[endu_name] 
269     perf_arm = bpy.data.armatures[perf_name] 
270     scene = bpy.context.scene
271     s_frame = scene.frame_start
272     e_frame = scene.frame_end
273     bonemap, bonemapr, root = createDictionary(perf_arm)    
274     inter_obj, inter_arm = createIntermediate(performer_obj,endu_name,bonemap, bonemapr,root,s_frame,e_frame,scene)
275     retargetEnduser(inter_obj, enduser_obj,root,s_frame,e_frame,scene)
276     copyTranslation(performer_obj,enduser_obj,["RightFoot","LeftFoot"],bonemap,bonemapr,root,s_frame,e_frame,scene)
277     bpy.ops.object.mode_set(mode='OBJECT')
278     bpy.ops.object.select_name(name=inter_obj.name,extend=False)
279     bpy.ops.object.delete()