# SPDX-License-Identifier: GPL-2.0-or-later import bpy import collections import heapq import operator from mathutils import Vector from mathutils.kdtree import KDTree from .errors import MetarigError from ..base_rig import stage, GenerateCallbackHost from ..base_generate import GeneratorPlugin class NodeMerger(GeneratorPlugin): """ Utility that allows rigs to interact based on common points in space. Rigs can register node objects representing locations during the initialize stage, and at the end the plugin sorts them into buckets based on proximity. For each such bucket a group object is created and allowed to further process the nodes. Nodes chosen by the groups as being 'final' become sub-objects of the plugin and receive stage callbacks. The domain parameter allows potentially having multiple completely separate layers of nodes with different purpose. """ epsilon = 1e-5 def __init__(self, generator, domain): super().__init__(generator) assert domain is not None assert generator.stage == 'initialize' self.domain = domain self.nodes = [] self.final_nodes = [] self.groups = [] self.frozen = False def register_node(self, node): assert not self.frozen node.generator_plugin = self self.nodes.append(node) def initialize(self): self.frozen = True nodes = self.nodes tree = KDTree(len(nodes)) for i, node in enumerate(nodes): tree.insert(node.point, i) tree.balance() processed = set() final_nodes = [] groups = [] for i in range(len(nodes)): if i in processed: continue # Find points to merge pending = [i] merge_set = set(pending) while pending: added = set() for j in pending: point = nodes[j].point eps = max(1, point.length) * self.epsilon for co, idx, dist in tree.find_range(point, eps): added.add(idx) pending = added.difference(merge_set) merge_set.update(added) assert merge_set.isdisjoint(processed) processed.update(merge_set) # Group the points merge_list = [nodes[i] for i in merge_set] merge_list.sort(key=lambda x: x.name) group_class = merge_list[0].group_class for item in merge_list[1:]: cls = item.group_class if issubclass(cls, group_class): group_class = cls elif not issubclass(group_class, cls): raise MetarigError( 'Group class conflict: {} and {} from {} of {}'.format( group_class, cls, item.name, item.rig.base_bone, ) ) group = group_class(merge_list) group.build(final_nodes) groups.append(group) self.final_nodes = self.rigify_sub_objects = final_nodes self.groups = groups class MergeGroup(object): """ Standard node group, merges nodes based on certain rules and priorities. 1. Nodes are classified into main and query nodes; query nodes are not merged. 2. Nodes owned by the same rig cannot merge with each other. 3. Node can only merge into target if node.can_merge_into(target) is true. 4. Among multiple candidates in one rig, node.get_merge_priority(target) is used. 5. The largest clusters of nodes that can merge are picked until none are left. The master nodes of the chosen clusters, plus query nodes, become 'final'. """ def __init__(self, nodes): self.nodes = nodes for node in nodes: node.group = self def is_main(node): return isinstance(node, MainMergeNode) self.main_nodes = [n for n in nodes if is_main(n)] self.query_nodes = [n for n in nodes if not is_main(n)] def build(self, final_nodes): main_nodes = self.main_nodes # Sort nodes into rig buckets - can't merge within the same rig rig_table = collections.defaultdict(list) for node in main_nodes: rig_table[node.rig].append(node) # Build a 'can merge' table merge_table = {n: set() for n in main_nodes} for node in main_nodes: for rig, tgt_nodes in rig_table.items(): if rig is not node.rig: nodes = [n for n in tgt_nodes if node.can_merge_into(n)] if nodes: best_node = max(nodes, key=node.get_merge_priority) merge_table[best_node].add(node) # Output groups starting with largest self.final_nodes = [] pending = set(main_nodes) while pending: # Find largest group nodes = [n for n in main_nodes if n in pending] max_len = max(len(merge_table[n]) for n in nodes) nodes = [n for n in nodes if len(merge_table[n]) == max_len] # If a tie, try to resolve using comparison if len(nodes) > 1: weighted_nodes = [ (n, sum( 1 if (n.is_better_cluster(n2) and not n2.is_better_cluster(n)) else 0 for n2 in nodes )) for n in nodes ] max_weight = max(wn[1] for wn in weighted_nodes) nodes = [wn[0] for wn in weighted_nodes if wn[1] == max_weight] # Final tie breaker is the name best = min(nodes, key=lambda n: n.name) child_set = merge_table[best] # Link children best.point = sum((c.point for c in child_set), best.point) / (len(child_set) + 1) for child in [n for n in main_nodes if n in child_set]: child.point = best.point best.merge_from(child) child.merge_into(best) final_nodes.append(best) self.final_nodes.append(best) best.merge_done() # Remove merged nodes from the table pending.remove(best) pending -= child_set for children in merge_table.values(): children &= pending for node in self.query_nodes: node.merge_done() final_nodes += self.query_nodes class BaseMergeNode(GenerateCallbackHost): """Base class of mergeable nodes.""" merge_domain = None merger = NodeMerger group_class = MergeGroup def __init__(self, rig, name, point, *, domain=None): self.rig = rig self.obj = rig.obj self.name = name self.point = Vector(point) merger = self.merger(rig.generator, domain or self.merge_domain) merger.register_node(self) def register_new_bone(self, new_name, old_name=None): self.generator_plugin.register_new_bone(new_name, old_name) def can_merge_into(self, other): raise NotImplementedError() def get_merge_priority(self, other): "Rank candidates to merge into." return 0 class MainMergeNode(BaseMergeNode): """ Base class of standard mergeable nodes. Each node can either be a master of its cluster or a merged child node. Children become sub-objects of their master to receive callbacks in defined order. """ def __init__(self, rig, name, point, *, domain=None): super().__init__(rig, name, point, domain=domain) self.merged_into = None self.merged = [] def get_merged_siblings(self): master = self.merged_master return [master, *master.merged] def is_better_cluster(self, other): "Compare with the other node to choose between cluster masters." return False def can_merge_from(self, other): return True def can_merge_into(self, other): return other.can_merge_from(self) def merge_into(self, other): self.merged_into = other def merge_from(self, other): self.merged.append(other) @property def is_master_node(self): return not self.merged_into def merge_done(self): self.merged_master = self.merged_into or self self.rigify_sub_objects = list(self.merged) for child in self.merged: child.merge_done() class QueryMergeNode(BaseMergeNode): """Base class for special nodes used only to query which nodes are at a certain location.""" is_master_node = False require_match = True def merge_done(self): self.matched_nodes = [ n for n in self.group.final_nodes if self.can_merge_into(n) ] self.matched_nodes.sort(key=self.get_merge_priority, reverse=True) if self.require_match and not self.matched_nodes: self.rig.raise_error( 'Could not match control node for {}', self.name)