rigify/utils/widgets.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 math
  23
  24 from mathutils import Matrix
  25
  26 from .errors import MetarigError
  27 from .collections import ensure_widget_collection
  28
  29 WGT_PREFIX = "WGT-"  # Prefix for widget objects
  30
  31 #=============================================
  32 # Widget creation
  33 #=============================================
  34
  35
  36 def obj_to_bone(obj, rig, bone_name, bone_transform_name=None):
  37     """ Places an object at the location/rotation/scale of the given bone.
  38     """
  39     if bpy.context.mode == 'EDIT_ARMATURE':
  40         raise MetarigError("obj_to_bone(): does not work while in edit mode")
  41
  42     bone = rig.pose.bones[bone_name]
  43     scale = bone.custom_shape_scale
  44
  45     if bone.use_custom_shape_bone_size:
  46         scale *= bone.length
  47
  48     if bone_transform_name is not None:
  49         bone = rig.pose.bones[bone_transform_name]
  50     elif bone.custom_shape_transform:
  51         bone = bone.custom_shape_transform
  52
  53     obj.rotation_mode = 'XYZ'
  54     obj.matrix_basis = rig.matrix_world @ bone.bone.matrix_local @ Matrix.Scale(scale, 4)
  55
  56
  57 def create_widget(rig, bone_name, bone_transform_name=None):
  58     """ Creates an empty widget object for a bone, and returns the object.
  59     """
  60     obj_name = WGT_PREFIX + rig.name + '_' + bone_name
  61     scene = bpy.context.scene
  62     collection = ensure_widget_collection(bpy.context)
  63
  64     # Check if it already exists in the scene
  65     if obj_name in scene.objects:
  66         # Move object to bone position, in case it changed
  67         obj = scene.objects[obj_name]
  68         obj_to_bone(obj, rig, bone_name, bone_transform_name)
  69
  70         return None
  71     else:
  72         # Delete object if it exists in blend data but not scene data.
  73         # This is necessary so we can then create the object without
  74         # name conflicts.
  75         if obj_name in bpy.data.objects:
  76             bpy.data.objects[obj_name].user_clear()
  77             bpy.data.objects.remove(bpy.data.objects[obj_name])
  78
  79         # Create mesh object
  80         mesh = bpy.data.meshes.new(obj_name)
  81         obj = bpy.data.objects.new(obj_name, mesh)
  82         collection.objects.link(obj)
  83
  84         # Move object to bone position and set layers
  85         obj_to_bone(obj, rig, bone_name, bone_transform_name)
  86         wgts_group_name = 'WGTS_' + rig.name
  87         if wgts_group_name in bpy.data.objects.keys():
  88             obj.parent = bpy.data.objects[wgts_group_name]
  89
  90         return obj
  91
  92
  93 def create_circle_polygon(number_verts, axis, radius=1.0, head_tail=0.0):
  94     """ Creates a basic circle around of an axis selected.
  95         number_verts: number of vertices of the polygon
  96         axis: axis normal to the circle
  97         radius: the radius of the circle
  98         head_tail: where along the length of the bone the circle is (0.0=head, 1.0=tail)
  99     """
 100     verts = []
 101     edges = []
 102     angle = 2 * math.pi / number_verts
 103     i = 0
 104
 105     assert(axis in 'XYZ')
 106
 107     while i < (number_verts):
 108         a = math.cos(i * angle)
 109         b = math.sin(i * angle)
 110
 111         if axis == 'X':
 112             verts.append((head_tail, a * radius, b * radius))
 113         elif axis == 'Y':
 114             verts.append((a * radius, head_tail, b * radius))
 115         elif axis == 'Z':
 116             verts.append((a * radius, b * radius, head_tail))
 117
 118         if i < (number_verts - 1):
 119             edges.append((i , i + 1))
 120
 121         i += 1
 122
 123     edges.append((0, number_verts - 1))
 124
 125     return verts, edges
 126
 127
 128 def adjust_widget_axis(obj, axis='y', offset=0.0):
 129     mesh = obj.data
 130
 131     if axis[0] == '-':
 132         s = -1.0
 133         axis = axis[1]
 134     else:
 135         s = 1.0
 136
 137     trans_matrix = Matrix.Translation((0.0, offset, 0.0))
 138     rot_matrix = Matrix.Diagonal((1.0, s, 1.0, 1.0))
 139
 140     if axis == "x":
 141         rot_matrix = Matrix.Rotation(-s*math.pi/2, 4, 'Z')
 142         trans_matrix = Matrix.Translation((offset, 0.0, 0.0))
 143
 144     elif axis == "z":
 145         rot_matrix = Matrix.Rotation(s*math.pi/2, 4, 'X')
 146         trans_matrix = Matrix.Translation((0.0, 0.0, offset))
 147
 148     matrix = trans_matrix @ rot_matrix
 149
 150     for vert in mesh.vertices:
 151         vert.co = matrix @ vert.co
 152
 153
 154 def adjust_widget_transform_mesh(obj, matrix, local=None):
 155     """Adjust the generated widget by applying a correction matrix to the mesh.
 156        If local is false, the matrix is in world space.
 157        If local is True, it's in the local space of the widget.
 158        If local is a bone, it's in the local space of the bone.
 159     """
 160     if obj:
 161         if local is not True:
 162             if local:
 163                 assert isinstance(local, bpy.types.PoseBone)
 164                 bonemat = local.id_data.matrix_world @ local.bone.matrix_local
 165                 matrix = bonemat @ matrix @ bonemat.inverted()
 166
 167             obmat = obj.matrix_basis
 168             matrix = obmat.inverted() @ matrix @ obmat
 169
 170         obj.data.transform(matrix)
 171
 172
 173 def write_widget(obj):
 174     """ Write a mesh object as a python script for widget use.
 175     """
 176     script = ""
 177     script += "def create_thing_widget(rig, bone_name, size=1.0, bone_transform_name=None):\n"
 178     script += "    obj = create_widget(rig, bone_name, bone_transform_name)\n"
 179     script += "    if obj is not None:\n"
 180
 181     # Vertices
 182     script += "        verts = ["
 183     for i, v in enumerate(obj.data.vertices):
 184         script += "({:g}*size, {:g}*size, {:g}*size),".format(v.co[0], v.co[1], v.co[2])
 185         script += "\n                 " if i % 2 == 1 else " "
 186     script += "]\n"
 187
 188     # Edges
 189     script += "        edges = ["
 190     for i, e in enumerate(obj.data.edges):
 191         script += "(" + str(e.vertices[0]) + ", " + str(e.vertices[1]) + "),"
 192         script += "\n                 " if i % 10 == 9 else " "
 193     script += "]\n"
 194
 195     # Faces
 196     script += "        faces = ["
 197     for i, f in enumerate(obj.data.polygons):
 198         script += "("
 199         for v in f.vertices:
 200             script += str(v) + ", "
 201         script += "),"
 202         script += "\n                 " if i % 10 == 9 else " "
 203     script += "]\n"
 204
 205     # Build mesh
 206     script += "\n        mesh = obj.data\n"
 207     script += "        mesh.from_pydata(verts, edges, faces)\n"
 208     script += "        mesh.update()\n"
 209     script += "        return obj\n"
 210     script += "    else:\n"
 211     script += "        return None\n"
 212
 213     return script