Skip to content
Snippets Groups Projects
  • Alexander Gavrilov's avatar
    3423174b
    Rigify: redesign generate.py and introduce a base rig class. · 3423174b
    Alexander Gavrilov authored
    The main goals are to provide an official way for rigs to
    interact in a structured way, and to remove mode switching
    within rigs.
    
    This involves introducing a base class for rigs that holds
    rig-to-rig and rig-to-bone references, converting the main
    generator into a class and passing it to rigs, and splitting
    the single generate method into multiple passes.
    
    For backward compatibility, old rigs are automatically handled
    via a wrapper that translates between old and new API.
    
    In addition, a way to create objects that receive the generate
    callbacks that aren't rigs is introduced via the GeneratorPlugin
    class. The UI script generation code is converted into a plugin.
    
    Making generic rig 'template' classes that are intended to be
    subclassed in specific rigs involves splitting operations done
    in each stage into multiple methods that can be overridden
    separately. The main callback thus ends up simply calling a
    sequence of other methods.
    
    To make such code cleaner it's better to allow registering
    those methods as new callbacks that would be automatically
    called by the system. This can be done via decorators.
    
    A new metaclass used for all rig and generate plugin classes
    builds and validates a table of all decorated methods, and
    allows calling them all together with the main callback.
    
    A new way to switch parents for IK bones based on the new
    features is introduced, and used in the existing limb rigs.
    
    Reviewers: icappiello campbellbarton
    
    Differential Revision: https://developer.blender.org/D4624
    3423174b
    History
    Rigify: redesign generate.py and introduce a base rig class.
    Alexander Gavrilov authored
    The main goals are to provide an official way for rigs to
    interact in a structured way, and to remove mode switching
    within rigs.
    
    This involves introducing a base class for rigs that holds
    rig-to-rig and rig-to-bone references, converting the main
    generator into a class and passing it to rigs, and splitting
    the single generate method into multiple passes.
    
    For backward compatibility, old rigs are automatically handled
    via a wrapper that translates between old and new API.
    
    In addition, a way to create objects that receive the generate
    callbacks that aren't rigs is introduced via the GeneratorPlugin
    class. The UI script generation code is converted into a plugin.
    
    Making generic rig 'template' classes that are intended to be
    subclassed in specific rigs involves splitting operations done
    in each stage into multiple methods that can be overridden
    separately. The main callback thus ends up simply calling a
    sequence of other methods.
    
    To make such code cleaner it's better to allow registering
    those methods as new callbacks that would be automatically
    called by the system. This can be done via decorators.
    
    A new metaclass used for all rig and generate plugin classes
    builds and validates a table of all decorated methods, and
    allows calling them all together with the main callback.
    
    A new way to switch parents for IK bones based on the new
    features is introduced, and used in the existing limb rigs.
    
    Reviewers: icappiello campbellbarton
    
    Differential Revision: https://developer.blender.org/D4624
base_generate.py 14.06 KiB
#====================== BEGIN GPL LICENSE BLOCK ======================
#
#  This program is free software; you can redistribute it and/or
#  modify it under the terms of the GNU General Public License
#  as published by the Free Software Foundation; either version 2
#  of the License, or (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software Foundation,
#  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
#======================= END GPL LICENSE BLOCK ========================

# <pep8 compliant>

import bpy
import sys
import traceback

from .utils.errors import MetarigError, RaiseErrorMixin
from .utils.naming import random_id
from .utils.metaclass import SingletonPluginMetaclass
from .utils.rig import list_bone_names_depth_first_sorted, get_rigify_type
from .utils.misc import assign_parameters

from . import base_rig


#=============================================
# Generator Plugin
#=============================================


class GeneratorPlugin(base_rig.GenerateCallbackHost, metaclass=SingletonPluginMetaclass):
    """
    Base class for generator plugins.

    Generator plugins are per-Generator singleton utility
    classes that receive the same stage callbacks as rigs.

    Useful for building entities shared by multiple rigs
    (e.g. the python script), or for making fire-and-forget
    utilities that actually require multiple stages to
    complete.

    This will create only one instance per set of args:

      instance = PluginClass(generator, ...init args)
    """

    priority = 0

    def __init__(self, generator):
        self.generator = generator
        self.obj = generator.obj

    def register_new_bone(self, new_name, old_name=None):
        self.generator.bone_owners[new_name] = None


#=============================================
# Rig Substitution Mechanism
#=============================================


class SubstitutionRig(RaiseErrorMixin):
    """A proxy rig that replaces itself with one or more different rigs."""

    def __init__(self, generator, pose_bone):
        self.generator = generator

        self.obj = generator.obj
        self.base_bone = pose_bone.name
        self.params = pose_bone.rigify_parameters

    def substitute(self):
        # return [rig1, rig2...]
        raise NotImplementedException()

    # Utility methods
    def register_new_bone(self, new_name, old_name=None):
        pass

    def get_params(self, bone_name):
        return self.obj.pose.bones[bone_name].rigify_parameters

    def assign_params(self, bone_name, param_dict=None, **params):
        assign_parameters(self.get_params(bone_name), param_dict, **params)

    def instantiate_rig(self, rig_class, bone_name):
        if isinstance(rig_class, str):
            rig_class = self.generator.find_rig_class(rig_class)

        return self.generator.instantiate_rig(rig_class, self.obj.pose.bones[bone_name])


