Files
blender-portable-repo/scripts/addons/bone_gizmos/bone_gizmo.py
T
2026-03-17 14:58:51 -06:00

383 lines
13 KiB
Python

# SPDX-FileCopyrightText: 2025 Blender Studio Tools Authors
#
# SPDX-License-Identifier: GPL-3.0-or-later
import bpy
from mathutils import Matrix, Euler
from bpy.types import Gizmo
import numpy as np
import gpu
from .shapes import MeshShape3D
from .utils import get_addon_prefs
# Let's fucking do it.
is_interacting = False
class MoveBoneGizmo(Gizmo):
"""In order to avoid re-implementing logic for transforming bones with
mouse movements, this gizmo instead binds its offset value to the
bpy.ops.transform.translate operator, giving us all that behaviour for free.
(Important behaviours like auto-keying, precision, snapping, axis locking, etc)
The downside of this is that we can't customize that behaviour very well,
for example we can't get the gizmo to draw during mouse interaction,
we can't hide the mouse cursor, etc. Minor sacrifices.
"""
bl_idname = "GIZMO_GT_bone_gizmo"
# The id must be "offset"
bl_target_properties = (
{"id": "offset", "type": 'FLOAT', "array_length": 3},
)
__slots__ = (
# This __slots__ thing allows us to use arbitrary Python variable
# assignments on instances of this gizmo.
"bone_name" # Name of the bone that owns this gizmo.
,"custom_shape" # Currently drawn shape, being passed into self.new_custom_shape().
,"meshshape" # Cache of vertex indicies. Can fall out of sync if the mesh is modified; Only re-calculated when gizmo properties are changed by user.
# Gizmos, like bones, have 3 states.
,"color_selected"
,"color_unselected"
,"alpha_selected"
,"alpha_unselected"
# The 3rd one is "highlighted".
# color_highlight and alpha_highlight are already provided by the API.
# We currently don't visually distinguish between selected and active gizmos.
)
def setup(self):
"""Called by Blender when the Gizmo is created."""
self.meshshape = None
self.custom_shape = None
def init_shape(self, context):
"""Should be called by the GizmoGroup, after it assigns the neccessary
__slots__ properties to properly initialize this Gizmo."""
if not self.poll(context):
return
props = self.get_props(context)
if self.is_using_vgroup(context):
self.load_shape_vertex_group(self.get_shape_object(context), props.vertex_group_name)
dg = context.evaluated_depsgraph_get()
ob = self.get_shape_object(context)
self.refresh_shape_vgroup(context, ob.evaluated_get(dg).to_mesh())
elif self.is_using_facemap(context):
# We use the built-in function to draw face maps, so we don't need to do any extra processing.
pass
else:
self.load_shape_entire_object(context)
def init_properties(self, context):
self.refresh_colors(context)
def refresh_colors(self, context):
prefs = get_addon_prefs(context)
self.line_width = prefs.line_width
props = self.get_props(context)
if self.is_using_bone_group_colors(context):
pb = self.get_pose_bone(context)
self.color_unselected = pb.bone_group.colors.normal[:]
self.color_selected = pb.bone_group.colors.select[:]
self.color_highlight = pb.bone_group.colors.select[:]
else:
self.color_unselected = props.color[:]
self.color_selected = props.color_highlight[:]
self.color_highlight = props.color_highlight[:]
if self.is_using_facemap(context) or self.is_using_vgroup(context):
self.alpha_unselected = prefs.mesh_alpha
self.alpha_selected = prefs.mesh_alpha + prefs.delta_alpha_select
self.alpha_highlight = min(0.999, self.alpha_selected + prefs.delta_alpha_highlight)
else:
self.alpha_unselected = prefs.widget_alpha
self.alpha_selected = prefs.widget_alpha + prefs.delta_alpha_select
self.alpha_highlight = min(0.999, self.alpha_selected + prefs.delta_alpha_highlight)
def poll(self, context):
"""Whether any gizmo logic should be executed or not. This function is not
from the API! Call this manually for early exists.
"""
global is_interacting
if is_interacting:
return False
pb = self.get_pose_bone(context)
if not pb or pb.bone.hide: return False
any_visible_layer = any(bl and al for bl, al in zip(pb.bone.layers[:], pb.id_data.data.layers[:]))
bone_visible = not pb.bone.hide and any_visible_layer
ret = self.get_shape_object(context) and bone_visible and pb.enable_bone_gizmo
return ret
def load_shape_vertex_group(self, obj, v_grp: str, weight_threshold=0.2):
"""Update the vertex indicies that the gizmo shape corresponds to when using
vertex group masking.
This is very expensive, should only be called on initial Gizmo creation,
manual updates, and changing of gizmo display object or mask group.
"""
self.meshshape = MeshShape3D(obj, vertex_groups=[v_grp], weight_threshold=weight_threshold)
def refresh_shape_vgroup(self, context, eval_mesh):
"""Update the custom shape based on the stored vertex indices."""
if not self.meshshape:
self.init_shape(context)
draw_style = 'TRIS'
if len(self.meshshape._indices) < 3:
draw_style = 'LINES'
if len(self.meshshape._indices) < 2:
return
self.custom_shape = self.new_custom_shape(draw_style, self.meshshape.get_vertices(eval_mesh))
return True
def load_shape_entire_object(self, context):
"""Update the custom shape to an entire object. This is somewhat expensive,
should only be called when Gizmo display object is changed or mask
facemap/vgroup is cleared.
"""
mesh = self.get_shape_object(context).data
vertices = np.zeros((len(mesh.vertices), 3), 'f')
mesh.vertices.foreach_get("co", vertices.ravel())
if len(mesh.polygons) > 0:
draw_style = 'TRIS'
mesh.calc_loop_triangles()
tris = np.zeros((len(mesh.loop_triangles), 3), 'i')
mesh.loop_triangles.foreach_get("vertices", tris.ravel())
custom_shape_verts = vertices[tris].reshape(-1,3)
else:
draw_style = 'LINES'
edges = np.zeros((len(mesh.edges), 2), 'i')
mesh.edges.foreach_get("vertices", edges.ravel())
custom_shape_verts = vertices[edges].reshape(-1,3)
self.custom_shape = self.new_custom_shape(draw_style, custom_shape_verts)
def draw_shape(self, context, select_id=None):
"""Shared drawing logic for selection and color.
We do not pass color here; The C functions read the
colors from self.color and self.color_highlight.
"""
ob = self.get_shape_object(context)
props = self.get_props(context)
face_map = ob.face_maps.get(props.face_map_name)
if face_map and props.use_face_map:
self.draw_preset_facemap(ob, face_map.index, select_id=select_id or 0)
elif self.custom_shape:
self.draw_custom_shape(self.custom_shape, select_id=select_id)
else:
# This can happen if the specified vertex group is empty.
return
def draw_shared(self, context, select_id=None):
if not self.poll(context):
return
if not self.get_shape_object(context):
return
self.update_basis_and_offset_matrix(context)
gpu.state.line_width_set(self.line_width)
gpu.state.blend_set('MULTIPLY')
self.draw_shape(context, select_id)
gpu.state.blend_set('NONE')
gpu.state.line_width_set(1)
def get_opacity(self, context):
"""Based factors of whether the bone corresponding to this gizmo
is currently selected, and whether the gizmo is being mouse hovered,
return the opacity value that is expected to be used for drawing this gizmo.
"""
prefs = get_addon_prefs(context)
is_selected = self.get_pose_bone(context).bone.select
opacity = prefs.widget_alpha
if self.is_using_facemap(context) or self.is_using_vgroup(context):
opacity = prefs.mesh_alpha
if self.is_highlight:
opacity += prefs.delta_alpha_highlight
elif is_selected:
opacity += prefs.delta_alpha_select
return opacity
def draw(self, context):
"""Called by Blender on every viewport update (including mouse moves).
Drawing functions called at this time will draw into the color pass.
"""
if not self.poll(context):
return
if self.use_draw_hover and not self.is_highlight:
return
if self.get_opacity(context) == 0:
return
pb = self.get_pose_bone(context)
if pb.bone.select:
self.color = self.color_selected
self.alpha = min(0.999, self.alpha_selected) # An alpha value of 1.0 or greater results in glitched drawing.
else:
self.color = self.color_unselected
self.alpha = min(0.999, self.alpha_unselected)
self.draw_shared(context)
def draw_select(self, context, select_id):
"""Called by Blender on every viewport update (including mouse moves).
Drawing functions called at this time will draw into an invisible pass
that is used for mouse interaction.
"""
if not self.poll(context):
return
self.draw_shared(context, select_id)
def is_using_vgroup(self, context):
ob = self.get_shape_object(context)
if not ob: return False
props = self.get_pose_bone(context).bone_gizmo
vgroup_exists = props.vertex_group_name in ob.vertex_groups
ret = ob and not props.use_face_map and vgroup_exists
return ret
def is_using_facemap(self, context):
props = self.get_props(context)
ob = self.get_shape_object(context)
if not ob: return False
return props.use_face_map and props.face_map_name in ob.face_maps
def is_using_bone_group_colors(self, context):
pb = self.get_pose_bone(context)
props = self.get_props(context)
return pb and pb.bone_group and pb.bone_group.color_set != 'DEFAULT' and props.gizmo_color_source == 'GROUP'
def get_pose_bone(self, context):
arm_ob = context.object
if not arm_ob or arm_ob.type != 'ARMATURE':
return
ret = arm_ob.pose.bones.get(self.bone_name)
return ret
def get_shape_object(self, context):
"""Get the shape object selected by the user in the Custom Gizmo panel.
"""
pb = self.get_pose_bone(context)
if not pb: return
ob = pb.bone_gizmo.shape_object
if not ob:
ob = pb.custom_shape
return ob
def get_props(self, context):
"""Use this context-based getter rather than any direct mean of referencing
the gizmo properties, because that would result in a crash on undo.
"""
pb = self.get_pose_bone(context)
return pb.bone_gizmo
def update_basis_and_offset_matrix(self, context):
"""Set the gizmo matrices self.matrix_basis and self.matrix_offset,
to position the gizmo correctly."""
pb = self.get_pose_bone(context)
armature = context.object
if self.is_using_facemap(context) or self.is_using_vgroup(context):
# If there is a face map or vertex group specified:
# The gizmo should stick strictly to the vertex group or face map of the shape object.
self.matrix_basis = self.get_shape_object(context).matrix_world.copy()
self.matrix_offset = Matrix.Identity(4)
else:
# If there is NO face map or vertex group specified:
# The gizmo should function as a replacement for the bone's Custom Shape
# properties. That means applying the custom shape transformation offsets
# and using the custom shape transform bone, if there is one specified.
self.matrix_basis = armature.matrix_world.copy()
display_bone = pb
if pb.custom_shape_transform:
display_bone = pb.custom_shape_transform
bone_mat = display_bone.matrix.copy()
trans_mat = Matrix.Translation(pb.custom_shape_translation)
rot = Euler((pb.custom_shape_rotation_euler.x, pb.custom_shape_rotation_euler.y, pb.custom_shape_rotation_euler.z), 'XYZ')
rot_mat = rot.to_matrix().to_4x4()
display_scale = pb.custom_shape_scale_xyz.copy()
if pb.use_custom_shape_bone_size:
display_scale *= pb.bone.length
scale_mat_x = Matrix.Scale(display_scale.x, 4, (1, 0, 0))
scale_mat_y = Matrix.Scale(display_scale.y, 4, (0, 1, 0))
scale_mat_z = Matrix.Scale(display_scale.z, 4, (0, 0, 1))
scale_mat = scale_mat_x @ scale_mat_y @ scale_mat_z
final_mat = bone_mat @ trans_mat @ rot_mat @ scale_mat
self.matrix_offset = final_mat
def invoke(self, context, event):
armature = context.object
if not event.shift:
for pb in armature.pose.bones:
pb.bone.select = False
pb = self.get_pose_bone(context)
if event.shift and pb.bone.select:
pb.bone.select = False
return {'FINISHED'}
if event.shift and not pb.bone.select:
pb.bone.select = True
armature.data.bones.active = pb.bone
return {'FINISHED'}
global is_interacting
is_interacting = True
pb.bone.select = True
armature.data.bones.active = pb.bone
# Allow executing an operator on bone interaction,
# based on data stored in the armatures 'gizmo_interaction' custom property.
# This should be a dictionary structured:
# op_bl_idname : [ ( [list of bone names], {op_kwargs} ) ]
# Whenever any of the bones from one list of bone names is interacted,
# the operator with op_bl_idname is executed, with op_kwargs passed in to it.
if 'gizmo_interactions' in armature.data:
interaction_data = armature.data['gizmo_interactions'].to_dict()
for op_name, op_datas in interaction_data.items():
op_category, op_name = op_name.split(".")
op_callable = getattr(getattr(bpy.ops, op_category), op_name)
for op_data in op_datas:
bone_names, op_kwargs = op_data
if pb.name in bone_names:
op_callable(**op_kwargs)
return {'RUNNING_MODAL'}
def exit(self, context, cancel):
global is_interacting
is_interacting = False
return
def modal(self, context, event, tweak):
return {'RUNNING_MODAL'}
classes = (
MoveBoneGizmo,
)
register, unregister = bpy.utils.register_classes_factory(classes)