File size: 9,927 Bytes
6adbbb1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | import { Bone, Group, Quaternion, Skeleton, Vector3 } from 'three';
// Three.js integration; the model and dense inference runtime are engine-independent.
export function validateRig(rig) {
const n = rig.names?.length;
if (!Number.isInteger(n) || n < 1 || new Set(rig.names).size !== n ||
rig.parents?.length !== n || rig.offsets?.length !== n || rig.end_sites?.length !== n) {
throw new Error('Rig requires unique names and equally sized hierarchy, offset and end-site arrays.');
}
if (rig.units !== 'metres' || rig.up_axis !== 'Y' || rig.handedness !== 'right') {
throw new Error('Convert the rig to right-handed Y-up metres before binding.');
}
for (let i = 0; i < n; i++) {
const p = rig.parents[i];
if (!Number.isInteger(p) || (i === 0 ? p !== -1 : p < 0 || p >= i)) {
throw new Error(`Invalid parent for joint ${rig.names[i]}.`);
}
if (typeof rig.end_sites[i] !== 'boolean' || rig.offsets[i]?.length !== 3 ||
!rig.offsets[i].every(Number.isFinite)) throw new Error(`Invalid joint ${rig.names[i]}.`);
}
return rig;
}
export function createSourceRig(rig) {
validateRig(rig);
const group = new Group();
group.name = 'SuperHQMotion';
const bones = rig.names.map((name, i) => {
const bone = new Bone();
bone.name = name;
bone.position.fromArray(rig.offsets[i]);
return bone;
});
bones.forEach((bone, i) => (rig.parents[i] < 0 ? group : bones[rig.parents[i]]).add(bone));
group.updateMatrixWorld(true);
const skeleton = new Skeleton(bones);
return { group, skeleton, adapter: bindSkeleton(skeleton, rig) };
}
/** Apply local rotations to a compatible Three.js Skeleton. Retargeting is a separate operation. */
export function bindSkeleton(skeleton, rig, nameMap = {}) {
validateRig(rig);
if (!(skeleton instanceof Skeleton)) throw new Error('Expected a Three.js Skeleton.');
const byName = new Map();
for (const bone of skeleton.bones) {
if (byName.has(bone.name)) throw new Error(`Duplicate target bone name: ${bone.name}`);
byName.set(bone.name, bone);
}
const mapped = rig.names.map((name, i) => {
const bone = byName.get(nameMap[name] ?? name);
if (!bone && !rig.end_sites[i]) throw new Error(`Missing articulated joint: ${name}`);
return bone;
});
if (new Set(mapped.filter(Boolean)).size !== mapped.filter(Boolean).length) {
throw new Error('Multiple source joints map to the same target bone.');
}
mapped.forEach((bone, i) => {
if (!bone) return;
if (rig.parents[i] >= 0 && bone.parent !== mapped[rig.parents[i]]) {
throw new Error(`Hierarchy mismatch: ${rig.names[i]}`);
}
// This initial adapter requires the canonical bind pose. It does not silently retarget proportions.
if (bone.position.distanceTo(new Vector3(...rig.offsets[i])) > 1e-4) {
throw new Error(`Rest offset mismatch: ${rig.names[i]}. Retarget before binding.`);
}
});
const root = mapped[0];
const work = new Quaternion();
return {
skeleton,
rig,
applyPose(pose) {
if (pose.rootPosition?.length !== 3 || !pose.rootPosition.every(Number.isFinite) ||
pose.localQuaternions?.length !== rig.names.length) throw new Error('Invalid pose dimensions.');
// Validate the whole pose before mutating any bone.
for (const q of pose.localQuaternions) {
if (q?.length !== 4 || !q.every(Number.isFinite) || Math.hypot(...q) < 1e-8) {
throw new Error('Nonfinite or degenerate pose rotation.');
}
}
root.position.fromArray(pose.rootPosition);
mapped.forEach((bone, i) => {
if (bone) bone.quaternion.copy(work.fromArray(pose.localQuaternions[i]).normalize());
});
root.updateWorldMatrix(true, true);
},
};
}
/** Playback of saved source or prediction sequences. This does not run inference. */
export function samplePoseSequence(clip, seconds) {
if (!['recorded_training_motion', 'saved_model_rollout'].includes(clip.kind) || clip.partition !== 'train' ||
!Number.isFinite(seconds) || seconds < 0 || !(clip.dt > 0) || clip.root_positions.length < 1) {
throw new Error('Expected a training pose sequence and a finite nonnegative time.');
}
const t = Math.min(seconds / clip.dt, clip.root_positions.length - 1);
const a = Math.floor(t), b = Math.min(a + 1, clip.root_positions.length - 1), alpha = t - a;
const rootPosition = new Vector3(...clip.root_positions[a]).lerp(new Vector3(...clip.root_positions[b]), alpha).toArray();
const localQuaternions = clip.local_quaternions[a].map((q, i) => new Quaternion(...q)
.slerp(new Quaternion(...clip.local_quaternions[b][i]), alpha).toArray());
return { rootPosition, localQuaternions };
}
export function sampleRecordedMotion(clip, seconds) {
if (clip.kind !== 'recorded_training_motion') throw new Error('Expected recorded source motion.');
return samplePoseSequence(clip, seconds);
}
/** Rest-aligned world-rotation retargeting. This does not solve foot contacts. */
export function createRetargeter(skeleton, rig, {nameMap, directionMap, restCorrectionMap, scale=1, rotation=[0,0,0,1], offset=[0,0,0]}={}) {
validateRig(rig);
const scaleVector=typeof scale==='number'?[scale,scale,scale]:scale;
if(!(skeleton instanceof Skeleton)||!nameMap||scaleVector?.length!==3||!scaleVector.every(v=>Number.isFinite(v)&&v>0)||
rotation?.length!==4||!rotation.every(Number.isFinite)||Math.hypot(...rotation)<1e-8||
offset?.length!==3||!offset.every(Number.isFinite))throw Error('Invalid retargeting configuration.');
const byName=new Map();
for(const bone of skeleton.bones){if(byName.has(bone.name))throw Error('Duplicate target bone name.');byName.set(bone.name,bone);}
const alignment=new Quaternion(...rotation).normalize(),inverseAlignment=alignment.clone().invert();
const entries=Object.entries(nameMap).map(([source,target])=>{
const index=rig.names.indexOf(source),bone=byName.get(target);
if(index<0||rig.end_sites[index]||!bone)throw Error(`Invalid retarget mapping: ${source} -> ${target}`);
bone.updateWorldMatrix(true,false);
return {index,bone,rest:bone.getWorldQuaternion(new Quaternion())};
});
if(!entries.some(e=>e.index===0)||new Set(entries.map(e=>e.bone)).size!==entries.length)throw Error('Map the source root and use unique target bones.');
function depth(bone){let n=0;for(let p=bone.parent;p;p=p.parent)n++;return n;}
entries.sort((a,b)=>depth(a.bone)-depth(b.bone));
const root=entries.find(e=>e.index===0).bone;
if(entries.some(e=>e.bone!==root&&!isBelow(e.bone,root)))throw Error('Mapped target bones must descend from the mapped root.');
function isBelow(bone,parent){for(let p=bone.parent;p;p=p.parent)if(p===parent)return true;return false;}
if(directionMap||restCorrectionMap){
const sourceRest=[];
rig.offsets.forEach((v,i)=>sourceRest.push(new Vector3(...v).add(rig.parents[i]<0?new Vector3():sourceRest[rig.parents[i]])));
const corrections=new Map();
for(const entry of [...entries].sort((a,b)=>a.index-b.index)){
const spec=directionMap?.[rig.names[entry.index]],explicit=restCorrectionMap?.[rig.names[entry.index]];let correction;
if(explicit){
if(explicit.length!==4||!explicit.every(Number.isFinite)||Math.hypot(...explicit)<1e-8)throw Error('Invalid rest correction quaternion.');
correction=new Quaternion(...explicit).normalize();
}else if(spec){
const child=rig.names.indexOf(spec.sourceChild),targetChild=byName.get(spec.targetChild);
if(child<0||!targetChild||targetChild.parent!==entry.bone)throw Error('Direction calibration requires a source descendant and direct target child.');
let ancestor=rig.parents[child];while(ancestor>=0&&ancestor!==entry.index)ancestor=rig.parents[ancestor];
if(ancestor!==entry.index)throw Error('Source calibration child must descend from its mapped bone.');
const from=targetChild.getWorldPosition(new Vector3()).sub(entry.bone.getWorldPosition(new Vector3()));
const to=sourceRest[child].clone().sub(sourceRest[entry.index]).applyQuaternion(alignment);
if(from.length()<1e-8||to.length()<1e-8)throw Error('Degenerate calibration direction.');
correction=new Quaternion().setFromUnitVectors(from.normalize(),to.normalize());
}else{
let ancestor=rig.parents[entry.index];while(ancestor>=0&&!corrections.has(ancestor))ancestor=rig.parents[ancestor];
correction=(corrections.get(ancestor)??new Quaternion()).clone();
}
corrections.set(entry.index,correction);entry.rest.premultiply(correction);
}
}
return {skeleton,rig,
applyPose(pose) {
if(pose.rootPosition?.length!==3||!pose.rootPosition.every(Number.isFinite)||pose.localQuaternions?.length!==rig.names.length||
pose.localQuaternions.some(q=>q?.length!==4||!q.every(Number.isFinite)||Math.hypot(...q)<1e-8))throw Error('Invalid source pose.');
const sourceWorld=[];
pose.localQuaternions.forEach((q,i)=>{
const local=new Quaternion(...q).normalize();
sourceWorld.push(rig.parents[i]<0?local:local.premultiply(sourceWorld[rig.parents[i]]));
});
const rootWorld=new Vector3(...pose.rootPosition).multiply(new Vector3(...scaleVector)).applyQuaternion(alignment).add(new Vector3(...offset));
if(root.parent){root.parent.updateWorldMatrix(true,false);root.position.copy(root.parent.worldToLocal(rootWorld));}
else root.position.copy(rootWorld);
for(const {index,bone,rest} of entries){
const world=alignment.clone().multiply(sourceWorld[index]).multiply(inverseAlignment).multiply(rest);
if(bone.parent){bone.parent.updateWorldMatrix(true,false);world.premultiply(bone.parent.getWorldQuaternion(new Quaternion()).invert());}
bone.quaternion.copy(world).normalize();bone.updateWorldMatrix(false,false);
}
root.updateWorldMatrix(true,true);
},
};
}
|