scripts-blender/addons/asset_pipeline_2/merge/transfer_data/transfer_functions.py

@@ -1,5 +1,5 @@
 import bpy
-from bpy import context
+from typing import Dict, Tuple, List
 from ..naming import get_basename, task_layer_prefix_name_get
 from ..drivers import find_drivers, copy_driver
 from ..visibility import override_obj_visability
@@ -15,8 +15,7 @@ from ... import constants
 import mathutils
 import bmesh
 import numpy as np
-import time
-
+from mathutils import Vector, kdtree
 
 ## FUNCTIONS SPECFIC TO TRANSFER DATA TYPES
 
@@ -55,35 +54,160 @@ def transfer_vertex_group(
     if target_obj == source_obj:
         return
 
-    if target_obj.vertex_groups.get(vertex_group_name):
-        target_obj.vertex_groups.remove(target_obj.vertex_groups.get(vertex_group_name))
-
     if not source_obj.vertex_groups.get(vertex_group_name):
         print(f"ERROR Vertex Group {vertex_group_name} not found in {source_obj.name}")
         return
+    vgroups = source_obj.vertex_groups
     
-    source_obj.vertex_groups.active = source_obj.vertex_groups.get(vertex_group_name)
-    # HACK without this sleep function Push will crash when transferring large amount of vertex groups
-    time.sleep(0.00000000000001)
+    kd_tree = build_kdtree(source_obj.data)
+    tgt_vg = target_obj.vertex_groups.get(vertex_group_name)
+    if tgt_vg:
+        target_obj.vertex_groups.remove(tgt_vg)
 
-    # DEBUG WHY THIS FAILS TO TRANSFER VERTEX GROUPS IN 4.0
-    with context.temp_override(
-        object=source_obj, selected_editable_objects=[target_obj, source_obj]
-    ):
-        bpy.ops.object.data_transfer(
-            data_type="VGROUP_WEIGHTS",
-            use_create=True,
-            vert_mapping='POLYINTERP_NEAREST',
-            layers_select_src="ACTIVE",
-            layers_select_dst="NAME",
-            mix_mode="REPLACE",
+    vert_influence_map = build_vert_influence_map(
+        source_obj, target_obj, kd_tree, 2
+    )
+    transfer_vertex_groups(source_obj, target_obj, vert_influence_map, vgroups)
+
+
+def precalc_and_transfer_single_group(source_obj, target_obj, vgroup_name, expand=2):
+    """Convenience function to transfer a single group. For transferring multiple groups,
+    this is very inefficient and shouldn't be used.
+
+    Instead, you should:
+    - build_kd_tree ONCE per source mesh.
+    - build_vert_influence_map and transfer_vertex_groups ONCE per object pair.
+    """
+
+    # Remove group from the target obj if it already exists.
+    tgt_vg = target_obj.vertex_groups.get(vgroup_name)
+    if tgt_vg:
+        target_obj.vertex_groups.remove(tgt_vg)
+
+    kd_tree = build_kdtree(source_obj.data)
+    vert_influence_map = build_vert_influence_map(
+        source_obj, target_obj, kd_tree, expand
+    )
+    transfer_vertex_groups(
+        source_obj,
+        target_obj,
+        vert_influence_map,
+        vgroups=[source_obj.vertex_groups[vgroup_name]],
+    )
+
+
+def build_kdtree(mesh):
+    kd = kdtree.KDTree(len(mesh.vertices))
+    for i, v in enumerate(mesh.vertices):
+        kd.insert(v.co, i)
+    kd.balance()
+    return kd
+
+
+def build_vert_influence_map(obj_from, obj_to, kd_tree, expand=2):
+    verts_of_edge = {
+        i: (e.vertices[0], e.vertices[1]) for i, e in enumerate(obj_from.data.edges)
+    }
+
+    edges_of_vert: Dict[int, List[int]] = {}
+    for edge_idx, edge in enumerate(obj_from.data.edges):
+        for vert_idx in edge.vertices:
+            if vert_idx not in edges_of_vert:
+                edges_of_vert[vert_idx] = []
+            edges_of_vert[vert_idx].append(edge_idx)
+
+    # A mapping from target vertex index to a list of source vertex indicies and
+    # their influence.
+    # This can be pre-calculated once per object pair, to minimize re-calculations
+    # of subsequent transferring of individual vertex groups.
+    vert_influence_map: List[int, List[Tuple[int, float]]] = {}
+    for i, dest_vert in enumerate(obj_to.data.vertices):
+        vert_influence_map[i] = get_source_vert_influences(
+            dest_vert, obj_from, kd_tree, expand, edges_of_vert, verts_of_edge
         )
-    if not target_obj.vertex_groups.get(vertex_group_name):
-        print(
-            f"FAILED to Transfer Vertex Group {vertex_group_name} to {target_obj.name}"
-        )
-        return
 
+    return vert_influence_map
+
+
+def get_source_vert_influences(
+    target_vert, obj_from, kd_tree, expand=2, edges_of_vert={}, verts_of_edge={}
+) -> List[Tuple[int, float]]:
+    _coord, idx, dist = get_nearest_vert(target_vert.co, kd_tree)
+    source_vert_indices = [idx]
+
+    if dist == 0:
+        # If the vertex position is a perfect match, just use that one vertex with max influence.
+        return [(idx, 1)]
+
+    for i in range(0, expand):
+        new_indices = []
+        for vert_idx in source_vert_indices:
+            for edge in edges_of_vert[vert_idx]:
+                vert_other = other_vert_of_edge(edge, vert_idx, verts_of_edge)
+                if vert_other not in source_vert_indices:
+                    new_indices.append(vert_other)
+        source_vert_indices.extend(new_indices)
+
+    distances: List[Tuple[int, float]] = []
+    distance_total = 0
+    for src_vert_idx in source_vert_indices:
+        distance = (target_vert.co - obj_from.data.vertices[src_vert_idx].co).length
+        distance_total += distance
+        distances.append((src_vert_idx, distance))
+
+    # Calculate influences such that the total of all influences adds up to 1.0,
+    # and the influence is inversely correlated with the distance.
+    parts = [1 / (dist / distance_total) for idx, dist in distances]
+    parts_sum = sum(parts)
+
+    influences = [
+        (idx, 1 if dist == 0 else part / parts_sum)
+        for part, dist in zip(parts, distances)
+    ]
+
+    return influences
+
+
+def get_nearest_vert(
+    coords: Vector, kd_tree: kdtree.KDTree
+) -> Tuple[Vector, int, float]:
+    """Return coordinate, index, and distance of nearest vert to coords in kd_tree."""
+    return kd_tree.find(coords)
+
+
+def other_vert_of_edge(
+    edge: int, vert: int, verts_of_edge: Dict[int, Tuple[int, int]]
+) -> int:
+    verts = verts_of_edge[edge]
+    assert vert in verts, f"Vert {vert} not part of edge {edge}."
+    return verts[0] if vert == verts[1] else verts[1]
+
+
+def transfer_vertex_groups(obj_from, obj_to, vert_influence_map, src_vgroups):
+    """Transfer src_vgroups in obj_from to obj_to using a pre-calculated vert_influence_map."""
+
+    for i, dest_vert in enumerate(obj_to.data.vertices):
+        source_verts = vert_influence_map[i]
+
+        # Vertex Group Name : Weight
+        vgroup_weights = {}
+
+        for src_vert_idx, influence in source_verts:
+            for group in obj_from.data.vertices[src_vert_idx].groups:
+                group_idx = group.group
+                vg = obj_from.vertex_groups[group_idx]
+                if vg.name not in src_vgroups:
+                    continue
+                if vg.name not in vgroup_weights:
+                    vgroup_weights[vg.name] = 0
+                vgroup_weights[vg.name] += vg.weight(src_vert_idx) * influence
+
+        # Assign final weights of this vertex in the vertex groups.
+        for vg_name in vgroup_weights.keys():
+            target_vg = obj_to.vertex_groups.get(vg_name)
+            if target_vg == None:
+                target_vg = obj_to.vertex_groups.new(name=vg_name)
+            target_vg.add([dest_vert.index], vgroup_weights[vg_name], 'REPLACE')
 
 # MODIFIERS
 def modifiers_clean(obj):