* added an armature submenu where python defined armatures can go.
[blender.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 SPECIAL_TYPES = "root",
27
28 class RigifyError(Exception):
29     """Exception raised for errors in the metarig.
30     """
31     def __init__(self, message):
32         self.message = message
33     def __str__(self):
34         return repr(self.message)
35
36 def submodule_func_from_type(bone_type):
37     type_pair = bone_type.split(".")
38
39     # 'leg.ik' will look for an ik function in the leg module
40     # 'leg' will look up leg.main
41     if len(type_pair) == 1:
42         type_pair = type_pair[0], "main"
43
44     submod_name, func_name = type_pair
45
46     # from rigify import leg
47     try:
48         submod = __import__(name="%s.%s" % (__package__, submod_name), fromlist=[submod_name])
49     except ImportError:
50         raise RigifyError("python module for type '%s' not found" % submod_name)
51         
52     reload(submod)
53     return submod, getattr(submod, func_name)
54
55
56 def submodule_types():
57     import os
58     submodules = []
59     files = os.listdir(os.path.dirname(__file__))
60     for f in files:
61         if not f.startswith("_") and f.endswith(".py"):
62             submodules.append(f[:-3])
63     
64     return sorted(submodules)
65
66
67 def validate_rig(context, obj):
68     '''
69     Makes no changes
70     only runs the metarig definitions and reports errors
71     '''
72     type_found = False
73     
74     for pbone in obj.pose.bones:
75         bone_name = pbone.name
76         bone_type = pbone.get("type", "")
77
78         if bone_type:
79             bone_type_list = [bt for bt in bone_type.replace(",", " ").split()]
80         else:
81             bone_type_list = []
82
83         for bone_type in bone_type_list:
84             if bone_type.split(".")[0] in SPECIAL_TYPES:
85                 continue
86
87             submod, type_func = submodule_func_from_type(bone_type)
88             reload(submod)
89             submod.metarig_definition(obj, bone_name)
90             type_found = True
91         
92         # missing, - check for duplicate root bone.
93     
94     if not type_found:
95         raise RigifyError("This rig has no 'type' properties defined on any pose bones, nothing to do")
96
97
98 def generate_rig(context, obj_orig, prefix="ORG-", META_DEF=True):
99     '''
100     Main function for generating 
101     '''
102     from collections import OrderedDict
103     import rigify_utils
104     reload(rigify_utils)
105     
106     # Not needed but catches any errors before duplicating
107     validate_rig(context, obj_orig)
108
109     global_undo = context.user_preferences.edit.global_undo
110     context.user_preferences.edit.global_undo = False
111     mode_orig = context.mode
112     rest_backup = obj_orig.data.pose_position
113     obj_orig.data.pose_position = 'REST'
114     
115     
116     bpy.ops.object.mode_set(mode='OBJECT')
117
118     scene = context.scene
119
120     # copy object and data
121     obj_orig.selected = False
122     obj = obj_orig.copy()
123     obj.data = obj_orig.data.copy()
124     scene.objects.link(obj)
125     scene.objects.active = obj
126     obj.selected = True
127
128     if META_DEF:
129         obj_def = obj_orig.copy()
130         obj_def.data = obj_orig.data.copy()
131         scene.objects.link(obj_def)
132
133     arm = obj.data
134
135     # original name mapping
136     base_names = {}
137     
138     # add all new parentless children to this bone
139     root_bone = None
140
141     bpy.ops.object.mode_set(mode='EDIT')
142     for bone in arm.edit_bones:
143         bone_name = bone.name
144         bone.name = prefix + bone_name
145         base_names[bone.name] = bone_name # new -> old mapping
146     bpy.ops.object.mode_set(mode='OBJECT')
147
148     # key: bone name
149     # value: {type:definition, ...}
150     #    where type is the submodule name - leg, arm etc
151     #    and definition is a list of bone names
152     bone_definitions = {}
153
154     # key: bone name
155     # value: [functions, ...]
156     #    each function is from the module. eg leg.ik, arm.main
157     bone_typeinfos = {}
158     
159     # key: bone name
160     # value: [new_bone_name, ...]
161     #   where each bone with a 'type' stores a list of bones that it created
162     #   ...needed so we can override the root parent
163     bone_genesis = {}
164
165     # inspect all bones and assign their definitions before modifying
166     for pbone in obj.pose.bones:
167         bone_name = pbone.name
168         bone_type = pbone.get("type", "")
169         if bone_type:
170             bone_type_list = [bt for bt in bone_type.replace(",", " ").split()]
171             
172             # not essential but means running autorig again wont do anything
173             del pbone["type"]
174         else:
175             bone_type_list = []
176             
177         if bone_type_list == ["root"]: # special case!
178             if root_bone:
179                 raise Exception("cant have more then 1 root bone, found '%s' and '%s' to have type==root" % (root_bone, bone_name))
180             root_bone = bone_name
181             bone_type_list[:] = []
182
183         for bone_type in bone_type_list:
184             submod, type_func = submodule_func_from_type(bone_type)
185             reload(submod)
186             submod_name = submod.__name__
187             
188             bone_def_dict = bone_definitions.setdefault(bone_name, {})
189
190             # Only calculate bone definitions once
191             if submod_name not in bone_def_dict:
192                 bone_def_dict[submod_name] = submod.metarig_definition(obj, bone_name)
193
194             bone_typeinfo = bone_typeinfos.setdefault(bone_name, [])
195             bone_typeinfo.append((submod_name, type_func))
196
197
198     # sort bones, not needed but gives more pradictable execution which may be useful in rare cases
199     bones_sorted = obj.pose.bones.values()
200     bones_sorted.sort(key=lambda pbone: pbone.name) # first sort by names
201     bones_sorted.sort(key=lambda pbone: - len(pbone.parent_recursive)) # children before parents
202
203     # now we have all the info about bones we can start operating on them
204     # for pbone in obj.pose.bones:
205     for pbone in bones_sorted:
206         bone_name = pbone.name
207
208         if bone_name not in bone_typeinfos:
209             continue
210
211         bone_def_dict = bone_definitions[bone_name]
212
213         # Only blend results from the same submodule, eg.
214         #    leg.ik and arm.fk could not be blended.
215         results = OrderedDict()
216         
217         bone_names_pre = set([bone.name for bone in arm.bones])
218
219         for submod_name, type_func in bone_typeinfos[bone_name]:
220             # this bones definition of the current typeinfo
221             definition = bone_def_dict[submod_name]
222
223             bpy.ops.object.mode_set(mode='EDIT')
224             ret = type_func(obj, definition, base_names)
225             bpy.ops.object.mode_set(mode='OBJECT')
226
227             if ret:
228                 result_submod = results.setdefault(submod_name, [])
229
230                 if result_submod and len(result_submod[-1]) != len(ret):
231                     raise Exception("bone lists not compatible: %s, %s" % (result_submod[-1], ret))
232
233                 result_submod.append(ret)
234
235         for result_submod in results.values():
236             # blend 2 chains
237             definition = bone_def_dict[submod_name]
238
239             if len(result_submod) == 2:
240                 blend_bone_list(obj, definition, result_submod[0], result_submod[1], target_bone=bone_name)
241
242
243         bone_names_post = set([bone.name for bone in arm.bones])
244         
245         # Store which bones were created from this one
246         bone_genesis[bone_name] = list(bone_names_post - bone_names_pre)
247     
248     # need a reverse lookup on bone_genesis so as to know immediately
249     # where a bone comes from
250     bone_genesis_reverse = {}
251     for bone_name, bone_children in bone_genesis.items():
252         for bone_child_name in bone_children:
253             bone_genesis_reverse[bone_child_name] = bone_name
254     
255
256     if root_bone:
257         # assign all new parentless bones to this
258         
259         bpy.ops.object.mode_set(mode='EDIT')
260         root_ebone = arm.edit_bones[root_bone]
261         for ebone in arm.edit_bones:
262             bone_name = ebone.name
263             if ebone.parent is None and bone_name not in base_names:
264                 # check for override
265                 bone_creator = bone_genesis_reverse[bone_name]
266                 pbone_creator = obj.pose.bones[bone_creator]
267                 root_bone_override = pbone_creator.get("root", "")
268
269                 if root_bone_override:
270                     root_ebone_tmp = arm.edit_bones[root_bone_override]
271                 else:
272                     root_ebone_tmp = root_ebone
273                 
274                 ebone.connected = False
275                 ebone.parent = root_ebone_tmp
276
277         bpy.ops.object.mode_set(mode='OBJECT')
278         
279
280     if META_DEF:
281         # for pbone in obj_def.pose.bones:
282         for bone_name, bone_name_new in base_names.items():
283             #pbone_from = bone_name
284             pbone = obj_def.pose.bones[bone_name_new]
285             
286             con = pbone.constraints.new('COPY_ROTATION')
287             con.target = obj
288             con.subtarget = bone_name
289
290             if not pbone.bone.connected:
291                 con = pbone.constraints.new('COPY_LOCATION')
292                 con.target = obj
293                 con.subtarget = bone_name
294
295         # would be 'REST' from when copied
296         obj_def.data.pose_position = 'POSE'
297
298     # Only for demo'ing
299
300     # obj.restrict_view = True
301     obj.data.draw_axes = False
302
303     bpy.ops.object.mode_set(mode=mode_orig)
304     obj_orig.data.pose_position = rest_backup
305     obj.data.pose_position = 'POSE'
306     context.user_preferences.edit.global_undo = global_undo
307     
308     
309     return obj
310
311
312 def generate_test(context, metarig_type="", GENERATE_FINAL=True):
313     import os
314     new_objects = []
315
316     scene = context.scene
317
318     def create_empty_armature(name):
319         obj_new = bpy.data.add_object('ARMATURE', name)
320         armature = bpy.data.add_armature(name)
321         obj_new.data = armature
322         scene.objects.link(obj_new)
323         scene.objects.active = obj_new
324         for obj in scene.objects:
325             obj.selected = False
326         obj_new.selected = True
327
328     for module_name in submodule_types():
329         if (metarig_type and module_name != metarig_type):
330             continue
331         
332         submodule, func = submodule_func_from_type(module_name)
333
334         metarig_template = getattr(submodule, "metarig_template", None)
335
336         if metarig_template:
337             create_empty_armature("meta_" + module_name) # sets active
338             metarig_template()
339             obj = context.active_object
340             obj.location = scene.cursor_location
341             
342             if GENERATE_FINAL:
343                 obj_new = generate_rig(context, obj)
344                 new_objects.append((obj, obj_new))
345             else:
346                 new_objects.append((obj, None))
347         else:
348             print("note: rig type '%s' has no metarig_template(), can't test this", module_name)
349
350     return new_objects
351
352
353 def generate_test_all(context, GRAPH=False):
354     import rigify
355     import rigify_utils
356     import graphviz_export
357     import os
358     reload(rigify)
359     reload(rigify_utils)
360     reload(graphviz_export)
361
362     new_objects = rigify.generate_test(context)
363     
364     if GRAPH:
365         base_name = os.path.splitext(bpy.data.filename)[0]
366         for obj, obj_new in new_objects:
367             for obj in (obj, obj_new):
368                 fn = base_name + "-" + bpy.utils.clean_name(obj.name)
369
370                 path_dot = fn + ".dot"
371                 path_png = fn + ".png"
372                 saved = graphviz_export.graph_armature(obj, path_dot, CONSTRAINTS=True, DRIVERS=True)
373
374                 #if saved:
375                 #    os.system("dot -Tpng %s > %s; eog %s" % (path_dot, path_png, path_png))
376
377     i = 0
378     for obj, obj_new in new_objects:
379         obj.data.drawtype = 'STICK'
380         obj.location[1] += i
381         obj_new.location[1] += i
382         obj_new.selected = False
383         obj.selected = True
384         i += 4
385
386
387 if __name__ == "__main__":
388     generate_rig(bpy.context, bpy.context.active_object)