Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: fix removeUnnecessaryJoints and combineSkeletons #1557

Merged
merged 4 commits into from
Dec 18, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
239 changes: 171 additions & 68 deletions packages/three-vrm/src/VRMUtils/combineSkeletons.ts
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);
ke456-png marked this conversation as resolved.
Show resolved Hide resolved
}

// 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
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]);
Comment on lines +113 to +114
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Considering cases multiple meshes share the same skin index attribute but not the same bones (which is unlikely) or boneInverses (which is more likely), we moved these lines here

}

const skeleton = new THREE.Skeleton(bones, boneInverses);
Expand Down
Loading