126 lines
6.0 KiB
Python
126 lines
6.0 KiB
Python
"""Read TokenRig GLB weights and match quantized joint heads to supplied bones.
|
|
|
|
Writes JSON only; does not rewrite binary assets. A global one-to-one assignment
|
|
avoids greedy nearest-bone matching errors around closely spaced fingers.
|
|
"""
|
|
import json
|
|
import struct
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
from scipy.optimize import linear_sum_assignment
|
|
from scipy.spatial import cKDTree
|
|
|
|
|
|
class GLB:
|
|
def __init__(self,path):
|
|
data=Path(path).read_bytes()
|
|
assert data[:4]==b'glTF'
|
|
self.bin=b''
|
|
offset=12
|
|
while offset<len(data):
|
|
size,kind=struct.unpack_from('<II',data,offset)
|
|
chunk=data[offset+8:offset+8+size]
|
|
if kind==0x4E4F534A:
|
|
self.doc=json.loads(chunk)
|
|
elif kind==0x004E4942:
|
|
self.bin=chunk
|
|
offset+=8+size
|
|
self.world={}
|
|
def visit(index,parent):
|
|
node=self.doc['nodes'][index]
|
|
if 'matrix' in node:
|
|
local=np.array(node['matrix']).reshape(4,4,order='F')
|
|
else:
|
|
x,y,z,w=node.get('rotation',[0,0,0,1])
|
|
rotation=np.array([[1-2*(y*y+z*z),2*(x*y-z*w),2*(x*z+y*w)],
|
|
[2*(x*y+z*w),1-2*(x*x+z*z),2*(y*z-x*w)],
|
|
[2*(x*z-y*w),2*(y*z+x*w),1-2*(x*x+y*y)]])
|
|
local=np.eye(4)
|
|
local[:3,:3]=rotation @ np.diag(node.get('scale',[1,1,1]))
|
|
local[:3,3]=node.get('translation',[0,0,0])
|
|
self.world[index]=parent@local
|
|
for child in node.get('children',[]):
|
|
visit(child,self.world[index])
|
|
for node in self.doc['scenes'][self.doc.get('scene',0)]['nodes']:
|
|
visit(node,np.eye(4))
|
|
|
|
def accessor(self,index):
|
|
a=self.doc['accessors'][index]
|
|
v=self.doc['bufferViews'][a['bufferView']]
|
|
dtype=np.dtype({5120:'i1',5121:'u1',5122:'<i2',5123:'<u2',5125:'<u4',5126:'<f4'}[a['componentType']])
|
|
count={'SCALAR':1,'VEC2':2,'VEC3':3,'VEC4':4,'MAT4':16}[a['type']]
|
|
stride=v.get('byteStride',dtype.itemsize*count)
|
|
result=np.ndarray((a['count'],count),dtype=dtype,buffer=self.bin,
|
|
offset=v.get('byteOffset',0)+a.get('byteOffset',0),strides=(stride,dtype.itemsize)).copy()
|
|
if a.get('normalized'):
|
|
result=result.astype(float)/np.iinfo(dtype).max
|
|
return result
|
|
|
|
def mesh(self):
|
|
index,node=next((i,n) for i,n in enumerate(self.doc['nodes']) if 'mesh' in n)
|
|
primitive=self.doc['meshes'][node['mesh']]['primitives'][0]
|
|
p=self.accessor(primitive['attributes']['POSITION'])
|
|
p4=np.column_stack((p,np.ones(len(p))))
|
|
if 'skin' in node:
|
|
# For skinned glTF nodes the mesh-node world transform cancels out
|
|
# of the skin transform. Evaluate the supplied inverse bind matrices
|
|
# rather than applying the armature object's scale a second time.
|
|
skin=self.doc['skins'][node['skin']]
|
|
inverse=self.accessor(skin['inverseBindMatrices'])
|
|
transforms=np.array([self.world[j] @ inverse[k].reshape(4,4,order='F') for k,j in enumerate(skin['joints'])])
|
|
ids=self.accessor(primitive['attributes']['JOINTS_0'])
|
|
weights=self.accessor(primitive['attributes']['WEIGHTS_0'])
|
|
p=np.zeros_like(p4,dtype=float)
|
|
for slot in range(ids.shape[1]):
|
|
p+=weights[:,slot,None]*np.einsum('nij,nj->ni',transforms[ids[:,slot]],p4)
|
|
p=p[:,:3]
|
|
else:
|
|
p=(p4 @ self.world[index].T)[:,:3]
|
|
return p,primitive,node
|
|
|
|
|
|
folder=Path(sys.argv[1])
|
|
raw_name=sys.argv[2] if len(sys.argv)>2 else 'rigged_raw.glb'
|
|
prefix=sys.argv[3] if len(sys.argv)>3 else ''
|
|
source=GLB(folder/'source.glb')
|
|
raw=GLB(folder/raw_name)
|
|
src_skin=source.doc['skins'][0]
|
|
raw_skin=raw.doc['skins'][0]
|
|
assert len(src_skin['joints'])==len(raw_skin['joints'])==41
|
|
src_heads=np.array([source.world[i][:3,3] for i in src_skin['joints']])
|
|
raw_heads=np.array([raw.world[i][:3,3] for i in raw_skin['joints']])
|
|
cost=np.linalg.norm(raw_heads[:,None,:]-src_heads[None,:,:],axis=2)
|
|
rows,cols=linear_sum_assignment(cost)
|
|
names=[source.doc['nodes'][i]['name'] for i in src_skin['joints']]
|
|
mapping={int(i):names[j] for i,j in zip(rows,cols)}
|
|
report={'joint_count':41,'max_joint_quantization_drift_m':float(cost[rows,cols].max()),
|
|
'joint_mapping':[{'output_skin_index':int(i),'output_name':raw.doc['nodes'][raw_skin['joints'][i]]['name'],
|
|
'input_name':names[j],'distance_m':float(cost[i,j]),
|
|
'row_second_best_distance_m':float(np.sort(cost[i])[1])} for i,j in zip(rows,cols)]}
|
|
assert report['max_joint_quantization_drift_m']<.015, report
|
|
positions,primitive,_=raw.mesh()
|
|
source_positions,_,_=source.mesh()
|
|
distances,_=cKDTree(positions).query(source_positions)
|
|
report['max_vertex_nearest_distance_m']=float(distances.max())
|
|
assert distances.max()<1e-5
|
|
reverse_distances,_=cKDTree(source_positions).query(positions)
|
|
report['max_reverse_vertex_nearest_distance_m']=float(reverse_distances.max())
|
|
report['source_render_vertices']=len(source_positions)
|
|
assert reverse_distances.max()<1e-5
|
|
attributes=primitive['attributes']
|
|
joints=raw.accessor(attributes['JOINTS_0'])
|
|
weights=raw.accessor(attributes['WEIGHTS_0'])
|
|
# Transfer may split render vertices at normal/UV corners. Both-direction
|
|
# geometric correspondence above is authoritative; render counts may differ.
|
|
decoded=[]
|
|
for p,ids,w in zip(positions,joints,weights):
|
|
named={mapping[int(i)]:float(value) for i,value in zip(ids,w) if value>0}
|
|
decoded.append({'position_y_up':p.tolist(),'weights':named})
|
|
report['vertices']=len(decoded)
|
|
report['weight_sum_range']=[float(weights.sum(1).min()),float(weights.sum(1).max())]
|
|
(folder/(prefix+'joint_mapping.json')).write_text(json.dumps(report,indent=2),encoding='utf-8')
|
|
(folder/(prefix+'decoded_weights.json')).write_text(json.dumps(decoded,separators=(',',':')),encoding='utf-8')
|
|
print(json.dumps({k:v for k,v in report.items() if k!='joint_mapping'},indent=2))
|