generate.py

#====================== 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 re
import time
from rna_prop_ui import rna_idprop_ui_prop_get

from .utils.bones import new_bone
from .utils.layers import ORG_LAYER, MCH_LAYER, DEF_LAYER, ROOT_LAYER
from .utils.naming import ORG_PREFIX, MCH_PREFIX, DEF_PREFIX, ROOT_NAME, make_original_name
from .utils.widgets import WGT_PREFIX
from .utils.widgets_special import create_root_widget
from .utils.misc import gamma_correct, select_object
from .utils.collections import ensure_widget_collection, list_layer_collections, filter_layer_collections_by_object
from .utils.rig import get_rigify_type

from . import base_generate
from . import rig_ui_template
from . import rig_lists

RIG_MODULE = "rigs"

class Timer:
    def __init__(self):
        self.timez = time.time()

    def tick(self, string):
        t = time.time()
        print(string + "%.3f" % (t - self.timez))
        self.timez = t


class Generator(base_generate.BaseGenerator):
    def __init__(self, context, metarig):
        super().__init__(context, metarig)

        self.id_store = context.window_manager

        self.rig_new_name = ""
        self.rig_old_name = ""


    def find_rig_class(self, rig_type):
        rig_module = rig_lists.rigs[rig_type]["module"]

        return rig_module.Rig


    def __create_rig_object(self):
        scene = self.scene
        id_store = self.id_store
        meta_data = self.metarig.data

        # Check if the generated rig already exists, so we can
        # regenerate in the same object.  If not, create a new
        # object to generate the rig in.
        print("Fetch rig.")

        self.rig_new_name = name = meta_data.rigify_rig_basename or "rig"

        obj = None

        if meta_data.rigify_generate_mode == 'overwrite':
            obj = meta_data.rigify_target_rig

            if not obj and name in scene.objects:
                obj = scene.objects[name]

            if obj:
                self.rig_old_name = obj.name

                obj.name = name
                obj.data.name = obj.name

                rig_collections = filter_layer_collections_by_object(self.usable_collections, obj)
                self.layer_collection = (rig_collections + [self.layer_collection])[0]
                self.collection = self.layer_collection.collection

            elif name in bpy.data.objects:
                obj = bpy.data.objects[name]

        if not obj:
            obj = bpy.data.objects.new(name, bpy.data.armatures.new(name))
            obj.display_type = 'WIRE'
            self.collection.objects.link(obj)

        elif obj.name not in self.collection.objects:  # rig exists but was deleted
            self.collection.objects.link(obj)

        meta_data.rigify_target_rig = obj
        obj.data.pose_position = 'POSE'

        self.obj = obj
        return obj


    def __create_widget_group(self, new_group_name):
        context = self.context
        scene = self.scene
        id_store = self.id_store

        # Create/find widge collection
        self.widget_collection = ensure_widget_collection(context)

        # Remove wgts if force update is set
        wgts_group_name = "WGTS_" + (self.rig_old_name or self.obj.name)
        if wgts_group_name in scene.objects and self.metarig.data.rigify_force_widget_update:
            bpy.ops.object.mode_set(mode='OBJECT')
            bpy.ops.object.select_all(action='DESELECT')
            for wgt in bpy.data.objects[wgts_group_name].children:
                wgt.select_set(True)
            bpy.ops.object.delete(use_global=False)
            if self.rig_old_name:
                bpy.data.objects[wgts_group_name].name = new_group_name

        # Create Group widget
        wgts_group_name = new_group_name
        if wgts_group_name not in scene.objects:
            if wgts_group_name in bpy.data.objects:
                bpy.data.objects[wgts_group_name].user_clear()
                bpy.data.objects.remove(bpy.data.objects[wgts_group_name])
            mesh = bpy.data.meshes.new(wgts_group_name)
            wgts_obj = bpy.data.objects.new(wgts_group_name, mesh)
            self.widget_collection.objects.link(wgts_obj)

        self.wgts_group_name = new_group_name


    def __duplicate_rig(self):
        obj = self.obj
        metarig = self.metarig
        context = self.context

        # Remove all bones from the generated rig armature.
        bpy.ops.object.mode_set(mode='EDIT')
        for bone in obj.data.edit_bones:
            obj.data.edit_bones.remove(bone)
        bpy.ops.object.mode_set(mode='OBJECT')

        # Select and duplicate metarig
        select_object(context, metarig, deselect_all=True)

        bpy.ops.object.duplicate()

        # Select the target rig and join
        select_object(context, obj)

        saved_matrix = obj.matrix_world.copy()
        obj.matrix_world = metarig.matrix_world

        bpy.ops.object.join()

        obj.matrix_world = saved_matrix

        # Select the generated rig
        select_object(context, obj, deselect_all=True)

        # Clean up animation data
        if obj.animation_data:
            obj.animation_data.action = None

            for track in obj.animation_data.nla_tracks:
                obj.animation_data.nla_tracks.remove(track)

            # Freeze drivers referring to custom properties
            for d in obj.animation_data.drivers:
                for var in d.driver.variables:
                    for tar in var.targets:
                        # If a custom property
                        if var.type == 'SINGLE_PROP' \
                        and re.match('^pose.bones\["[^"\]]*"\]\["[^"\]]*"\]$', tar.data_path):
                            tar.data_path = "RIGIFY-" + tar.data_path


    def __rename_org_bones(self):
        obj = self.obj

        #----------------------------------
        # Make a list of the original bones so we can keep track of them.
        original_bones = [bone.name for bone in obj.data.bones]

        # Add the ORG_PREFIX to the original bones.
        for i in range(0, len(original_bones)):
            bone = obj.pose.bones[original_bones[i]]

            # This rig type is special in that it preserves the name of the bone.
            if get_rigify_type(bone) != 'basic.raw_copy':
                bone.name = make_original_name(original_bones[i])
                original_bones[i] = bone.name

        self.original_bones = original_bones


    def __create_root_bone(self):
        obj = self.obj
        metarig = self.metarig

        #----------------------------------
        # Create the root bone.
        root_bone = new_bone(obj, ROOT_NAME)
        spread = get_xy_spread(metarig.data.bones) or metarig.data.bones[0].length
        spread = float('%.3g' % spread)
        scale = spread/0.589
        obj.data.edit_bones[root_bone].head = (0, 0, 0)
        obj.data.edit_bones[root_bone].tail = (0, scale, 0)
        obj.data.edit_bones[root_bone].roll = 0
        self.root_bone = root_bone
        self.bone_owners[root_bone] = None


    def __parent_bones_to_root(self):
        eb = self.obj.data.edit_bones

        # Parent loose bones to root
        for bone in eb:
            if bone.name in self.noparent_bones:
                continue
            elif bone.parent is None:
                bone.use_connect = False
                bone.parent = eb[self.root_bone]


    def __lock_transforms(self):
        # Lock transforms on all non-control bones
        r = re.compile("[A-Z][A-Z][A-Z]-")
        for pb in self.obj.pose.bones:
            if r.match(pb.name):
                pb.lock_location = (True, True, True)
                pb.lock_rotation = (True, True, True)
                pb.lock_rotation_w = True
                pb.lock_scale = (True, True, True)


    def __assign_layers(self):
        bones = self.obj.data.bones

        bones[self.root_bone].layers = ROOT_LAYER

        # Every bone that has a name starting with "DEF-" make deforming.  All the
        # others make non-deforming.
        for bone in bones:
            name = bone.name

            bone.use_deform = name.startswith(DEF_PREFIX)

            # Move all the original bones to their layer.
            if name.startswith(ORG_PREFIX):
                bone.layers = ORG_LAYER
            # Move all the bones with names starting with "MCH-" to their layer.
            elif name.startswith(MCH_PREFIX):
                bone.layers = MCH_LAYER
            # Move all the bones with names starting with "DEF-" to their layer.
            elif name.startswith(DEF_PREFIX):
                bone.layers = DEF_LAYER

            bone.bbone_x = bone.bbone_z = bone.length * 0.05


    def __restore_driver_vars(self):
        obj = self.obj

        # Alter marked driver targets
        if obj.animation_data:
            for d in obj.animation_data.drivers:
                for v in d.driver.variables:
                    for tar in v.targets:
                        if tar.data_path.startswith("RIGIFY-"):
                            temp, bone, prop = tuple([x.strip('"]') for x in tar.data_path.split('["')])
                            if bone in obj.data.bones \
                            and prop in obj.pose.bones[bone].keys():
                                tar.data_path = tar.data_path[7:]
                            else:
                                org_name = make_original_name(bone)
                                org_name = self.org_rename_table.get(org_name, org_name)
                                tar.data_path = 'pose.bones["%s"]["%s"]' % (org_name, prop)


    def __assign_widgets(self):
        obj_table = {obj.name: obj for obj in self.scene.objects}

        # Assign shapes to bones
        # Object's with name WGT-<bone_name> get used as that bone's shape.
        for bone in self.obj.pose.bones:
            # Object names are limited to 63 characters... arg
            wgt_name = (WGT_PREFIX + self.obj.name + '_' + bone.name)[:63]

            if wgt_name in obj_table:
                bone.custom_shape = obj_table[wgt_name]


    def __compute_visible_layers(self):
        # Reveal all the layers with control bones on them
        vis_layers = [False for n in range(0, 32)]

        for bone in self.obj.data.bones:
            for i in range(0, 32):
                vis_layers[i] = vis_layers[i] or bone.layers[i]

        for i in range(0, 32):
            vis_layers[i] = vis_layers[i] and not (ORG_LAYER[i] or MCH_LAYER[i] or DEF_LAYER[i])

        self.obj.data.layers = vis_layers


    def generate(self):
        context = self.context
        metarig = self.metarig
        scene = self.scene
        id_store = self.id_store
        view_layer = self.view_layer
        t = Timer()

        self.usable_collections = list_layer_collections(view_layer.layer_collection, selectable=True)

        if self.layer_collection not in self.usable_collections:
            metarig_collections = filter_layer_collections_by_object(self.usable_collections, self.metarig)
            self.layer_collection = (metarig_collections + [view_layer.layer_collection])[0]
            self.collection = self.layer_collection.collection

        bpy.ops.object.mode_set(mode='OBJECT')

        #------------------------------------------
        # Create/find the rig object and set it up
        obj = self.__create_rig_object()

        # Get rid of anim data in case the rig already existed
        print("Clear rig animation data.")

        obj.animation_data_clear()
        obj.data.animation_data_clear()

        select_object(context, obj, deselect_all=True)

        #------------------------------------------
        # Create Group widget
        self.__create_widget_group("WGTS_" + obj.name)

        t.tick("Create main WGTS: ")

        #------------------------------------------
        # Get parented objects to restore later
        childs = {}  # {object: bone}
        for child in obj.children:
            childs[child] = child.parent_bone

        #------------------------------------------
        # Copy bones from metarig to obj
        self.__duplicate_rig()

        obj.data.use_mirror_x = False

        t.tick("Duplicate rig: ")

        #------------------------------------------
        # Add the ORG_PREFIX to the original bones.
        bpy.ops.object.mode_set(mode='OBJECT')

        self.__rename_org_bones()

        t.tick("Make list of org bones: ")

        #------------------------------------------
        # Put the rig_name in the armature custom properties
        rna_idprop_ui_prop_get(obj.data, "rig_id", create=True)
        obj.data["rig_id"] = self.rig_id

        self.script = rig_ui_template.ScriptGenerator(self)

        #------------------------------------------
        bpy.ops.object.mode_set(mode='OBJECT')

        self.instantiate_rig_tree()

        t.tick("Instantiate rigs: ")

        #------------------------------------------
        bpy.ops.object.mode_set(mode='OBJECT')

        self.invoke_initialize()

        t.tick("Initialize rigs: ")

        #------------------------------------------
        bpy.ops.object.mode_set(mode='EDIT')

        self.invoke_prepare_bones()

        t.tick("Prepare bones: ")

        #------------------------------------------
        bpy.ops.object.mode_set(mode='OBJECT')
        bpy.ops.object.mode_set(mode='EDIT')

        self.__create_root_bone()

        self.invoke_generate_bones()

        t.tick("Generate bones: ")

        #------------------------------------------
        bpy.ops.object.mode_set(mode='OBJECT')
        bpy.ops.object.mode_set(mode='EDIT')

        self.invoke_parent_bones()

        self.__parent_bones_to_root()

        t.tick("Parent bones: ")

        #------------------------------------------
        bpy.ops.object.mode_set(mode='OBJECT')

        self.invoke_configure_bones()

        t.tick("Configure bones: ")

        #------------------------------------------
        bpy.ops.object.mode_set(mode='OBJECT')

        self.invoke_preapply_bones()

        t.tick("Preapply bones: ")

        #------------------------------------------
        bpy.ops.object.mode_set(mode='EDIT')

        self.invoke_apply_bones()

        t.tick("Apply bones: ")

        #------------------------------------------
        bpy.ops.object.mode_set(mode='OBJECT')

        self.invoke_rig_bones()

        t.tick("Rig bones: ")

        #------------------------------------------
        bpy.ops.object.mode_set(mode='OBJECT')

        create_root_widget(obj, "root")

        self.invoke_generate_widgets()

        t.tick("Generate widgets: ")

        #------------------------------------------
        bpy.ops.object.mode_set(mode='OBJECT')

        self.__lock_transforms()
        self.__assign_layers()
        self.__compute_visible_layers()
        self.__restore_driver_vars()

        t.tick("Assign layers: ")

        #------------------------------------------
        bpy.ops.object.mode_set(mode='OBJECT')

        self.invoke_finalize()

        t.tick("Finalize: ")

        #------------------------------------------
        bpy.ops.object.mode_set(mode='OBJECT')

        self.__assign_widgets()

        # Create Selection Sets
        create_selection_sets(obj, metarig)

        # Create Bone Groups
        create_bone_groups(obj, metarig, self.layer_group_priorities)

        t.tick("The rest: ")

        #----------------------------------
        # Deconfigure
        bpy.ops.object.mode_set(mode='OBJECT')
        obj.data.pose_position = 'POSE'

        # Restore parent to bones
        for child, sub_parent in childs.items():
            if sub_parent in obj.pose.bones:
                mat = child.matrix_world.copy()
                child.parent_bone = sub_parent
                child.matrix_world = mat

        #----------------------------------
        # Restore active collection
        view_layer.active_layer_collection = self.layer_collection


