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