Merge pull request #1557 from pixiv/fix-combine-meshes

0b5vr · web-flow · commit 81b699174823 · 2024-12-18T19:10:34.000+09:00
fix: fix removeUnnecessaryJoints and combineSkeletons
diff --git a/packages/three-vrm/src/VRMUtils/combineSkeletons.ts b/packages/three-vrm/src/VRMUtils/combineSkeletons.ts
@@ -9,98 +9,201 @@ import * as THREE from 'three';
  * @param root Root object that will be traversed
  */
 export function combineSkeletons(root: THREE.Object3D): void {
+  const skinnedMeshes = collectSkinnedMeshes(root);
+
+  // List all used skin indices for each skin index attribute
+  const attributeUsedIndexSetMap = new Map<THREE.BufferAttribute | THREE.InterleavedBufferAttribute, Set<number>>();
+  for (const mesh of skinnedMeshes) {
+    const geometry = mesh.geometry;
+    const skinIndexAttr = geometry.getAttribute('skinIndex');
+    const skinWeightAttr = geometry.getAttribute('skinWeight');
+    const usedIndicesSet = listUsedIndices(skinIndexAttr, skinWeightAttr);
+    attributeUsedIndexSetMap.set(skinIndexAttr, usedIndicesSet);
+  }
+
+  // List all bones and boneInverses for each meshes
+  const meshBoneInverseMapMap = new Map<THREE.SkinnedMesh, Map<THREE.Bone, THREE.Matrix4>>();
+  for (const mesh of skinnedMeshes) {
+    const boneInverseMap = listUsedBones(mesh, attributeUsedIndexSetMap);
+    meshBoneInverseMapMap.set(mesh, boneInverseMap);
+  }
+
+  // Group meshes by bone sets
+  const groups: { boneInverseMap: Map<THREE.Bone, THREE.Matrix4>; meshes: Set<THREE.SkinnedMesh> }[] = [];
+  for (const [mesh, boneInverseMap] of meshBoneInverseMapMap) {
+    let foundMergeableGroup = false;
+    for (const candidate of groups) {
+      // check if the candidate group is mergeable
+      const isMergeable = boneInverseMapIsMergeable(boneInverseMap, candidate.boneInverseMap);
+
+      // if we found a mergeable group, add the mesh to the group
+      if (isMergeable) {
+        foundMergeableGroup = true;
+        candidate.meshes.add(mesh);
+
+        // add lacking bones to the group
+        for (const [bone, boneInverse] of boneInverseMap) {
+          candidate.boneInverseMap.set(bone, boneInverse);
+        }
+
+        break;
+      }
+    }
+
+    // if we couldn't find a mergeable group, create a new group
+    if (!foundMergeableGroup) {
+      groups.push({ boneInverseMap, meshes: new Set([mesh]) });
+    }
+  }
+
+  // prepare new skeletons for each group, and bind them to the meshes
+  for (const { boneInverseMap, meshes } of groups) {
+    // create a new skeleton
+    const newBones = Array.from(boneInverseMap.keys());
+    const newBoneInverses = Array.from(boneInverseMap.values());
+    const newSkeleton = new THREE.Skeleton(newBones, newBoneInverses);
+
+    const attributeProcessedSet = new Set<THREE.BufferAttribute | THREE.InterleavedBufferAttribute>();
+
+    for (const mesh of meshes) {
+      const attribute = mesh.geometry.getAttribute('skinIndex');
+
+      if (!attributeProcessedSet.has(attribute)) {
+        // remap skin index attribute
+        remapSkinIndexAttribute(attribute, mesh.skeleton.bones, newBones);
+        attributeProcessedSet.add(attribute);
+      }
+
+      // bind the new skeleton to the mesh
+      mesh.bind(newSkeleton, new THREE.Matrix4());
+    }
+  }
+}
+
+/**
+ * Traverse an entire tree and collect skinned meshes.
+ */
+function collectSkinnedMeshes(scene: THREE.Object3D): Set<THREE.SkinnedMesh> {
   const skinnedMeshes = new Set<THREE.SkinnedMesh>();
-  const geometryToSkinnedMesh = new Map<THREE.BufferGeometry, THREE.SkinnedMesh>();
 
-  // Traverse entire tree and collect skinned meshes
-  root.traverse((obj) => {
-    if (obj.type !== 'SkinnedMesh') {
+  scene.traverse((obj) => {
+    if (!(obj as any).isSkinnedMesh) {
       return;
     }
 
     const skinnedMesh = obj as THREE.SkinnedMesh;
-
-    // Check if the geometry has already been encountered
-    const previousSkinnedMesh = geometryToSkinnedMesh.get(skinnedMesh.geometry);
-    if (previousSkinnedMesh) {
-      // Skinned meshes that share their geometry with other skinned meshes can't be processed.
-      // The skinnedMeshes already contain previousSkinnedMesh, so remove it now.
-      skinnedMeshes.delete(previousSkinnedMesh);
-    } else {
-      geometryToSkinnedMesh.set(skinnedMesh.geometry, skinnedMesh);
-      skinnedMeshes.add(skinnedMesh);
-    }
+    skinnedMeshes.add(skinnedMesh);
   });
 
-  // Prepare new skeletons for the skinned meshes
-  const newSkeletons: Array<{ bones: THREE.Bone[]; boneInverses: THREE.Matrix4[]; meshes: THREE.SkinnedMesh[] }> = [];
-  skinnedMeshes.forEach((skinnedMesh) => {
-    const skeleton = skinnedMesh.skeleton;
+  return skinnedMeshes;
+}
 
-    // Find suitable skeleton
-    let newSkeleton = newSkeletons.find((candidate) => skeletonMatches(skeleton, candidate));
-    if (!newSkeleton) {
-      newSkeleton = { bones: [], boneInverses: [], meshes: [] };
-      newSkeletons.push(newSkeleton);
+/**
+ * List all skin indices used by the given geometry.
+ * If the skin weight is 0, the index won't be considered as used.
+ * @param skinIndexAttr The skin index attribute to list used indices
+ * @param skinWeightAttr The skin weight attribute corresponding to the skin index attribute
+ */
+function listUsedIndices(
+  skinIndexAttr: THREE.BufferAttribute | THREE.InterleavedBufferAttribute,
+  skinWeightAttr: THREE.BufferAttribute | THREE.InterleavedBufferAttribute,
+): Set<number> {
+  const usedIndices = new Set<number>();
+
+  for (let i = 0; i < skinIndexAttr.count; i++) {
+    for (let j = 0; j < skinIndexAttr.itemSize; j++) {
+      const index = skinIndexAttr.getComponent(i, j);
+      const weight = skinWeightAttr.getComponent(i, j);
+
+      if (weight !== 0) {
+        usedIndices.add(index);
+      }
     }
+  }
 
-    // Add skinned mesh to the new skeleton
-    newSkeleton.meshes.push(skinnedMesh);
+  return usedIndices;
+}
 
-    // Determine bone index mapping from skeleton -> newSkeleton
-    const boneIndexMap: number[] = skeleton.bones.map((bone) => newSkeleton.bones.indexOf(bone));
+/**
+ * List all bones used by the given skinned mesh.
+ * @param mesh The skinned mesh to list used bones
+ * @param attributeUsedIndexSetMap A map from skin index attribute to the set of used skin indices
+ * @returns A map from used bone to the corresponding bone inverse matrix
+ */
+function listUsedBones(
+  mesh: THREE.SkinnedMesh,
+  attributeUsedIndexSetMap: Map<THREE.BufferAttribute | THREE.InterleavedBufferAttribute, Set<number>>,
+): Map<THREE.Bone, THREE.Matrix4> {
+  const boneInverseMap = new Map<THREE.Bone, THREE.Matrix4>();
 
-    // Update skinIndex attribute
-    const geometry = skinnedMesh.geometry;
-    const attribute = geometry.getAttribute('skinIndex');
-    const weightAttribute = geometry.getAttribute('skinWeight');
+  const skeleton = mesh.skeleton;
 
-    for (let i = 0; i < attribute.count; i++) {
-      for (let j = 0; j < attribute.itemSize; j++) {
-        // check bone weight
-        const weight = weightAttribute.getComponent(i, j);
-        if (weight === 0) {
-          continue;
-        }
+  const geometry = mesh.geometry;
+  const skinIndexAttr = geometry.getAttribute('skinIndex');
+  const usedIndicesSet = attributeUsedIndexSetMap.get(skinIndexAttr);
 
-        const index = attribute.getComponent(i, j);
+  if (!usedIndicesSet) {
+    throw new Error('Unreachable. attributeUsedIndexSetMap does not know the skin index attribute');
+  }
 
-        // new skinIndex buffer
-        if (boneIndexMap[index] === -1) {
-          boneIndexMap[index] = newSkeleton.bones.length;
-          newSkeleton.bones.push(skeleton.bones[index]);
-          newSkeleton.boneInverses.push(skeleton.boneInverses[index]);
-        }
+  for (const index of usedIndicesSet) {
+    boneInverseMap.set(skeleton.bones[index], skeleton.boneInverses[index]);
+  }
 
-        attribute.setComponent(i, j, boneIndexMap[index]);
+  return boneInverseMap;
+}
+
+/**
+ * Check if the given bone inverse map is mergeable to the candidate bone inverse map.
+ * @param toCheck The bone inverse map to check
+ * @param candidate The candidate bone inverse map
+ * @returns True if the bone inverse map is mergeable to the candidate bone inverse map
+ */
+function boneInverseMapIsMergeable(
+  toCheck: Map<THREE.Bone, THREE.Matrix4>,
+  candidate: Map<THREE.Bone, THREE.Matrix4>,
+): boolean {
+  for (const [bone, boneInverse] of toCheck.entries()) {
+    // if the bone is in the candidate group and the boneInverse is different, it's not mergeable
+    const candidateBoneInverse = candidate.get(bone);
+    if (candidateBoneInverse != null) {
+      if (!matrixEquals(boneInverse, candidateBoneInverse)) {
+        return false;
       }
     }
+  }
 
-    attribute.needsUpdate = true;
-  });
+  return true;
+}
 
-  // Bind new skeleton to the meshes
-  for (const { bones, boneInverses, meshes } of newSkeletons) {
-    const newSkeleton = new THREE.Skeleton(bones, boneInverses);
-    meshes.forEach((mesh) => mesh.bind(newSkeleton, new THREE.Matrix4()));
+function remapSkinIndexAttribute(
+  attribute: THREE.BufferAttribute | THREE.InterleavedBufferAttribute,
+  oldBones: THREE.Bone[],
+  newBones: THREE.Bone[],
+): void {
+  // a map from bone to old index
+  const boneOldIndexMap = new Map<THREE.Bone, number>();
+  for (const bone of oldBones) {
+    boneOldIndexMap.set(bone, boneOldIndexMap.size);
   }
-}
 
-/**
- * Checks if a given skeleton matches a candidate skeleton. For the skeletons to match,
- * all bones must either be in the candidate skeleton with the same boneInverse OR
- * not part of the candidate skeleton (as it can be added to it).
- * @param skeleton The skeleton to check.
- * @param candidate The candidate skeleton to match against.
- */
-function skeletonMatches(skeleton: THREE.Skeleton, candidate: { bones: THREE.Bone[]; boneInverses: THREE.Matrix4[] }) {
-  return skeleton.bones.every((bone, index) => {
-    const candidateIndex = candidate.bones.indexOf(bone);
-    if (candidateIndex !== -1) {
-      return matrixEquals(skeleton.boneInverses[index], candidate.boneInverses[candidateIndex]);
+  // a map from old skin index to new skin index
+  const oldToNew = new Map<number, number>();
+  for (const [i, bone] of newBones.entries()) {
+    const oldIndex = boneOldIndexMap.get(bone)!;
+    oldToNew.set(oldIndex, i);
+  }
+
+  // replace the skin index attribute with new indices
+  for (let i = 0; i < attribute.count; i++) {
+    for (let j = 0; j < attribute.itemSize; j++) {
+      const oldIndex = attribute.getComponent(i, j);
+      const newIndex = oldToNew.get(oldIndex)!;
+      attribute.setComponent(i, j, newIndex);
     }
-    return true;
-  });
+  }
+
+  attribute.needsUpdate = true;
 }
 
 // https://github.com/mrdoob/three.js/blob/r170/test/unit/src/math/Matrix4.tests.js#L12
diff --git a/packages/three-vrm/src/VRMUtils/removeUnnecessaryJoints.ts b/packages/three-vrm/src/VRMUtils/removeUnnecessaryJoints.ts
@@ -46,64 +46,72 @@ export function removeUnnecessaryJoints(
     skinnedMeshes.push(obj as THREE.SkinnedMesh);
   });
 
-  // A map from meshes to bones and boneInverses
+  // A map from meshes to new-to-old bone index map
   // some meshes might share a same skinIndex attribute, and this map also prevents to convert the attribute twice
-  const bonesList: Map<
-    THREE.SkinnedMesh,
-    {
-      bones: THREE.Bone[];
-      boneInverses: THREE.Matrix4[];
-    }
+  const attributeToBoneIndexMapMap: Map<
+    THREE.BufferAttribute | THREE.InterleavedBufferAttribute,
+    Map<number, number>
   > = new Map();
 
   // A maximum number of bones
   let maxBones = 0;
 
-  // Iterate over all skinned meshes and collect bones and boneInverses
+  // Iterate over all skinned meshes and remap bones for each skin index attribute
   for (const mesh of skinnedMeshes) {
     const geometry = mesh.geometry;
     const attribute = geometry.getAttribute('skinIndex');
 
-    const bones: THREE.Bone[] = []; // new list of bone
-    const boneInverses: THREE.Matrix4[] = []; // new list of boneInverse
-    const boneIndexMap: { [index: number]: number } = {}; // map of old bone index vs. new bone index
+    if (attributeToBoneIndexMapMap.has(attribute)) {
+      continue;
+    }
+
+    const oldToNew = new Map<number, number>(); // map of old bone index vs. new bone index
+    const newToOld = new Map<number, number>(); // map of new bone index vs. old bone index
 
     // create a new bone map
     for (let i = 0; i < attribute.count; i++) {
       for (let j = 0; j < attribute.itemSize; j++) {
-        const index = attribute.getComponent(i, j);
+        const oldIndex = attribute.getComponent(i, j);
+        let newIndex = oldToNew.get(oldIndex);
 
         // new skinIndex buffer
-        if (boneIndexMap[index] == null) {
-          boneIndexMap[index] = bones.length;
-          bones.push(mesh.skeleton.bones[index]);
-          boneInverses.push(mesh.skeleton.boneInverses[index]);
+        if (newIndex == null) {
+          newIndex = oldToNew.size;
+          oldToNew.set(oldIndex, newIndex);
+          newToOld.set(newIndex, oldIndex);
         }
 
-        attribute.setComponent(i, j, boneIndexMap[index]);
+        attribute.setComponent(i, j, newIndex);
       }
     }
 
     // replace with new indices
     attribute.needsUpdate = true;
 
     // update boneList
-    bonesList.set(mesh, { bones, boneInverses });
+    attributeToBoneIndexMapMap.set(attribute, newToOld);
 
     // update max bones count
-    maxBones = Math.max(maxBones, bones.length);
+    maxBones = Math.max(maxBones, oldToNew.size);
   }
 
   // Let's actually set the skeletons
   for (const mesh of skinnedMeshes) {
-    const { bones, boneInverses } = bonesList.get(mesh)!;
+    const geometry = mesh.geometry;
+    const attribute = geometry.getAttribute('skinIndex');
+    const newToOld = attributeToBoneIndexMapMap.get(attribute)!;
+
+    const bones: THREE.Bone[] = [];
+    const boneInverses: THREE.Matrix4[] = [];
 
     // if `experimentalSameBoneCounts` is `true`, compensate skeletons with dummy bones to keep the bone count same between skeletons
-    if (experimentalSameBoneCounts) {
-      for (let i = bones.length; i < maxBones; i++) {
-        bones[i] = bones[0];
-        boneInverses[i] = boneInverses[0];
-      }
+    const nBones = experimentalSameBoneCounts ? maxBones : newToOld.size;
+
+    for (let newIndex = 0; newIndex < nBones; newIndex++) {
+      const oldIndex = newToOld.get(newIndex) ?? 0;
+
+      bones.push(mesh.skeleton.bones[oldIndex]);
+      boneInverses.push(mesh.skeleton.boneInverses[oldIndex]);
     }
 
     const skeleton = new THREE.Skeleton(bones, boneInverses);
