Skip to content

Commit

Permalink
Merge pull request #1557 from pixiv/fix-combine-meshes
Browse files Browse the repository at this point in the history 
fix: fix removeUnnecessaryJoints and combineSkeletons
  • Loading branch information
@0b5vr
0b5vr authored Dec 18, 2024
2 parents d21d2b4 + 78e3695 commit 81b6991
Show file tree
Hide file tree
Showing 2 changed files with 204 additions and 93 deletions.
239 changes: 171 additions & 68 deletions packages/three-vrm/src/VRMUtils/combineSkeletons.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
58 changes: 33 additions & 25 deletions packages/three-vrm/src/VRMUtils/removeUnnecessaryJoints.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down

0 comments on commit 81b6991

Please sign in to comment.