139 lines
6.7 KiB
Python
139 lines
6.7 KiB
Python
"""Prepare a weight-blind fixed Mixamo-41 TokenRig input from the static FBX.
|
|
|
|
No website skin weights are retained. Non-anatomical placeholder weights keep
|
|
all 41 joints from being trimmed by TokenRig. TokenRig predicts new weights.
|
|
"""
|
|
import itertools
|
|
import json
|
|
import math
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import bpy
|
|
import numpy as np
|
|
from mathutils import Matrix, Vector
|
|
from mathutils.kdtree import KDTree
|
|
|
|
|
|
def world_points(mesh):
|
|
return [mesh.matrix_world @ v.co for v in mesh.data.vertices]
|
|
|
|
|
|
def prepare(static, reference, out):
|
|
bpy.ops.wm.read_factory_settings(use_empty=True)
|
|
bpy.ops.import_scene.fbx(filepath=str(reference))
|
|
arm = next(o for o in bpy.context.scene.objects if o.type == 'ARMATURE')
|
|
assert len(arm.data.bones) == 41
|
|
refmesh = next(o for o in bpy.context.scene.objects if o.type == 'MESH')
|
|
refpoints = world_points(refmesh)
|
|
refarray = np.array(refpoints)
|
|
tree = KDTree(len(refpoints))
|
|
for i,p in enumerate(refpoints):
|
|
tree.insert(p,i)
|
|
tree.balance()
|
|
for obj in list(bpy.data.objects):
|
|
if obj != arm:
|
|
bpy.data.objects.remove(obj,do_unlink=True)
|
|
arm.animation_data_clear()
|
|
bpy.ops.import_scene.fbx(filepath=str(static))
|
|
mesh = next(o for o in bpy.context.scene.objects if o.type == 'MESH')
|
|
assert len(mesh.data.vertices) == len(refpoints)
|
|
rawpoints = world_points(mesh)
|
|
center = (refarray.min(0)+refarray.max(0))*.5
|
|
candidates=[]
|
|
for perm in itertools.permutations(range(3)):
|
|
for signs in itertools.product((-1,1),repeat=3):
|
|
m = np.zeros((3,3))
|
|
for i,j in enumerate(perm):
|
|
m[i,j] = signs[i]
|
|
if np.linalg.det(m) < .5:
|
|
continue
|
|
points = np.array(rawpoints) @ m.T
|
|
if np.max(np.abs(np.ptp(points,axis=0)-np.ptp(refarray,axis=0))) > 1e-4:
|
|
continue
|
|
shift = center-(points.min(0)+points.max(0))*.5
|
|
distances=[tree.find(Vector(p+shift))[2] for p in points]
|
|
candidates.append((max(distances),float(np.mean(distances)),m,shift))
|
|
worst,mean,rotation,shift = min(candidates,key=lambda x:x[1])
|
|
assert worst < 1e-5, f'static/reference geometry mismatch {worst}'
|
|
alignment = Matrix(rotation.tolist()).to_4x4()
|
|
alignment.translation = Vector(shift)
|
|
mesh.matrix_world = alignment @ mesh.matrix_world
|
|
for modifier in list(mesh.modifiers):
|
|
if modifier.type == 'ARMATURE':
|
|
mesh.modifiers.remove(modifier)
|
|
mesh.vertex_groups.clear()
|
|
mesh_world = mesh.matrix_world.copy()
|
|
mesh.parent = arm
|
|
mesh.matrix_world = mesh_world
|
|
for bone in arm.data.bones:
|
|
mesh.vertex_groups.new(name=bone.name)
|
|
mesh.vertex_groups['mixamorig:Hips'].add(list(range(len(mesh.data.vertices))),1.0,'REPLACE')
|
|
# TokenRig's Asset.trim_skeleton() removes unweighted joints. Give each bone
|
|
# one arbitrary vertex, without supplying any anatomical weighting prior.
|
|
for i,bone in enumerate(arm.data.bones):
|
|
if bone.name != 'mixamorig:Hips':
|
|
mesh.vertex_groups['mixamorig:Hips'].remove([i])
|
|
mesh.vertex_groups[bone.name].add([i],1.0,'REPLACE')
|
|
modifier = mesh.modifiers.new('TokenRigInputSkeleton','ARMATURE')
|
|
modifier.object = arm
|
|
arm.matrix_world = Matrix.Rotation(math.pi/2,4,'X') @ arm.matrix_world
|
|
bpy.context.view_layer.update()
|
|
out.mkdir(parents=True,exist_ok=True)
|
|
manifest = {'static_mesh_source': str(static),'skeleton_source':str(reference),
|
|
'bones':{b.name:{'parent':b.parent.name if b.parent else None,
|
|
'world_matrix':[list(row) for row in arm.matrix_world @ b.matrix_local],
|
|
'head':list(arm.matrix_world @ b.head_local)} for b in arm.data.bones},
|
|
'vertices':len(mesh.data.vertices),'triangles':sum(len(p.vertices)-2 for p in mesh.data.polygons),
|
|
'alignment_max_nearest_vertex_distance_m':worst,
|
|
'input_skin':'Non-anatomical keepalive placeholders: one arbitrary vertex per bone, remaining vertices Hips=1; all website weights removed',
|
|
'blender_up':'Z','blender_front':'-Y', 'alignment':[list(row) for row in alignment]}
|
|
(out/'input_manifest.json').write_text(json.dumps(manifest,indent=2),encoding='utf-8')
|
|
bpy.ops.export_scene.gltf(filepath=str(out/'source.glb'),export_format='GLB',export_yup=True,
|
|
export_skins=True,export_animations=False,export_def_bones=False)
|
|
print('PREPARED',len(arm.data.bones),len(mesh.data.vertices),manifest['triangles'],'alignment',worst)
|
|
return arm,mesh
|
|
|
|
|
|
def restore(static,reference,out):
|
|
arm,mesh=prepare(static,reference,out)
|
|
decoded=json.loads((out/'decoded_weights.json').read_text(encoding='utf-8'))
|
|
tree=KDTree(len(decoded))
|
|
for i,item in enumerate(decoded):
|
|
x,y,z=item['position_y_up']
|
|
tree.insert(Vector((x,-z,y)),i)
|
|
tree.balance()
|
|
mesh.vertex_groups.clear()
|
|
for bone in arm.data.bones:
|
|
mesh.vertex_groups.new(name=bone.name)
|
|
distances=[]
|
|
for v in mesh.data.vertices:
|
|
_,index,distance=tree.find(mesh.matrix_world @ v.co)
|
|
distances.append(distance)
|
|
weights=decoded[index]['weights']
|
|
total=sum(weights.values())
|
|
assert total>.9999
|
|
for name,weight in weights.items():
|
|
mesh.vertex_groups[name].add([v.index],weight/total,'REPLACE')
|
|
print('RESTORE_DISTANCE',max(distances),'TARGET_BOUNDS',np.array(world_points(mesh)).min(0),np.array(world_points(mesh)).max(0),
|
|
'DECODED_BOUNDS',np.array([item['position_y_up'] for item in decoded]).min(0),np.array([item['position_y_up'] for item in decoded]).max(0))
|
|
assert max(distances)<1e-5
|
|
(out/'restore_manifest.json').write_text(json.dumps({'vertex_nearest_max_distance_m':max(distances),
|
|
'bones':len(arm.data.bones),'vertices':len(mesh.data.vertices),
|
|
'skeleton':'Original website FBX skeleton, exact original rest matrices after common world-axis normalization',
|
|
'weights':'Only TokenRig predicted weights; no filtering, locality correction or website weight transfer',
|
|
'blender_up':'Z','blender_front':'-Y'},indent=2),encoding='utf-8')
|
|
bpy.ops.wm.save_as_mainfile(filepath=str(out/'tokenrig_mixamo41.blend'),compress=True)
|
|
bpy.ops.export_scene.gltf(filepath=str(out/'tokenrig_mixamo41.glb'),export_format='GLB',export_yup=True,
|
|
export_skins=True,export_animations=False,export_def_bones=False)
|
|
print('RESTORED',len(arm.data.bones),max(distances))
|
|
|
|
|
|
args=sys.argv[sys.argv.index('--')+1:]
|
|
if args[0]=='prepare':
|
|
prepare(*map(Path,args[1:]))
|
|
elif args[0]=='restore':
|
|
restore(*map(Path,args[1:]))
|
|
else:
|
|
raise SystemExit('use prepare|restore <static.fbx> <Mixamo.fbx> <out_dir>')
|