#=============================================
# Legacy Rig Wrapper
#=============================================


class LegacyRig(base_rig.BaseRig):
    """Wrapper around legacy style rigs without a common base class"""

    def __init__(self, generator, pose_bone, wrapped_class):
        self.wrapped_rig = None
        self.wrapped_class = wrapped_class

        super().__init__(generator, pose_bone)

    def find_org_bones(self, pose_bone):
        bone_name = pose_bone.name

        if not self.wrapped_rig:
            self.wrapped_rig = self.wrapped_class(self.obj, self.base_bone, self.params)

            # Switch back to OBJECT mode if the rig changed it
            if self.obj.mode != 'OBJECT':
                bpy.ops.object.mode_set(mode='OBJECT')

        # Try to extract the main list of bones - old rigs often have it.
        # This is not actually strictly necessary, so failing is OK.
        if hasattr(self.wrapped_rig, 'org_bones'):
            bones = self.wrapped_rig.org_bones
            if isinstance(bones, list):
                return bones

        return [bone_name]

    def generate_bones(self):
        # Inject references into the rig if it won't cause conflict
        if not hasattr(self.wrapped_rig, 'rigify_generator'):
            self.wrapped_rig.rigify_generator = self.generator
        if not hasattr(self.wrapped_rig, 'rigify_wrapper'):
            self.wrapped_rig.rigify_wrapper = self

        # Old rigs only have one generate method, so call it from
        # generate_bones, which is the only stage allowed to add bones.
        scripts = self.wrapped_rig.generate()

        # Switch back to EDIT mode if the rig changed it
        if self.obj.mode != 'EDIT':
            bpy.ops.object.mode_set(mode='EDIT')

        if isinstance(scripts, dict):
            if 'script' in scripts:
                self.script.add_panel_code(scripts['script'])
            if 'imports' in scripts:
                self.script.add_imports(scripts['imports'])
            if 'utilities' in scripts:
                self.script.add_utilities(scripts['utilities'])
            if 'register' in scripts:
                self.script.register_classes(scripts['register'])
            if 'register_drivers' in scripts:
                self.script.register_driver_functions(scripts['register_drivers'])
            if 'register_props' in scripts:
                for prop, val in scripts['register_props']:
                    self.script.register_property(prop, val)
            if 'noparent_bones' in scripts:
                for bone_name in scripts['noparent_bones']:
                    self.generator.disable_auto_parent(bone_name)
        elif scripts is not None:
            self.script.add_panel_code([scripts[0]])

    def finalize(self):
        if hasattr(self.wrapped_rig, 'glue'):
            self.wrapped_rig.glue()

            # Switch back to OBJECT mode if the rig changed it
            if self.obj.mode != 'OBJECT':
                bpy.ops.object.mode_set(mode='OBJECT')


#=============================================
# Base Generate Engine
#=============================================