def generate_rig(context, metarig):
    """ Generates a rig from a metarig.

    """
    # Initial configuration
    rest_backup = metarig.data.pose_position
    metarig.data.pose_position = 'REST'

    try:
        Generator(context, metarig).generate()

        metarig.data.pose_position = rest_backup

    except Exception as e:
        # Cleanup if something goes wrong
        print("Rigify: failed to generate rig.")

        bpy.ops.object.mode_set(mode='OBJECT')
        metarig.data.pose_position = rest_backup

        # Continue the exception
        raise e


def create_selection_set_for_rig_layer(
        rig: bpy.types.Object,
        set_name: str,
        layer_idx: int
    ) -> None:
    """Create a single selection set on a rig.

    The set will contain all bones on the rig layer with the given index.
    """
    selset = rig.selection_sets.add()
    selset.name = set_name

    for b in rig.pose.bones:
        if not b.bone.layers[layer_idx] or b.name in selset.bone_ids:
            continue

        bone_id = selset.bone_ids.add()
        bone_id.name = b.name

def create_selection_sets(obj, metarig):
    """Create selection sets if the Selection Sets addon is enabled.

    Whether a selection set for a rig layer is created is controlled in the
    Rigify Layer Names panel.
    """
    # Check if selection sets addon is installed
    if 'bone_selection_groups' not in bpy.context.preferences.addons \
            and 'bone_selection_sets' not in bpy.context.preferences.addons:
        return

    obj.selection_sets.clear()

    for i, name in enumerate(metarig.data.rigify_layers.keys()):
        if name == '' or not metarig.data.rigify_layers[i].selset:
            continue

        create_selection_set_for_rig_layer(obj, name, i)


