2eed5bbb74e694af57347830ca1a8049356c819e
[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., 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 Vector
23
24 # TODO, have these in a more general module
25 from rna_prop_ui import rna_idprop_ui_prop_get
26 SPECIAL_TYPES = "root",
27 LAYER_TYPES = "main", "extra", "ik", "fk"
28
29 ORG_LAYERS = [n == 31 for n in range(0, 32)]
30 MCH_LAYERS = [n == 30 for n in range(0, 32)]
31 DEF_LAYERS = [n == 29 for n in range(0, 32)]
32 ROOT_LAYERS = [n == 28 for n in range(0, 32)]
33
34 ORG_PREFIX = "ORG-"
35 MCH_PREFIX = "MCH-"
36 DEF_PREFIX = "DEF-"
37
38 WGT_PREFIX = "WGT-"
39
40
41 class RigifyError(Exception):
42     """Exception raised for errors in the metarig.
43     """
44
45     def __init__(self, message):
46         self.message = message
47
48     def __str__(self):
49         return repr(self.message)
50
51
52 def submodule_func_from_type(bone_type):
53     type_pair = bone_type.split(".")
54
55     # 'leg.ik' will look for an ik function in the leg module
56     # 'leg' will look up leg.main
57     if len(type_pair) == 1:
58         type_pair = type_pair[0], "main"
59
60     type_name, func_name = type_pair
61
62     # from rigify import leg
63     try:
64         submod = __import__(name="%s.%s" % (__package__, type_name), fromlist=[type_name])
65     except ImportError:
66         raise RigifyError("python module for type '%s' not found" % type_name)
67
68     reload(submod)
69     return type_name, submod, getattr(submod, func_name)
70
71
72 def get_submodule_types():
73     import os
74     submodules = []
75     files = os.listdir(os.path.dirname(__file__))
76     for f in files:
77         if not f.startswith("_") and f.endswith(".py"):
78             submodules.append(f[:-3])
79
80     return sorted(submodules)
81
82
83 def get_bone_type_options(pbone, type_name):
84     options = {}
85     bone_name = pbone.name
86     for key, value in pbone.items():
87         key_pair = key.rsplit(".")
88         # get all bone properties
89         """"
90         if key_pair[0] == type_name:
91             if len(key_pair) != 2:
92                 raise RigifyError("option error for bone '%s', property name was not a pair '%s'" % (bone_name, key_pair))
93             options[key_pair[1]] = value
94         """
95         options[key] = value
96
97     return options
98
99
100 def get_layer_dict(options):
101     '''
102     Extracts layer info from a bone options dict
103     defaulting to the layer index if not set.
104     '''
105     layer_default = [False] * 32
106     result = {}
107     for i, layer_type in enumerate(LAYER_TYPES):
108         # no matter if its not defined
109         layer_index = options.get("layer_" + layer_type, i + 2)
110         layer = layer_default[:]
111         layer[layer_index-1] = True
112         result[layer_type] = layer
113     return result
114
115
116 def validate_rig(context, obj):
117     '''
118     Makes no changes
119     only runs the metarig definitions and reports errors
120     '''
121     type_found = False
122
123     for pbone in obj.pose.bones:
124         bone_name = pbone.name
125         bone_type = pbone.get("type", "")
126
127         if bone_type:
128             bone_type_list = [bt for bt in bone_type.replace(",", " ").split()]
129         else:
130             bone_type_list = []
131
132         for bone_type in bone_type_list:
133             if bone_type.split(".")[0] in SPECIAL_TYPES:
134                 continue
135
136             type_name, submod, type_func = submodule_func_from_type(bone_type)
137             reload(submod)
138             submod.metarig_definition(obj, bone_name)
139             type_found = True
140
141             get_bone_type_options(pbone, bone_type)
142
143         # missing, - check for duplicate root bone.
144
145     if not type_found:
146         raise RigifyError("This rig has no 'type' properties defined on any pose bones, nothing to do")
147
148
149 def generate_rig(context, obj_orig, prefix="ORG-", META_DEF=True):
150     '''
151     Main function for generating
152     '''
153     from collections import OrderedDict
154     import rigify_utils
155     reload(rigify_utils)
156
157     print("Begin...")
158
159     # Not needed but catches any errors before duplicating
160     validate_rig(context, obj_orig)
161
162     use_global_undo = context.user_preferences.edit.use_global_undo
163     context.user_preferences.edit.use_global_undo = False
164     mode_orig = context.mode
165     rest_backup = obj_orig.data.pose_position
166     obj_orig.data.pose_position = 'REST'
167
168     bpy.ops.object.mode_set(mode='OBJECT')
169
170     scene = context.scene
171
172     # Check if the generated rig already exists, so we can
173     # regenerate in the same object.  If not, create a new
174     # object to generate the rig in.
175     print("Fetch rig.")
176     try:
177         name = obj_orig["rig_object_name"]
178     except KeyError:
179         name = "rig"
180
181     try:
182         obj = scene.objects[name]
183     except KeyError:
184         obj = bpy.data.objects.new(name, bpy.data.armatures.new(name))
185         scene.objects.link(obj)
186
187     obj.data.pose_position = 'POSE'
188
189     # Get rid of anim data in case the rig already existed
190     print("Clear rig animation data.")
191     obj.animation_data_clear()
192
193     # Select generated rig object
194     obj_orig.select = False
195     obj.select = True
196     scene.objects.active = obj
197
198     # Remove all bones from the generated rig armature.
199     bpy.ops.object.mode_set(mode='EDIT')
200     for bone in obj.data.edit_bones:
201         obj.data.edit_bones.remove(bone)
202     bpy.ops.object.mode_set(mode='OBJECT')
203
204     # Create temporary duplicates for merging
205     temp_rig_1 = obj_orig.copy()
206     temp_rig_1.data = obj_orig.data.copy()
207     scene.objects.link(temp_rig_1)
208
209     temp_rig_2 = obj_orig.copy()
210     temp_rig_2.data = obj.data
211     scene.objects.link(temp_rig_2)
212
213     # Select the temp rigs for merging
214     for objt in scene.objects:
215         objt.select = False # deselect all objects
216     temp_rig_1.select = True
217     temp_rig_2.select = True
218     scene.objects.active = temp_rig_2
219
220     # Merge the temporary rigs
221     bpy.ops.object.join(context)
222
223     # Delete the second temp rig
224     bpy.ops.object.delete()
225
226     # Select the generated rig
227     for objt in scene.objects:
228         objt.select = False # deselect all objects
229     obj.select = True
230     scene.objects.active = obj
231
232     # Copy over the pose_bone properties
233     for bone in obj_orig.pose.bones:
234         bone_gen = obj.pose.bones[bone.name]
235
236         # Rotation mode and transform locks
237         bone_gen.rotation_mode = bone.rotation_mode
238         bone_gen.lock_rotation = tuple(bone.lock_rotation)
239         bone_gen.lock_rotation_w = bone.lock_rotation_w
240         bone_gen.lock_rotations_4d = bone.lock_rotations_4d
241         bone_gen.lock_location = tuple(bone.lock_location)
242         bone_gen.lock_scale = tuple(bone.lock_scale)
243
244         # Custom properties
245         for prop in bone.keys():
246             bone_gen[prop] = bone[prop]
247
248     # Copy over bone properties
249     for bone in obj_orig.data.bones:
250         bone_gen = obj.data.bones[bone.name]
251
252         # B-bone stuff
253         bone_gen.bbone_segments = bone.bbone_segments
254         bone_gen.bbone_in = bone.bbone_in
255         bone_gen.bbone_out = bone.bbone_out
256
257
258     # Create proxy deformation rig
259     # TODO: remove this
260     if META_DEF:
261         obj_def = obj_orig.copy()
262         obj_def.data = obj_orig.data.copy()
263         scene.objects.link(obj_def)
264
265     scene.update()
266     print("On to the real work.")
267
268     arm = obj.data
269
270     # prepend the ORG prefix to the bones, and create the base_names mapping
271     base_names = {}
272     bpy.ops.object.mode_set(mode='EDIT')
273     for bone in arm.edit_bones:
274         bone_name = bone.name
275         bone.name = ORG_PREFIX + bone_name
276         base_names[bone.name] = bone_name
277
278     # create root_bone
279     bpy.ops.object.mode_set(mode='EDIT')
280     edit_bone = obj.data.edit_bones.new("root")
281     root_bone = edit_bone.name
282     edit_bone.head = (0.0, 0.0, 0.0)
283     edit_bone.tail = (0.0, 1.0, 0.0)
284     edit_bone.roll = 0.0
285     edit_bone.layers = ROOT_LAYERS
286     bpy.ops.object.mode_set(mode='OBJECT')
287
288     # key: bone name
289     # value: {type:definition, ...}
290     #    where type is the submodule name - leg, arm etc
291     #    and definition is a list of bone names
292     bone_definitions = {}
293
294     # key: bone name
295     # value: [functions, ...]
296     #    each function is from the module. eg leg.ik, arm.main
297     bone_typeinfos = {}
298
299     # key: bone name
300     # value: [new_bone_name, ...]
301     #   where each bone with a 'type' stores a list of bones that it created
302     #   ...needed so we can override the root parent
303     bone_genesis = {}
304
305
306     # inspect all bones and assign their definitions before modifying
307     for pbone in obj.pose.bones:
308         bone_name = pbone.name
309         bone_type = pbone.get("type", "")
310         if bone_type:
311             bone_type_list = [bt for bt in bone_type.replace(",", " ").split()]
312
313             # not essential but means running autorig again wont do anything
314             del pbone["type"]
315         else:
316             bone_type_list = []
317
318         for bone_type in bone_type_list:
319             type_name, submod, type_func = submodule_func_from_type(bone_type)
320             reload(submod)
321
322             bone_def_dict = bone_definitions.setdefault(bone_name, {})
323
324             # Only calculate bone definitions once
325             if type_name not in bone_def_dict:
326                 bone_def_dict[type_name] = submod.metarig_definition(obj, bone_name)
327
328             bone_typeinfo = bone_typeinfos.setdefault(bone_name, [])
329             bone_typeinfo.append((type_name, type_func))
330
331
332     # sort bones, not needed but gives more pradictable execution which may be useful in rare cases
333     bones_sorted = obj.pose.bones.values()
334     bones_sorted.sort(key=lambda pbone: pbone.name) # first sort by names
335     bones_sorted.sort(key=lambda pbone: len(pbone.parent_recursive)) # parents before children
336
337     # now we have all the info about bones we can start operating on them
338     # for pbone in obj.pose.bones:
339     for pbone in bones_sorted:
340         bone_name = pbone.name
341         print(bone_name)
342         if bone_name not in bone_typeinfos:
343             continue
344
345         bone_def_dict = bone_definitions[bone_name]
346
347         # Only blend results from the same submodule, eg.
348         #    leg.ik and arm.fk could not be blended.
349         results = OrderedDict()
350
351         bone_names_pre = {bone.name for bone in arm.bones}
352
353         for type_name, type_func in bone_typeinfos[bone_name]:
354             print("    " + type_name)
355             # this bones definition of the current typeinfo
356             definition = bone_def_dict[type_name]
357             options = get_bone_type_options(pbone, type_name)
358
359             bpy.ops.object.mode_set(mode='EDIT')
360             ret = type_func(obj, definition, base_names, options)
361             bpy.ops.object.mode_set(mode='OBJECT')
362
363             if ret:
364                 result_submod = results.setdefault(type_name, [])
365
366                 if result_submod and len(result_submod[-1]) != len(ret):
367                     raise RigifyError("bone lists not compatible: %s, %s" % (result_submod[-1], ret))
368
369                 result_submod.append(ret)
370
371         for result_submod in results.values():
372             # blend 2 chains
373             definition = bone_def_dict[type_name]
374
375             if len(result_submod) == 2:
376                 blend_bone_list(obj, definition, result_submod[0], result_submod[1], target_bone=bone_name)
377
378
379         bone_names_post = {bone.name for bone in arm.bones}
380
381         # Store which bones were created from this one
382         bone_genesis[bone_name] = list(bone_names_post - bone_names_pre)
383
384     # need a reverse lookup on bone_genesis so as to know immediately
385     # where a bone comes from
386     bone_genesis_reverse = {}
387     '''
388     for bone_name, bone_children in bone_genesis.items():
389         for bone_child_name in bone_children:
390             bone_genesis_reverse[bone_child_name] = bone_name
391     '''
392
393
394     if root_bone:
395         # assign all new parentless bones to this
396
397         bpy.ops.object.mode_set(mode='EDIT')
398         root_ebone = arm.edit_bones[root_bone]
399         for ebone in arm.edit_bones:
400             bone_name = ebone.name
401             if ebone.parent is None:
402                 ebone.parent = root_ebone
403             '''
404             if ebone.parent is None and bone_name not in base_names:
405                 # check for override
406                 bone_creator = bone_genesis_reverse[bone_name]
407                 pbone_creator = obj.pose.bones[bone_creator]
408                 root_bone_override = pbone_creator.get("root", "")
409
410                 if root_bone_override:
411                     root_ebone_tmp = arm.edit_bones[root_bone_override]
412                 else:
413                     root_ebone_tmp = root_ebone
414
415                 ebone.use_connect = False
416                 ebone.parent = root_ebone_tmp
417             '''
418
419         bpy.ops.object.mode_set(mode='OBJECT')
420
421
422     if META_DEF:
423         # for pbone in obj_def.pose.bones:
424         for bone_name, bone_name_new in base_names.items():
425             #pbone_from = bone_name
426             pbone = obj_def.pose.bones[bone_name_new]
427
428             con = pbone.constraints.new('COPY_ROTATION')
429             con.target = obj
430             con.subtarget = bone_name
431
432             if not pbone.bone.use_connect:
433                 con = pbone.constraints.new('COPY_LOCATION')
434                 con.target = obj
435                 con.subtarget = bone_name
436
437         # would be 'REST' from when copied
438         obj_def.data.pose_position = 'POSE'
439
440     # todo - make a more generic system?
441     layer_tot = [False] * 32
442     layer_last = layer_tot[:]
443     layer_last[31] = True
444     layer_second_last = layer_tot[:]
445     layer_second_last[30] = True
446
447     for bone_name, bone in arm.bones.items():
448         bone.use_deform = False  # Non DEF bones shouldn't deform
449         if bone_name.startswith(ORG_PREFIX):
450             bone.layers = ORG_LAYERS
451         elif bone_name.startswith(MCH_PREFIX): # XXX fixme
452             bone.layers = MCH_LAYERS
453         elif bone_name.startswith(DEF_PREFIX): # XXX fixme
454             bone.layers = DEF_LAYERS
455             bone.use_deform = True
456         else:
457             # Assign bone appearance if there is a widget for it
458             obj.pose.bones[bone_name].custom_shape = context.scene.objects.get(WGT_PREFIX + bone_name)
459
460         layer_tot[:] = [max(lay) for lay in zip(layer_tot, bone.layers)]
461
462     # Only for demo'ing
463     layer_show = [a and not (b or c or d) for a, b, c, d in zip(layer_tot, ORG_LAYERS, MCH_LAYERS, DEF_LAYERS)]
464     arm.layers = layer_show
465
466
467     # obj.hide = True
468     obj.data.show_axes = False
469
470     bpy.ops.object.mode_set(mode=mode_orig)
471     obj_orig.data.pose_position = rest_backup
472     obj.data.pose_position = 'POSE'
473     obj_orig.data.pose_position = 'POSE'
474     context.user_preferences.edit.use_global_undo = use_global_undo
475
476     print("Done.\n")
477
478     return obj
479
480
481 def generate_test(context, metarig_type="", GENERATE_FINAL=True):
482     import os
483     new_objects = []
484
485     scene = context.scene
486
487     def create_empty_armature(name):
488         armature = bpy.data.armatures.new(name)
489         obj_new = bpy.data.objects.new(name, armature)
490         scene.objects.link(obj_new)
491         scene.objects.active = obj_new
492         for obj in scene.objects:
493             obj.select = False
494         obj_new.select = True
495
496     for module_name in get_submodule_types():
497         if (metarig_type and module_name != metarig_type):
498             continue
499
500         # XXX workaround!, problem with updating the pose matrix.
501         if module_name == "delta":
502             continue
503
504         type_name, submodule, func = submodule_func_from_type(module_name)
505
506         metarig_template = getattr(submodule, "metarig_template", None)
507
508         if metarig_template:
509             create_empty_armature("meta_" + module_name) # sets active
510             metarig_template()
511             obj = context.active_object
512             obj.location = scene.cursor_location
513
514             if GENERATE_FINAL:
515                 obj_new = generate_rig(context, obj)
516                 new_objects.append((obj, obj_new))
517             else:
518                 new_objects.append((obj, None))
519         else:
520             print("note: rig type '%s' has no metarig_template(), can't test this" % module_name)
521
522     return new_objects
523
524
525 def generate_test_all(context, GRAPH=False):
526     import rigify
527     import rigify_utils
528     import graphviz_export
529     import os
530     reload(rigify)
531     reload(rigify_utils)
532     reload(graphviz_export)
533
534     new_objects = rigify.generate_test(context)
535
536     if GRAPH:
537         base_name = os.path.splitext(bpy.data.filepath)[0]
538         for obj, obj_new in new_objects:
539             for obj in (obj, obj_new):
540                 fn = base_name + "-" + bpy.path.clean_name(obj.name)
541
542                 path_dot = fn + ".dot"
543                 path_png = fn + ".png"
544                 saved = graphviz_export.graph_armature(obj, path_dot, CONSTRAINTS=True, DRIVERS=True)
545
546                 #if saved:
547                 #    os.system("dot -Tpng %s > %s; eog %s" % (path_dot, path_png, path_png))
548
549     i = 0
550     for obj, obj_new in new_objects:
551         obj.data.drawtype = 'STICK'
552         obj.location[1] += i
553         obj_new.location[1] += i
554         obj_new.select = False
555         obj.select = True
556         i += 4
557
558
559 if __name__ == "__main__":
560     generate_rig(bpy.context, bpy.context.active_object)