class BaseGenerator:
    """Base class for the main generator object. Contains rig and plugin management code."""

    def __init__(self, context, metarig):
        self.context = context
        self.scene = context.scene
        self.view_layer = context.view_layer
        self.layer_collection = context.layer_collection
        self.collection = self.layer_collection.collection
        self.metarig = metarig
        self.obj = None

        # List of all rig instances
        self.rig_list = []
        # List of rigs that don't have a parent
        self.root_rigs = []
        # Map from bone names to their rigs
        self.bone_owners = {}

        # Set of plugins
        self.plugin_list = []
        self.plugin_map = {}

        # Current execution stage so plugins could check they are used correctly
        self.stage = None

        # Set of bones that should be left without parent
        self.noparent_bones = set()

        # Random string with time appended so that
        # different rigs don't collide id's
        self.rig_id = random_id(16)


    def disable_auto_parent(self, bone_name):
        """Prevent automatically parenting the bone to root if parentless."""
        self.noparent_bones.add(bone_name)


    def __run_object_stage(self, method_name):
        assert(self.context.active_object == self.obj)
        assert(self.obj.mode == 'OBJECT')
        num_bones = len(self.obj.data.bones)

        self.stage = method_name

        for rig in [*self.rig_list, *self.plugin_list]:
            rig.rigify_invoke_stage(method_name)

            assert(self.context.active_object == self.obj)
            assert(self.obj.mode == 'OBJECT')
            assert(num_bones == len(self.obj.data.bones))


    def __run_edit_stage(self, method_name):
        assert(self.context.active_object == self.obj)
        assert(self.obj.mode == 'EDIT')
        num_bones = len(self.obj.data.edit_bones)

        self.stage = method_name

        for rig in [*self.rig_list, *self.plugin_list]:
            rig.rigify_invoke_stage(method_name)

            assert(self.context.active_object == self.obj)
            assert(self.obj.mode == 'EDIT')
            assert(num_bones == len(self.obj.data.edit_bones))


    def invoke_initialize(self):
        self.__run_object_stage('initialize')


    def invoke_prepare_bones(self):
        self.__run_edit_stage('prepare_bones')


    def __auto_register_bones(self, bones, rig):
        """Find bones just added and not registered by this rig."""
        for bone in bones:
            name = bone.name
            if name not in self.bone_owners:
                self.bone_owners[name] = rig
                if rig:
                    rig.rigify_new_bones[name] = None

                if not isinstance(rig, LegacyRig):
                    print("WARNING: rig %s didn't register bone %s\n" % (self.describe_rig(rig), name))


    def invoke_generate_bones(self):
        assert(self.context.active_object == self.obj)
        assert(self.obj.mode == 'EDIT')

        self.stage = 'generate_bones'

        for rig in self.rig_list:
            rig.rigify_invoke_stage('generate_bones')

            assert(self.context.active_object == self.obj)
            assert(self.obj.mode == 'EDIT')

            self.__auto_register_bones(self.obj.data.edit_bones, rig)

        for plugin in self.plugin_list:
            plugin.rigify_invoke_stage('generate_bones')

            assert(self.context.active_object == self.obj)
            assert(self.obj.mode == 'EDIT')

            self.__auto_register_bones(self.obj.data.edit_bones, None)


    def invoke_parent_bones(self):
        self.__run_edit_stage('parent_bones')


    def invoke_configure_bones(self):
        self.__run_object_stage('configure_bones')


    def invoke_apply_bones(self):
        self.__run_edit_stage('apply_bones')


    def invoke_rig_bones(self):
        self.__run_object_stage('rig_bones')


    def invoke_generate_widgets(self):
        self.__run_object_stage('generate_widgets')


    def invoke_finalize(self):
        self.__run_object_stage('finalize')


    def instantiate_rig(self, rig_class, pose_bone):
        assert not issubclass(rig_class, SubstitutionRig)

        if issubclass(rig_class, base_rig.BaseRig):
            return rig_class(self, pose_bone)
        else:
            return LegacyRig(self, pose_bone, rig_class)


    def instantiate_rig_by_type(self, rig_type, pose_bone):
        return self.instantiate_rig(self.find_rig_class(rig_type), pose_bone)


    def describe_rig(self, rig):
        base_bone = rig.base_bone

        if isinstance(rig, LegacyRig):
            rig = rig.wrapped_rig

        return "%s (%s)" % (rig.__class__, base_bone)


    def __create_rigs(self, bone_name, halt_on_missing):
        """Recursively walk bones and create rig instances."""

        pose_bone = self.obj.pose.bones[bone_name]

        rig_type = get_rigify_type(pose_bone)

        if rig_type != "":
            try:
                rig_class = self.find_rig_class(rig_type)

                if issubclass(rig_class, SubstitutionRig):
                    rigs = rig_class(self, pose_bone).substitute()
                else:
                    rigs = [self.instantiate_rig(rig_class, pose_bone)]

                assert(self.context.active_object == self.obj)
                assert(self.obj.mode == 'OBJECT')

                for rig in rigs:
                    self.rig_list.append(rig)

                    for org_name in rig.rigify_org_bones:
                        if org_name in self.bone_owners:
                            old_rig = self.describe_rig(self.bone_owners[org_name])
                            new_rig = self.describe_rig(rig)
                            print("CONFLICT: bone %s is claimed by rigs %s and %s\n" % (org_name, old_rig, new_rig))

                        self.bone_owners[org_name] = rig

            except ImportError:
                message = "Rig Type Missing: python module for type '%s' not found (bone: %s)" % (rig_type, bone_name)
                if halt_on_missing:
                    raise MetarigError(message)
                else:
                    print(message)
                    print('print_exc():')
                    traceback.print_exc(file=sys.stdout)


    def __build_rig_tree_rec(self, bone, current_rig, handled):
        """Recursively walk bones and connect rig instances into a tree."""

        rig = self.bone_owners.get(bone.name)

        if rig:
            if rig is current_rig:
                pass

            elif rig not in handled:
                rig.rigify_parent = current_rig

                if current_rig:
                    current_rig.rigify_children.append(rig)
                else:
                    self.root_rigs.append(rig)

                handled[rig] = bone.name

            elif rig.rigify_parent is not current_rig:
                raise MetarigError("CONFLICT: bone %s owned by rig %s has different parent rig from %s\n" %
                                   (bone.name, rig.base_bone, handled[rig]))

            current_rig = rig
        else:
            if current_rig:
                current_rig.rigify_child_bones.add(bone.name)

            self.bone_owners[bone.name] = current_rig

        for child in bone.children:
            self.__build_rig_tree_rec(child, current_rig, handled)


    def instantiate_rig_tree(self, halt_on_missing=False):
        """Create rig instances and connect them into a tree."""

        assert(self.context.active_object == self.obj)
        assert(self.obj.mode == 'OBJECT')

        # Construct the rig instances
        for name in list_bone_names_depth_first_sorted(self.obj):
            self.__create_rigs(name, halt_on_missing)

        # Connect rigs and bones into a tree
        handled = {}

        for bone in self.obj.data.bones:
            if bone.parent is None:
                self.__build_rig_tree_rec(bone, None, handled)