def create_bone_groups(obj, metarig, priorities={}):

    bpy.ops.object.mode_set(mode='OBJECT')
    pb = obj.pose.bones
    layers = metarig.data.rigify_layers
    groups = metarig.data.rigify_colors
    dummy = {}

    # Create BGs
    for l in layers:
        if l.group == 0:
            continue
        g_id = l.group - 1
        name = groups[g_id].name
        if name not in obj.pose.bone_groups.keys():
            bg = obj.pose.bone_groups.new(name=name)
            bg.color_set = 'CUSTOM'
            bg.colors.normal = gamma_correct(groups[g_id].normal)
            bg.colors.select = gamma_correct(groups[g_id].select)
            bg.colors.active = gamma_correct(groups[g_id].active)

    for b in pb:
        try:
            prios = priorities.get(b.name, dummy)
            enabled = [ i for i, v in enumerate(b.bone.layers) if v ]
            layer_index = max(enabled, key=lambda i: prios.get(i, 0))
        except ValueError:
            continue
        if layer_index > len(layers) - 1:   # bone is on reserved layers
            continue
        g_id = layers[layer_index].group - 1
        if g_id >= 0:
            name = groups[g_id].name
            b.bone_group = obj.pose.bone_groups[name]


def get_xy_spread(bones):
    x_max = 0
    y_max = 0
    for b in bones:
        x_max = max((x_max, abs(b.head[0]), abs(b.tail[0])))
        y_max = max((y_max, abs(b.head[1]), abs(b.tail[1])))

    return max((x_max, y_max))


def param_matches_type(param_name, rig_type):
    """ Returns True if the parameter name is consistent with the rig type.
    """
    if param_name.rsplit(".", 1)[0] == rig_type:
        return True
    else:
        return False


def param_name(param_name, rig_type):
    """ Get the actual parameter name, sans-rig-type.
    """
    return param_name[len(rig_type) + 1:]