rigify/rigs/skin/skin_nodes.py

   1 # SPDX-License-Identifier: GPL-2.0-or-later
   2
   3 # <pep8 compliant>
   4
   5 import bpy
   6 import enum
   7
   8 from mathutils import Vector, Quaternion
   9
  10 from ...utils.layers import set_bone_layers
  11 from ...utils.naming import NameSides, make_derived_name, get_name_base_and_sides, change_name_side, Side, SideZ
  12 from ...utils.bones import BoneUtilityMixin, set_bone_widget_transform
  13 from ...utils.widgets_basic import create_cube_widget, create_sphere_widget
  14 from ...utils.mechanism import MechanismUtilityMixin
  15 from ...utils.rig import get_parent_rigs
  16
  17 from ...utils.node_merger import MainMergeNode, QueryMergeNode
  18
  19 from .skin_parents import ControlBoneParentLayer, ControlBoneWeakParentLayer, ControlBoneParentMix
  20 from .skin_rigs import BaseSkinRig, BaseSkinChainRig
  21
  22
  23 class ControlNodeLayer(enum.IntEnum):
  24     FREE = 0
  25     MIDDLE_PIVOT = 10
  26     TWEAK = 20
  27
  28
  29 class ControlNodeIcon(enum.IntEnum):
  30     TWEAK = 0
  31     MIDDLE_PIVOT = 1
  32     FREE = 2
  33     CUSTOM = 3
  34
  35
  36 class ControlNodeEnd(enum.IntEnum):
  37     START = -1
  38     MIDDLE = 0
  39     END = 1
  40
  41
  42 class BaseSkinNode(MechanismUtilityMixin, BoneUtilityMixin):
  43     """Base class for skin control and query nodes."""
  44
  45     node_parent_built = False
  46
  47     def do_build_parent(self):
  48         """Create and intern the parent mechanism generator."""
  49         assert self.rig.generator.stage == 'initialize'
  50
  51         result = self.rig.build_own_control_node_parent(self)
  52         parents = self.rig.get_all_parent_skin_rigs()
  53
  54         for rig in reversed(parents):
  55             result = rig.extend_control_node_parent(result, self)
  56
  57         for rig in parents:
  58             result = rig.extend_control_node_parent_post(result, self)
  59
  60         result = self.merged_master.intern_parent(self, result)
  61         result.is_parent_frozen = True
  62         return result
  63
  64     def build_parent(self, use=True, reparent=False):
  65         """Create and activate if needed the parent mechanism for this node."""
  66         if not self.node_parent_built:
  67             self.node_parent = self.do_build_parent()
  68             self.node_parent_built = True
  69
  70         if use:
  71             self.merged_master.register_use_parent(self.node_parent)
  72
  73         if reparent:
  74             self.merged_master.request_reparent(self.node_parent)
  75
  76         return self.node_parent
  77
  78     @property
  79     def control_bone(self):
  80         """The generated control bone."""
  81         return self.merged_master._control_bone
  82
  83     @property
  84     def reparent_bone(self):
  85         """The generated reparent bone for this node's parent mechanism."""
  86         return self.merged_master.get_reparent_bone(self.node_parent)
  87
  88
  89 class ControlBoneNode(MainMergeNode, BaseSkinNode):
  90     """Node representing controls of skin chain rigs."""
  91
  92     merge_domain = 'ControlNetNode'
  93
  94     def __init__(
  95         self, rig, org, name, *, point=None, size=None,
  96         needs_parent=False, needs_reparent=False, allow_scale=False,
  97         chain_end=ControlNodeEnd.MIDDLE,
  98         layer=ControlNodeLayer.FREE, index=None, icon=ControlNodeIcon.TWEAK,
  99     ):
 100         assert isinstance(rig, BaseSkinChainRig)
 101
 102         super().__init__(rig, name, point or rig.get_bone(org).head)
 103
 104         self.org = org
 105
 106         self.name_split = get_name_base_and_sides(name)
 107
 108         self.name_merged = None
 109         self.name_merged_split = None
 110
 111         self.size = size or rig.get_bone(org).length
 112         self.layer = layer
 113         self.icon = icon
 114         self.rotation = None
 115         self.chain_end = chain_end
 116
 117         # Create the parent mechanism even if not master
 118         self.node_needs_parent = needs_parent
 119         # If this node's own parent mechanism differs from master, generate a conversion bone
 120         self.node_needs_reparent = needs_reparent
 121
 122         # Generate the control as a MCH bone to hide it from the user
 123         self.hide_control = False
 124         # Unlock scale channels
 125         self.allow_scale = allow_scale
 126
 127         # For use by the owner rig: index in chain
 128         self.index = index
 129         # If this node is the end of a chain, points to the next one
 130         self.chain_end_neighbor = None
 131
 132     def can_merge_into(self, other):
 133         # Only merge up the layers (towards more mechanism)
 134         dprio = self.rig.chain_priority - other.rig.chain_priority
 135         return (
 136             dprio <= 0 and
 137             (self.layer <= other.layer or dprio < 0) and
 138             super().can_merge_into(other)
 139         )
 140
 141     def get_merge_priority(self, other):
 142         # Prefer higher and closest layer
 143         if self.layer <= other.layer:
 144             return -abs(self.layer - other.layer)
 145         else:
 146             return -abs(self.layer - other.layer) - 100
 147
 148     def is_better_cluster(self, other):
 149         """Check if the current bone is preferable as master when choosing of same sized groups."""
 150
 151         # Prefer bones that have strictly more parents
 152         my_parents = list(reversed(get_parent_rigs(self.rig.rigify_parent)))
 153         other_parents = list(reversed(get_parent_rigs(other.rig.rigify_parent)))
 154
 155         if len(my_parents) > len(other_parents) and my_parents[0:len(other_parents)] == other_parents:
 156             return True
 157         if len(other_parents) > len(my_parents) and other_parents[0:len(other_parents)] == my_parents:
 158             return False
 159
 160         # Prefer side chains
 161         side_x_my, side_z_my = map(abs, self.name_split[1:])
 162         side_x_other, side_z_other = map(abs, other.name_split[1:])
 163
 164         if ((side_x_my < side_x_other and side_z_my <= side_z_other) or
 165                 (side_x_my <= side_x_other and side_z_my < side_z_other)):
 166             return False
 167         if ((side_x_my > side_x_other and side_z_my >= side_z_other) or
 168                 (side_x_my >= side_x_other and side_z_my > side_z_other)):
 169             return True
 170
 171         return False
 172
 173     def merge_done(self):
 174         if self.is_master_node:
 175             self.parent_subrig_cache = []
 176             self.parent_subrig_names = {}
 177             self.reparent_requests = []
 178             self.used_parents = {}
 179
 180         super().merge_done()
 181
 182         self.find_mirror_siblings()
 183
 184     def find_mirror_siblings(self):
 185         """Find merged nodes that have their names in mirror symmetry with this one."""
 186
 187         self.mirror_siblings = {}
 188         self.mirror_sides_x = set()
 189         self.mirror_sides_z = set()
 190
 191         for node in self.get_merged_siblings():
 192             if node.name_split.base == self.name_split.base:
 193                 self.mirror_siblings[node.name_split] = node
 194                 self.mirror_sides_x.add(node.name_split.side)
 195                 self.mirror_sides_z.add(node.name_split.side_z)
 196
 197         assert self.mirror_siblings[self.name_split] is self
 198
 199         # Remove sides that merged with a mirror from the name
 200         side_x = Side.MIDDLE if len(self.mirror_sides_x) > 1 else self.name_split.side
 201         side_z = SideZ.MIDDLE if len(self.mirror_sides_z) > 1 else self.name_split.side_z
 202
 203         self.name_merged = change_name_side(self.name, side=side_x, side_z=side_z)
 204         self.name_merged_split = NameSides(self.name_split.base, side_x, side_z)
 205
 206     def get_best_mirror(self):
 207         """Find best mirror sibling for connecting via mirror."""
 208
 209         base, side, sidez = self.name_split
 210
 211         for flip in [(base, -side, -sidez), (base, -side, sidez), (base, side, -sidez)]:
 212             mirror = self.mirror_siblings.get(flip, None)
 213             if mirror and mirror is not self:
 214                 return mirror
 215
 216         return None
 217
 218     def intern_parent(self, node, parent):
 219         """De-duplicate the parent layer chain within this merge group."""
 220
 221         # Quick check for the same object
 222         if id(parent) in self.parent_subrig_names:
 223             return parent
 224
 225         # Find if an identical parent is already in the cache
 226         cache = self.parent_subrig_cache
 227
 228         for previous in cache:
 229             if previous == parent:
 230                 previous.is_parent_frozen = True
 231                 return previous
 232
 233         # Add to cache and intern the layer parent if exists
 234         cache.append(parent)
 235
 236         self.parent_subrig_names[id(parent)] = node.name
 237
 238         if isinstance(parent, ControlBoneParentLayer):
 239             parent.parent = self.intern_parent(node, parent.parent)
 240         elif isinstance(parent, ControlBoneParentMix):
 241             parent.parents = [self.intern_parent(node, p) for p in parent.parents]
 242
 243         return parent
 244
 245     def register_use_parent(self, parent):
 246         """Activate this parent mechanism generator."""
 247         self.used_parents[id(parent)] = parent
 248
 249     def request_reparent(self, parent):
 250         """Request a reparent bone to be generated for this parent mechanism."""
 251         requests = self.reparent_requests
 252
 253         if parent not in requests:
 254             # If the actual reparent would be generated, weak parent will be needed.
 255             if self.has_weak_parent and not self.use_weak_parent:
 256                 if parent is not self.node_parent:
 257                     self.use_weak_parent = True
 258                     self.register_use_parent(self.node_parent_weak)
 259
 260             self.register_use_parent(parent)
 261             requests.append(parent)
 262
 263     def get_reparent_bone(self, parent):
 264         """Returns the generated reparent bone for this parent mechanism."""
 265         return self.reparent_bones[id(parent)]
 266
 267     def get_rotation(self):
 268         """Returns the orientation quaternion provided for this node by parents."""
 269         if self.rotation is None:
 270             self.rotation = self.rig.get_final_control_node_rotation(self)
 271
 272         return self.rotation
 273
 274     def initialize(self):
 275         if self.is_master_node:
 276             sibling_list = self.get_merged_siblings()
 277             mirror_sibling_list = self.mirror_siblings.values()
 278
 279             # Compute size
 280             best = max(sibling_list, key=lambda n: n.icon)
 281             best_mirror = best.mirror_siblings.values()
 282
 283             self.size = sum(node.size for node in best_mirror) / len(best_mirror)
 284
 285             # Compute orientation
 286             self.rotation = sum(
 287                 (node.get_rotation() for node in mirror_sibling_list),
 288                 Quaternion((0, 0, 0, 0))
 289             ).normalized()
 290
 291             self.matrix = self.rotation.to_matrix().to_4x4()
 292             self.matrix.translation = self.point
 293
 294             # Create parents and decide if mix would be needed
 295             weak_parent_list = [node.build_parent(use=False) for node in mirror_sibling_list]
 296
 297             if all(parent == self.node_parent for parent in weak_parent_list):
 298                 weak_parent_list = [self.node_parent]
 299                 self.node_parent_weak = self.node_parent
 300             else:
 301                 self.node_parent_weak = ControlBoneParentMix(self.rig, self, weak_parent_list)
 302
 303             # Prepare parenting without weak layers
 304             parent_list = [ControlBoneWeakParentLayer.strip(p) for p in weak_parent_list]
 305
 306             self.use_weak_parent = False
 307             self.has_weak_parent = any((p is not pw)
 308                                        for p, pw in zip(weak_parent_list, parent_list))
 309
 310             if not self.has_weak_parent:
 311                 self.node_parent = self.node_parent_weak
 312             elif len(parent_list) > 1:
 313                 self.node_parent = ControlBoneParentMix(
 314                     self.rig, self, parent_list, suffix='_mix_base')
 315             else:
 316                 self.node_parent = parent_list[0]
 317
 318             # Mirror siblings share the mixed parent for reparent
 319             self.register_use_parent(self.node_parent)
 320
 321             for node in mirror_sibling_list:
 322                 node.node_parent = self.node_parent
 323
 324         # All nodes
 325         if self.node_needs_parent or self.node_needs_reparent:
 326             self.build_parent(reparent=self.node_needs_reparent)
 327
 328     def prepare_bones(self):
 329         # Activate parent components once all reparents are registered
 330         if self.is_master_node:
 331             for parent in self.used_parents.values():
 332                 parent.enable_component()
 333
 334             self.used_parents = None
 335
 336     def make_bone(self, name, scale, *, rig=None, orientation=None):
 337         """
 338         Creates a bone associated with this node, using the appropriate
 339         orientation, location and size.
 340         """
 341         name = (rig or self).copy_bone(self.org, name)
 342
 343         if orientation is not None:
 344             matrix = orientation.to_matrix().to_4x4()
 345             matrix.translation = self.merged_master.point
 346         else:
 347             matrix = self.merged_master.matrix
 348
 349         bone = self.get_bone(name)
 350         bone.matrix = matrix
 351         bone.length = self.merged_master.size * scale
 352
 353         return name
 354
 355     def find_master_name_node(self):
 356         """Find which node to name the control bone from."""
 357
 358         # Chain end nodes have sub-par names, so try to find another chain
 359         if self.chain_end == ControlNodeEnd.END:
 360             # Choose possible other nodes so that it doesn't lose mirror tags
 361             siblings = [
 362                 node for node in self.get_merged_siblings()
 363                 if self.mirror_sides_x.issubset(node.mirror_sides_x)
 364                 and self.mirror_sides_z.issubset(node.mirror_sides_z)
 365             ]
 366
 367             # Prefer chain start, then middle nodes
 368             candidates = [node for node in siblings if node.chain_end == ControlNodeEnd.START]
 369
 370             if not candidates:
 371                 candidates = [node for node in siblings if node.chain_end == ControlNodeEnd.MIDDLE]
 372
 373             # Choose based on priority and name alphabetical order
 374             if candidates:
 375                 return min(candidates, key=lambda c: (-c.rig.chain_priority, c.name_merged))
 376
 377         return self
 378
 379     def generate_bones(self):
 380         if self.is_master_node:
 381             # Make control bone
 382             self._control_bone = self.make_master_bone()
 383
 384             # Make weak parent bone
 385             if self.use_weak_parent:
 386                 self.weak_parent_bone = self.make_bone(
 387                     make_derived_name(self._control_bone, 'mch', '_weak_parent'), 1/2)
 388
 389             # Make requested reparents
 390             self.reparent_bones = {id(self.node_parent): self._control_bone}
 391             self.reparent_bones_fake = set(self.reparent_bones.values())
 392
 393             for parent in self.reparent_requests:
 394                 if id(parent) not in self.reparent_bones:
 395                     parent_name = self.parent_subrig_names[id(parent)]
 396                     bone = self.make_bone(make_derived_name(parent_name, 'mch', '_reparent'), 1/3)
 397                     self.reparent_bones[id(parent)] = bone
 398
 399     def make_master_bone(self):
 400         choice = self.find_master_name_node()
 401         name = choice.name_merged
 402
 403         if self.hide_control:
 404             name = make_derived_name(name, 'mch')
 405
 406         return choice.make_bone(name, 1)
 407
 408     def parent_bones(self):
 409         if self.is_master_node:
 410             self.set_bone_parent(
 411                 self._control_bone, self.node_parent.output_bone, inherit_scale='AVERAGE')
 412
 413             if self.use_weak_parent:
 414                 self.set_bone_parent(
 415                     self.weak_parent_bone, self.node_parent_weak.output_bone, inherit_scale='FULL')
 416
 417             for parent in self.reparent_requests:
 418                 bone = self.reparent_bones[id(parent)]
 419                 if bone not in self.reparent_bones_fake:
 420                     self.set_bone_parent(bone, parent.output_bone, inherit_scale='AVERAGE')
 421
 422     def configure_bones(self):
 423         if self.is_master_node:
 424             if not any(node.allow_scale for node in self.get_merged_siblings()):
 425                 self.get_bone(self.control_bone).lock_scale = (True, True, True)
 426
 427         layers = self.rig.get_control_node_layers(self)
 428         if layers:
 429             bone = self.get_bone(self.control_bone).bone
 430             set_bone_layers(bone, layers, not self.is_master_node)
 431
 432     def rig_bones(self):
 433         if self.is_master_node:
 434             # Invoke parent rig callbacks
 435             for rig in reversed(self.rig.get_all_parent_skin_rigs()):
 436                 rig.extend_control_node_rig(self)
 437
 438             # Rig reparent bones
 439             reparent_source = self.control_bone
 440
 441             if self.use_weak_parent:
 442                 reparent_source = self.weak_parent_bone
 443
 444                 self.make_constraint(reparent_source, 'COPY_TRANSFORMS',
 445                                      self.control_bone, space='LOCAL')
 446
 447                 set_bone_widget_transform(self.obj, self.control_bone, reparent_source)
 448
 449             for parent in self.reparent_requests:
 450                 bone = self.reparent_bones[id(parent)]
 451                 if bone not in self.reparent_bones_fake:
 452                     self.make_constraint(bone, 'COPY_TRANSFORMS', reparent_source)
 453
 454     def generate_widgets(self):
 455         if self.is_master_node:
 456             best = max(self.get_merged_siblings(), key=lambda n: n.icon)
 457
 458             if best.icon == ControlNodeIcon.TWEAK:
 459                 create_sphere_widget(self.obj, self.control_bone)
 460             elif best.icon in (ControlNodeIcon.MIDDLE_PIVOT, ControlNodeIcon.FREE):
 461                 create_cube_widget(self.obj, self.control_bone)
 462             else:
 463                 best.rig.make_control_node_widget(best)
 464
 465
 466 class ControlQueryNode(QueryMergeNode, BaseSkinNode):
 467     """Node representing controls of skin chain rigs."""
 468
 469     merge_domain = 'ControlNetNode'
 470
 471     def __init__(self, rig, org, *, name=None, point=None, find_highest_layer=False):
 472         assert isinstance(rig, BaseSkinRig)
 473
 474         super().__init__(rig, name or org, point or rig.get_bone(org).head)
 475
 476         self.org = org
 477         self.find_highest_layer = find_highest_layer
 478
 479     def can_merge_into(self, other):
 480         return True
 481
 482     def get_merge_priority(self, other):
 483         return other.layer if self.find_highest_layer else -other.layer
 484
 485     @property
 486     def merged_master(self):
 487         return self.matched_nodes[0]