if bone type root is defined, all new parentless bones become children of it
[blender-staging.git] / release / scripts / modules / rigify / __init__.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 import bpy
22 from Mathutils import Vector
23
24 # TODO, have these in a more general module
25 from rna_prop_ui import rna_idprop_ui_prop_get
26
27 empty_layer = [False] * 32
28
29
30 def auto_class(slots, name="ContainerClass", class_dict=None):
31
32     if class_dict:
33         class_dict = class_dict.copy()
34     else:
35         class_dict = {}
36
37     class_dict["__slots__"] = tuple(slots)
38
39     return type(name, (object,), class_dict)
40
41
42 def auto_class_instance(slots, name="ContainerClass", class_dict=None):
43     return auto_class(slots, name, class_dict)()
44
45
46 def _bone_class_instance_update(self):
47     ''' Re-Assigns bones from the blender data
48     '''
49     arm = self.obj.data
50     bbones = arm.bones
51     pbones = self.obj.pose.bones
52     ebones = arm.edit_bones
53
54     for member in self.attr_names:
55         name = getattr(self, member, None)
56         if name is not None:
57             setattr(self, member + "_b", bbones.get(name, None))
58             setattr(self, member + "_p", pbones.get(name, None))
59             setattr(self, member + "_e", ebones.get(name, None))
60
61
62 def _bone_class_instance_rename(self, attr, new_name):
63     ''' Rename bones, editmode only
64     '''
65
66     if self.obj.mode != 'EDIT':
67         raise Exception("Only rename in editmode supported")
68
69     ebone = getattr(self, attr + "_e")
70     ebone.name = new_name
71
72     # we may not get what is asked for so get the name from the editbone
73     setattr(self, attr, ebone.name)
74
75
76 def _bone_class_instance_copy(self, from_fmt="%s", to_fmt="%s", exclude_attrs=(), base_names=None):
77     from_name_ls = []
78     new_name_ls = []
79     new_slot_ls = []
80
81     for attr in self.attr_names:
82         
83         if attr in exclude_attrs:
84             continue
85         
86         bone_name_orig = getattr(self, attr)
87         ebone = getattr(self, attr + "_e")
88         # orig_names[attr] = bone_name_orig
89
90         # insert formatting
91         if from_fmt != "%s":
92             bone_name = from_fmt % bone_name_orig
93             ebone.name = bone_name
94             bone_name = ebone.name # cant be sure we get what we ask for
95         else:
96             bone_name = bone_name_orig
97
98         setattr(self, attr, bone_name)
99
100         new_slot_ls.append(attr)
101         from_name_ls.append(bone_name)
102         if base_names:
103             bone_name_orig = base_names[bone_name_orig]
104         new_name_ls.append(to_fmt % bone_name_orig)
105
106     new_bones = copy_bone_simple_list(self.obj.data, from_name_ls, new_name_ls, True)
107     new_bc = bone_class_instance(self.obj, new_slot_ls)
108
109     for i, attr in enumerate(new_slot_ls):
110         ebone = new_bones[i]
111         setattr(new_bc, attr + "_e", ebone)
112         setattr(new_bc, attr, ebone.name)
113
114     return new_bc
115
116
117 def _bone_class_instance_names(self):
118     return [getattr(self, attr) for attr in self.attr_names]
119
120
121 def _bone_class_instance_blend(self, from_bc, to_bc, target_bone=None, target_prop="blend"):
122     '''
123     Use for blending bone chains.
124
125     blend_target = (bone_name, bone_property)
126     default to the last bone, blend prop
127
128     XXX - toggles editmode, need to re-validate all editbones :(
129     '''
130
131     if self.attr_names != from_bc.attr_names or self.attr_names != to_bc.attr_names:
132         raise Exception("can only blend between matching chains")
133
134     apply_bones = [getattr(self, attr) for attr in self.attr_names]
135     from_bones = [getattr(from_bc, attr) for attr in from_bc.attr_names]
136     to_bones = [getattr(to_bc, attr) for attr in to_bc.attr_names]
137
138     blend_bone_list(self.obj, apply_bones, from_bones, to_bones, target_bone, target_prop)
139
140
141 def bone_class_instance(obj, slots, name="BoneContainer"):
142     
143     if len(slots) != len(set(slots)):
144         raise Exception("duplicate entries found %s" % attr_names)
145
146     attr_names = tuple(slots) # dont modify the original
147     slots = list(slots) # dont modify the original
148     for i in range(len(slots)):
149         member = slots[i]
150         slots.append(member + "_b") # bone bone
151         slots.append(member + "_p") # pose bone
152         slots.append(member + "_e") # edit bone
153
154     class_dict = { \
155         "obj": obj, \
156         "attr_names": attr_names, \
157         "update": _bone_class_instance_update, \
158         "rename": _bone_class_instance_rename, \
159         "names": _bone_class_instance_names, \
160         "copy": _bone_class_instance_copy, \
161         "blend": _bone_class_instance_blend, \
162     }
163
164     instance = auto_class_instance(slots, name, class_dict)
165     return instance
166
167
168 def get_bone_data(obj, bone_name):
169     arm = obj.data
170     pbone = obj.pose.bones[bone_name]
171     if obj.mode == 'EDIT':
172         bone = arm.edit_bones[bone_name]
173     else:
174         bone = arm.bones[bone_name]
175
176     return arm, pbone, bone
177
178
179 def copy_bone_simple(arm, from_bone, name, parent=False):
180     ebone = arm.edit_bones[from_bone]
181     ebone_new = arm.edit_bones.new(name)
182
183     if parent:
184         ebone_new.connected = ebone.connected
185         ebone_new.parent = ebone.parent
186
187     ebone_new.head = ebone.head
188     ebone_new.tail = ebone.tail
189     ebone_new.roll = ebone.roll
190     return ebone_new
191
192
193 def copy_bone_simple_list(arm, from_bones, to_bones, parent=False):
194
195     if len(from_bones) != len(to_bones):
196         raise Exception("bone list sizes must match")
197
198     copy_bones = [copy_bone_simple(arm, bone_name, to_bones[i], True) for i, bone_name in enumerate(from_bones)]
199
200     # now we need to re-parent
201     for ebone in copy_bones:
202         parent = ebone.parent
203         if parent:
204             try:
205                 i = from_bones.index(parent.name)
206             except:
207                 i = -1
208
209             if i == -1:
210                 ebone.parent = None
211             else:
212                 ebone.parent = copy_bones[i]
213
214     return copy_bones
215
216
217 def blend_bone_list(obj, apply_bones, from_bones, to_bones, target_bone=None, target_prop="blend"):
218
219     if obj.mode == 'EDIT':
220         raise Exception("blending cant be called in editmode")
221
222     if len(apply_bones) != len(from_bones):
223         raise Exception("lists differ in length (from -> apply): \n\t%s\n\t%s" % (from_bones, apply_bones))
224     if len(apply_bones) != len(to_bones):
225         raise Exception("lists differ in length (to -> apply): \n\t%s\n\t%s" % (to_bones, apply_bones))
226
227     # setup the blend property
228     if target_bone is None:
229         target_bone = apply_bones[-1] # default to the last bone
230
231     prop_pbone = obj.pose.bones[target_bone]
232     if prop_pbone.get(target_bone, None) is None:
233         prop = rna_idprop_ui_prop_get(prop_pbone, target_prop, create=True)
234         prop_pbone[target_prop] = 0.5
235         prop["soft_min"] = 0.0
236         prop["soft_max"] = 1.0
237
238     driver_path = prop_pbone.path_to_id() + ('["%s"]' % target_prop)
239
240     def blend_target(driver):
241         tar = driver.targets.new()
242         tar.name = target_bone
243         tar.id_type = 'OBJECT'
244         tar.id = obj
245         tar.rna_path = driver_path
246
247     def blend_location(new_pbone, from_bone_name, to_bone_name):
248         con = new_pbone.constraints.new('COPY_LOCATION')
249         con.target = obj
250         con.subtarget = from_bone_name
251
252         con = new_pbone.constraints.new('COPY_LOCATION')
253         con.target = obj
254         con.subtarget = to_bone_name
255
256         fcurve = con.driver_add("influence", 0)
257         driver = fcurve.driver
258         driver.type = 'AVERAGE'
259         fcurve.modifiers.remove(0) # grr dont need a modifier
260
261         blend_target(driver)
262
263     def blend_rotation(new_pbone, from_bone_name, to_bone_name):
264         con = new_pbone.constraints.new('COPY_ROTATION')
265         con.target = obj
266         con.subtarget = from_bone_name
267
268         con = new_pbone.constraints.new('COPY_ROTATION')
269         con.target = obj
270         con.subtarget = to_bone_name
271
272         fcurve = con.driver_add("influence", 0)
273         driver = fcurve.driver
274         driver.type = 'AVERAGE'
275         fcurve.modifiers.remove(0) # grr dont need a modifier
276
277         blend_target(driver)
278
279     for i, new_bone_name in enumerate(apply_bones):
280         from_bone_name = from_bones[i]
281         to_bone_name = to_bones[i]
282
283         # allow skipping some bones by having None in the list
284         if None in (new_bone_name, from_bone_name, to_bone_name):
285             continue
286
287         new_pbone = obj.pose.bones[new_bone_name]
288
289         # if the bone is connected or its location is totally locked then dont add location blending.
290         if not (new_pbone.bone.connected or (False not in new_pbone.lock_location)):
291             blend_location(new_pbone, from_bone_name, to_bone_name)
292
293         if not (False not in new_pbone.lock_rotation): # TODO. 4D chech?
294             blend_rotation(new_pbone, from_bone_name, to_bone_name)
295
296
297 def add_stretch_to(obj, from_name, to_name, name):
298     '''
299     Adds a bone that stretches from one to another
300     '''
301
302     mode_orig = obj.mode
303     bpy.ops.object.mode_set(mode='EDIT')
304
305     arm = obj.data
306     stretch_ebone = arm.edit_bones.new(name)
307     stretch_name = stretch_ebone.name
308     del name
309
310     head = stretch_ebone.head = arm.edit_bones[from_name].head.copy()
311     #tail = stretch_ebone.tail = arm.edit_bones[to_name].head.copy()
312
313     # annoying exception for zero length bones, since its using stretch_to the rest pose doesnt really matter
314     #if (head - tail).length < 0.1:
315     if 1:
316         tail = stretch_ebone.tail = arm.edit_bones[from_name].tail.copy()
317
318
319     # Now for the constraint
320     bpy.ops.object.mode_set(mode='OBJECT')
321
322     stretch_pbone = obj.pose.bones[stretch_name]
323
324     con = stretch_pbone.constraints.new('COPY_LOCATION')
325     con.target = obj
326     con.subtarget = from_name
327
328     con = stretch_pbone.constraints.new('STRETCH_TO')
329     con.target = obj
330     con.subtarget = to_name
331     con.original_length = (head - tail).length
332     con.keep_axis = 'PLANE_X'
333     con.volume = 'NO_VOLUME'
334
335     bpy.ops.object.mode_set(mode=mode_orig)
336     
337     return stretch_name
338
339 def add_pole_target_bone(obj, base_name, name, mode='CROSS'):
340     '''
341     Does not actually create a poll target, just the bone to use as a poll target
342     '''
343     mode_orig = obj.mode
344     bpy.ops.object.mode_set(mode='EDIT')
345
346     arm = obj.data
347
348     poll_ebone = arm.edit_bones.new(base_name + "_poll")
349     base_ebone = arm.edit_bones[base_name]
350     poll_name = poll_ebone.name
351     parent_ebone = base_ebone.parent
352
353     base_head = base_ebone.head.copy()
354     base_tail = base_ebone.tail.copy()
355     base_dir = base_head - base_tail
356
357     parent_head = parent_ebone.head.copy()
358     parent_tail = parent_ebone.tail.copy()
359     parent_dir = parent_head - parent_tail
360
361     distance = (base_dir.length + parent_dir.length)
362     
363     if mode == 'CROSS':
364         # direction from the angle of the joint
365         offset = base_dir.copy().normalize() - parent_dir.copy().normalize()
366         offset.length = distance
367     elif mode == 'ZAVERAGE':
368         # between both bones Z axis
369         z_axis_a = base_ebone.matrix.copy().rotationPart() * Vector(0.0, 0.0, -1.0)
370         z_axis_b = parent_ebone.matrix.copy().rotationPart() * Vector(0.0, 0.0, -1.0)
371         offset = (z_axis_a + z_axis_b).normalize() * distance
372     else:
373         # preset axis
374         offset = Vector(0, 0, 0)
375         if mode[0] == "+":
376             val = distance
377         else:
378             val = - distance
379
380         setattr(offset, mode[1].lower(), val)
381
382     poll_ebone.head = base_head + offset
383     poll_ebone.tail = base_head + (offset * (1.0 - (1.0 / 4.0)))
384
385     bpy.ops.object.mode_set(mode=mode_orig)
386
387     return poll_name
388
389
390 def generate_rig(context, obj_orig, prefix="ORG-", META_DEF=True):
391     from collections import OrderedDict
392
393     global_undo = context.user_preferences.edit.global_undo
394     context.user_preferences.edit.global_undo = False
395
396     bpy.ops.object.mode_set(mode='OBJECT')
397
398     scene = context.scene
399
400     # copy object and data
401     obj_orig.selected = False
402     obj = obj_orig.copy()
403     obj.data = obj_orig.data.copy()
404     scene.objects.link(obj)
405     scene.objects.active = obj
406     obj.selected = True
407     
408     if META_DEF:
409         obj_def = obj_orig.copy()
410         obj_def.data = obj_orig.data.copy()
411         scene.objects.link(obj_def)
412
413     arm = obj.data
414
415     # original name mapping
416     base_names = {}
417     
418     # add all new parentless children to this bone
419     root_bone = None
420
421     bpy.ops.object.mode_set(mode='EDIT')
422     for bone in arm.edit_bones:
423         bone_name = bone.name
424         bone.name = prefix + bone_name
425         base_names[bone.name] = bone_name # new -> old mapping
426     bpy.ops.object.mode_set(mode='OBJECT')
427
428     # key: bone name
429     # value: {type:definition, ...}
430     #    where type is the submodule name - leg, arm etc
431     #    and definition is a list of bone names
432     bone_definitions = {}
433
434     # key: bone name
435     # value: [functions, ...]
436     #    each function is from the module. eg leg.ik, arm.main
437     bone_typeinfos = {}
438
439     # inspect all bones and assign their definitions before modifying
440     for pbone in obj.pose.bones:
441         bone_name = pbone.name
442         bone_type = pbone.get("type", "")
443         if bone_type:
444             bone_type_list = [bt for bt in bone_type.replace(",", " ").split()]
445             
446             # not essential but means running autorig again wont do anything
447             del pbone["type"]
448         else:
449             bone_type_list = []
450             
451         if bone_type_list == ["root"]: # special case!
452             if root_bone:
453                 raise Exception("cant have more then 1 root bone, found '%s' and '%s' to have type==root" % (root_bone, bone_name))
454             root_bone = bone_name
455             bone_type_list[:] = []
456
457         for bone_type in bone_type_list:
458             
459             type_pair = bone_type.split(".")
460
461             # 'leg.ik' will look for an ik function in the leg module
462             # 'leg' will look up leg.main
463             if len(type_pair) == 1:
464                 type_pair = type_pair[0], "main"
465
466             submod_name, func_name = type_pair
467
468             # from rigify import leg
469             submod = __import__(name="%s.%s" % (__package__, submod_name), fromlist=[submod_name])
470             reload(submod)
471
472             bone_def_dict = bone_definitions.setdefault(bone_name, {})
473
474             # Only calculate bone definitions once
475             if submod_name not in bone_def_dict:
476                 metarig_definition_func = getattr(submod, "metarig_definition")
477                 bone_def_dict[submod_name] = metarig_definition_func(obj, bone_name)
478
479
480             bone_typeinfo = bone_typeinfos.setdefault(bone_name, [])
481
482             type_func = getattr(submod, func_name)
483             bone_typeinfo.append((submod_name, type_func))
484
485
486     # sort bones, not needed but gives more pradictable execution which may be useful in rare cases
487     bones_sorted = obj.pose.bones.values()
488     bones_sorted.sort(key=lambda pbone: pbone.name) # first sort by names
489     bones_sorted.sort(key=lambda pbone: - len(pbone.parent_recursive)) # children before parents
490
491     # now we have all the info about bones we can start operating on them
492     # for pbone in obj.pose.bones:
493     for pbone in bones_sorted:
494         bone_name = pbone.name
495
496         if bone_name not in bone_typeinfos:
497             continue
498
499         bone_def_dict = bone_definitions[bone_name]
500
501         # Only blend results from the same submodule, eg.
502         #    leg.ik and arm.fk could not be blended.
503         results = OrderedDict()
504
505         for submod_name, type_func in bone_typeinfos[bone_name]:
506             # this bones definition of the current typeinfo
507             definition = bone_def_dict[submod_name]
508
509             bpy.ops.object.mode_set(mode='EDIT')
510             ret = type_func(obj, definition, base_names)
511             bpy.ops.object.mode_set(mode='OBJECT')
512
513             if ret:
514                 result_submod = results.setdefault(submod_name, [])
515
516                 if result_submod and len(result_submod[-1]) != len(ret):
517                     raise Exception("bone lists not compatible: %s, %s" % (result_submod[-1], ret))
518
519                 result_submod.append(ret)
520
521         for result_submod in results.values():
522             # blend 2 chains
523             definition = bone_def_dict[submod_name]
524
525             if len(result_submod) == 2:
526                 blend_bone_list(obj, definition, result_submod[0], result_submod[1], target_bone=bone_name)
527
528     if root_bone:
529         # assign all new parentless bones to this
530         
531         bpy.ops.object.mode_set(mode='EDIT')
532         root_ebone = arm.edit_bones[root_bone]
533         for ebone in arm.edit_bones:
534             if ebone.parent is None and ebone.name not in base_names:
535                 ebone.connected = False
536                 ebone.parent = root_ebone
537         bpy.ops.object.mode_set(mode='OBJECT')
538         
539
540     if META_DEF:
541         # for pbone in obj_def.pose.bones:
542         for bone_name, bone_name_new in base_names.items():
543             #pbone_from = bone_name
544             pbone = obj_def.pose.bones[bone_name_new]
545             
546             con = pbone.constraints.new('COPY_ROTATION')
547             con.target = obj
548             con.subtarget = bone_name
549
550             if not pbone.bone.connected:
551                 con = pbone.constraints.new('COPY_LOCATION')
552                 con.target = obj
553                 con.subtarget = bone_name
554
555     # Only for demo'ing
556
557     # obj.restrict_view = True
558     obj.data.draw_axes = False
559
560     context.user_preferences.edit.global_undo = global_undo
561
562     return obj
563
564
565 def write_meta_rig(obj, func_name="metarig_template"):
566     ''' Must be in editmode
567     '''
568     code = []
569
570     code.append("def %s():" % func_name)
571     code.append("    # generated by rigify.write_meta_rig")
572     bpy.ops.object.mode_set(mode='EDIT')
573     code.append("    bpy.ops.object.mode_set(mode='EDIT')")
574
575     code.append("    obj = bpy.context.object")
576     code.append("    arm = obj.data")
577
578     arm = obj.data
579     # write parents first
580     bones = [(len(bone.parent_recursive), bone.name) for bone in arm.edit_bones]
581     bones.sort(key=lambda item: item[0])
582     bones = [item[1] for item in bones]
583
584
585     for bone_name in bones:
586         bone = arm.edit_bones[bone_name]
587         code.append("    bone = arm.edit_bones.new('%s')" % bone.name)
588         code.append("    bone.head[:] = %.4f, %.4f, %.4f" % bone.head.toTuple(4))
589         code.append("    bone.tail[:] = %.4f, %.4f, %.4f" % bone.tail.toTuple(4))
590         code.append("    bone.roll = %.4f" % bone.roll)
591         code.append("    bone.connected = %s" % str(bone.connected))
592         if bone.parent:
593             code.append("    bone.parent = arm.edit_bones['%s']" % bone.parent.name)
594
595     bpy.ops.object.mode_set(mode='OBJECT')
596     code.append("")
597     code.append("    bpy.ops.object.mode_set(mode='OBJECT')")
598
599     for bone_name in bones:
600         pbone = obj.pose.bones[bone_name]
601         pbone_written = False
602
603         # Only 1 level of props, simple types supported
604         for key, value in pbone.items():
605             if key.startswith("_"):
606                 continue
607
608             if type(value) not in (float, str, int):
609                 print("Unsupported ID Prop:", str((key, value)))
610                 continue
611
612             if type(value) == str:
613                 value = "'" + value + "'"
614
615             if not pbone_written: # only write bones we need
616                 code.append("    pbone = obj.pose.bones['%s']" % bone_name)
617
618             code.append("    pbone['%s'] = %s" % (key, value))
619
620     return "\n".join(code)
621
622
623 def generate_test(context):
624     import os
625     new_objects = []
626
627     scene = context.scene
628
629     def create_empty_armature(name):
630         obj_new = bpy.data.add_object('ARMATURE', name)
631         armature = bpy.data.add_armature(name)
632         obj_new.data = armature
633         scene.objects.link(obj_new)
634         scene.objects.active = obj_new
635
636     files = os.listdir(os.path.dirname(__file__))
637     for f in files:
638         if f.startswith("_"):
639             continue
640
641         if not f.endswith(".py"):
642             continue
643
644         module_name = f[:-3]
645         submodule = __import__(name="%s.%s" % (__package__, module_name), fromlist=[module_name])
646
647         metarig_template = getattr(submodule, "metarig_template", None)
648
649         if metarig_template:
650             create_empty_armature("meta_" + module_name) # sets active
651             metarig_template()
652             obj = context.object
653             obj_new = generate_rig(context, obj)
654
655             new_objects.append((obj, obj_new))
656         else:
657             print("note: rig type '%s' has no metarig_template(), can't test this", module_name)
658
659     return new_objects
660
661
662 def generate_test_all(context):
663     import rigify
664     import graphviz_export
665     import os
666     reload(rigify)
667     reload(graphviz_export)
668
669     new_objects = rigify.generate_test(context)
670
671     base_name = os.path.splitext(bpy.data.filename)[0]
672     for obj, obj_new in new_objects:
673         for obj in (obj, obj_new):
674             fn = base_name + "-" + bpy.utils.clean_name(obj.name)
675
676             path_dot = fn + ".dot"
677             path_png = fn + ".png"
678             saved = graphviz_export.graph_armature(obj, path_dot, CONSTRAINTS=True, DRIVERS=True)
679
680             #if saved:
681             #    os.system("dot -Tpng %s > %s; eog %s" % (path_dot, path_png, path_png))
682
683     i = 0
684     for obj, obj_new in new_objects:
685         obj.data.drawtype = 'STICK'
686         obj.location[1] += i
687         obj_new.location[1] += i
688         obj_new.selected = False
689         obj.selected = True
690         i += 4
691
692
693 if __name__ == "__main__":
694     generate_rig(bpy.context, bpy.context.object)