rigify/base_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 sys
  23 import traceback
  24 import collections
  25
  26 from .utils.errors import MetarigError, RaiseErrorMixin
  27 from .utils.naming import random_id
  28 from .utils.metaclass import SingletonPluginMetaclass
  29 from .utils.rig import list_bone_names_depth_first_sorted, get_rigify_type
  30 from .utils.misc import clone_parameters, assign_parameters
  31
  32 from . import base_rig
  33
  34
  35 #=============================================
  36 # Generator Plugin
  37 #=============================================
  38
  39
  40 class GeneratorPlugin(base_rig.GenerateCallbackHost, metaclass=SingletonPluginMetaclass):
  41     """
  42     Base class for generator plugins.
  43
  44     Generator plugins are per-Generator singleton utility
  45     classes that receive the same stage callbacks as rigs.
  46
  47     Useful for building entities shared by multiple rigs
  48     (e.g. the python script), or for making fire-and-forget
  49     utilities that actually require multiple stages to
  50     complete.
  51
  52     This will create only one instance per set of args:
  53
  54       instance = PluginClass(generator, ...init args)
  55     """
  56
  57     priority = 0
  58
  59     def __init__(self, generator):
  60         self.generator = generator
  61         self.obj = generator.obj
  62
  63     def register_new_bone(self, new_name, old_name=None):
  64         self.generator.bone_owners[new_name] = None
  65
  66
  67 #=============================================
  68 # Rig Substitution Mechanism
  69 #=============================================
  70
  71
  72 class SubstitutionRig(RaiseErrorMixin):
  73     """A proxy rig that replaces itself with one or more different rigs."""
  74
  75     def __init__(self, generator, pose_bone):
  76         self.generator = generator
  77
  78         self.obj = generator.obj
  79         self.base_bone = pose_bone.name
  80         self.params = pose_bone.rigify_parameters
  81         self.params_copy = clone_parameters(self.params)
  82
  83     def substitute(self):
  84         # return [rig1, rig2...]
  85         raise NotImplementedException()
  86
  87     # Utility methods
  88     def register_new_bone(self, new_name, old_name=None):
  89         pass
  90
  91     def get_params(self, bone_name):
  92         return self.obj.pose.bones[bone_name].rigify_parameters
  93
  94     def assign_params(self, bone_name, param_dict=None, **params):
  95         assign_parameters(self.get_params(bone_name), param_dict, **params)
  96
  97     def instantiate_rig(self, rig_class, bone_name):
  98         if isinstance(rig_class, str):
  99             rig_class = self.generator.find_rig_class(rig_class)
 100
 101         return self.generator.instantiate_rig(rig_class, self.obj.pose.bones[bone_name])
 102
 103
 104 #=============================================
 105 # Legacy Rig Wrapper
 106 #=============================================
 107
 108
 109 class LegacyRig(base_rig.BaseRig):
 110     """Wrapper around legacy style rigs without a common base class"""
 111
 112     def __init__(self, generator, pose_bone, wrapped_class):
 113         self.wrapped_rig = None
 114         self.wrapped_class = wrapped_class
 115
 116         super().__init__(generator, pose_bone)
 117
 118     def find_org_bones(self, pose_bone):
 119         bone_name = pose_bone.name
 120
 121         if not self.wrapped_rig:
 122             self.wrapped_rig = self.wrapped_class(self.obj, self.base_bone, self.params)
 123
 124             # Switch back to OBJECT mode if the rig changed it
 125             if self.obj.mode != 'OBJECT':
 126                 bpy.ops.object.mode_set(mode='OBJECT')
 127
 128         # Try to extract the main list of bones - old rigs often have it.
 129         # This is not actually strictly necessary, so failing is OK.
 130         if hasattr(self.wrapped_rig, 'org_bones'):
 131             bones = self.wrapped_rig.org_bones
 132             if isinstance(bones, list):
 133                 return bones
 134
 135         return [bone_name]
 136
 137     def generate_bones(self):
 138         # Inject references into the rig if it won't cause conflict
 139         if not hasattr(self.wrapped_rig, 'rigify_generator'):
 140             self.wrapped_rig.rigify_generator = self.generator
 141         if not hasattr(self.wrapped_rig, 'rigify_wrapper'):
 142             self.wrapped_rig.rigify_wrapper = self
 143
 144         # Old rigs only have one generate method, so call it from
 145         # generate_bones, which is the only stage allowed to add bones.
 146         scripts = self.wrapped_rig.generate()
 147
 148         # Switch back to EDIT mode if the rig changed it
 149         if self.obj.mode != 'EDIT':
 150             bpy.ops.object.mode_set(mode='EDIT')
 151
 152         if isinstance(scripts, dict):
 153             if 'script' in scripts:
 154                 self.script.add_panel_code(scripts['script'])
 155             if 'imports' in scripts:
 156                 self.script.add_imports(scripts['imports'])
 157             if 'utilities' in scripts:
 158                 self.script.add_utilities(scripts['utilities'])
 159             if 'register' in scripts:
 160                 self.script.register_classes(scripts['register'])
 161             if 'register_drivers' in scripts:
 162                 self.script.register_driver_functions(scripts['register_drivers'])
 163             if 'register_props' in scripts:
 164                 for prop, val in scripts['register_props']:
 165                     self.script.register_property(prop, val)
 166             if 'noparent_bones' in scripts:
 167                 for bone_name in scripts['noparent_bones']:
 168                     self.generator.disable_auto_parent(bone_name)
 169         elif scripts is not None:
 170             self.script.add_panel_code([scripts[0]])
 171
 172     def finalize(self):
 173         if hasattr(self.wrapped_rig, 'glue'):
 174             self.wrapped_rig.glue()
 175
 176             # Switch back to OBJECT mode if the rig changed it
 177             if self.obj.mode != 'OBJECT':
 178                 bpy.ops.object.mode_set(mode='OBJECT')
 179
 180
 181 #=============================================
 182 # Base Generate Engine
 183 #=============================================
 184
 185
 186 class BaseGenerator:
 187     """Base class for the main generator object. Contains rig and plugin management code."""
 188
 189     def __init__(self, context, metarig):
 190         self.context = context
 191         self.scene = context.scene
 192         self.view_layer = context.view_layer
 193         self.layer_collection = context.layer_collection
 194         self.collection = self.layer_collection.collection
 195         self.metarig = metarig
 196         self.obj = None
 197
 198         # List of all rig instances
 199         self.rig_list = []
 200         # List of rigs that don't have a parent
 201         self.root_rigs = []
 202         # Map from bone names to their rigs
 203         self.bone_owners = {}
 204
 205         # Set of plugins
 206         self.plugin_list = []
 207         self.plugin_map = {}
 208
 209         # Current execution stage so plugins could check they are used correctly
 210         self.stage = None
 211
 212         # Set of bones that should be left without parent
 213         self.noparent_bones = set()
 214
 215         # Table of layer priorities for defining bone groups
 216         self.layer_group_priorities = collections.defaultdict(dict)
 217
 218         # Random string with time appended so that
 219         # different rigs don't collide id's
 220         self.rig_id = random_id(16)
 221
 222         # Table of renamed ORG bones
 223         self.org_rename_table = dict()
 224
 225
 226     def disable_auto_parent(self, bone_name):
 227         """Prevent automatically parenting the bone to root if parentless."""
 228         self.noparent_bones.add(bone_name)
 229
 230
 231     def set_layer_group_priority(self, bone_name, layers, priority):
 232         for i, val in enumerate(layers):
 233             if val:
 234                 self.layer_group_priorities[bone_name][i] = priority
 235
 236
 237     def rename_org_bone(self, old_name, new_name):
 238         assert self.stage == 'instantiate'
 239         assert old_name == self.org_rename_table.get(old_name, None)
 240         assert old_name not in self.bone_owners
 241
 242         bone = self.obj.data.bones[old_name]
 243
 244         bone.name = new_name
 245         new_name = bone.name
 246
 247         self.org_rename_table[old_name] = new_name
 248         return new_name
 249
 250
 251     def __run_object_stage(self, method_name):
 252         assert(self.context.active_object == self.obj)
 253         assert(self.obj.mode == 'OBJECT')
 254         num_bones = len(self.obj.data.bones)
 255
 256         self.stage = method_name
 257
 258         for rig in [*self.rig_list, *self.plugin_list]:
 259             rig.rigify_invoke_stage(method_name)
 260
 261             assert(self.context.active_object == self.obj)
 262             assert(self.obj.mode == 'OBJECT')
 263             assert(num_bones == len(self.obj.data.bones))
 264
 265
 266     def __run_edit_stage(self, method_name):
 267         assert(self.context.active_object == self.obj)
 268         assert(self.obj.mode == 'EDIT')
 269         num_bones = len(self.obj.data.edit_bones)
 270
 271         self.stage = method_name
 272
 273         for rig in [*self.rig_list, *self.plugin_list]:
 274             rig.rigify_invoke_stage(method_name)
 275
 276             assert(self.context.active_object == self.obj)
 277             assert(self.obj.mode == 'EDIT')
 278             assert(num_bones == len(self.obj.data.edit_bones))
 279
 280
 281     def invoke_initialize(self):
 282         self.__run_object_stage('initialize')
 283
 284
 285     def invoke_prepare_bones(self):
 286         self.__run_edit_stage('prepare_bones')
 287
 288
 289     def __auto_register_bones(self, bones, rig):
 290         """Find bones just added and not registered by this rig."""
 291         for bone in bones:
 292             name = bone.name
 293             if name not in self.bone_owners:
 294                 self.bone_owners[name] = rig
 295                 if rig:
 296                     rig.rigify_new_bones[name] = None
 297
 298                 if not isinstance(rig, LegacyRig):
 299                     print("WARNING: rig %s didn't register bone %s\n" % (self.describe_rig(rig), name))
 300
 301
 302     def invoke_generate_bones(self):
 303         assert(self.context.active_object == self.obj)
 304         assert(self.obj.mode == 'EDIT')
 305
 306         self.stage = 'generate_bones'
 307
 308         for rig in self.rig_list:
 309             rig.rigify_invoke_stage('generate_bones')
 310
 311             assert(self.context.active_object == self.obj)
 312             assert(self.obj.mode == 'EDIT')
 313
 314             self.__auto_register_bones(self.obj.data.edit_bones, rig)
 315
 316         for plugin in self.plugin_list:
 317             plugin.rigify_invoke_stage('generate_bones')
 318
 319             assert(self.context.active_object == self.obj)
 320             assert(self.obj.mode == 'EDIT')
 321
 322             self.__auto_register_bones(self.obj.data.edit_bones, None)
 323
 324
 325     def invoke_parent_bones(self):
 326         self.__run_edit_stage('parent_bones')
 327
 328
 329     def invoke_configure_bones(self):
 330         self.__run_object_stage('configure_bones')
 331
 332
 333     def invoke_apply_bones(self):
 334         self.__run_edit_stage('apply_bones')
 335
 336
 337     def invoke_rig_bones(self):
 338         self.__run_object_stage('rig_bones')
 339
 340
 341     def invoke_generate_widgets(self):
 342         self.__run_object_stage('generate_widgets')
 343
 344
 345     def invoke_finalize(self):
 346         self.__run_object_stage('finalize')
 347
 348
 349     def instantiate_rig(self, rig_class, pose_bone):
 350         assert not issubclass(rig_class, SubstitutionRig)
 351
 352         if issubclass(rig_class, base_rig.BaseRig):
 353             return rig_class(self, pose_bone)
 354         else:
 355             return LegacyRig(self, pose_bone, rig_class)
 356
 357
 358     def instantiate_rig_by_type(self, rig_type, pose_bone):
 359         return self.instantiate_rig(self.find_rig_class(rig_type), pose_bone)
 360
 361
 362     def describe_rig(self, rig):
 363         base_bone = rig.base_bone
 364
 365         if isinstance(rig, LegacyRig):
 366             rig = rig.wrapped_rig
 367
 368         return "%s (%s)" % (rig.__class__, base_bone)
 369
 370
 371     def __create_rigs(self, bone_name, halt_on_missing):
 372         """Recursively walk bones and create rig instances."""
 373
 374         pose_bone = self.obj.pose.bones[bone_name]
 375
 376         rig_type = get_rigify_type(pose_bone)
 377
 378         if rig_type != "":
 379             try:
 380                 rig_class = self.find_rig_class(rig_type)
 381
 382                 if issubclass(rig_class, SubstitutionRig):
 383                     rigs = rig_class(self, pose_bone).substitute()
 384                 else:
 385                     rigs = [self.instantiate_rig(rig_class, pose_bone)]
 386
 387                 assert(self.context.active_object == self.obj)
 388                 assert(self.obj.mode == 'OBJECT')
 389
 390                 for rig in rigs:
 391                     self.rig_list.append(rig)
 392
 393                     for org_name in rig.rigify_org_bones:
 394                         if org_name in self.bone_owners:
 395                             old_rig = self.describe_rig(self.bone_owners[org_name])
 396                             new_rig = self.describe_rig(rig)
 397                             print("CONFLICT: bone %s is claimed by rigs %s and %s\n" % (org_name, old_rig, new_rig))
 398
 399                         self.bone_owners[org_name] = rig
 400
 401             except ImportError:
 402                 message = "Rig Type Missing: python module for type '%s' not found (bone: %s)" % (rig_type, bone_name)
 403                 if halt_on_missing:
 404                     raise MetarigError(message)
 405                 else:
 406                     print(message)
 407                     print('print_exc():')
 408                     traceback.print_exc(file=sys.stdout)
 409
 410
 411     def __build_rig_tree_rec(self, bone, current_rig, handled):
 412         """Recursively walk bones and connect rig instances into a tree."""
 413
 414         rig = self.bone_owners.get(bone.name)
 415
 416         if rig:
 417             if rig is current_rig:
 418                 pass
 419
 420             elif rig not in handled:
 421                 rig.rigify_parent = current_rig
 422
 423                 if current_rig:
 424                     current_rig.rigify_children.append(rig)
 425                 else:
 426                     self.root_rigs.append(rig)
 427
 428                 handled[rig] = bone.name
 429
 430             elif rig.rigify_parent is not current_rig:
 431                 raise MetarigError("CONFLICT: bone %s owned by rig %s has different parent rig from %s\n" %
 432                                    (bone.name, rig.base_bone, handled[rig]))
 433
 434             current_rig = rig
 435         else:
 436             if current_rig:
 437                 current_rig.rigify_child_bones.add(bone.name)
 438
 439             self.bone_owners[bone.name] = current_rig
 440
 441         for child in bone.children:
 442             self.__build_rig_tree_rec(child, current_rig, handled)
 443
 444
 445     def instantiate_rig_tree(self, halt_on_missing=False):
 446         """Create rig instances and connect them into a tree."""
 447
 448         assert(self.context.active_object == self.obj)
 449         assert(self.obj.mode == 'OBJECT')
 450
 451         self.stage = 'instantiate'
 452
 453         # Compute the list of bones
 454         bone_list = list_bone_names_depth_first_sorted(self.obj)
 455
 456         self.org_rename_table = {n: n for n in bone_list}
 457
 458         # Construct the rig instances
 459         for name in bone_list:
 460             self.__create_rigs(self.org_rename_table[name], halt_on_missing)
 461
 462         # Connect rigs and bones into a tree
 463         handled = {}
 464
 465         for bone in self.obj.data.bones:
 466             if bone.parent is None:
 467                 self.__build_rig_tree_rec(bone, None, handled)