Post Retarget fixes - added an Update Constraints button, that recalculates all fixes...
[blender.git] / release / scripts / modules / mocap_constraints.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 bl_operators import nla
24 from retarget import hasIKConstraint
25
26 ### Utility Functions
27
28
29 def getConsObj(bone):
30     #utility function - returns related IK target if bone has IK
31     ik = [constraint for constraint in bone.constraints if constraint.type == "IK"]
32     if ik:
33         ik = ik[0]
34         cons_obj = ik.target
35         if ik.subtarget:
36             cons_obj = ik.target.pose.bones[ik.subtarget]
37     else:
38         cons_obj = bone
39     return cons_obj
40
41
42 def consObjToBone(cons_obj):
43     #Utility function - returns related bone from ik object
44     if cons_obj.name[-3:] == "Org":
45         return cons_obj.name[:-3]
46     else:
47         return cons_obj.name
48
49 ### And and Remove Constraints (called from operators)
50
51
52 def addNewConstraint(m_constraint, cons_obj):
53      #Decide the correct Blender constraint according to the Mocap constraint type
54     if m_constraint.type == "point" or m_constraint.type == "freeze":
55         c_type = "LIMIT_LOCATION"
56     if m_constraint.type == "distance":
57         c_type = "LIMIT_DISTANCE"
58     if m_constraint.type == "floor":
59         c_type = "LIMIT_LOCATION"
60         #create and store the new constraint within m_constraint
61     real_constraint = cons_obj.constraints.new(c_type)
62     real_constraint.name = "Mocap fix " + str(len(cons_obj.constraints))
63     m_constraint.real_constraint_bone = consObjToBone(cons_obj)
64     m_constraint.real_constraint = real_constraint.name
65     #set the rest of the constraint properties
66     setConstraint(m_constraint, bpy.context)
67
68
69 def removeConstraint(m_constraint, cons_obj):
70     #remove the influence fcurve and Blender constraint
71     oldConstraint = cons_obj.constraints[m_constraint.real_constraint]
72     removeFcurves(cons_obj, bpy.context.active_object, oldConstraint, m_constraint)
73     cons_obj.constraints.remove(oldConstraint)
74
75 ### Update functions. There are 3: UpdateType/Bone
76 ### update framing (deals with changes in the desired frame range)
77 ### And setConstraint which deals with the rest
78
79
80 def updateConstraintBoneType(m_constraint, context):
81     #If the constraint exists, we need to remove it
82     #from the old bone
83     obj = context.active_object
84     bones = obj.pose.bones
85     if m_constraint.real_constraint:
86         bone = bones[m_constraint.real_constraint_bone]
87         cons_obj = getConsObj(bone)
88         removeConstraint(m_constraint, cons_obj)
89     #Regardless, after that we create a new constraint
90     if m_constraint.constrained_bone:
91         bone = bones[m_constraint.constrained_bone]
92         cons_obj = getConsObj(bone)
93         addNewConstraint(m_constraint, cons_obj)
94
95
96 def setConstraintFraming(m_constraint, context):
97     obj = context.active_object
98     bones = obj.pose.bones
99     bone = bones[m_constraint.constrained_bone]
100     cons_obj = getConsObj(bone)
101     real_constraint = cons_obj.constraints[m_constraint.real_constraint]
102     #remove the old keyframes
103     removeFcurves(cons_obj, obj, real_constraint, m_constraint)
104     #set the new ones according to the m_constraint properties
105     s, e = m_constraint.s_frame, m_constraint.e_frame
106     s_in, s_out = m_constraint.smooth_in, m_constraint.smooth_out
107     real_constraint.influence = 1
108     real_constraint.keyframe_insert(data_path="influence", frame=s)
109     real_constraint.keyframe_insert(data_path="influence", frame=e)
110     real_constraint.influence = 0
111     real_constraint.keyframe_insert(data_path="influence", frame=s - s_in)
112     real_constraint.keyframe_insert(data_path="influence", frame=e + s_out)
113
114
115 def removeFcurves(cons_obj, obj, real_constraint, m_constraint):
116     #Determine if the constrained object is a bone or an empty
117     if isinstance(cons_obj, bpy.types.PoseBone):
118         fcurves = obj.animation_data.action.fcurves
119     else:
120         fcurves = cons_obj.animation_data.action.fcurves
121     #Find the RNA data path of the constraint's influence
122     RNA_paths = []
123     RNA_paths.append(real_constraint.path_from_id("influence"))
124     if m_constraint.type == "floor" or m_constraint.type == "point":
125         RNA_paths += [real_constraint.path_from_id("max_x"), real_constraint.path_from_id("min_x")]
126         RNA_paths += [real_constraint.path_from_id("max_y"), real_constraint.path_from_id("min_y")]
127         RNA_paths += [real_constraint.path_from_id("max_z"), real_constraint.path_from_id("min_z")]
128     #Retrieve the correct fcurve via the RNA data path and remove it
129     fcurves_del = [fcurve for fcurve in fcurves if fcurve.data_path in RNA_paths]
130     #clear the fcurve and set the frames.
131     if fcurves_del:
132         for fcurve in fcurves_del:
133             fcurves.remove(fcurve)
134     #remove armature fcurves (if user keyframed m_constraint properties)
135     if obj.data.animation_data and m_constraint.type == "point":
136         if obj.data.animation_data.action:
137             path = m_constraint.path_from_id("targetPoint")
138             m_fcurves = [fcurve for fcurve in obj.data.animation_data.action.fcurves if fcurve.data_path == path]
139             for curve in m_fcurves:
140                 obj.data.animation_data.action.fcurves.remove(curve)
141
142 #Utility function for copying property fcurves over
143
144
145 def copyFCurve(newCurve, oldCurve):
146     for point in oldCurve.keyframe_points:
147         newCurve.keyframe_points.insert(frame=point.co.x, value=point.co.y)
148
149 #Creates new fcurves for the constraint properties (for floor and point)
150
151
152 def createConstraintFCurves(cons_obj, obj, real_constraint):
153     if isinstance(cons_obj, bpy.types.PoseBone):
154         c_fcurves = obj.animation_data.action.fcurves
155     else:
156         c_fcurves = cons_obj.animation_data.action.fcurves
157     c_x_path = [real_constraint.path_from_id("max_x"), real_constraint.path_from_id("min_x")]
158     c_y_path = [real_constraint.path_from_id("max_y"), real_constraint.path_from_id("min_y")]
159     c_z_path = [real_constraint.path_from_id("max_z"), real_constraint.path_from_id("min_z")]
160     c_constraints_path = c_x_path + c_y_path + c_z_path
161     existing_curves = [fcurve for fcurve in c_fcurves if fcurve.data_path in c_constraints_path]
162     if existing_curves:
163         for curve in existing_curves:
164             c_fcurves.remove(curve)
165     xCurves, yCurves, zCurves = [], [], []
166     for path in c_constraints_path:
167         newCurve = c_fcurves.new(path)
168         if path in c_x_path:
169             xCurves.append(newCurve)
170         elif path in c_y_path:
171             yCurves.append(newCurve)
172         else:
173             zCurves.append(newCurve)
174     return xCurves, yCurves, zCurves
175
176
177 # Function that copies all settings from m_constraint to the real Blender constraints
178 # Is only called when blender constraint already exists
179
180
181 def setConstraint(m_constraint, context):
182     if not m_constraint.constrained_bone:
183         return
184     obj = context.active_object
185     bones = obj.pose.bones
186     bone = bones[m_constraint.constrained_bone]
187     cons_obj = getConsObj(bone)
188     real_constraint = cons_obj.constraints[m_constraint.real_constraint]
189
190     #frame changing section
191     setConstraintFraming(m_constraint, context)
192     s, e = m_constraint.s_frame, m_constraint.e_frame
193     s_in, s_out = m_constraint.smooth_in, m_constraint.smooth_out
194     s -= s_in
195     e += s_out
196     #Set the blender constraint parameters
197     if m_constraint.type == "point":
198         constraint_settings = False  # are fix settings keyframed?
199         if not m_constraint.targetSpace == "constrained_boneB":
200             real_constraint.owner_space = m_constraint.targetSpace
201         else:
202             real_constraint.owner_space = "LOCAL"
203         if obj.data.animation_data:
204             if obj.data.animation_data.action:
205                 path = m_constraint.path_from_id("targetPoint")
206                 m_fcurves = [fcurve for fcurve in obj.data.animation_data.action.fcurves if fcurve.data_path == path]
207                 if m_fcurves:
208                     constraint_settings = True
209                     xCurves, yCurves, zCurves = createConstraintFCurves(cons_obj, obj, real_constraint)
210                     for curve in xCurves:
211                         copyFCurve(curve, m_fcurves[0])
212                     for curve in yCurves:
213                         copyFCurve(curve, m_fcurves[1])
214                     for curve in zCurves:
215                         copyFCurve(curve, m_fcurves[2])
216         if m_constraint.targetSpace == "constrained_boneB" and m_constraint.constrained_boneB:
217             c_frame = context.scene.frame_current
218             bakedPos = {}
219             src_bone = bones[m_constraint.constrained_boneB]
220             if not constraint_settings:
221                 xCurves, yCurves, zCurves = createConstraintFCurves(cons_obj, obj, real_constraint)
222             print("please wait a moment, calculating fix")
223             for t in range(s, e):
224                 context.scene.frame_set(t)
225                 src_bone_pos = src_bone.matrix.to_translation()
226                 bakedPos[t] = src_bone_pos + m_constraint.targetPoint  # final position for constrained bone in object space
227             context.scene.frame_set(c_frame)
228             for frame in bakedPos.keys():
229                 pos = bakedPos[frame]
230                 for xCurve in xCurves:
231                     xCurve.keyframe_points.insert(frame=frame, value=pos.x)
232                 for yCurve in yCurves:
233                     yCurve.keyframe_points.insert(frame=frame, value=pos.y)
234                 for zCurve in zCurves:
235                     zCurve.keyframe_points.insert(frame=frame, value=pos.z)
236
237         if not constraint_settings:
238             x, y, z = m_constraint.targetPoint
239             real_constraint.max_x = x
240             real_constraint.max_y = y
241             real_constraint.max_z = z
242             real_constraint.min_x = x
243             real_constraint.min_y = y
244             real_constraint.min_z = z
245             real_constraint.use_max_x = True
246             real_constraint.use_max_y = True
247             real_constraint.use_max_z = True
248             real_constraint.use_min_x = True
249             real_constraint.use_min_y = True
250             real_constraint.use_min_z = True
251
252     if m_constraint.type == "freeze":
253         real_constraint.owner_space = m_constraint.targetSpace
254         bpy.context.scene.frame_set(m_constraint.s_frame)
255         if isinstance(cons_obj, bpy.types.PoseBone):
256             x, y, z = cons_obj.bone.center + (cons_obj.bone.vector / 2) + obj.matrix_world.to_translation()
257         else:
258             x, y, z = cons_obj.matrix_world.to_translation()
259
260         real_constraint.max_x = x
261         real_constraint.max_y = y
262         real_constraint.max_z = z
263         real_constraint.min_x = x
264         real_constraint.min_y = y
265         real_constraint.min_z = z
266         real_constraint.use_max_x = True
267         real_constraint.use_max_y = True
268         real_constraint.use_max_z = True
269         real_constraint.use_min_x = True
270         real_constraint.use_min_y = True
271         real_constraint.use_min_z = True
272
273     if m_constraint.type == "distance" and m_constraint.constrained_boneB:
274         real_constraint.owner_space = "WORLD"
275         real_constraint.target = getConsObj(bones[m_constraint.constrained_boneB])
276         real_constraint.limit_mode = "LIMITDIST_ONSURFACE"
277         real_constraint.distance = m_constraint.targetDist
278
279     if m_constraint.type == "floor" and m_constraint.targetMesh:
280         real_constraint.mute = True
281         real_constraint.owner_space = "WORLD"
282         #calculate the positions thoughout the range
283         s, e = m_constraint.s_frame, m_constraint.e_frame
284         s_in, s_out = m_constraint.smooth_in, m_constraint.smooth_out
285         s -= s_in
286         e += s_out
287         bakedPos = {}
288         floor = bpy.data.objects[m_constraint.targetMesh]
289         c_frame = context.scene.frame_current
290         print("please wait a moment, calculating fix")
291         for t in range(s, e):
292             context.scene.frame_set(t)
293             axis = Vector((0, 0, 100)) * obj.matrix_world.to_3x3()
294             offset = Vector((0, 0, m_constraint.targetDist)) * obj.matrix_world.to_3x3()
295             ray_origin = cons_obj.matrix_world.to_translation() - offset  # world position of constrained bone
296             ray_target = ray_origin + axis
297             #convert ray points to floor's object space
298             ray_origin *= floor.matrix_world.inverted()
299             ray_target *= floor.matrix_world.inverted()
300             hit, nor, ind = floor.ray_cast(ray_origin, ray_target)
301             if hit != Vector((0, 0, 0)):
302                 bakedPos[t] = (hit * floor.matrix_world)
303                 bakedPos[t] += Vector((0, 0, m_constraint.targetDist))
304             else:
305                 bakedPos[t] = cons_obj.matrix_world.to_translation()
306         context.scene.frame_set(c_frame)
307         #create keyframes for real constraint
308         xCurves, yCurves, zCurves = createConstraintFCurves(cons_obj, obj, real_constraint)
309         for frame in bakedPos.keys():
310             pos = bakedPos[frame]
311             for xCurve in xCurves:
312                 xCurve.keyframe_points.insert(frame=frame, value=pos.x)
313             for yCurve in yCurves:
314                 yCurve.keyframe_points.insert(frame=frame, value=pos.y)
315             for zCurve in zCurves:
316                 zCurve.keyframe_points.insert(frame=frame, value=pos.z)
317         real_constraint.use_max_x = True
318         real_constraint.use_max_y = True
319         real_constraint.use_max_z = True
320         real_constraint.use_min_x = True
321         real_constraint.use_min_y = True
322         real_constraint.use_min_z = True
323
324     # active/baked check
325     real_constraint.mute = (not m_constraint.active)
326
327
328 def locBake(s_frame, e_frame, bones):
329     scene = bpy.context.scene
330     bakeDict = {}
331     for bone in bones:
332         bakeDict[bone.name] = {}
333     for t in range(s_frame, e_frame):
334         scene.frame_set(t)
335         for bone in bones:
336             bakeDict[bone.name][t] = bone.matrix.copy()
337     for t in range(s_frame, e_frame):
338         for bone in bones:
339             print(bone.bone.matrix_local.to_translation())
340             bone.matrix = bakeDict[bone.name][t]
341             bone.keyframe_insert("location", frame=t)
342
343
344 # Baking function which bakes all bones effected by the constraint
345 def bakeAllConstraints(obj, s_frame, e_frame, bones):
346     for bone in bones:
347         bone.bone.select = False
348     selectedBones = []  # Marks bones that need a full bake
349     simpleBake = []  # Marks bones that need only a location bake
350     for end_bone in bones:
351         if end_bone.name in [m_constraint.real_constraint_bone for m_constraint in obj.data.mocap_constraints]:
352             #For all bones that have a constraint:
353             ik = hasIKConstraint(end_bone)
354             cons_obj = getConsObj(end_bone)
355             if ik:
356                     #If it's an auto generated IK:
357                     if ik.chain_count == 0:
358                         selectedBones += bones  # Chain len 0, bake everything
359                     else:
360                         selectedBones += [end_bone] + end_bone.parent_recursive[:ik.chain_count - 1]  # Bake the chain
361             else:
362                 #It's either an FK bone which we should just bake
363                 #OR a user created IK target bone
364                 simpleBake += [end_bone]
365     for bone in selectedBones:
366         bone.bone.select = True
367     constraintTrack = obj.animation_data.nla_tracks["Mocap fixes"]
368     constraintStrip = constraintTrack.strips[0]
369     constraintStrip.action_frame_start = s_frame
370     constraintStrip.action_frame_end = e_frame
371     constraintStrip.frame_start = s_frame
372     constraintStrip.frame_end = e_frame
373     if selectedBones:
374         #Use bake function from NLA Bake Action operator
375         nla.bake(s_frame, e_frame, action=constraintStrip.action, only_selected=True, do_pose=True, do_object=False)
376     if simpleBake:
377         #Do a "simple" bake, location only, world space only.
378         locBake(s_frame, e_frame, simpleBake)
379
380
381 #Calls the baking function and decativates releveant constraints
382 def bakeConstraints(context):
383     obj = context.active_object
384     bones = obj.pose.bones
385     s_frame, e_frame = context.scene.frame_start, context.scene.frame_end
386     #Bake relevant bones
387     bakeAllConstraints(obj, s_frame, e_frame, bones)
388     for m_constraint in obj.data.mocap_constraints:
389         end_bone = bones[m_constraint.real_constraint_bone]
390         cons_obj = getConsObj(end_bone)
391         # It's a control empty: turn the ik off
392         if not isinstance(cons_obj, bpy.types.PoseBone):
393             ik_con = hasIKConstraint(end_bone)
394             if ik_con:
395                 ik_con.mute = True
396         # Deactivate related Blender Constraint
397         m_constraint.active = False
398
399
400 #Deletes the baked fcurves and reactivates relevant constraints
401 def unbakeConstraints(context):
402     # to unbake constraints we delete the whole strip
403     obj = context.active_object
404     bones = obj.pose.bones
405     scene = bpy.context.scene
406     constraintTrack = obj.animation_data.nla_tracks["Mocap fixes"]
407     constraintStrip = constraintTrack.strips[0]
408     action = constraintStrip.action
409     # delete the fcurves on the strip
410     for fcurve in action.fcurves:
411         action.fcurves.remove(fcurve)
412     # reactivate relevant constraints
413     for m_constraint in obj.data.mocap_constraints:
414         end_bone = bones[m_constraint.real_constraint_bone]
415         cons_obj = getConsObj(end_bone)
416         # It's a control empty: turn the ik back on
417         if not isinstance(cons_obj, bpy.types.PoseBone):
418             ik_con = hasIKConstraint(end_bone)
419             if ik_con:
420                 ik_con.mute = False
421         m_constraint.active = True
422
423
424 def updateConstraints(obj, context):
425     fixes = obj.data.mocap_constraints
426     for fix in fixes:
427         fix.active = False
428         fix.active = True