rigify/generate.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 import re
  23 import time
  24 import traceback
  25 import sys
  26 from rna_prop_ui import rna_idprop_ui_prop_get
  27
  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 .rig_ui_template import UI_SLIDERS, layers_ui, UI_REGISTER
  35
  36 RIG_MODULE = "rigs"
  37 ORG_LAYER = [n == 31 for n in range(0, 32)]  # Armature layer that original bones should be moved to.
  38 MCH_LAYER = [n == 30 for n in range(0, 32)]  # Armature layer that mechanism bones should be moved to.
  39 DEF_LAYER = [n == 29 for n in range(0, 32)]  # Armature layer that deformation bones should be moved to.
  40 ROOT_LAYER = [n == 28 for n in range(0, 32)]  # Armature layer that root bone should be moved to.
  41
  42
  43 class Timer:
  44     def __init__(self):
  45         self.timez = time.time()
  46
  47     def tick(self, string):
  48         t = time.time()
  49         print(string + "%.3f" % (t - self.timez))
  50         self.timez = t
  51
  52
  53 # TODO: generalize to take a group as input instead of an armature.
  54 def generate_rig(context, metarig):
  55     """ Generates a rig from a metarig.
  56
  57     """
  58     t = Timer()
  59
  60     # Random string with time appended so that
  61     # different rigs don't collide id's
  62     rig_id = random_id(16)
  63
  64     # Initial configuration
  65     # mode_orig = context.mode  # UNUSED
  66     rest_backup = metarig.data.pose_position
  67     metarig.data.pose_position = 'REST'
  68
  69     bpy.ops.object.mode_set(mode='OBJECT')
  70
  71     scene = context.scene
  72
  73     #------------------------------------------
  74     # Create/find the rig object and set it up
  75
  76     # Check if the generated rig already exists, so we can
  77     # regenerate in the same object.  If not, create a new
  78     # object to generate the rig in.
  79     print("Fetch rig.")
  80     try:
  81         name = metarig["rig_object_name"]
  82     except KeyError:
  83         name = "rig"
  84
  85     try:
  86         obj = scene.objects[name]
  87     except KeyError:
  88         obj = bpy.data.objects.new(name, bpy.data.armatures.new(name))
  89         obj.draw_type = 'WIRE'
  90         scene.objects.link(obj)
  91
  92     obj.data.pose_position = 'POSE'
  93
  94     # Get rid of anim data in case the rig already existed
  95     print("Clear rig animation data.")
  96     obj.animation_data_clear()
  97
  98     # Select generated rig object
  99     metarig.select = False
 100     obj.select = True
 101     scene.objects.active = obj
 102
 103     # Remove all bones from the generated rig armature.
 104     bpy.ops.object.mode_set(mode='EDIT')
 105     for bone in obj.data.edit_bones:
 106         obj.data.edit_bones.remove(bone)
 107     bpy.ops.object.mode_set(mode='OBJECT')
 108
 109     # Create temporary duplicates for merging
 110     temp_rig_1 = metarig.copy()
 111     temp_rig_1.data = metarig.data.copy()
 112     scene.objects.link(temp_rig_1)
 113
 114     temp_rig_2 = metarig.copy()
 115     temp_rig_2.data = obj.data
 116     scene.objects.link(temp_rig_2)
 117
 118     # Select the temp rigs for merging
 119     for objt in scene.objects:
 120         objt.select = False  # deselect all objects
 121     temp_rig_1.select = True
 122     temp_rig_2.select = True
 123     scene.objects.active = temp_rig_2
 124
 125     # Merge the temporary rigs
 126     bpy.ops.object.join()
 127
 128     # Delete the second temp rig
 129     bpy.ops.object.delete()
 130
 131     # Select the generated rig
 132     for objt in scene.objects:
 133         objt.select = False  # deselect all objects
 134     obj.select = True
 135     scene.objects.active = obj
 136
 137     # Copy over bone properties
 138     for bone in metarig.data.bones:
 139         bone_gen = obj.data.bones[bone.name]
 140
 141         # B-bone stuff
 142         bone_gen.bbone_segments = bone.bbone_segments
 143         bone_gen.bbone_in = bone.bbone_in
 144         bone_gen.bbone_out = bone.bbone_out
 145
 146     # Copy over the pose_bone properties
 147     for bone in metarig.pose.bones:
 148         bone_gen = obj.pose.bones[bone.name]
 149
 150         # Rotation mode and transform locks
 151         bone_gen.rotation_mode = bone.rotation_mode
 152         bone_gen.lock_rotation = tuple(bone.lock_rotation)
 153         bone_gen.lock_rotation_w = bone.lock_rotation_w
 154         bone_gen.lock_rotations_4d = bone.lock_rotations_4d
 155         bone_gen.lock_location = tuple(bone.lock_location)
 156         bone_gen.lock_scale = tuple(bone.lock_scale)
 157
 158         # rigify_type and rigify_parameters
 159         bone_gen.rigify_type = bone.rigify_type
 160         for prop in dir(bone_gen.rigify_parameters):
 161             if (not prop.startswith("_")) \
 162             and (not prop.startswith("bl_")) \
 163             and (prop != "rna_type"):
 164                 try:
 165                     setattr(bone_gen.rigify_parameters, prop, \
 166                             getattr(bone.rigify_parameters, prop))
 167                 except AttributeError:
 168                     print("FAILED TO COPY PARAMETER: " + str(prop))
 169
 170         # Custom properties
 171         for prop in bone.keys():
 172             try:
 173                 bone_gen[prop] = bone[prop]
 174             except KeyError:
 175                 pass
 176
 177         # Constraints
 178         for con1 in bone.constraints:
 179             con2 = bone_gen.constraints.new(type=con1.type)
 180             copy_attributes(con1, con2)
 181
 182             # Set metarig target to rig target
 183             if "target" in dir(con2):
 184                 if con2.target == metarig:
 185                     con2.target = obj
 186
 187     # Copy drivers
 188     if metarig.animation_data:
 189         for d1 in metarig.animation_data.drivers:
 190             d2 = obj.driver_add(d1.data_path)
 191             copy_attributes(d1, d2)
 192             copy_attributes(d1.driver, d2.driver)
 193
 194             # Remove default modifiers, variables, etc.
 195             for m in d2.modifiers:
 196                 d2.modifiers.remove(m)
 197             for v in d2.driver.variables:
 198                 d2.driver.variables.remove(v)
 199
 200             # Copy modifiers
 201             for m1 in d1.modifiers:
 202                 m2 = d2.modifiers.new(type=m1.type)
 203                 copy_attributes(m1, m2)
 204
 205             # Copy variables
 206             for v1 in d1.driver.variables:
 207                 v2 = d2.driver.variables.new()
 208                 copy_attributes(v1, v2)
 209                 for i in range(len(v1.targets)):
 210                     copy_attributes(v1.targets[i], v2.targets[i])
 211                     # Switch metarig targets to rig targets
 212                     if v2.targets[i].id == metarig:
 213                         v2.targets[i].id = obj
 214
 215                     # Mark targets that may need to be altered after rig generation
 216                     tar = v2.targets[i]
 217                     # If a custom property
 218                     if v2.type == 'SINGLE_PROP' \
 219                     and re.match('^pose.bones\["[^"\]]*"\]\["[^"\]]*"\]$', tar.data_path):
 220                         tar.data_path = "RIGIFY-" + tar.data_path
 221
 222             # Copy key frames
 223             for i in range(len(d1.keyframe_points)):
 224                 d2.keyframe_points.add()
 225                 k1 = d1.keyframe_points[i]
 226                 k2 = d2.keyframe_points[i]
 227                 copy_attributes(k1, k2)
 228
 229     t.tick("Duplicate rig: ")
 230     #----------------------------------
 231     # Make a list of the original bones so we can keep track of them.
 232     original_bones = [bone.name for bone in obj.data.bones]
 233
 234     # Add the ORG_PREFIX to the original bones.
 235     bpy.ops.object.mode_set(mode='OBJECT')
 236     for i in range(0, len(original_bones)):
 237         obj.data.bones[original_bones[i]].name = make_original_name(original_bones[i])
 238         original_bones[i] = make_original_name(original_bones[i])
 239
 240     # Create a sorted list of the original bones, sorted in the order we're
 241     # going to traverse them for rigging.
 242     # (root-most -> leaf-most, alphabetical)
 243     bones_sorted = []
 244     for name in original_bones:
 245         bones_sorted += [name]
 246     bones_sorted.sort()  # first sort by names
 247     bones_sorted.sort(key=lambda bone: len(obj.pose.bones[bone].parent_recursive))  # then parents before children
 248
 249     t.tick("Make list of org bones: ")
 250     #----------------------------------
 251     # Create the root bone.
 252     bpy.ops.object.mode_set(mode='EDIT')
 253     root_bone = new_bone(obj, ROOT_NAME)
 254     obj.data.edit_bones[root_bone].head = (0, 0, 0)
 255     obj.data.edit_bones[root_bone].tail = (0, 1, 0)
 256     obj.data.edit_bones[root_bone].roll = 0
 257     bpy.ops.object.mode_set(mode='OBJECT')
 258     obj.data.bones[root_bone].layers = ROOT_LAYER
 259     # Put the rig_name in the armature custom properties
 260     rna_idprop_ui_prop_get(obj.data, "rig_id", create=True)
 261     obj.data["rig_id"] = rig_id
 262
 263     t.tick("Create root bone: ")
 264     #----------------------------------
 265     try:
 266         # Collect/initialize all the rigs.
 267         rigs = []
 268         for bone in bones_sorted:
 269             bpy.ops.object.mode_set(mode='EDIT')
 270             rigs += get_bone_rigs(obj, bone)
 271         t.tick("Initialize rigs: ")
 272
 273         # Generate all the rigs.
 274         ui_scripts = []
 275         for rig in rigs:
 276             # Go into editmode in the rig armature
 277             bpy.ops.object.mode_set(mode='OBJECT')
 278             context.scene.objects.active = obj
 279             obj.select = True
 280             bpy.ops.object.mode_set(mode='EDIT')
 281             scripts = rig.generate()
 282             if scripts != None:
 283                 ui_scripts += [scripts[0]]
 284         t.tick("Generate rigs: ")
 285     except Exception as e:
 286         # Cleanup if something goes wrong
 287         print("Rigify: failed to generate rig.")
 288         metarig.data.pose_position = rest_backup
 289         obj.data.pose_position = 'POSE'
 290         bpy.ops.object.mode_set(mode='OBJECT')
 291
 292         # Continue the exception
 293         raise e
 294
 295     #----------------------------------
 296     bpy.ops.object.mode_set(mode='OBJECT')
 297
 298     # Get a list of all the bones in the armature
 299     bones = [bone.name for bone in obj.data.bones]
 300
 301     # Parent any free-floating bones to the root.
 302     bpy.ops.object.mode_set(mode='EDIT')
 303     for bone in bones:
 304         if obj.data.edit_bones[bone].parent is None:
 305             obj.data.edit_bones[bone].use_connect = False
 306             obj.data.edit_bones[bone].parent = obj.data.edit_bones[root_bone]
 307     bpy.ops.object.mode_set(mode='OBJECT')
 308
 309     # Lock transforms on all non-control bones
 310     r = re.compile("[A-Z][A-Z][A-Z]-")
 311     for bone in bones:
 312         if r.match(bone):
 313             pb = obj.pose.bones[bone]
 314             pb.lock_location = (True, True, True)
 315             pb.lock_rotation = (True, True, True)
 316             pb.lock_rotation_w = True
 317             pb.lock_scale = (True, True, True)
 318
 319     # Every bone that has a name starting with "DEF-" make deforming.  All the
 320     # others make non-deforming.
 321     for bone in bones:
 322         if obj.data.bones[bone].name.startswith(DEF_PREFIX):
 323             obj.data.bones[bone].use_deform = True
 324         else:
 325             obj.data.bones[bone].use_deform = False
 326
 327     # Alter marked driver targets
 328     if obj.animation_data:
 329         for d in obj.animation_data.drivers:
 330             for v in d.driver.variables:
 331                 for tar in v.targets:
 332                     if tar.data_path.startswith("RIGIFY-"):
 333                         temp, bone, prop = tuple([x.strip('"]') for x in tar.data_path.split('["')])
 334                         if bone in obj.data.bones \
 335                         and prop in obj.pose.bones[bone].keys():
 336                             tar.data_path = tar.data_path[7:]
 337                         else:
 338                             tar.data_path = 'pose.bones["%s"]["%s"]' % (make_original_name(bone), prop)
 339
 340     # Move all the original bones to their layer.
 341     for bone in original_bones:
 342         obj.data.bones[bone].layers = ORG_LAYER
 343
 344     # Move all the bones with names starting with "MCH-" to their layer.
 345     for bone in bones:
 346         if obj.data.bones[bone].name.startswith(MCH_PREFIX):
 347             obj.data.bones[bone].layers = MCH_LAYER
 348
 349     # Move all the bones with names starting with "DEF-" to their layer.
 350     for bone in bones:
 351         if obj.data.bones[bone].name.startswith(DEF_PREFIX):
 352             obj.data.bones[bone].layers = DEF_LAYER
 353
 354     # Create root bone widget
 355     create_root_widget(obj, "root")
 356
 357     # Assign shapes to bones
 358     # Object's with name WGT-<bone_name> get used as that bone's shape.
 359     for bone in bones:
 360         wgt_name = (WGT_PREFIX + obj.data.bones[bone].name)[:63]  # Object names are limited to 63 characters... arg
 361         if wgt_name in context.scene.objects:
 362             # Weird temp thing because it won't let me index by object name
 363             for ob in context.scene.objects:
 364                 if ob.name == wgt_name:
 365                     obj.pose.bones[bone].custom_shape = ob
 366                     break
 367             # This is what it should do:
 368             # obj.pose.bones[bone].custom_shape = context.scene.objects[wgt_name]
 369     # Reveal all the layers with control bones on them
 370     vis_layers = [False for n in range(0, 32)]
 371     for bone in bones:
 372         for i in range(0, 32):
 373             vis_layers[i] = vis_layers[i] or obj.data.bones[bone].layers[i]
 374     for i in range(0, 32):
 375         vis_layers[i] = vis_layers[i] and not (ORG_LAYER[i] or MCH_LAYER[i] or DEF_LAYER[i])
 376     obj.data.layers = vis_layers
 377
 378     # Ensure the collection of layer names exists
 379     for i in range(1 + len(metarig.data.rigify_layers), 29):
 380         metarig.data.rigify_layers.add()
 381
 382     # Create list of layer name/row pairs
 383     layer_layout = []
 384     for l in metarig.data.rigify_layers:
 385         layer_layout += [(l.name, l.row)]
 386
 387     # Generate the UI script
 388     if "rig_ui.py" in bpy.data.texts:
 389         script = bpy.data.texts["rig_ui.py"]
 390         script.clear()
 391     else:
 392         script = bpy.data.texts.new("rig_ui.py")
 393     script.write(UI_SLIDERS % rig_id)
 394     for s in ui_scripts:
 395         script.write("\n        " + s.replace("\n", "\n        ") + "\n")
 396     script.write(layers_ui(vis_layers, layer_layout))
 397     script.write(UI_REGISTER)
 398     script.use_module = True
 399
 400     # Run UI script
 401     exec(script.as_string(), {})
 402
 403     t.tick("The rest: ")
 404     #----------------------------------
 405     # Deconfigure
 406     bpy.ops.object.mode_set(mode='OBJECT')
 407     metarig.data.pose_position = rest_backup
 408     obj.data.pose_position = 'POSE'
 409
 410
 411 def get_bone_rigs(obj, bone_name, halt_on_missing=False):
 412     """ Fetch all the rigs specified on a bone.
 413     """
 414     rigs = []
 415     rig_type = obj.pose.bones[bone_name].rigify_type
 416     rig_type = rig_type.replace(" ", "")
 417
 418     if rig_type == "":
 419         pass
 420     else:
 421         # Gather parameters
 422         params = obj.pose.bones[bone_name].rigify_parameters
 423
 424         # Get the rig
 425         try:
 426             rig = get_rig_type(rig_type).Rig(obj, bone_name, params)
 427         except ImportError:
 428             message = "Rig Type Missing: python module for type '%s' not found (bone: %s)" % (rig_type, bone_name)
 429             if halt_on_missing:
 430                 raise MetarigError(message)
 431             else:
 432                 print(message)
 433                 print('print_exc():')
 434                 traceback.print_exc(file=sys.stdout)
 435         else:
 436             rigs += [rig]
 437     return rigs
 438
 439
 440 def param_matches_type(param_name, rig_type):
 441     """ Returns True if the parameter name is consistent with the rig type.
 442     """
 443     if param_name.rsplit(".", 1)[0] == rig_type:
 444         return True
 445     else:
 446         return False
 447
 448
 449 def param_name(param_name, rig_type):
 450     """ Get the actual parameter name, sans-rig-type.
 451     """
 452     return param_name[len(rig_type) + 1:]