Update for changes in Blender's API
[blender-addons.git] / rigify / generate.py
blobe2492708a00ae4735c11b84ff6bfc0dee6b89f6b
1 #====================== BEGIN GPL LICENSE BLOCK ======================
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.
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.
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.
17 #======================= END GPL LICENSE BLOCK ========================
19 # <pep8 compliant>
21 import bpy
22 import re
23 import time
24 import traceback
25 import sys
26 from rna_prop_ui import rna_idprop_ui_prop_get
28 from .utils import MetarigError, new_bone, get_rig_type
29 from .utils import ORG_PREFIX, MCH_PREFIX, DEF_PREFIX, WGT_PREFIX, ROOT_NAME, make_original_name
30 from .utils import RIG_DIR
31 from .utils import create_root_widget
32 from .utils import random_id
33 from .utils import copy_attributes
34 from .utils import gamma_correct
35 from .rig_ui_template import UI_SLIDERS, layers_ui, UI_REGISTER
38 RIG_MODULE = "rigs"
39 ORG_LAYER = [n == 31 for n in range(0, 32)] # Armature layer that original bones should be moved to.
40 MCH_LAYER = [n == 30 for n in range(0, 32)] # Armature layer that mechanism bones should be moved to.
41 DEF_LAYER = [n == 29 for n in range(0, 32)] # Armature layer that deformation bones should be moved to.
42 ROOT_LAYER = [n == 28 for n in range(0, 32)] # Armature layer that root bone should be moved to.
43 WGT_LAYERS = [x == 19 for x in range(0, 20)] # Widgets go on the last scene layer.
46 class Timer:
47 def __init__(self):
48 self.timez = time.time()
50 def tick(self, string):
51 t = time.time()
52 print(string + "%.3f" % (t - self.timez))
53 self.timez = t
56 # TODO: generalize to take a group as input instead of an armature.
57 def generate_rig(context, metarig):
58 """ Generates a rig from a metarig.
60 """
61 t = Timer()
63 # Random string with time appended so that
64 # different rigs don't collide id's
65 rig_id = random_id(16)
67 # Initial configuration
68 # mode_orig = context.mode # UNUSED
69 rest_backup = metarig.data.pose_position
70 metarig.data.pose_position = 'REST'
72 bpy.ops.object.mode_set(mode='OBJECT')
74 scene = context.scene
75 id_store = context.window_manager
76 #------------------------------------------
77 # Create/find the rig object and set it up
79 # Check if the generated rig already exists, so we can
80 # regenerate in the same object. If not, create a new
81 # object to generate the rig in.
82 print("Fetch rig.")
84 rig_new_name = ""
85 rig_old_name = ""
86 if id_store.rigify_rig_basename:
87 rig_new_name = id_store.rigify_rig_basename + "_rig"
89 if id_store.rigify_generate_mode == 'overwrite':
90 name = id_store.rigify_target_rig or "rig"
91 try:
92 obj = scene.objects[name]
93 rig_old_name = name
94 obj.name = rig_new_name or name
95 except KeyError:
96 rig_old_name = name
97 name = rig_new_name or name
98 obj = bpy.data.objects.new(name, bpy.data.armatures.new(name))
99 obj.draw_type = 'WIRE'
100 scene.objects.link(obj)
101 else:
102 name = rig_new_name or "rig"
103 obj = bpy.data.objects.new(name, bpy.data.armatures.new(name)) # in case name 'rig' exists it will be rig.001
104 obj.draw_type = 'WIRE'
105 scene.objects.link(obj)
107 id_store.rigify_target_rig = obj.name
108 obj.data.pose_position = 'POSE'
110 # Get rid of anim data in case the rig already existed
111 print("Clear rig animation data.")
112 obj.animation_data_clear()
114 # Select generated rig object
115 metarig.select = False
116 obj.select = True
117 scene.objects.active = obj
119 # Remove wgts if force update is set
120 wgts_group_name = "WGTS_" + (rig_old_name or obj.name)
121 if wgts_group_name in scene.objects and id_store.rigify_force_widget_update:
122 bpy.ops.object.select_all(action='DESELECT')
123 for i, lyr in enumerate(WGT_LAYERS):
124 if lyr:
125 context.scene.layers[i] = True
126 for wgt in bpy.data.objects[wgts_group_name].children:
127 wgt.select = True
128 bpy.ops.object.delete(use_global=False)
129 for i, lyr in enumerate(WGT_LAYERS):
130 if lyr:
131 context.scene.layers[i] = False
132 if rig_old_name:
133 bpy.data.objects[wgts_group_name].name = "WGTS_" + obj.name
135 wgts_group_name = "WGTS_" + obj.name
137 # Get parented objects to restore later
138 childs = {} # {object: bone}
139 for child in obj.children:
140 childs[child] = child.parent_bone
142 # Remove all bones from the generated rig armature.
143 bpy.ops.object.mode_set(mode='EDIT')
144 for bone in obj.data.edit_bones:
145 obj.data.edit_bones.remove(bone)
146 bpy.ops.object.mode_set(mode='OBJECT')
148 # Create temporary duplicates for merging
149 temp_rig_1 = metarig.copy()
150 temp_rig_1.data = metarig.data.copy()
151 scene.objects.link(temp_rig_1)
153 temp_rig_2 = metarig.copy()
154 temp_rig_2.data = obj.data
155 scene.objects.link(temp_rig_2)
157 # Select the temp rigs for merging
158 for objt in scene.objects:
159 objt.select = False # deselect all objects
160 temp_rig_1.select = True
161 temp_rig_2.select = True
162 scene.objects.active = temp_rig_2
164 # Merge the temporary rigs
165 bpy.ops.object.join()
167 # Delete the second temp rig
168 bpy.ops.object.delete()
170 # Select the generated rig
171 for objt in scene.objects:
172 objt.select = False # deselect all objects
173 obj.select = True
174 scene.objects.active = obj
176 # Copy over bone properties
177 for bone in metarig.data.bones:
178 bone_gen = obj.data.bones[bone.name]
180 # B-bone stuff
181 bone_gen.bbone_segments = bone.bbone_segments
182 bone_gen.bbone_easein = bone.bbone_easein
183 bone_gen.bbone_easeout = bone.bbone_easeout
185 # Copy over the pose_bone properties
186 for bone in metarig.pose.bones:
187 bone_gen = obj.pose.bones[bone.name]
189 # Rotation mode and transform locks
190 bone_gen.rotation_mode = bone.rotation_mode
191 bone_gen.lock_rotation = tuple(bone.lock_rotation)
192 bone_gen.lock_rotation_w = bone.lock_rotation_w
193 bone_gen.lock_rotations_4d = bone.lock_rotations_4d
194 bone_gen.lock_location = tuple(bone.lock_location)
195 bone_gen.lock_scale = tuple(bone.lock_scale)
197 # rigify_type and rigify_parameters
198 bone_gen.rigify_type = bone.rigify_type
199 for prop in dir(bone_gen.rigify_parameters):
200 if (not prop.startswith("_")) \
201 and (not prop.startswith("bl_")) \
202 and (prop != "rna_type"):
203 try:
204 setattr(bone_gen.rigify_parameters, prop, \
205 getattr(bone.rigify_parameters, prop))
206 except AttributeError:
207 print("FAILED TO COPY PARAMETER: " + str(prop))
209 # Custom properties
210 for prop in bone.keys():
211 try:
212 bone_gen[prop] = bone[prop]
213 except KeyError:
214 pass
216 # Constraints
217 for con1 in bone.constraints:
218 con2 = bone_gen.constraints.new(type=con1.type)
219 copy_attributes(con1, con2)
221 # Set metarig target to rig target
222 if "target" in dir(con2):
223 if con2.target == metarig:
224 con2.target = obj
226 # Copy drivers
227 if metarig.animation_data:
228 for d1 in metarig.animation_data.drivers:
229 d2 = obj.driver_add(d1.data_path)
230 copy_attributes(d1, d2)
231 copy_attributes(d1.driver, d2.driver)
233 # Remove default modifiers, variables, etc.
234 for m in d2.modifiers:
235 d2.modifiers.remove(m)
236 for v in d2.driver.variables:
237 d2.driver.variables.remove(v)
239 # Copy modifiers
240 for m1 in d1.modifiers:
241 m2 = d2.modifiers.new(type=m1.type)
242 copy_attributes(m1, m2)
244 # Copy variables
245 for v1 in d1.driver.variables:
246 v2 = d2.driver.variables.new()
247 copy_attributes(v1, v2)
248 for i in range(len(v1.targets)):
249 copy_attributes(v1.targets[i], v2.targets[i])
250 # Switch metarig targets to rig targets
251 if v2.targets[i].id == metarig:
252 v2.targets[i].id = obj
254 # Mark targets that may need to be altered after rig generation
255 tar = v2.targets[i]
256 # If a custom property
257 if v2.type == 'SINGLE_PROP' \
258 and re.match('^pose.bones\["[^"\]]*"\]\["[^"\]]*"\]$', tar.data_path):
259 tar.data_path = "RIGIFY-" + tar.data_path
261 # Copy key frames
262 for i in range(len(d1.keyframe_points)):
263 d2.keyframe_points.add()
264 k1 = d1.keyframe_points[i]
265 k2 = d2.keyframe_points[i]
266 copy_attributes(k1, k2)
268 t.tick("Duplicate rig: ")
269 #----------------------------------
270 # Make a list of the original bones so we can keep track of them.
271 original_bones = [bone.name for bone in obj.data.bones]
273 # Add the ORG_PREFIX to the original bones.
274 bpy.ops.object.mode_set(mode='OBJECT')
275 for i in range(0, len(original_bones)):
276 obj.data.bones[original_bones[i]].name = make_original_name(original_bones[i])
277 original_bones[i] = make_original_name(original_bones[i])
279 # Create a sorted list of the original bones, sorted in the order we're
280 # going to traverse them for rigging.
281 # (root-most -> leaf-most, alphabetical)
282 bones_sorted = []
283 for name in original_bones:
284 bones_sorted += [name]
285 bones_sorted.sort() # first sort by names
286 bones_sorted.sort(key=lambda bone: len(obj.pose.bones[bone].parent_recursive)) # then parents before children
288 t.tick("Make list of org bones: ")
289 #----------------------------------
290 # Create the root bone.
291 bpy.ops.object.mode_set(mode='EDIT')
292 root_bone = new_bone(obj, ROOT_NAME)
293 spread = get_xy_spread(metarig.data.bones) or metarig.data.bones[0].length
294 spread = float('%.3g' % spread)
295 scale = spread/0.589
296 obj.data.edit_bones[root_bone].head = (0, 0, 0)
297 obj.data.edit_bones[root_bone].tail = (0, scale, 0)
298 obj.data.edit_bones[root_bone].roll = 0
299 bpy.ops.object.mode_set(mode='OBJECT')
300 obj.data.bones[root_bone].layers = ROOT_LAYER
302 # Put the rig_name in the armature custom properties
303 rna_idprop_ui_prop_get(obj.data, "rig_id", create=True)
304 obj.data["rig_id"] = rig_id
306 t.tick("Create root bone: ")
308 # Create Group widget
309 # wgts_group_name = "WGTS"
310 if wgts_group_name not in scene.objects:
311 if wgts_group_name in bpy.data.objects:
312 bpy.data.objects[wgts_group_name].user_clear()
313 bpy.data.objects.remove(bpy.data.objects[wgts_group_name])
314 mesh = bpy.data.meshes.new(wgts_group_name)
315 wgts_obj = bpy.data.objects.new(wgts_group_name, mesh)
316 scene.objects.link(wgts_obj)
317 wgts_obj.layers = WGT_LAYERS
318 t.tick("Create main WGTS: ")
320 # if id_store.rigify_generate_mode == 'new':
321 # bpy.ops.object.select_all(action='DESELECT')
322 # for wgt in bpy.data.objects[wgts_group_name].children:
323 # wgt.select = True
324 # for i, lyr in enumerate(WGT_LAYERS):
325 # if lyr:
326 # context.scene.layers[i] = True
327 # bpy.ops.object.make_single_user(obdata=True)
328 # for i, lyr in enumerate(WGT_LAYERS):
329 # if lyr:
330 # context.scene.layers[i] = False
332 #----------------------------------
333 try:
334 # Collect/initialize all the rigs.
335 rigs = []
336 for bone in bones_sorted:
337 bpy.ops.object.mode_set(mode='EDIT')
338 rigs += get_bone_rigs(obj, bone)
339 t.tick("Initialize rigs: ")
341 # Generate all the rigs.
342 ui_scripts = []
343 for rig in rigs:
344 # Go into editmode in the rig armature
345 bpy.ops.object.mode_set(mode='OBJECT')
346 context.scene.objects.active = obj
347 obj.select = True
348 bpy.ops.object.mode_set(mode='EDIT')
349 scripts = rig.generate()
350 if scripts is not None:
351 ui_scripts += [scripts[0]]
352 t.tick("Generate rigs: ")
353 except Exception as e:
354 # Cleanup if something goes wrong
355 print("Rigify: failed to generate rig.")
356 metarig.data.pose_position = rest_backup
357 obj.data.pose_position = 'POSE'
358 bpy.ops.object.mode_set(mode='OBJECT')
360 # Continue the exception
361 raise e
363 #----------------------------------
364 bpy.ops.object.mode_set(mode='OBJECT')
366 # Get a list of all the bones in the armature
367 bones = [bone.name for bone in obj.data.bones]
369 # Parent any free-floating bones to the root excluding bones with child of constraint.
370 pbones = obj.pose.bones
373 ik_follow_drivers = []
375 if obj.animation_data:
376 for drv in obj.animation_data.drivers:
377 for var in drv.driver.variables:
378 if 'IK_follow' == var.name:
379 ik_follow_drivers.append(drv.data_path)
381 noparent_bones = []
382 for bone in bones:
383 # if 'IK_follow' in pbones[bone].keys():
384 # noparent_bones += [bone]
385 for d in ik_follow_drivers:
386 if bone in d:
387 noparent_bones += [bone]
389 bpy.ops.object.mode_set(mode='EDIT')
390 for bone in bones:
391 if bone in noparent_bones:
392 continue
393 elif obj.data.edit_bones[bone].parent is None:
394 obj.data.edit_bones[bone].use_connect = False
395 obj.data.edit_bones[bone].parent = obj.data.edit_bones[root_bone]
397 bpy.ops.object.mode_set(mode='OBJECT')
399 # Lock transforms on all non-control bones
400 r = re.compile("[A-Z][A-Z][A-Z]-")
401 for bone in bones:
402 if r.match(bone):
403 pb = obj.pose.bones[bone]
404 pb.lock_location = (True, True, True)
405 pb.lock_rotation = (True, True, True)
406 pb.lock_rotation_w = True
407 pb.lock_scale = (True, True, True)
409 # Every bone that has a name starting with "DEF-" make deforming. All the
410 # others make non-deforming.
411 for bone in bones:
412 if obj.data.bones[bone].name.startswith(DEF_PREFIX):
413 obj.data.bones[bone].use_deform = True
414 else:
415 obj.data.bones[bone].use_deform = False
417 # Alter marked driver targets
418 if obj.animation_data:
419 for d in obj.animation_data.drivers:
420 for v in d.driver.variables:
421 for tar in v.targets:
422 if tar.data_path.startswith("RIGIFY-"):
423 temp, bone, prop = tuple([x.strip('"]') for x in tar.data_path.split('["')])
424 if bone in obj.data.bones \
425 and prop in obj.pose.bones[bone].keys():
426 tar.data_path = tar.data_path[7:]
427 else:
428 tar.data_path = 'pose.bones["%s"]["%s"]' % (make_original_name(bone), prop)
430 # Move all the original bones to their layer.
431 for bone in original_bones:
432 obj.data.bones[bone].layers = ORG_LAYER
434 # Move all the bones with names starting with "MCH-" to their layer.
435 for bone in bones:
436 if obj.data.bones[bone].name.startswith(MCH_PREFIX):
437 obj.data.bones[bone].layers = MCH_LAYER
439 # Move all the bones with names starting with "DEF-" to their layer.
440 for bone in bones:
441 if obj.data.bones[bone].name.startswith(DEF_PREFIX):
442 obj.data.bones[bone].layers = DEF_LAYER
444 # Create root bone widget
445 create_root_widget(obj, "root")
447 # Assign shapes to bones
448 # Object's with name WGT-<bone_name> get used as that bone's shape.
449 for bone in bones:
450 wgt_name = (WGT_PREFIX + obj.name + '_' + obj.data.bones[bone].name)[:63] # Object names are limited to 63 characters... arg
451 if wgt_name in context.scene.objects:
452 # Weird temp thing because it won't let me index by object name
453 for ob in context.scene.objects:
454 if ob.name == wgt_name:
455 obj.pose.bones[bone].custom_shape = ob
456 break
457 # This is what it should do:
458 # obj.pose.bones[bone].custom_shape = context.scene.objects[wgt_name]
459 # Reveal all the layers with control bones on them
460 vis_layers = [False for n in range(0, 32)]
461 for bone in bones:
462 for i in range(0, 32):
463 vis_layers[i] = vis_layers[i] or obj.data.bones[bone].layers[i]
464 for i in range(0, 32):
465 vis_layers[i] = vis_layers[i] and not (ORG_LAYER[i] or MCH_LAYER[i] or DEF_LAYER[i])
466 obj.data.layers = vis_layers
468 # Ensure the collection of layer names exists
469 for i in range(1 + len(metarig.data.rigify_layers), 29):
470 metarig.data.rigify_layers.add()
472 # Create list of layer name/row pairs
473 layer_layout = []
474 for l in metarig.data.rigify_layers:
475 print(l.name)
476 layer_layout += [(l.name, l.row)]
478 # Generate the UI script
479 if id_store.rigify_generate_mode == 'overwrite':
480 rig_ui_name = id_store.rigify_rig_ui or 'rig_ui.py'
481 else:
482 rig_ui_name = 'rig_ui.py'
484 if id_store.rigify_generate_mode == 'overwrite' and rig_ui_name in bpy.data.texts.keys():
485 script = bpy.data.texts[rig_ui_name]
486 script.clear()
487 else:
488 script = bpy.data.texts.new("rig_ui.py")
490 rig_ui_old_name = ""
491 if id_store.rigify_rig_basename:
492 rig_ui_old_name = script.name
493 script.name = id_store.rigify_rig_basename + "_rig_ui.py"
495 id_store.rigify_rig_ui = script.name
497 script.write(UI_SLIDERS % rig_id)
498 for s in ui_scripts:
499 script.write("\n " + s.replace("\n", "\n ") + "\n")
500 script.write(layers_ui(vis_layers, layer_layout))
501 script.write(UI_REGISTER)
502 script.use_module = True
504 # Run UI script
505 exec(script.as_string(), {})
507 # Create Selection Sets
508 create_selection_sets(obj, metarig)
510 # Create Bone Groups
511 create_bone_groups(obj, metarig)
513 # Add rig_ui to logic
514 skip = False
515 ctrls = obj.game.controllers
517 for c in ctrls:
518 if 'Python' in c.name and c.text.name == script.name:
519 skip = True
520 break
521 if not skip:
522 bpy.ops.logic.controller_add(type='PYTHON', object=obj.name)
523 ctrl = obj.game.controllers[-1]
524 ctrl.text = bpy.data.texts[script.name]
527 t.tick("The rest: ")
528 #----------------------------------
529 # Deconfigure
530 bpy.ops.object.mode_set(mode='OBJECT')
531 metarig.data.pose_position = rest_backup
532 obj.data.pose_position = 'POSE'
534 # Restore parent to bones
535 for child, sub_parent in childs.items():
536 if sub_parent in obj.pose.bones:
537 mat = child.matrix_world.copy()
538 child.parent_bone = sub_parent
539 child.matrix_world = mat
541 def create_selection_sets(obj, metarig):
543 # Check if selection sets addon is installed
544 if 'bone_selection_groups' not in bpy.context.user_preferences.addons \
545 and 'bone_selection_sets' not in bpy.context.user_preferences.addons:
546 return
548 bpy.ops.object.mode_set(mode='POSE')
550 bpy.context.scene.objects.active = obj
551 obj.select = True
552 metarig.select = False
553 pbones = obj.pose.bones
555 for i, name in enumerate(metarig.data.rigify_layers.keys()):
556 if name == '' or not metarig.data.rigify_layers[i].set:
557 continue
559 bpy.ops.pose.select_all(action='DESELECT')
560 for b in pbones:
561 if b.bone.layers[i]:
562 b.bone.select = True
564 #bpy.ops.pose.selection_set_add()
565 obj.selection_sets.add()
566 obj.selection_sets[-1].name = name
567 if 'bone_selection_sets' in bpy.context.user_preferences.addons:
568 act_sel_set = obj.selection_sets[-1]
570 # iterate only the selected bones in current pose that are not hidden
571 for bone in bpy.context.selected_pose_bones:
572 if bone.name not in act_sel_set.bone_ids:
573 bone_id = act_sel_set.bone_ids.add()
574 bone_id.name = bone.name
577 def create_bone_groups(obj, metarig):
579 bpy.ops.object.mode_set(mode='OBJECT')
580 pb = obj.pose.bones
581 layers = metarig.data.rigify_layers
582 groups = metarig.data.rigify_colors
584 # Create BGs
585 for l in layers:
586 if l.group == 0:
587 continue
588 g_id = l.group - 1
589 name = groups[g_id].name
590 if name not in obj.pose.bone_groups.keys():
591 bg = obj.pose.bone_groups.new(name)
592 bg.color_set = 'CUSTOM'
593 bg.colors.normal = gamma_correct(groups[g_id].normal)
594 bg.colors.select = gamma_correct(groups[g_id].select)
595 bg.colors.active = gamma_correct(groups[g_id].active)
597 for b in pb:
598 try:
599 layer_index = b.bone.layers[:].index(True)
600 except ValueError:
601 continue
602 if layer_index > len(layers) - 1: # bone is on reserved layers
603 continue
604 g_id = layers[layer_index].group - 1
605 if g_id >= 0:
606 name = groups[g_id].name
607 b.bone_group = obj.pose.bone_groups[name]
610 def get_bone_rigs(obj, bone_name, halt_on_missing=False):
611 """ Fetch all the rigs specified on a bone.
613 rigs = []
614 rig_type = obj.pose.bones[bone_name].rigify_type
615 rig_type = rig_type.replace(" ", "")
617 if rig_type == "":
618 pass
619 else:
620 # Gather parameters
621 params = obj.pose.bones[bone_name].rigify_parameters
623 # Get the rig
624 try:
625 rig = get_rig_type(rig_type).Rig(obj, bone_name, params)
626 except ImportError:
627 message = "Rig Type Missing: python module for type '%s' not found (bone: %s)" % (rig_type, bone_name)
628 if halt_on_missing:
629 raise MetarigError(message)
630 else:
631 print(message)
632 print('print_exc():')
633 traceback.print_exc(file=sys.stdout)
634 else:
635 rigs += [rig]
636 return rigs
639 def get_xy_spread(bones):
640 x_max = 0
641 y_max = 0
642 for b in bones:
643 x_max = max((x_max, abs(b.head[0]), abs(b.tail[0])))
644 y_max = max((y_max, abs(b.head[1]), abs(b.tail[1])))
646 return max((x_max, y_max))
649 def param_matches_type(param_name, rig_type):
650 """ Returns True if the parameter name is consistent with the rig type.
652 if param_name.rsplit(".", 1)[0] == rig_type:
653 return True
654 else:
655 return False
658 def param_name(param_name, rig_type):
659 """ Get the actual parameter name, sans-rig-type.
661 return param_name[len(rig_type) + 1:]