rigify/utils/rig.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 importlib
  23 import importlib.util
  24 import os
  25
  26 RIG_DIR = "rigs"  # Name of the directory where rig types are kept
  27 METARIG_DIR = "metarigs"  # Name of the directory where metarigs are kept
  28 TEMPLATE_DIR = "ui_templates"  # Name of the directory where ui templates are kept
  29
  30 MODULE_NAME = "rigify"  # Windows/Mac blender is weird, so __package__ doesn't work
  31
  32 outdated_types = {"pitchipoy.limbs.super_limb": "limbs.super_limb",
  33                   "pitchipoy.limbs.super_arm": "limbs.super_limb",
  34                   "pitchipoy.limbs.super_leg": "limbs.super_limb",
  35                   "pitchipoy.limbs.super_front_paw": "limbs.super_limb",
  36                   "pitchipoy.limbs.super_rear_paw": "limbs.super_limb",
  37                   "pitchipoy.limbs.super_finger": "limbs.super_finger",
  38                   "pitchipoy.super_torso_turbo": "spines.super_spine",
  39                   "pitchipoy.simple_tentacle": "limbs.simple_tentacle",
  40                   "pitchipoy.super_face": "faces.super_face",
  41                   "pitchipoy.super_palm": "limbs.super_palm",
  42                   "pitchipoy.super_copy": "basic.super_copy",
  43                   "pitchipoy.tentacle": "",
  44                   "palm": "limbs.super_palm",
  45                   "basic.copy": "basic.super_copy",
  46                   "biped.arm": "",
  47                   "biped.leg": "",
  48                   "finger": "",
  49                   "neck_short": "",
  50                   "misc.delta": "",
  51                   "spine": ""
  52                   }
  53
  54 def upgradeMetarigTypes(metarig, revert=False):
  55     """Replaces rigify_type properties from old versions with their current names
  56
  57     :param revert: revert types to previous version (if old type available)
  58     """
  59
  60     if revert:
  61         vals = list(outdated_types.values())
  62         rig_defs = {v: k for k, v in outdated_types.items() if vals.count(v) == 1}
  63     else:
  64         rig_defs = outdated_types
  65
  66     for bone in metarig.pose.bones:
  67         rig_type = bone.rigify_type
  68         if rig_type in rig_defs:
  69             bone.rigify_type = rig_defs[rig_type]
  70             if 'leg' in rig_type:
  71                 bone.rigfy_parameters.limb_type = 'leg'
  72             if 'arm' in rig_type:
  73                 bone.rigfy_parameters.limb_type = 'arm'
  74             if 'paw' in rig_type:
  75                 bone.rigfy_parameters.limb_type = 'paw'
  76             if rig_type == "basic.copy":
  77                 bone.rigify_parameters.make_widget = False
  78
  79
  80 #=============================================
  81 # Misc
  82 #=============================================
  83
  84 def get_rig_type(rig_type, base_path=''):
  85     return get_resource(rig_type, base_path=base_path)
  86
  87 def get_resource(resource_name, base_path=''):
  88     """ Fetches a rig module by name, and returns it.
  89     """
  90
  91     if '.' in resource_name:
  92         module_subpath = str.join(os.sep, resource_name.split('.'))
  93         package = resource_name.split('.')[0]
  94         for sub in resource_name.split('.')[1:]:
  95             package = '.'.join([package, sub])
  96             submod = importlib.import_module(package)
  97     else:
  98         module_subpath = resource_name
  99
 100     spec = importlib.util.spec_from_file_location(resource_name, os.path.join(base_path, module_subpath + '.py'))
 101     submod = importlib.util.module_from_spec(spec)
 102     spec.loader.exec_module(submod)
 103     return submod
 104
 105
 106 def get_metarig_module(metarig_name, path=METARIG_DIR):
 107     """ Fetches a rig module by name, and returns it.
 108     """
 109
 110     name = ".%s.%s" % (path, metarig_name)
 111     submod = importlib.import_module(name, package=MODULE_NAME)
 112     importlib.reload(submod)
 113     return submod
 114
 115
 116 def connected_children_names(obj, bone_name):
 117     """ Returns a list of bone names (in order) of the bones that form a single
 118         connected chain starting with the given bone as a parent.
 119         If there is a connected branch, the list stops there.
 120     """
 121     bone = obj.data.bones[bone_name]
 122     names = []
 123
 124     while True:
 125         connects = 0
 126         con_name = ""
 127
 128         for child in bone.children:
 129             if child.use_connect:
 130                 connects += 1
 131                 con_name = child.name
 132
 133         if connects == 1:
 134             names += [con_name]
 135             bone = obj.data.bones[con_name]
 136         else:
 137             break
 138
 139     return names
 140
 141
 142 def has_connected_children(bone):
 143     """ Returns true/false whether a bone has connected children or not.
 144     """
 145     t = False
 146     for b in bone.children:
 147         t = t or b.use_connect
 148     return t
 149
 150
 151 def write_metarig(obj, layers=False, func_name="create", groups=False):
 152     """
 153     Write a metarig as a python script, this rig is to have all info needed for
 154     generating the real rig with rigify.
 155     """
 156     code = []
 157
 158     code.append("import bpy\n\n")
 159     code.append("from mathutils import Color\n\n")
 160
 161     code.append("def %s(obj):" % func_name)
 162     code.append("    # generated by rigify.utils.write_metarig")
 163     bpy.ops.object.mode_set(mode='EDIT')
 164     code.append("    bpy.ops.object.mode_set(mode='EDIT')")
 165     code.append("    arm = obj.data")
 166
 167     arm = obj.data
 168
 169     # Rigify bone group colors info
 170     if groups and len(arm.rigify_colors) > 0:
 171         code.append("\n    for i in range(" + str(len(arm.rigify_colors)) + "):")
 172         code.append("        arm.rigify_colors.add()\n")
 173
 174         for i in range(len(arm.rigify_colors)):
 175             name = arm.rigify_colors[i].name
 176             active = arm.rigify_colors[i].active
 177             normal = arm.rigify_colors[i].normal
 178             select = arm.rigify_colors[i].select
 179             standard_colors_lock = arm.rigify_colors[i].standard_colors_lock
 180             code.append('    arm.rigify_colors[' + str(i) + '].name = "' + name + '"')
 181             code.append('    arm.rigify_colors[' + str(i) + '].active = Color(' + str(active[:]) + ')')
 182             code.append('    arm.rigify_colors[' + str(i) + '].normal = Color(' + str(normal[:]) + ')')
 183             code.append('    arm.rigify_colors[' + str(i) + '].select = Color(' + str(select[:]) + ')')
 184             code.append('    arm.rigify_colors[' + str(i) + '].standard_colors_lock = ' + str(standard_colors_lock))
 185
 186     # Rigify layer layout info
 187     if layers and len(arm.rigify_layers) > 0:
 188         code.append("\n    for i in range(" + str(len(arm.rigify_layers)) + "):")
 189         code.append("        arm.rigify_layers.add()\n")
 190
 191         for i in range(len(arm.rigify_layers)):
 192             name = arm.rigify_layers[i].name
 193             row = arm.rigify_layers[i].row
 194             selset = arm.rigify_layers[i].selset
 195             group = arm.rigify_layers[i].group
 196             code.append('    arm.rigify_layers[' + str(i) + '].name = "' + name + '"')
 197             code.append('    arm.rigify_layers[' + str(i) + '].row = ' + str(row))
 198             code.append('    arm.rigify_layers[' + str(i) + '].selset = ' + str(selset))
 199             code.append('    arm.rigify_layers[' + str(i) + '].group = ' + str(group))
 200
 201     # write parents first
 202     bones = [(len(bone.parent_recursive), bone.name) for bone in arm.edit_bones]
 203     bones.sort(key=lambda item: item[0])
 204     bones = [item[1] for item in bones]
 205
 206     code.append("\n    bones = {}\n")
 207
 208     for bone_name in bones:
 209         bone = arm.edit_bones[bone_name]
 210         code.append("    bone = arm.edit_bones.new(%r)" % bone.name)
 211         code.append("    bone.head[:] = %.4f, %.4f, %.4f" % bone.head.to_tuple(4))
 212         code.append("    bone.tail[:] = %.4f, %.4f, %.4f" % bone.tail.to_tuple(4))
 213         code.append("    bone.roll = %.4f" % bone.roll)
 214         code.append("    bone.use_connect = %s" % str(bone.use_connect))
 215         if bone.parent:
 216             code.append("    bone.parent = arm.edit_bones[bones[%r]]" % bone.parent.name)
 217         code.append("    bones[%r] = bone.name" % bone.name)
 218
 219     bpy.ops.object.mode_set(mode='OBJECT')
 220     code.append("")
 221     code.append("    bpy.ops.object.mode_set(mode='OBJECT')")
 222
 223     # Rig type and other pose properties
 224     for bone_name in bones:
 225         pbone = obj.pose.bones[bone_name]
 226
 227         code.append("    pbone = obj.pose.bones[bones[%r]]" % bone_name)
 228         code.append("    pbone.rigify_type = %r" % pbone.rigify_type)
 229         code.append("    pbone.lock_location = %s" % str(tuple(pbone.lock_location)))
 230         code.append("    pbone.lock_rotation = %s" % str(tuple(pbone.lock_rotation)))
 231         code.append("    pbone.lock_rotation_w = %s" % str(pbone.lock_rotation_w))
 232         code.append("    pbone.lock_scale = %s" % str(tuple(pbone.lock_scale)))
 233         code.append("    pbone.rotation_mode = %r" % pbone.rotation_mode)
 234         if layers:
 235             code.append("    pbone.bone.layers = %s" % str(list(pbone.bone.layers)))
 236         # Rig type parameters
 237         for param_name in pbone.rigify_parameters.keys():
 238             param = getattr(pbone.rigify_parameters, param_name, '')
 239             if str(type(param)) == "<class 'bpy_prop_array'>":
 240                 param = list(param)
 241             if type(param) == str:
 242                 param = '"' + param + '"'
 243             code.append("    try:")
 244             code.append("        pbone.rigify_parameters.%s = %s" % (param_name, str(param)))
 245             code.append("    except AttributeError:")
 246             code.append("        pass")
 247
 248     code.append("\n    bpy.ops.object.mode_set(mode='EDIT')")
 249     code.append("    for bone in arm.edit_bones:")
 250     code.append("        bone.select = False")
 251     code.append("        bone.select_head = False")
 252     code.append("        bone.select_tail = False")
 253
 254     code.append("    for b in bones:")
 255     code.append("        bone = arm.edit_bones[bones[b]]")
 256     code.append("        bone.select = True")
 257     code.append("        bone.select_head = True")
 258     code.append("        bone.select_tail = True")
 259     code.append("        arm.edit_bones.active = bone")
 260
 261     # Set appropriate layers visible
 262     if layers:
 263         # Find what layers have bones on them
 264         active_layers = []
 265         for bone_name in bones:
 266             bone = obj.data.bones[bone_name]
 267             for i in range(len(bone.layers)):
 268                 if bone.layers[i]:
 269                     if i not in active_layers:
 270                         active_layers.append(i)
 271         active_layers.sort()
 272
 273         code.append("\n    arm.layers = [(x in " + str(active_layers) + ") for x in range(" + str(len(arm.layers)) + ")]")
 274
 275     code.append('\nif __name__ == "__main__":')
 276     code.append("    " + func_name + "(bpy.context.active_object)\n")
 277
 278     return "\n".join(code)