Small bugfix for prior commit - Removing constraints no longer causes an error
[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 constraint " + 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
193     #Set the blender constraint parameters
194     if m_constraint.type == "point":
195         constraint_settings = False
196         real_constraint.owner_space = m_constraint.targetSpace
197         if obj.data.animation_data:
198             if obj.data.animation_data.action:
199                 path = m_constraint.path_from_id("targetPoint")
200                 m_fcurves = [fcurve for fcurve in obj.data.animation_data.action.fcurves if fcurve.data_path == path]
201                 if m_fcurves:
202                     constraint_settings = True
203                     xCurves, yCurves, zCurves = createConstraintFCurves(cons_obj, obj, real_constraint)
204                     for curve in xCurves:
205                         copyFCurve(curve, m_fcurves[0])
206                     for curve in yCurves:
207                         copyFCurve(curve, m_fcurves[1])
208                     for curve in zCurves:
209                         copyFCurve(curve, m_fcurves[2])
210         if not constraint_settings:
211             x, y, z = m_constraint.targetPoint
212             real_constraint.max_x = x
213             real_constraint.max_y = y
214             real_constraint.max_z = z
215             real_constraint.min_x = x
216             real_constraint.min_y = y
217             real_constraint.min_z = z
218             real_constraint.use_max_x = True
219             real_constraint.use_max_y = True
220             real_constraint.use_max_z = True
221             real_constraint.use_min_x = True
222             real_constraint.use_min_y = True
223             real_constraint.use_min_z = True
224
225     if m_constraint.type == "freeze":
226         real_constraint.owner_space = m_constraint.targetSpace
227         bpy.context.scene.frame_set(m_constraint.s_frame)
228         if isinstance(cons_obj, bpy.types.PoseBone):
229             x, y, z = cons_obj.bone.center + (cons_obj.bone.vector / 2) + obj.matrix_world.to_translation()
230         else:
231             x, y, z = cons_obj.matrix_world.to_translation()
232
233         real_constraint.max_x = x
234         real_constraint.max_y = y
235         real_constraint.max_z = z
236         real_constraint.min_x = x
237         real_constraint.min_y = y
238         real_constraint.min_z = z
239         real_constraint.use_max_x = True
240         real_constraint.use_max_y = True
241         real_constraint.use_max_z = True
242         real_constraint.use_min_x = True
243         real_constraint.use_min_y = True
244         real_constraint.use_min_z = True
245
246     if m_constraint.type == "distance" and m_constraint.constrained_boneB:
247         real_constraint.owner_space = "WORLD"
248         real_constraint.target = getConsObj(bones[m_constraint.constrained_boneB])
249         real_constraint.limit_mode = "LIMITDIST_ONSURFACE"
250         real_constraint.distance = m_constraint.targetDist
251
252     if m_constraint.type == "floor" and m_constraint.targetMesh:
253         real_constraint.mute = True
254         real_constraint.owner_space = "WORLD"
255         #calculate the positions thoughout the range
256         s, e = m_constraint.s_frame, m_constraint.e_frame
257         s_in, s_out = m_constraint.smooth_in, m_constraint.smooth_out
258         s -= s_in
259         e += s_out
260         bakedPos = {}
261         floor = bpy.data.objects[m_constraint.targetMesh]
262         c_frame = context.scene.frame_current
263         for t in range(s, e):
264             context.scene.frame_set(t)
265             axis = Vector((0, 0, 100)) * obj.matrix_world.to_3x3()
266             offset = Vector((0, 0, m_constraint.targetDist)) * obj.matrix_world.to_3x3()
267             ray_origin = cons_obj.matrix_world.to_translation() - offset  # world position of constrained bone
268             ray_target = ray_origin + axis
269             #convert ray points to floor's object space
270             ray_origin *= floor.matrix_world.inverted()
271             ray_target *= floor.matrix_world.inverted()
272             hit, nor, ind = floor.ray_cast(ray_origin, ray_target)
273             if hit != Vector((0, 0, 0)):
274                 bakedPos[t] = (hit * floor.matrix_world)
275                 bakedPos[t] += Vector((0, 0, m_constraint.targetDist))
276             else:
277                 bakedPos[t] = cons_obj.matrix_world.to_translation()
278         context.scene.frame_set(c_frame)
279         #create keyframes for real constraint
280         xCurves, yCurves, zCurves = createConstraintFCurves(cons_obj, obj, real_constraint)
281         for frame in bakedPos.keys():
282             pos = bakedPos[frame]
283             for xCurve in xCurves:
284                 xCurve.keyframe_points.insert(frame=frame, value=pos.x)
285             for yCurve in yCurves:
286                 yCurve.keyframe_points.insert(frame=frame, value=pos.y)
287             for zCurve in zCurves:
288                 zCurve.keyframe_points.insert(frame=frame, value=pos.z)
289         real_constraint.use_max_x = True
290         real_constraint.use_max_y = True
291         real_constraint.use_max_z = True
292         real_constraint.use_min_x = True
293         real_constraint.use_min_y = True
294         real_constraint.use_min_z = True
295
296     # active/baked check
297     real_constraint.mute = (not m_constraint.active)
298
299
300 def locBake(s_frame, e_frame, bones):
301     scene = bpy.context.scene
302     bakeDict = {}
303     for bone in bones:
304         bakeDict[bone.name] = {}
305     for t in range(s_frame, e_frame):
306         scene.frame_set(t)
307         for bone in bones:
308             bakeDict[bone.name][t] = bone.matrix.copy()
309     for t in range(s_frame, e_frame):
310         for bone in bones:
311             print(bone.bone.matrix_local.to_translation())
312             bone.matrix = bakeDict[bone.name][t]
313             bone.keyframe_insert("location", frame=t)
314
315
316 # Baking function which bakes all bones effected by the constraint
317 def bakeAllConstraints(obj, s_frame, e_frame, bones):
318     for bone in bones:
319         bone.bone.select = False
320     selectedBones = []  # Marks bones that need a full bake
321     simpleBake = []  # Marks bones that need only a location bake
322     for end_bone in bones:
323         if end_bone.name in [m_constraint.real_constraint_bone for m_constraint in obj.data.mocap_constraints]:
324             #For all bones that have a constraint:
325             ik = hasIKConstraint(end_bone)
326             cons_obj = getConsObj(end_bone)
327             if ik:
328                     #If it's an auto generated IK:
329                     if ik.chain_count == 0:
330                         selectedBones += bones  # Chain len 0, bake everything
331                     else:
332                         selectedBones += [end_bone] + end_bone.parent_recursive[:ik.chain_count - 1]  # Bake the chain
333             else:
334                 #It's either an FK bone which we should just bake
335                 #OR a user created IK target bone
336                 simpleBake += [end_bone]
337     for bone in selectedBones:
338         bone.bone.select = True
339     constraintTrack = obj.animation_data.nla_tracks["Mocap constraints"]
340     constraintStrip = constraintTrack.strips[0]
341     constraintStrip.action_frame_start = s_frame
342     constraintStrip.action_frame_end = e_frame
343     constraintStrip.frame_start = s_frame
344     constraintStrip.frame_end = e_frame
345     if selectedBones:
346         #Use bake function from NLA Bake Action operator
347         nla.bake(s_frame, e_frame, action=constraintStrip.action, only_selected=True, do_pose=True, do_object=False)
348     if simpleBake:
349         #Do a "simple" bake, location only, world space only.
350         locBake(s_frame, e_frame, simpleBake)
351
352
353 #Calls the baking function and decativates releveant constraints
354 def bakeConstraints(context):
355     obj = context.active_object
356     bones = obj.pose.bones
357     s_frame, e_frame = context.scene.frame_start, context.scene.frame_end
358     #Bake relevant bones
359     bakeAllConstraints(obj, s_frame, e_frame, bones)
360     for m_constraint in obj.data.mocap_constraints:
361         end_bone = bones[m_constraint.real_constraint_bone]
362         cons_obj = getConsObj(end_bone)
363         # It's a control empty: turn the ik off
364         if not isinstance(cons_obj, bpy.types.PoseBone):
365             ik_con = hasIKConstraint(end_bone)
366             if ik_con:
367                 ik_con.mute = True
368         # Deactivate related Blender Constraint
369         m_constraint.active = False
370
371
372 #Deletes the baked fcurves and reactivates relevant constraints
373 def unbakeConstraints(context):
374     # to unbake constraints we delete the whole strip
375     obj = context.active_object
376     bones = obj.pose.bones
377     scene = bpy.context.scene
378     constraintTrack = obj.animation_data.nla_tracks["Mocap constraints"]
379     constraintStrip = constraintTrack.strips[0]
380     action = constraintStrip.action
381     # delete the fcurves on the strip
382     for fcurve in action.fcurves:
383         action.fcurves.remove(fcurve)
384     # reactivate relevant constraints
385     for m_constraint in obj.data.mocap_constraints:
386         end_bone = bones[m_constraint.real_constraint_bone]
387         cons_obj = getConsObj(end_bone)
388         # It's a control empty: turn the ik back on
389         if not isinstance(cons_obj, bpy.types.PoseBone):
390             ik_con = hasIKConstraint(end_bone)
391             if ik_con:
392                 ik_con.mute = False
393         m_constraint.active = True