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 from bpy.types import bpy_struct, bpy_prop_array, Constraint
  27
  28 RIG_DIR = "rigs"  # Name of the directory where rig types are kept
  29 METARIG_DIR = "metarigs"  # Name of the directory where metarigs are kept
  30 TEMPLATE_DIR = "ui_templates"  # Name of the directory where ui templates are kept
  31
  32
  33 outdated_types = {"pitchipoy.limbs.super_limb": "limbs.super_limb",
  34                   "pitchipoy.limbs.super_arm": "limbs.super_limb",
  35                   "pitchipoy.limbs.super_leg": "limbs.super_limb",
  36                   "pitchipoy.limbs.super_front_paw": "limbs.super_limb",
  37                   "pitchipoy.limbs.super_rear_paw": "limbs.super_limb",
  38                   "pitchipoy.limbs.super_finger": "limbs.super_finger",
  39                   "pitchipoy.super_torso_turbo": "spines.super_spine",
  40                   "pitchipoy.simple_tentacle": "limbs.simple_tentacle",
  41                   "pitchipoy.super_face": "faces.super_face",
  42                   "pitchipoy.super_palm": "limbs.super_palm",
  43                   "pitchipoy.super_copy": "basic.super_copy",
  44                   "pitchipoy.tentacle": "",
  45                   "palm": "limbs.super_palm",
  46                   "basic.copy": "basic.super_copy",
  47                   "biped.arm": "",
  48                   "biped.leg": "",
  49                   "finger": "",
  50                   "neck_short": "",
  51                   "misc.delta": "",
  52                   "spine": ""
  53                   }
  54
  55 def get_rigify_type(pose_bone):
  56     return pose_bone.rigify_type.replace(" ", "")
  57
  58 def is_rig_base_bone(obj, name):
  59     return bool(get_rigify_type(obj.pose.bones[name]))
  60
  61 def upgradeMetarigTypes(metarig, revert=False):
  62     """Replaces rigify_type properties from old versions with their current names
  63
  64     :param revert: revert types to previous version (if old type available)
  65     """
  66
  67     if revert:
  68         vals = list(outdated_types.values())
  69         rig_defs = {v: k for k, v in outdated_types.items() if vals.count(v) == 1}
  70     else:
  71         rig_defs = outdated_types
  72
  73     for bone in metarig.pose.bones:
  74         rig_type = bone.rigify_type
  75         if rig_type in rig_defs:
  76             bone.rigify_type = rig_defs[rig_type]
  77             if 'leg' in rig_type:
  78                 bone.rigfy_parameters.limb_type = 'leg'
  79             if 'arm' in rig_type:
  80                 bone.rigfy_parameters.limb_type = 'arm'
  81             if 'paw' in rig_type:
  82                 bone.rigfy_parameters.limb_type = 'paw'
  83             if rig_type == "basic.copy":
  84                 bone.rigify_parameters.make_widget = False
  85
  86
  87 #=============================================
  88 # Misc
  89 #=============================================
  90
  91 def get_resource(resource_name):
  92     """ Fetches a rig module by name, and returns it.
  93     """
  94
  95     module = importlib.import_module(resource_name)
  96     importlib.reload(module)
  97     return module
  98
  99
 100 def connected_children_names(obj, bone_name):
 101     """ Returns a list of bone names (in order) of the bones that form a single
 102         connected chain starting with the given bone as a parent.
 103         If there is a connected branch, the list stops there.
 104     """
 105     bone = obj.data.bones[bone_name]
 106     names = []
 107
 108     while True:
 109         connects = 0
 110         con_name = ""
 111
 112         for child in bone.children:
 113             if child.use_connect:
 114                 connects += 1
 115                 con_name = child.name
 116
 117         if connects == 1:
 118             names += [con_name]
 119             bone = obj.data.bones[con_name]
 120         else:
 121             break
 122
 123     return names
 124
 125
 126 def has_connected_children(bone):
 127     """ Returns true/false whether a bone has connected children or not.
 128     """
 129     t = False
 130     for b in bone.children:
 131         t = t or b.use_connect
 132     return t
 133
 134
 135 def _list_bone_names_depth_first_sorted_rec(result_list, bone):
 136     result_list.append(bone.name)
 137
 138     for child in sorted(list(bone.children), key=lambda b: b.name):
 139         _list_bone_names_depth_first_sorted_rec(result_list, child)
 140
 141 def list_bone_names_depth_first_sorted(obj):
 142     """Returns a list of bone names in depth first name sorted order."""
 143     result_list = []
 144
 145     for bone in sorted(list(obj.data.bones), key=lambda b: b.name):
 146         if bone.parent is None:
 147             _list_bone_names_depth_first_sorted_rec(result_list, bone)
 148
 149     return result_list
 150
 151
 152 def _get_property_value(obj, name):
 153     value = getattr(obj, name, None)
 154     if isinstance(value, bpy_prop_array):
 155         value = tuple(value)
 156     return value
 157
 158 def _generate_properties(lines, prefix, obj, base_class, *, defaults={}, objects={}):
 159     block_props = set(prop.identifier for prop in base_class.bl_rna.properties) - set(defaults.keys())
 160
 161     for prop in type(obj).bl_rna.properties:
 162         if prop.identifier not in block_props and not prop.is_readonly:
 163             cur_value = _get_property_value(obj, prop.identifier)
 164
 165             if prop.identifier in defaults:
 166                 if cur_value == defaults[prop.identifier]:
 167                     continue
 168
 169             if isinstance(cur_value, bpy_struct):
 170                 if cur_value in objects:
 171                     lines.append('%s.%s = %s' % (prefix, prop.identifier, objects[cur_value]))
 172             else:
 173                 lines.append('%s.%s = %r' % (prefix, prop.identifier, cur_value))
 174
 175
 176 def write_metarig(obj, layers=False, func_name="create", groups=False):
 177     """
 178     Write a metarig as a python script, this rig is to have all info needed for
 179     generating the real rig with rigify.
 180     """
 181     code = []
 182
 183     code.append("import bpy\n\n")
 184     code.append("from mathutils import Color\n\n")
 185
 186     code.append("def %s(obj):" % func_name)
 187     code.append("    # generated by rigify.utils.write_metarig")
 188     bpy.ops.object.mode_set(mode='EDIT')
 189     code.append("    bpy.ops.object.mode_set(mode='EDIT')")
 190     code.append("    arm = obj.data")
 191
 192     arm = obj.data
 193
 194     # Rigify bone group colors info
 195     if groups and len(arm.rigify_colors) > 0:
 196         code.append("\n    for i in range(" + str(len(arm.rigify_colors)) + "):")
 197         code.append("        arm.rigify_colors.add()\n")
 198
 199         for i in range(len(arm.rigify_colors)):
 200             name = arm.rigify_colors[i].name
 201             active = arm.rigify_colors[i].active
 202             normal = arm.rigify_colors[i].normal
 203             select = arm.rigify_colors[i].select
 204             standard_colors_lock = arm.rigify_colors[i].standard_colors_lock
 205             code.append('    arm.rigify_colors[' + str(i) + '].name = "' + name + '"')
 206             code.append('    arm.rigify_colors[' + str(i) + '].active = Color(' + str(active[:]) + ')')
 207             code.append('    arm.rigify_colors[' + str(i) + '].normal = Color(' + str(normal[:]) + ')')
 208             code.append('    arm.rigify_colors[' + str(i) + '].select = Color(' + str(select[:]) + ')')
 209             code.append('    arm.rigify_colors[' + str(i) + '].standard_colors_lock = ' + str(standard_colors_lock))
 210
 211     # Rigify layer layout info
 212     if layers and len(arm.rigify_layers) > 0:
 213         code.append("\n    for i in range(" + str(len(arm.rigify_layers)) + "):")
 214         code.append("        arm.rigify_layers.add()\n")
 215
 216         for i in range(len(arm.rigify_layers)):
 217             name = arm.rigify_layers[i].name
 218             row = arm.rigify_layers[i].row
 219             selset = arm.rigify_layers[i].selset
 220             group = arm.rigify_layers[i].group
 221             code.append('    arm.rigify_layers[' + str(i) + '].name = "' + name + '"')
 222             code.append('    arm.rigify_layers[' + str(i) + '].row = ' + str(row))
 223             code.append('    arm.rigify_layers[' + str(i) + '].selset = ' + str(selset))
 224             code.append('    arm.rigify_layers[' + str(i) + '].group = ' + str(group))
 225
 226     # write parents first
 227     bones = [(len(bone.parent_recursive), bone.name) for bone in arm.edit_bones]
 228     bones.sort(key=lambda item: item[0])
 229     bones = [item[1] for item in bones]
 230
 231     code.append("\n    bones = {}\n")
 232
 233     for bone_name in bones:
 234         bone = arm.edit_bones[bone_name]
 235         code.append("    bone = arm.edit_bones.new(%r)" % bone.name)
 236         code.append("    bone.head = %.4f, %.4f, %.4f" % bone.head.to_tuple(4))
 237         code.append("    bone.tail = %.4f, %.4f, %.4f" % bone.tail.to_tuple(4))
 238         code.append("    bone.roll = %.4f" % bone.roll)
 239         code.append("    bone.use_connect = %s" % str(bone.use_connect))
 240         if bone.inherit_scale != 'FULL':
 241             code.append("    bone.inherit_scale = %r" % str(bone.inherit_scale))
 242         if bone.parent:
 243             code.append("    bone.parent = arm.edit_bones[bones[%r]]" % bone.parent.name)
 244         code.append("    bones[%r] = bone.name" % bone.name)
 245
 246     bpy.ops.object.mode_set(mode='OBJECT')
 247     code.append("")
 248     code.append("    bpy.ops.object.mode_set(mode='OBJECT')")
 249
 250     # Rig type and other pose properties
 251     for bone_name in bones:
 252         pbone = obj.pose.bones[bone_name]
 253
 254         code.append("    pbone = obj.pose.bones[bones[%r]]" % bone_name)
 255         code.append("    pbone.rigify_type = %r" % pbone.rigify_type)
 256         code.append("    pbone.lock_location = %s" % str(tuple(pbone.lock_location)))
 257         code.append("    pbone.lock_rotation = %s" % str(tuple(pbone.lock_rotation)))
 258         code.append("    pbone.lock_rotation_w = %s" % str(pbone.lock_rotation_w))
 259         code.append("    pbone.lock_scale = %s" % str(tuple(pbone.lock_scale)))
 260         code.append("    pbone.rotation_mode = %r" % pbone.rotation_mode)
 261         if layers:
 262             code.append("    pbone.bone.layers = %s" % str(list(pbone.bone.layers)))
 263         # Rig type parameters
 264         for param_name in pbone.rigify_parameters.keys():
 265             param = getattr(pbone.rigify_parameters, param_name, '')
 266             if str(type(param)) == "<class 'bpy_prop_array'>":
 267                 param = list(param)
 268             if type(param) == str:
 269                 param = '"' + param + '"'
 270             code.append("    try:")
 271             code.append("        pbone.rigify_parameters.%s = %s" % (param_name, str(param)))
 272             code.append("    except AttributeError:")
 273             code.append("        pass")
 274         # Constraints
 275         for con in pbone.constraints:
 276             code.append("    con = pbone.constraints.new(%r)" % (con.type))
 277             code.append("    con.name = %r" % (con.name))
 278             # Add target first because of target_space handling
 279             if con.type == 'ARMATURE':
 280                 for tgt in con.targets:
 281                     code.append("    tgt = con.targets.new()")
 282                     code.append("    tgt.target = obj")
 283                     code.append("    tgt.subtarget = %r" % (tgt.subtarget))
 284                     code.append("    tgt.weight = %.3f" % (tgt.weight))
 285             elif getattr(con, 'target', None) == obj:
 286                 code.append("    con.target = obj")
 287             # Generic properties
 288             _generate_properties(
 289                 code, "    con", con, Constraint,
 290                 defaults={
 291                     'owner_space': 'WORLD', 'target_space': 'WORLD',
 292                     'mute': False, 'influence': 1.0,
 293                     'target': obj,
 294                 },
 295                 objects={obj: 'obj'},
 296             )
 297
 298     code.append("\n    bpy.ops.object.mode_set(mode='EDIT')")
 299     code.append("    for bone in arm.edit_bones:")
 300     code.append("        bone.select = False")
 301     code.append("        bone.select_head = False")
 302     code.append("        bone.select_tail = False")
 303
 304     code.append("    for b in bones:")
 305     code.append("        bone = arm.edit_bones[bones[b]]")
 306     code.append("        bone.select = True")
 307     code.append("        bone.select_head = True")
 308     code.append("        bone.select_tail = True")
 309     code.append("        bone.bbone_x = bone.bbone_z = bone.length * 0.05")
 310     code.append("        arm.edit_bones.active = bone")
 311
 312     # Set appropriate layers visible
 313     if layers:
 314         # Find what layers have bones on them
 315         active_layers = []
 316         for bone_name in bones:
 317             bone = obj.data.bones[bone_name]
 318             for i in range(len(bone.layers)):
 319                 if bone.layers[i]:
 320                     if i not in active_layers:
 321                         active_layers.append(i)
 322         active_layers.sort()
 323
 324         code.append("\n    arm.layers = [(x in " + str(active_layers) + ") for x in range(" + str(len(arm.layers)) + ")]")
 325
 326     code.append("\n    return bones")
 327
 328     code.append('\nif __name__ == "__main__":')
 329     code.append("    " + func_name + "(bpy.context.active_object)\n")
 330
 331     return "\n".join(code)