scripts/transfer_learned_rig.py

"""Transfer a provider's learned skin to the unchanged production GLB."""

import argparse
import json
import sys
from pathlib import Path

import bpy
from mathutils import Matrix
from mathutils.kdtree import KDTree


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--source", required=True)
    parser.add_argument("--rig", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--normalization")
    args = parser.parse_args(sys.argv[sys.argv.index("--") + 1 :])
    bpy.ops.wm.read_factory_settings(use_empty=True)
    if Path(args.rig).suffix.lower() == ".fbx":
        bpy.ops.import_scene.fbx(filepath=args.rig)
    else:
        bpy.ops.import_scene.gltf(filepath=args.rig)
    rigs = [o for o in bpy.context.scene.objects if o.type == "ARMATURE"]
    if len(rigs) != 1:
        raise RuntimeError("Learned rig must contain exactly one armature")
    rig = rigs[0]
    rig.animation_data_clear()
    for bone in rig.pose.bones:
        bone.matrix_basis.identity()
    donors = [o for o in bpy.context.scene.objects if o.type == "MESH"]
    if args.normalization:
        correction = json.loads(Path(args.normalization).read_text())
        donor_coords = [o.matrix_world @ v.co for o in donors for v in o.data.vertices]
        span = max(v.z for v in donor_coords) - min(v.z for v in donor_coords)
        scale = span / correction["normalized_height"]
        matrix = Matrix(correction["inverse_blender"]) @ Matrix.Diagonal((1 / scale, 1 / scale, 1 / scale, 1))
        worlds = {o: o.matrix_world.copy() for o in [rig, *donors]}
        rig.matrix_world = matrix @ worlds[rig]
        for obj in donors:
            obj.matrix_world = matrix @ worlds[obj]
        bpy.context.view_layer.update()
    points = []
    weights = []
    for donor in donors:
        names = {g.index: g.name for g in donor.vertex_groups}
        for vertex in donor.data.vertices:
            points.append(donor.matrix_world @ vertex.co)
            influences = [
                (names[g.group], g.weight)
                for g in vertex.groups
                if names[g.group] in rig.data.bones and g.weight > 0
            ]
            weights.append(sorted(influences, key=lambda p: -p[1])[:4])
    tree = KDTree(len(points))
    for index, point in enumerate(points):
        tree.insert(point, index)
    tree.balance()
    before = set(bpy.context.scene.objects)
    bpy.ops.import_scene.gltf(filepath=args.source)
    targets = [o for o in bpy.context.scene.objects if o not in before and o.type == "MESH"]
    coords = [o.matrix_world @ v.co for o in targets for v in o.data.vertices]
    height = max(p.z for p in coords) - min(p.z for p in coords)
    max_distance = 0.0
    for obj in targets:
        obj.vertex_groups.clear()
        groups = {b.name: obj.vertex_groups.new(name=b.name) for b in rig.data.bones}
        for vertex in obj.data.vertices:
            _, index, distance = tree.find(obj.matrix_world @ vertex.co)
            max_distance = max(max_distance, distance)
            influences = weights[index]
            total = sum(weight for _, weight in influences)
            if total <= 0:
                raise RuntimeError("Provider returned unweighted vertices")
            for name, weight in influences:
                groups[name].add([vertex.index], weight / total, "REPLACE")
        world = obj.matrix_world.copy()
        obj.parent = rig
        obj.matrix_world = world
        modifier = obj.modifiers.new("Learned skin", "ARMATURE")
        modifier.object = rig
    if max_distance > max(height * 0.005, 1e-5):
        raise RuntimeError(f"Learned rig geometry differs from source: {max_distance / height:.4%} height")
    for obj in donors:
        bpy.data.objects.remove(obj, do_unlink=True)
    for action in list(bpy.data.actions):
        bpy.data.actions.remove(action)
    bpy.ops.object.select_all(action="DESELECT")
    for obj in [rig, *targets]:
        obj.select_set(True)
    bpy.ops.export_scene.gltf(
        filepath=args.output,
        export_format="GLB",
        use_selection=True,
        export_animations=False,
        export_skins=True,
        export_yup=True,
    )
    Path(args.output).with_suffix(".transfer.json").write_text(
        json.dumps(
            {
                "components": len(targets),
                "bones": len(rig.data.bones),
                "vertices": len(coords),
                "max_transfer_distance": max_distance,
                "height": height,
                "method": "learned skin nearest original rest vertex; top four normalized",
                "materials": "original meshes, UVs and material nodes retained",
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()