Sculpt Branch:
[blender.git] / release / scripts / modules / rigify_utils.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., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
16 #
17 # ##### END GPL LICENSE BLOCK #####
18
19 # <pep8 compliant>
20
21 # rigify its self does not depend on this module, however some of the
22 # rigify templates use these utility functions.
23 #
24 # So even though this can be for general purpose use, this module was created
25 # for rigify so in some cases seemingly generic functions make assumptions
26 # that a generic function would need to check for.
27
28 import bpy
29 from Mathutils import Vector
30 from rna_prop_ui import rna_idprop_ui_prop_get
31
32 DELIMITER = '-._'
33 EMPTY_LAYER = [False] * 32
34
35 def add_stretch_to(obj, from_name, to_name, name):
36     '''
37     Adds a bone that stretches from one to another
38     '''
39
40     mode_orig = obj.mode
41     bpy.ops.object.mode_set(mode='EDIT')
42
43     arm = obj.data
44     stretch_ebone = arm.edit_bones.new(name)
45     stretch_name = stretch_ebone.name
46     del name
47
48     head = stretch_ebone.head = arm.edit_bones[from_name].head.copy()
49     #tail = stretch_ebone.tail = arm.edit_bones[to_name].head.copy()
50
51     # annoying exception for zero length bones, since its using stretch_to the rest pose doesnt really matter
52     #if (head - tail).length < 0.1:
53     if 1:
54         tail = stretch_ebone.tail = arm.edit_bones[from_name].tail.copy()
55
56
57     # Now for the constraint
58     bpy.ops.object.mode_set(mode='OBJECT')
59
60     stretch_pbone = obj.pose.bones[stretch_name]
61
62     con = stretch_pbone.constraints.new('COPY_LOCATION')
63     con.target = obj
64     con.subtarget = from_name
65
66     con = stretch_pbone.constraints.new('STRETCH_TO')
67     con.target = obj
68     con.subtarget = to_name
69     con.original_length = (head - tail).length
70     con.keep_axis = 'PLANE_X'
71     con.volume = 'NO_VOLUME'
72
73     bpy.ops.object.mode_set(mode=mode_orig)
74     
75     return stretch_name
76
77
78 def copy_bone_simple(arm, from_bone, name, parent=False):
79     ebone = arm.edit_bones[from_bone]
80     ebone_new = arm.edit_bones.new(name)
81
82     if parent:
83         ebone_new.connected = ebone.connected
84         ebone_new.parent = ebone.parent
85
86     ebone_new.head = ebone.head
87     ebone_new.tail = ebone.tail
88     ebone_new.roll = ebone.roll
89     return ebone_new
90
91
92 def copy_bone_simple_list(arm, from_bones, to_bones, parent=False):
93
94     if len(from_bones) != len(to_bones):
95         raise Exception("bone list sizes must match")
96
97     copy_bones = [copy_bone_simple(arm, bone_name, to_bones[i], True) for i, bone_name in enumerate(from_bones)]
98
99     # now we need to re-parent
100     for ebone in copy_bones:
101         parent = ebone.parent
102         if parent:
103             try:
104                 i = from_bones.index(parent.name)
105             except:
106                 i = -1
107
108             if i == -1:
109                 ebone.parent = None
110             else:
111                 ebone.parent = copy_bones[i]
112
113     return copy_bones
114
115
116 def blend_bone_list(obj, apply_bones, from_bones, to_bones, target_bone=None, target_prop="blend"):
117
118     if obj.mode == 'EDIT':
119         raise Exception("blending cant be called in editmode")
120
121     if len(apply_bones) != len(from_bones):
122         raise Exception("lists differ in length (from -> apply): \n\t%s\n\t%s" % (from_bones, apply_bones))
123     if len(apply_bones) != len(to_bones):
124         raise Exception("lists differ in length (to -> apply): \n\t%s\n\t%s" % (to_bones, apply_bones))
125
126     # setup the blend property
127     if target_bone is None:
128         target_bone = apply_bones[-1] # default to the last bone
129
130     prop_pbone = obj.pose.bones[target_bone]
131     if prop_pbone.get(target_bone, None) is None:
132         prop = rna_idprop_ui_prop_get(prop_pbone, target_prop, create=True)
133         prop_pbone[target_prop] = 0.5
134         prop["soft_min"] = 0.0
135         prop["soft_max"] = 1.0
136
137     driver_path = prop_pbone.path_to_id() + ('["%s"]' % target_prop)
138
139     def blend_target(driver):
140         tar = driver.targets.new()
141         tar.name = target_bone
142         tar.id_type = 'OBJECT'
143         tar.id = obj
144         tar.data_path = driver_path
145
146     def blend_location(new_pbone, from_bone_name, to_bone_name):
147         con = new_pbone.constraints.new('COPY_LOCATION')
148         con.target = obj
149         con.subtarget = from_bone_name
150
151         con = new_pbone.constraints.new('COPY_LOCATION')
152         con.target = obj
153         con.subtarget = to_bone_name
154
155         fcurve = con.driver_add("influence", 0)
156         driver = fcurve.driver
157         driver.type = 'AVERAGE'
158         fcurve.modifiers.remove(0) # grr dont need a modifier
159
160         blend_target(driver)
161
162     def blend_rotation(new_pbone, from_bone_name, to_bone_name):
163         con = new_pbone.constraints.new('COPY_ROTATION')
164         con.target = obj
165         con.subtarget = from_bone_name
166
167         con = new_pbone.constraints.new('COPY_ROTATION')
168         con.target = obj
169         con.subtarget = to_bone_name
170
171         fcurve = con.driver_add("influence", 0)
172         driver = fcurve.driver
173         driver.type = 'AVERAGE'
174         fcurve.modifiers.remove(0) # grr dont need a modifier
175
176         blend_target(driver)
177
178     for i, new_bone_name in enumerate(apply_bones):
179         from_bone_name = from_bones[i]
180         to_bone_name = to_bones[i]
181
182         # allow skipping some bones by having None in the list
183         if None in (new_bone_name, from_bone_name, to_bone_name):
184             continue
185
186         new_pbone = obj.pose.bones[new_bone_name]
187
188         # if the bone is connected or its location is totally locked then dont add location blending.
189         if not (new_pbone.bone.connected or (False not in new_pbone.lock_location)):
190             blend_location(new_pbone, from_bone_name, to_bone_name)
191
192         if not (False not in new_pbone.lock_rotation): # TODO. 4D chech?
193             blend_rotation(new_pbone, from_bone_name, to_bone_name)
194
195
196 def add_pole_target_bone(obj, base_bone_name, name, mode='CROSS'):
197     '''
198     Does not actually create a poll target, just the bone to use as a poll target
199     '''
200     mode_orig = obj.mode
201     bpy.ops.object.mode_set(mode='EDIT')
202
203     arm = obj.data
204
205     poll_ebone = arm.edit_bones.new(name)
206     base_ebone = arm.edit_bones[base_bone_name]
207     poll_name = poll_ebone.name
208     parent_ebone = base_ebone.parent
209
210     base_head = base_ebone.head.copy()
211     base_tail = base_ebone.tail.copy()
212     base_dir = base_head - base_tail
213
214     parent_head = parent_ebone.head.copy()
215     parent_tail = parent_ebone.tail.copy()
216     parent_dir = parent_head - parent_tail
217
218     distance = (base_dir.length + parent_dir.length)
219     
220     if mode == 'CROSS':
221         # direction from the angle of the joint
222         offset = base_dir.copy().normalize() - parent_dir.copy().normalize()
223         offset.length = distance
224     elif mode == 'ZAVERAGE':
225         # between both bones Z axis
226         z_axis_a = base_ebone.matrix.copy().rotationPart() * Vector(0.0, 0.0, -1.0)
227         z_axis_b = parent_ebone.matrix.copy().rotationPart() * Vector(0.0, 0.0, -1.0)
228         offset = (z_axis_a + z_axis_b).normalize() * distance
229     else:
230         # preset axis
231         offset = Vector(0, 0, 0)
232         if mode[0] == "+":
233             val = distance
234         else:
235             val = - distance
236
237         setattr(offset, mode[1].lower(), val)
238
239     poll_ebone.head = base_head + offset
240     poll_ebone.tail = base_head + (offset * (1.0 - (1.0 / 4.0)))
241
242     bpy.ops.object.mode_set(mode=mode_orig)
243
244     return poll_name
245
246
247
248 def get_side_name(name):
249     '''
250     Returns the last part of a string (typically a bone's name) indicating
251     whether it is a a left or right (or center, or whatever) bone.
252     Returns an empty string if nothing is found.
253     '''
254     if name[-2] in DELIMITER:
255         return name[-2:]
256     else:
257         return ""
258
259 def get_base_name(name):
260     '''
261     Returns the part of a string (typically a bone's name) corresponding to it's
262     base name (no sidedness, no ORG prefix).
263     '''
264     if name[-2] in DELIMITER:
265         return name[:-2]
266     else:
267         return name
268
269
270 def write_meta_rig(obj, func_name="metarig_template"):
271     '''
272     Write a metarig as a python script, this rig is to have all info needed for
273     generating the real rig with rigify.
274     '''
275     code = []
276
277     code.append("def %s():" % func_name)
278     code.append("    # generated by rigify.write_meta_rig")
279     bpy.ops.object.mode_set(mode='EDIT')
280     code.append("    bpy.ops.object.mode_set(mode='EDIT')")
281
282     code.append("    obj = bpy.context.active_object")
283     code.append("    arm = obj.data")
284
285     arm = obj.data
286     # write parents first
287     bones = [(len(bone.parent_recursive), bone.name) for bone in arm.edit_bones]
288     bones.sort(key=lambda item: item[0])
289     bones = [item[1] for item in bones]
290
291
292     for bone_name in bones:
293         bone = arm.edit_bones[bone_name]
294         code.append("    bone = arm.edit_bones.new('%s')" % bone.name)
295         code.append("    bone.head[:] = %.4f, %.4f, %.4f" % bone.head.toTuple(4))
296         code.append("    bone.tail[:] = %.4f, %.4f, %.4f" % bone.tail.toTuple(4))
297         code.append("    bone.roll = %.4f" % bone.roll)
298         code.append("    bone.connected = %s" % str(bone.connected))
299         if bone.parent:
300             code.append("    bone.parent = arm.edit_bones['%s']" % bone.parent.name)
301
302     bpy.ops.object.mode_set(mode='OBJECT')
303     code.append("")
304     code.append("    bpy.ops.object.mode_set(mode='OBJECT')")
305
306     for bone_name in bones:
307         pbone = obj.pose.bones[bone_name]
308         pbone_written = False
309
310         # Only 1 level of props, simple types supported
311         for key, value in pbone.items():
312             if key.startswith("_"):
313                 continue
314
315             if type(value) not in (float, str, int):
316                 print("Unsupported ID Prop:", str((key, value)))
317                 continue
318
319             if type(value) == str:
320                 value = "'" + value + "'"
321
322             if not pbone_written: # only write bones we need
323                 code.append("    pbone = obj.pose.bones['%s']" % bone_name)
324
325             code.append("    pbone['%s'] = %s" % (key, value))
326
327     return "\n".join(code)
328
329
330
331 # *** bone class collection ***
332 def bone_class_instance(obj, slots, name="BoneContainer"):
333     '''
334     bone collection utility class to help manage cases with
335     edit/pose/bone bones where switching modes can invalidate some of the members.
336     
337     there are also utility functions for manipulating all members.
338     '''
339     
340     if len(slots) != len(set(slots)):
341         raise Exception("duplicate entries found %s" % attr_names)
342
343     attr_names = tuple(slots) # dont modify the original
344     slots = list(slots) # dont modify the original
345     for i in range(len(slots)):
346         member = slots[i]
347         slots.append(member + "_b") # bone bone
348         slots.append(member + "_p") # pose bone
349         slots.append(member + "_e") # edit bone
350
351     class_dict = { \
352         "obj": obj, \
353         "attr_names": attr_names, \
354         "update": _bone_class_instance_update, \
355         "rename": _bone_class_instance_rename, \
356         "names": _bone_class_instance_names, \
357         "copy": _bone_class_instance_copy, \
358         "blend": _bone_class_instance_blend, \
359     }
360
361     instance = auto_class_instance(slots, name, class_dict)
362     return instance
363
364 def auto_class(slots, name="ContainerClass", class_dict=None):
365
366     if class_dict:
367         class_dict = class_dict.copy()
368     else:
369         class_dict = {}
370
371     class_dict["__slots__"] = tuple(slots)
372
373     return type(name, (object,), class_dict)
374
375
376 def auto_class_instance(slots, name="ContainerClass", class_dict=None):
377     return auto_class(slots, name, class_dict)()
378
379
380 def _bone_class_instance_update(self):
381     ''' Re-Assigns bones from the blender data
382     '''
383     arm = self.obj.data
384     bbones = arm.bones
385     pbones = self.obj.pose.bones
386     ebones = arm.edit_bones
387
388     for member in self.attr_names:
389         name = getattr(self, member, None)
390         if name is not None:
391             setattr(self, member + "_b", bbones.get(name, None))
392             setattr(self, member + "_p", pbones.get(name, None))
393             setattr(self, member + "_e", ebones.get(name, None))
394
395
396 def _bone_class_instance_rename(self, attr, new_name):
397     ''' Rename bones, editmode only
398     '''
399
400     if self.obj.mode != 'EDIT':
401         raise Exception("Only rename in editmode supported")
402
403     ebone = getattr(self, attr + "_e")
404     ebone.name = new_name
405
406     # we may not get what is asked for so get the name from the editbone
407     setattr(self, attr, ebone.name)
408
409
410 def _bone_class_instance_copy(self, from_fmt="%s", to_fmt="%s", exclude_attrs=(), base_names=None):
411     from_name_ls = []
412     new_name_ls = []
413     new_slot_ls = []
414
415     for attr in self.attr_names:
416         
417         if attr in exclude_attrs:
418             continue
419         
420         bone_name_orig = getattr(self, attr)
421         ebone = getattr(self, attr + "_e")
422         # orig_names[attr] = bone_name_orig
423
424         # insert formatting
425         if from_fmt != "%s":
426             bone_name = from_fmt % bone_name_orig
427             ebone.name = bone_name
428             bone_name = ebone.name # cant be sure we get what we ask for
429         else:
430             bone_name = bone_name_orig
431
432         setattr(self, attr, bone_name)
433
434         new_slot_ls.append(attr)
435         from_name_ls.append(bone_name)
436         if base_names:
437             bone_name_orig = base_names[bone_name_orig]
438         new_name_ls.append(to_fmt % bone_name_orig)
439
440     new_bones = copy_bone_simple_list(self.obj.data, from_name_ls, new_name_ls, True)
441     new_bc = bone_class_instance(self.obj, new_slot_ls)
442
443     for i, attr in enumerate(new_slot_ls):
444         ebone = new_bones[i]
445         setattr(new_bc, attr + "_e", ebone)
446         setattr(new_bc, attr, ebone.name)
447
448     return new_bc
449
450
451 def _bone_class_instance_names(self):
452     return [getattr(self, attr) for attr in self.attr_names]
453
454
455 def _bone_class_instance_blend(self, from_bc, to_bc, target_bone=None, target_prop="blend"):
456     '''
457     Use for blending bone chains.
458
459     blend_target = (bone_name, bone_property)
460     default to the last bone, blend prop
461
462     XXX - toggles editmode, need to re-validate all editbones :(
463     '''
464
465     if self.attr_names != from_bc.attr_names or self.attr_names != to_bc.attr_names:
466         raise Exception("can only blend between matching chains")
467
468     apply_bones = [getattr(self, attr) for attr in self.attr_names]
469     from_bones = [getattr(from_bc, attr) for attr in from_bc.attr_names]
470     to_bones = [getattr(to_bc, attr) for attr in to_bc.attr_names]
471
472     blend_bone_list(self.obj, apply_bones, from_bones, to_bones, target_bone, target_prop)