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