diff options
Diffstat (limited to 'rigify/utils/node_merger.py')
| -rw-r--r-- | rigify/utils/node_merger.py | 317 |
1 files changed, 317 insertions, 0 deletions
diff --git a/rigify/utils/node_merger.py b/rigify/utils/node_merger.py new file mode 100644 index 00000000..49ebaf27 --- /dev/null +++ b/rigify/utils/node_merger.py @@ -0,0 +1,317 @@ +# ====================== BEGIN GPL LICENSE BLOCK ====================== +# +# This program is free software; you can redistribute it and/or +# modify it under the terms of the GNU General Public License +# as published by the Free Software Foundation; either version 2 +# of the License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program; if not, write to the Free Software Foundation, +# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. +# +# ======================= END GPL LICENSE BLOCK ======================== + +# <pep8 compliant> + +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: + for co, idx, dist in tree.find_range(nodes[j].point, self.epsilon): + 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) |