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);
    },
  };
}