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