rigify/utils/widgets.py

   1 # SPDX-FileCopyrightText: 2019-2022 Blender Foundation
   2 #
   3 # SPDX-License-Identifier: GPL-2.0-or-later
   4
   5 import bpy
   6 import math
   7 import inspect
   8 import functools
   9
  10 from typing import Optional, Callable, TYPE_CHECKING
  11 from bpy.types import Mesh, Object, UILayout, WindowManager
  12 from mathutils import Matrix, Vector, Euler
  13 from itertools import count
  14
  15 from .errors import MetarigError
  16 from .collections import ensure_collection
  17 from .misc import ArmatureObject, MeshObject, AnyVector, verify_mesh_obj, IdPropSequence
  18 from .naming import change_name_side, get_name_side, Side
  19
  20 if TYPE_CHECKING:
  21     from .. import RigifyName
  22
  23
  24 WGT_PREFIX = "WGT-"  # Prefix for widget objects
  25 WGT_GROUP_PREFIX = "WGTS_"  # noqa; Prefix for the widget collection
  26
  27
  28 ##############################################
  29 # Widget creation
  30 ##############################################
  31
  32 def obj_to_bone(obj: Object, rig: ArmatureObject, bone_name: str,
  33                 bone_transform_name: Optional[str] = None):
  34     """ Places an object at the location/rotation/scale of the given bone.
  35     """
  36     if bpy.context.mode == 'EDIT_ARMATURE':
  37         raise MetarigError("obj_to_bone(): does not work while in edit mode")
  38
  39     bone = rig.pose.bones[bone_name]
  40
  41     loc = bone.custom_shape_translation
  42     rot = bone.custom_shape_rotation_euler
  43     scale = Vector(bone.custom_shape_scale_xyz)
  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     shape_mat = Matrix.LocRotScale(loc, Euler(rot), scale)
  54
  55     obj.rotation_mode = 'XYZ'
  56     obj.matrix_basis = rig.matrix_world @ bone.bone.matrix_local @ shape_mat
  57
  58
  59 def create_widget(rig: ArmatureObject, bone_name: str,
  60                   bone_transform_name: Optional[str] = None, *,
  61                   widget_name: Optional[str] = None,
  62                   widget_force_new=False, subsurf=0) -> Optional[MeshObject]:
  63     """
  64     Creates an empty widget object for a bone, and returns the object.
  65     If the object already existed, returns None.
  66     """
  67     assert rig.mode != 'EDIT'
  68
  69     from ..base_generate import BaseGenerator
  70
  71     scene = bpy.context.scene
  72     bone = rig.pose.bones[bone_name]
  73
  74     # Access the current generator instance when generating (ugh, globals)
  75     generator = BaseGenerator.instance
  76
  77     if generator:
  78         collection = generator.widget_collection
  79     else:
  80         collection = ensure_collection(bpy.context, WGT_GROUP_PREFIX + rig.name, hidden=True)
  81
  82     use_mirror = generator and generator.use_mirror_widgets
  83     bone_mid_name = change_name_side(bone_name, Side.MIDDLE) if use_mirror else bone_name
  84
  85     obj_name = widget_name or WGT_PREFIX + rig.name + '_' + bone_name
  86     reuse_mesh = None
  87
  88     obj: Optional[MeshObject]
  89
  90     # Check if it already exists in the scene
  91     if not widget_force_new:
  92         obj = None
  93
  94         if generator:
  95             # Check if the widget was already generated
  96             if bone_name in generator.new_widget_table:
  97                 return None
  98
  99             # If re-generating, check widgets used by the previous rig
 100             obj = generator.old_widget_table.get(bone_name)
 101
 102         if not obj:
 103             # Search the scene by name
 104             obj = scene.objects.get(obj_name)
 105             if obj and obj.library:
 106                 # Second brute force try if the first result is linked
 107                 local_objs = [obj for obj in scene.objects
 108                               if obj.name == obj_name and not obj.library]
 109                 obj = local_objs[0] if local_objs else None
 110
 111         if obj:
 112             # Record the generated widget
 113             if generator:
 114                 generator.new_widget_table[bone_name] = obj
 115
 116             # Re-add to the collection if not there for some reason
 117             if obj.name not in collection.objects:
 118                 collection.objects.link(obj)
 119
 120             # Flip scale for originally mirrored widgets
 121             if obj.scale.x < 0 < bone.custom_shape_scale_xyz.x:
 122                 bone.custom_shape_scale_xyz.x *= -1
 123
 124             # Move object to bone position, in case it changed
 125             obj_to_bone(obj, rig, bone_name, bone_transform_name)
 126
 127             return None
 128
 129         # Create a linked duplicate of the widget assigned in the metarig
 130         reuse_widget = rig.pose.bones[bone_name].custom_shape
 131         if reuse_widget:
 132             subsurf = 0
 133             reuse_mesh = reuse_widget.data
 134
 135         # Create a linked duplicate with the mirror widget
 136         if not reuse_mesh and use_mirror and bone_mid_name != bone_name:
 137             reuse_mesh = generator.widget_mirror_mesh.get(bone_mid_name)
 138
 139     # Create an empty mesh datablock if not linking
 140     if reuse_mesh:
 141         mesh = reuse_mesh
 142
 143     elif use_mirror and bone_mid_name != bone_name:
 144         # When mirroring, untag side from mesh name, and remember it
 145         mesh = bpy.data.meshes.new(change_name_side(obj_name, Side.MIDDLE))
 146
 147         generator.widget_mirror_mesh[bone_mid_name] = mesh
 148
 149     else:
 150         mesh = bpy.data.meshes.new(obj_name)
 151
 152     # Create the object
 153     obj = verify_mesh_obj(bpy.data.objects.new(obj_name, mesh))
 154     collection.objects.link(obj)
 155
 156     # Add the subdivision surface modifier
 157     if subsurf > 0:
 158         mod = obj.modifiers.new("subsurf", 'SUBSURF')
 159         mod.levels = subsurf
 160
 161     # Record the generated widget
 162     if generator:
 163         generator.new_widget_table[bone_name] = obj
 164
 165     # Flip scale for right side if mirroring widgets
 166     if use_mirror and get_name_side(bone_name) == Side.RIGHT:
 167         if bone.custom_shape_scale_xyz.x > 0:
 168             bone.custom_shape_scale_xyz.x *= -1
 169
 170     # Move object to bone position and set layers
 171     obj_to_bone(obj, rig, bone_name, bone_transform_name)
 172
 173     if reuse_mesh:
 174         return None
 175
 176     return obj
 177
 178
 179 ##############################################
 180 # Widget choice dropdown
 181 ##############################################
 182
 183 _registered_widgets = {}
 184
 185
 186 def _get_valid_args(callback, skip):
 187     spec = inspect.getfullargspec(callback)
 188     return set(spec.args[skip:] + spec.kwonlyargs)
 189
 190
 191 def register_widget(name: str, callback, **default_args):
 192     unwrapped = inspect.unwrap(callback)
 193     if unwrapped != callback:
 194         valid_args = _get_valid_args(unwrapped, 1)
 195     else:
 196         valid_args = _get_valid_args(callback, 2)
 197
 198     _registered_widgets[name] = (callback, valid_args, default_args)
 199
 200
 201 def get_rigify_widgets(id_store: WindowManager) -> IdPropSequence['RigifyName']:
 202     return id_store.rigify_widgets  # noqa
 203
 204
 205 def layout_widget_dropdown(layout: UILayout, props, prop_name: str, **kwargs):
 206     """Create a UI dropdown to select a widget from the known list."""
 207
 208     id_store = bpy.context.window_manager
 209     rigify_widgets = get_rigify_widgets(id_store)
 210
 211     rigify_widgets.clear()
 212
 213     for name in sorted(_registered_widgets):
 214         item = rigify_widgets.add()
 215         item.name = name
 216
 217     layout.prop_search(props, prop_name, id_store, "rigify_widgets", **kwargs)
 218
 219
 220 def create_registered_widget(obj: ArmatureObject, bone_name: str, widget_id: str, **kwargs):
 221     try:
 222         callback, valid_args, default_args = _registered_widgets[widget_id]
 223     except KeyError:
 224         raise MetarigError("Unknown widget name: " + widget_id)
 225
 226     # Convert between radius and size
 227     if kwargs.get('size') and 'size' not in valid_args:
 228         if 'radius' in valid_args and not kwargs.get('radius'):
 229             kwargs['radius'] = kwargs['size'] / 2
 230
 231     elif kwargs.get('radius') and 'radius' not in valid_args:
 232         if 'size' in valid_args and not kwargs.get('size'):
 233             kwargs['size'] = kwargs['radius'] * 2
 234
 235     args = {**default_args, **kwargs}
 236
 237     return callback(obj, bone_name, **{k: v for k, v in args.items() if k in valid_args})
 238
 239
 240 ##############################################
 241 # Widget geometry
 242 ##############################################
 243
 244 class GeometryData:
 245     verts: list[AnyVector]
 246     edges: list[tuple[int, int]]
 247     faces: list[tuple[int, ...]]
 248
 249     def __init__(self):
 250         self.verts = []
 251         self.edges = []
 252         self.faces = []
 253
 254
 255 def widget_generator(generate_func=None, *, register=None, subsurf=0) -> Callable:
 256     """
 257     Decorator that encapsulates a call to create_widget, and only requires
 258     the actual function to fill the provided vertex and edge lists.
 259
 260     Accepts parameters of create_widget, plus any keyword arguments the
 261     wrapped function has.
 262     """
 263     if generate_func is None:
 264         return functools.partial(widget_generator, register=register, subsurf=subsurf)
 265
 266     @functools.wraps(generate_func)
 267     def wrapper(rig: ArmatureObject, bone_name: str, bone_transform_name=None,
 268                 widget_name=None, widget_force_new=False, **kwargs):
 269         obj = create_widget(rig, bone_name, bone_transform_name,
 270                             widget_name=widget_name, widget_force_new=widget_force_new,
 271                             subsurf=subsurf)
 272         if obj is not None:
 273             geom = GeometryData()
 274
 275             generate_func(geom, **kwargs)
 276
 277             mesh: Mesh = obj.data
 278             mesh.from_pydata(geom.verts, geom.edges, geom.faces)
 279             mesh.update()
 280
 281             return obj
 282         else:
 283             return None
 284
 285     if register:
 286         register_widget(register, wrapper)
 287
 288     return wrapper
 289
 290
 291 def generate_lines_geometry(geom: GeometryData,
 292                             points: list[AnyVector], *,
 293                             matrix: Optional[Matrix] = None, closed_loop=False):
 294     """
 295     Generates a polyline using given points, optionally closing the loop.
 296     """
 297     assert len(points) >= 2
 298
 299     base = len(geom.verts)
 300
 301     for i, raw_point in enumerate(points):
 302         point = Vector(raw_point).to_3d()
 303
 304         if matrix:
 305             point = matrix @ point
 306
 307         geom.verts.append(point)
 308
 309         if i > 0:
 310             geom.edges.append((base + i - 1, base + i))
 311
 312     if closed_loop:
 313         geom.edges.append((len(geom.verts) - 1, base))
 314
 315
 316 def generate_circle_geometry(geom: GeometryData, center: AnyVector, radius: float, *,
 317                              matrix: Optional[Matrix] = None,
 318                              angle_range: Optional[tuple[float, float]] = None,
 319                              steps=24, radius_x: Optional[float] = None, depth_x=0):
 320     """
 321     Generates a circle, adding vertices and edges to the lists.
 322     center, radius: parameters of the circle
 323     matrix: transformation matrix (by default the circle is in the XY plane)
 324     angle_range: a pair of angles to generate an arc of the circle
 325     steps: number of edges to cover the whole circle (reduced for arcs)
 326     """
 327     assert steps >= 3
 328
 329     start = 0
 330     delta = math.pi * 2 / steps
 331
 332     if angle_range:
 333         start, end = angle_range
 334         if start == end:
 335             steps = 1
 336         else:
 337             steps = max(3, math.ceil(abs(end - start) / delta) + 1)
 338             delta = (end - start) / (steps - 1)
 339
 340     if radius_x is None:
 341         radius_x = radius
 342
 343     center = Vector(center).to_3d()  # allow 2d center
 344     points = []
 345
 346     for i in range(steps):
 347         angle = start + delta * i
 348         x = math.cos(angle)
 349         y = math.sin(angle)
 350         points.append(center + Vector((x * radius_x, y * radius, x * x * depth_x)))
 351
 352     generate_lines_geometry(geom, points, matrix=matrix, closed_loop=not angle_range)
 353
 354
 355 def generate_circle_hull_geometry(geom: GeometryData, points: list[AnyVector],
 356                                   radius: float, gap: float, *,
 357                                   matrix: Optional[Matrix] = None, steps=24):
 358     """
 359     Given a list of 2D points forming a convex hull, generate a contour around
 360     it, with each point being circumscribed with a circle arc of given radius,
 361     and keeping the given distance gap from the lines connecting the circles.
 362     """
 363     assert radius >= gap
 364
 365     if len(points) <= 1:
 366         if points:
 367             generate_circle_geometry(
 368                 geom, points[0], radius,
 369                 matrix=matrix, steps=steps
 370             )
 371         return
 372
 373     base = len(geom.verts)
 374     points_ex = [points[-1], *points, points[0]]
 375     angle_gap = math.asin(gap / radius)
 376
 377     for i, pt_prev, pt_cur, pt_next in zip(count(0), points_ex[0:], points_ex[1:], points_ex[2:]):
 378         vec_prev = pt_prev - pt_cur
 379         vec_next = pt_next - pt_cur
 380
 381         # Compute bearings to adjacent points
 382         angle_prev = math.atan2(vec_prev.y, vec_prev.x)
 383         angle_next = math.atan2(vec_next.y, vec_next.x)
 384         if angle_next <= angle_prev:
 385             angle_next += math.pi * 2
 386
 387         # Adjust gap for circles that are too close
 388         angle_prev += max(angle_gap, math.acos(min(1, vec_prev.length/radius/2)))
 389         angle_next -= max(angle_gap, math.acos(min(1, vec_next.length/radius/2)))
 390
 391         if angle_next > angle_prev:
 392             if len(geom.verts) > base:
 393                 geom.edges.append((len(geom.verts)-1, len(geom.verts)))
 394
 395             generate_circle_geometry(
 396                 geom, pt_cur, radius, angle_range=(angle_prev, angle_next),
 397                 matrix=matrix, steps=steps
 398             )
 399
 400     if len(geom.verts) > base:
 401         geom.edges.append((len(geom.verts)-1, base))
 402
 403
 404 def create_circle_polygon(number_verts: int, axis: str, radius=1.0, head_tail=0.0):
 405     """ Creates a basic circle around of an axis selected.
 406         number_verts: number of vertices of the polygon
 407         axis: axis normal to the circle
 408         radius: the radius of the circle
 409         head_tail: where along the length of the bone the circle is (0.0=head, 1.0=tail)
 410     """
 411     verts = []
 412     edges = []
 413     angle = 2 * math.pi / number_verts
 414     i = 0
 415
 416     assert(axis in 'XYZ')
 417
 418     while i < number_verts:
 419         a = math.cos(i * angle)
 420         b = math.sin(i * angle)
 421
 422         if axis == 'X':
 423             verts.append((head_tail, a * radius, b * radius))
 424         elif axis == 'Y':
 425             verts.append((a * radius, head_tail, b * radius))
 426         elif axis == 'Z':
 427             verts.append((a * radius, b * radius, head_tail))
 428
 429         if i < (number_verts - 1):
 430             edges.append((i, i + 1))
 431
 432         i += 1
 433
 434     edges.append((0, number_verts - 1))
 435
 436     return verts, edges
 437
 438
 439 ##############################################
 440 # Widget transformation
 441 ##############################################
 442
 443 def adjust_widget_axis(obj: Object, axis='y', offset=0.0):
 444     mesh = obj.data
 445     assert isinstance(mesh, Mesh)
 446
 447     if axis[0] == '-':
 448         s = -1.0
 449         axis = axis[1]
 450     else:
 451         s = 1.0
 452
 453     trans_matrix = Matrix.Translation((0.0, offset, 0.0))
 454     rot_matrix = Matrix.Diagonal((1.0, s, 1.0, 1.0))
 455
 456     if axis == "x":
 457         rot_matrix = Matrix.Rotation(-s*math.pi/2, 4, 'Z')
 458         trans_matrix = Matrix.Translation((offset, 0.0, 0.0))
 459
 460     elif axis == "z":
 461         rot_matrix = Matrix.Rotation(s*math.pi/2, 4, 'X')
 462         trans_matrix = Matrix.Translation((0.0, 0.0, offset))
 463
 464     matrix = trans_matrix @ rot_matrix
 465
 466     for vert in mesh.vertices:
 467         vert.co = matrix @ vert.co
 468
 469
 470 def adjust_widget_transform_mesh(obj: Optional[Object], matrix: Matrix,
 471                                  local: bool | None = None):
 472     """Adjust the generated widget by applying a correction matrix to the mesh.
 473        If local is false, the matrix is in world space.
 474        If local is True, it's in the local space of the widget.
 475        If local is a bone, it's in the local space of the bone.
 476     """
 477     if obj:
 478         mesh = obj.data
 479         assert isinstance(mesh, Mesh)
 480
 481         if local is not True:
 482             if local:
 483                 assert isinstance(local, bpy.types.PoseBone)
 484                 bone_mat = local.id_data.matrix_world @ local.bone.matrix_local
 485                 matrix = bone_mat @ matrix @ bone_mat.inverted()
 486
 487             obj_mat = obj.matrix_basis
 488             matrix = obj_mat.inverted() @ matrix @ obj_mat
 489
 490         mesh.transform(matrix)
 491
 492
 493 def write_widget(obj: Object, name='thing', use_size=True):
 494     """ Write a mesh object as a python script for widget use.
 495     """
 496     script = ""
 497     script += "@widget_generator\n"
 498     script += "def create_"+name+"_widget(geom"
 499     if use_size:
 500         script += ", *, size=1.0"
 501     script += "):\n"
 502
 503     # Vertices
 504     szs = "*size" if use_size else ""
 505     width = 2 if use_size else 3
 506
 507     mesh = obj.data
 508     assert isinstance(mesh, Mesh)
 509
 510     script += "    geom.verts = ["
 511     for i, v in enumerate(mesh.vertices):
 512         script += "({:g}{}, {:g}{}, {:g}{}),".format(v.co[0], szs, v.co[1], szs, v.co[2], szs)
 513         script += "\n                  " if i % width == (width - 1) else " "
 514     script += "]\n"
 515
 516     # Edges
 517     script += "    geom.edges = ["
 518     for i, e in enumerate(mesh.edges):
 519         script += "(" + str(e.vertices[0]) + ", " + str(e.vertices[1]) + "),"
 520         script += "\n                  " if i % 10 == 9 else " "
 521     script += "]\n"
 522
 523     # Faces
 524     if mesh.polygons:
 525         script += "    geom.faces = ["
 526         for i, f in enumerate(mesh.polygons):
 527             script += "(" + ", ".join(str(v) for v in f.vertices) + "),"
 528             script += "\n                  " if i % 10 == 9 else " "
 529         script += "]\n"
 530
 531     return script