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