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 from rna_prop_ui import rna_idprop_ui_prop_get
  25
  26 from .utils.bones import new_bone
  27 from .utils.layers import ORG_LAYER, MCH_LAYER, DEF_LAYER, ROOT_LAYER
  28 from .utils.naming import ORG_PREFIX, MCH_PREFIX, DEF_PREFIX, ROOT_NAME, make_original_name
  29 from .utils.widgets import WGT_PREFIX
  30 from .utils.widgets_special import create_root_widget
  31 from .utils.misc import gamma_correct, select_object
  32 from .utils.collections import ensure_widget_collection, list_layer_collections, filter_layer_collections_by_object
  33 from .utils.rig import get_rigify_type
  34
  35 from . import base_generate
  36 from . import rig_ui_template
  37 from . import rig_lists
  38
  39 RIG_MODULE = "rigs"
  40
  41 class Timer:
  42     def __init__(self):
  43         self.timez = time.time()
  44
  45     def tick(self, string):
  46         t = time.time()
  47         print(string + "%.3f" % (t - self.timez))
  48         self.timez = t
  49
  50
  51 class Generator(base_generate.BaseGenerator):
  52     def __init__(self, context, metarig):
  53         super().__init__(context, metarig)
  54
  55         self.id_store = context.window_manager
  56
  57         self.rig_new_name = ""
  58         self.rig_old_name = ""
  59
  60
  61     def find_rig_class(self, rig_type):
  62         rig_module = rig_lists.rigs[rig_type]["module"]
  63
  64         return rig_module.Rig
  65
  66
  67     def __create_rig_object(self):
  68         scene = self.scene
  69         id_store = self.id_store
  70         meta_data = self.metarig.data
  71
  72         # Check if the generated rig already exists, so we can
  73         # regenerate in the same object.  If not, create a new
  74         # object to generate the rig in.
  75         print("Fetch rig.")
  76
  77         self.rig_new_name = name = meta_data.rigify_rig_basename or "rig"
  78
  79         obj = None
  80
  81         if meta_data.rigify_generate_mode == 'overwrite':
  82             obj = meta_data.rigify_target_rig
  83
  84             if not obj and name in scene.objects:
  85                 obj = scene.objects[name]
  86
  87             if obj:
  88                 self.rig_old_name = obj.name
  89
  90                 obj.name = name
  91                 obj.data.name = obj.name
  92
  93                 rig_collections = filter_layer_collections_by_object(self.usable_collections, obj)
  94                 self.layer_collection = (rig_collections + [self.layer_collection])[0]
  95                 self.collection = self.layer_collection.collection
  96
  97             elif name in bpy.data.objects:
  98                 obj = bpy.data.objects[name]
  99
 100         if not obj:
 101             obj = bpy.data.objects.new(name, bpy.data.armatures.new(name))
 102             obj.display_type = 'WIRE'
 103             self.collection.objects.link(obj)
 104
 105         elif obj.name not in self.collection.objects:  # rig exists but was deleted
 106             self.collection.objects.link(obj)
 107
 108         meta_data.rigify_target_rig = obj
 109         obj.data.pose_position = 'POSE'
 110
 111         self.obj = obj
 112         return obj
 113
 114
 115     def __create_widget_group(self):
 116         new_group_name = "WGTS_" + self.obj.name
 117         wgts_group_name = "WGTS_" + (self.rig_old_name or self.obj.name)
 118
 119         # Find the old widgets collection
 120         old_collection = bpy.data.collections.get(wgts_group_name)
 121
 122         if not old_collection:
 123             # Update the old 'Widgets' collection
 124             legacy_collection = bpy.data.collections.get('Widgets')
 125
 126             if legacy_collection and wgts_group_name in legacy_collection.objects:
 127                 legacy_collection.name = wgts_group_name
 128                 old_collection = legacy_collection
 129
 130         if old_collection:
 131             # Remove widgets if force update is set
 132             if self.metarig.data.rigify_force_widget_update:
 133                 for obj in list(old_collection.objects):
 134                     bpy.data.objects.remove(obj)
 135
 136             # Rename widgets and collection if renaming
 137             if self.rig_old_name:
 138                 old_prefix = WGT_PREFIX + self.rig_old_name + "_"
 139                 new_prefix = WGT_PREFIX + self.obj.name + "_"
 140
 141                 for obj in list(old_collection.objects):
 142                     if obj.name.startswith(old_prefix):
 143                         new_name = new_prefix + obj.name[len(old_prefix):]
 144                     elif obj.name == wgts_group_name:
 145                         new_name = new_group_name
 146                     else:
 147                         continue
 148
 149                     obj.data.name = new_name
 150                     obj.name = new_name
 151
 152                 old_collection.name = new_group_name
 153
 154         # Create/find widget collection
 155         self.widget_collection = ensure_widget_collection(self.context, new_group_name)
 156         self.wgts_group_name = new_group_name
 157
 158
 159     def __duplicate_rig(self):
 160         obj = self.obj
 161         metarig = self.metarig
 162         context = self.context
 163
 164         # Remove all bones from the generated rig armature.
 165         bpy.ops.object.mode_set(mode='EDIT')
 166         for bone in obj.data.edit_bones:
 167             obj.data.edit_bones.remove(bone)
 168         bpy.ops.object.mode_set(mode='OBJECT')
 169
 170         # Select and duplicate metarig
 171         select_object(context, metarig, deselect_all=True)
 172
 173         bpy.ops.object.duplicate()
 174
 175         # Rename org bones in the temporary object
 176         temp_obj = context.view_layer.objects.active
 177
 178         assert temp_obj and temp_obj != metarig
 179
 180         self.__freeze_driver_vars(temp_obj)
 181         self.__rename_org_bones(temp_obj)
 182
 183         # Select the target rig and join
 184         select_object(context, obj)
 185
 186         saved_matrix = obj.matrix_world.copy()
 187         obj.matrix_world = metarig.matrix_world
 188
 189         bpy.ops.object.join()
 190
 191         obj.matrix_world = saved_matrix
 192
 193         # Select the generated rig
 194         select_object(context, obj, deselect_all=True)
 195
 196         # Clean up animation data
 197         if obj.animation_data:
 198             obj.animation_data.action = None
 199
 200             for track in obj.animation_data.nla_tracks:
 201                 obj.animation_data.nla_tracks.remove(track)
 202
 203
 204     def __freeze_driver_vars(self, obj):
 205         if obj.animation_data:
 206             # Freeze drivers referring to custom properties
 207             for d in obj.animation_data.drivers:
 208                 for var in d.driver.variables:
 209                     for tar in var.targets:
 210                         # If a custom property
 211                         if var.type == 'SINGLE_PROP' \
 212                         and re.match('^pose.bones\["[^"\]]*"\]\["[^"\]]*"\]$', tar.data_path):
 213                             tar.data_path = "RIGIFY-" + tar.data_path
 214
 215
 216     def __rename_org_bones(self, obj):
 217         #----------------------------------
 218         # Make a list of the original bones so we can keep track of them.
 219         original_bones = [bone.name for bone in obj.data.bones]
 220
 221         # Add the ORG_PREFIX to the original bones.
 222         for i in range(0, len(original_bones)):
 223             bone = obj.pose.bones[original_bones[i]]
 224
 225             # This rig type is special in that it preserves the name of the bone.
 226             if get_rigify_type(bone) != 'basic.raw_copy':
 227                 bone.name = make_original_name(original_bones[i])
 228                 original_bones[i] = bone.name
 229
 230         self.original_bones = original_bones
 231
 232
 233     def __create_root_bone(self):
 234         obj = self.obj
 235         metarig = self.metarig
 236
 237         #----------------------------------
 238         # Create the root bone.
 239         root_bone = new_bone(obj, ROOT_NAME)
 240         spread = get_xy_spread(metarig.data.bones) or metarig.data.bones[0].length
 241         spread = float('%.3g' % spread)
 242         scale = spread/0.589
 243         obj.data.edit_bones[root_bone].head = (0, 0, 0)
 244         obj.data.edit_bones[root_bone].tail = (0, scale, 0)
 245         obj.data.edit_bones[root_bone].roll = 0
 246         self.root_bone = root_bone
 247         self.bone_owners[root_bone] = None
 248
 249
 250     def __parent_bones_to_root(self):
 251         eb = self.obj.data.edit_bones
 252
 253         # Parent loose bones to root
 254         for bone in eb:
 255             if bone.name in self.noparent_bones:
 256                 continue
 257             elif bone.parent is None:
 258                 bone.use_connect = False
 259                 bone.parent = eb[self.root_bone]
 260
 261
 262     def __lock_transforms(self):
 263         # Lock transforms on all non-control bones
 264         r = re.compile("[A-Z][A-Z][A-Z]-")
 265         for pb in self.obj.pose.bones:
 266             if r.match(pb.name):
 267                 pb.lock_location = (True, True, True)
 268                 pb.lock_rotation = (True, True, True)
 269                 pb.lock_rotation_w = True
 270                 pb.lock_scale = (True, True, True)
 271
 272
 273     def __assign_layers(self):
 274         pbones = self.obj.pose.bones
 275
 276         pbones[self.root_bone].bone.layers = ROOT_LAYER
 277
 278         # Every bone that has a name starting with "DEF-" make deforming.  All the
 279         # others make non-deforming.
 280         for pbone in pbones:
 281             bone = pbone.bone
 282             name = bone.name
 283             layers = None
 284
 285             bone.use_deform = name.startswith(DEF_PREFIX)
 286
 287             # Move all the original bones to their layer.
 288             if name.startswith(ORG_PREFIX):
 289                 layers = ORG_LAYER
 290             # Move all the bones with names starting with "MCH-" to their layer.
 291             elif name.startswith(MCH_PREFIX):
 292                 layers = MCH_LAYER
 293             # Move all the bones with names starting with "DEF-" to their layer.
 294             elif name.startswith(DEF_PREFIX):
 295                 layers = DEF_LAYER
 296
 297             if layers is not None:
 298                 bone.layers = layers
 299
 300                 # Remove custom shapes from non-control bones
 301                 pbone.custom_shape = None
 302
 303             bone.bbone_x = bone.bbone_z = bone.length * 0.05
 304
 305
 306     def __restore_driver_vars(self):
 307         obj = self.obj
 308
 309         # Alter marked driver targets
 310         if obj.animation_data:
 311             for d in obj.animation_data.drivers:
 312                 for v in d.driver.variables:
 313                     for tar in v.targets:
 314                         if tar.data_path.startswith("RIGIFY-"):
 315                             temp, bone, prop = tuple([x.strip('"]') for x in tar.data_path.split('["')])
 316                             if bone in obj.data.bones \
 317                             and prop in obj.pose.bones[bone].keys():
 318                                 tar.data_path = tar.data_path[7:]
 319                             else:
 320                                 org_name = make_original_name(bone)
 321                                 org_name = self.org_rename_table.get(org_name, org_name)
 322                                 tar.data_path = 'pose.bones["%s"]["%s"]' % (org_name, prop)
 323
 324
 325     def __assign_widgets(self):
 326         obj_table = {obj.name: obj for obj in self.scene.objects}
 327
 328         # Assign shapes to bones
 329         # Object's with name WGT-<bone_name> get used as that bone's shape.
 330         for bone in self.obj.pose.bones:
 331             # Object names are limited to 63 characters... arg
 332             wgt_name = (WGT_PREFIX + self.obj.name + '_' + bone.name)[:63]
 333
 334             if wgt_name in obj_table:
 335                 bone.custom_shape = obj_table[wgt_name]
 336
 337
 338     def __compute_visible_layers(self):
 339         # Reveal all the layers with control bones on them
 340         vis_layers = [False for n in range(0, 32)]
 341
 342         for bone in self.obj.data.bones:
 343             for i in range(0, 32):
 344                 vis_layers[i] = vis_layers[i] or bone.layers[i]
 345
 346         for i in range(0, 32):
 347             vis_layers[i] = vis_layers[i] and not (ORG_LAYER[i] or MCH_LAYER[i] or DEF_LAYER[i])
 348
 349         self.obj.data.layers = vis_layers
 350
 351
 352     def generate(self):
 353         context = self.context
 354         metarig = self.metarig
 355         scene = self.scene
 356         id_store = self.id_store
 357         view_layer = self.view_layer
 358         t = Timer()
 359
 360         self.usable_collections = list_layer_collections(view_layer.layer_collection, selectable=True)
 361
 362         if self.layer_collection not in self.usable_collections:
 363             metarig_collections = filter_layer_collections_by_object(self.usable_collections, self.metarig)
 364             self.layer_collection = (metarig_collections + [view_layer.layer_collection])[0]
 365             self.collection = self.layer_collection.collection
 366
 367         bpy.ops.object.mode_set(mode='OBJECT')
 368
 369         #------------------------------------------
 370         # Create/find the rig object and set it up
 371         obj = self.__create_rig_object()
 372
 373         # Get rid of anim data in case the rig already existed
 374         print("Clear rig animation data.")
 375
 376         obj.animation_data_clear()
 377         obj.data.animation_data_clear()
 378
 379         select_object(context, obj, deselect_all=True)
 380
 381         #------------------------------------------
 382         # Create Group widget
 383         self.__create_widget_group()
 384
 385         t.tick("Create main WGTS: ")
 386
 387         #------------------------------------------
 388         # Get parented objects to restore later
 389         childs = {}  # {object: bone}
 390         for child in obj.children:
 391             childs[child] = child.parent_bone
 392
 393         #------------------------------------------
 394         # Copy bones from metarig to obj (adds ORG_PREFIX)
 395         self.__duplicate_rig()
 396
 397         obj.data.use_mirror_x = False
 398
 399         t.tick("Duplicate rig: ")
 400
 401         #------------------------------------------
 402         # Put the rig_name in the armature custom properties
 403         rna_idprop_ui_prop_get(obj.data, "rig_id", create=True)
 404         obj.data["rig_id"] = self.rig_id
 405
 406         self.script = rig_ui_template.ScriptGenerator(self)
 407
 408         #------------------------------------------
 409         bpy.ops.object.mode_set(mode='OBJECT')
 410
 411         self.instantiate_rig_tree()
 412
 413         t.tick("Instantiate rigs: ")
 414
 415         #------------------------------------------
 416         bpy.ops.object.mode_set(mode='OBJECT')
 417
 418         self.invoke_initialize()
 419
 420         t.tick("Initialize rigs: ")
 421
 422         #------------------------------------------
 423         bpy.ops.object.mode_set(mode='EDIT')
 424
 425         self.invoke_prepare_bones()
 426
 427         t.tick("Prepare bones: ")
 428
 429         #------------------------------------------
 430         bpy.ops.object.mode_set(mode='OBJECT')
 431         bpy.ops.object.mode_set(mode='EDIT')
 432
 433         self.__create_root_bone()
 434
 435         self.invoke_generate_bones()
 436
 437         t.tick("Generate bones: ")
 438
 439         #------------------------------------------
 440         bpy.ops.object.mode_set(mode='OBJECT')
 441         bpy.ops.object.mode_set(mode='EDIT')
 442
 443         self.invoke_parent_bones()
 444
 445         self.__parent_bones_to_root()
 446
 447         t.tick("Parent bones: ")
 448
 449         #------------------------------------------
 450         bpy.ops.object.mode_set(mode='OBJECT')
 451
 452         self.invoke_configure_bones()
 453
 454         t.tick("Configure bones: ")
 455
 456         #------------------------------------------
 457         bpy.ops.object.mode_set(mode='OBJECT')
 458
 459         self.invoke_preapply_bones()
 460
 461         t.tick("Preapply bones: ")
 462
 463         #------------------------------------------
 464         bpy.ops.object.mode_set(mode='EDIT')
 465
 466         self.invoke_apply_bones()
 467
 468         t.tick("Apply bones: ")
 469
 470         #------------------------------------------
 471         bpy.ops.object.mode_set(mode='OBJECT')
 472
 473         self.invoke_rig_bones()
 474
 475         t.tick("Rig bones: ")
 476
 477         #------------------------------------------
 478         bpy.ops.object.mode_set(mode='OBJECT')
 479
 480         create_root_widget(obj, "root")
 481
 482         self.invoke_generate_widgets()
 483
 484         t.tick("Generate widgets: ")
 485
 486         #------------------------------------------
 487         bpy.ops.object.mode_set(mode='OBJECT')
 488
 489         self.__lock_transforms()
 490         self.__assign_layers()
 491         self.__compute_visible_layers()
 492         self.__restore_driver_vars()
 493
 494         t.tick("Assign layers: ")
 495
 496         #------------------------------------------
 497         bpy.ops.object.mode_set(mode='OBJECT')
 498
 499         self.invoke_finalize()
 500
 501         t.tick("Finalize: ")
 502
 503         #------------------------------------------
 504         bpy.ops.object.mode_set(mode='OBJECT')
 505
 506         self.__assign_widgets()
 507
 508         # Create Selection Sets
 509         create_selection_sets(obj, metarig)
 510
 511         # Create Bone Groups
 512         create_bone_groups(obj, metarig, self.layer_group_priorities)
 513
 514         t.tick("The rest: ")
 515
 516         #----------------------------------
 517         # Deconfigure
 518         bpy.ops.object.mode_set(mode='OBJECT')
 519         obj.data.pose_position = 'POSE'
 520
 521         # Restore parent to bones
 522         for child, sub_parent in childs.items():
 523             if sub_parent in obj.pose.bones:
 524                 mat = child.matrix_world.copy()
 525                 child.parent_bone = sub_parent
 526                 child.matrix_world = mat
 527
 528         #----------------------------------
 529         # Restore active collection
 530         view_layer.active_layer_collection = self.layer_collection
 531
 532
 533 def generate_rig(context, metarig):
 534     """ Generates a rig from a metarig.
 535
 536     """
 537     # Initial configuration
 538     rest_backup = metarig.data.pose_position
 539     metarig.data.pose_position = 'REST'
 540
 541     try:
 542         Generator(context, metarig).generate()
 543
 544         metarig.data.pose_position = rest_backup
 545
 546     except Exception as e:
 547         # Cleanup if something goes wrong
 548         print("Rigify: failed to generate rig.")
 549
 550         bpy.ops.object.mode_set(mode='OBJECT')
 551         metarig.data.pose_position = rest_backup
 552
 553         # Continue the exception
 554         raise e
 555
 556
 557 def create_selection_set_for_rig_layer(
 558         rig: bpy.types.Object,
 559         set_name: str,
 560         layer_idx: int
 561     ) -> None:
 562     """Create a single selection set on a rig.
 563
 564     The set will contain all bones on the rig layer with the given index.
 565     """
 566     selset = rig.selection_sets.add()
 567     selset.name = set_name
 568
 569     for b in rig.pose.bones:
 570         if not b.bone.layers[layer_idx] or b.name in selset.bone_ids:
 571             continue
 572
 573         bone_id = selset.bone_ids.add()
 574         bone_id.name = b.name
 575
 576 def create_selection_sets(obj, metarig):
 577     """Create selection sets if the Selection Sets addon is enabled.
 578
 579     Whether a selection set for a rig layer is created is controlled in the
 580     Rigify Layer Names panel.
 581     """
 582     # Check if selection sets addon is installed
 583     if 'bone_selection_groups' not in bpy.context.preferences.addons \
 584             and 'bone_selection_sets' not in bpy.context.preferences.addons:
 585         return
 586
 587     obj.selection_sets.clear()
 588
 589     for i, name in enumerate(metarig.data.rigify_layers.keys()):
 590         if name == '' or not metarig.data.rigify_layers[i].selset:
 591             continue
 592
 593         create_selection_set_for_rig_layer(obj, name, i)
 594
 595
 596 def create_bone_groups(obj, metarig, priorities={}):
 597
 598     bpy.ops.object.mode_set(mode='OBJECT')
 599     pb = obj.pose.bones
 600     layers = metarig.data.rigify_layers
 601     groups = metarig.data.rigify_colors
 602     dummy = {}
 603
 604     # Create BGs
 605     for l in layers:
 606         if l.group == 0:
 607             continue
 608         g_id = l.group - 1
 609         name = groups[g_id].name
 610         if name not in obj.pose.bone_groups.keys():
 611             bg = obj.pose.bone_groups.new(name=name)
 612             bg.color_set = 'CUSTOM'
 613             bg.colors.normal = gamma_correct(groups[g_id].normal)
 614             bg.colors.select = gamma_correct(groups[g_id].select)
 615             bg.colors.active = gamma_correct(groups[g_id].active)
 616
 617     for b in pb:
 618         try:
 619             prios = priorities.get(b.name, dummy)
 620             enabled = [ i for i, v in enumerate(b.bone.layers) if v ]
 621             layer_index = max(enabled, key=lambda i: prios.get(i, 0))
 622         except ValueError:
 623             continue
 624         if layer_index > len(layers) - 1:   # bone is on reserved layers
 625             continue
 626         g_id = layers[layer_index].group - 1
 627         if g_id >= 0:
 628             name = groups[g_id].name
 629             b.bone_group = obj.pose.bone_groups[name]
 630
 631
 632 def get_xy_spread(bones):
 633     x_max = 0
 634     y_max = 0
 635     for b in bones:
 636         x_max = max((x_max, abs(b.head[0]), abs(b.tail[0])))
 637         y_max = max((y_max, abs(b.head[1]), abs(b.tail[1])))
 638
 639     return max((x_max, y_max))
 640
 641
 642 def param_matches_type(param_name, rig_type):
 643     """ Returns True if the parameter name is consistent with the rig type.
 644     """
 645     if param_name.rsplit(".", 1)[0] == rig_type:
 646         return True
 647     else:
 648         return False
 649
 650
 651 def param_name(param_name, rig_type):
 652     """ Get the actual parameter name, sans-rig-type.
 653     """
 654     return param_name[len(rig_type) + 1:]