From 227fbf7a2e0d1523efacf4b96644b00b5fa4ea6a Mon Sep 17 00:00:00 2001 From: Salad Dais Date: Tue, 18 Oct 2022 19:30:23 +0000 Subject: [PATCH] Improve avatar skeleton implementation --- hippolyzer/lib/base/gltftools.py | 23 +++++---- hippolyzer/lib/base/mesh_skeleton.py | 76 ++++++++++++++++------------ tests/base/test_skeleton.py | 31 ++++++++++++ 3 files changed, 87 insertions(+), 43 deletions(-) create mode 100644 tests/base/test_skeleton.py diff --git a/hippolyzer/lib/base/gltftools.py b/hippolyzer/lib/base/gltftools.py index 9163d7a..d37b9bd 100644 --- a/hippolyzer/lib/base/gltftools.py +++ b/hippolyzer/lib/base/gltftools.py @@ -385,9 +385,12 @@ class GLTFBuilder: return buffer_view def add_joints(self, skin: SkinSegmentDict) -> JOINT_CONTEXT_DICT: - joints: JOINT_CONTEXT_DICT = {} # There may be some joints not present in the mesh that we need to add to reach the mPelvis root - required_joints = AVATAR_SKELETON.get_required_joints(skin['joint_names']) + required_joints = set() + for joint_name in skin['joint_names']: + joint_node = AVATAR_SKELETON[joint_name] + required_joints.add(joint_node) + required_joints.update(joint_node.ancestors) # If this is present, it may override the joint positions from the skeleton definition if 'alt_inverse_bind_matrix' in skin: @@ -395,12 +398,12 @@ class GLTFBuilder: else: joint_overrides = {} - for joint_name in required_joints: - joint = AVATAR_SKELETON[joint_name] + built_joints: JOINT_CONTEXT_DICT = {} + for joint in required_joints: joint_matrix = joint.matrix # Do we have a joint position override that would affect joint_matrix? - override = joint_overrides.get(joint_name) + override = joint_overrides.get(joint.name) if override: decomp = list(transformations.decompose_matrix(joint_matrix)) # We specifically only want the translation from the override! @@ -419,16 +422,16 @@ class GLTFBuilder: # TODO: populate "extras" here with the metadata the Blender collada stuff uses to store # "bind_mat" and "rest_mat" so we can go back to our original matrices when exporting # from blender to .dae! - node = self.add_node(joint_name, transform=joint_matrix) + gltf_joint = self.add_node(joint.name, transform=joint_matrix) # Store the node along with any fixups we may need to apply to the bind matrices later - joints[joint_name] = JointContext(node, orig_matrix, fixup_matrix) + built_joints[joint.name] = JointContext(gltf_joint, orig_matrix, fixup_matrix) # Add each joint to the child list of their respective parent - for joint_name, joint_ctx in joints.items(): + for joint_name, joint_ctx in built_joints.items(): if parent := AVATAR_SKELETON[joint_name].parent: - joints[parent().name].node.children.append(self.model.nodes.index(joint_ctx.node)) - return joints + built_joints[parent().name].node.children.append(self.model.nodes.index(joint_ctx.node)) + return built_joints def _fix_blender_joint(self, joint_matrix: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: """ diff --git a/hippolyzer/lib/base/mesh_skeleton.py b/hippolyzer/lib/base/mesh_skeleton.py index e2d8e57..00506af 100644 --- a/hippolyzer/lib/base/mesh_skeleton.py +++ b/hippolyzer/lib/base/mesh_skeleton.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import dataclasses import weakref from typing import * @@ -10,18 +12,23 @@ from hippolyzer.lib.base.helpers import get_resource_filename MAYBE_JOINT_REF = Optional[Callable[[], "JointNode"]] +SKELETON_REF = Optional[Callable[[], "Skeleton"]] -@dataclasses.dataclass(unsafe_hash=True) +@dataclasses.dataclass class JointNode: name: str parent: MAYBE_JOINT_REF + skeleton: SKELETON_REF translation: Vector3 pivot: Vector3 # pivot point for the joint, generally the same as translation rotation: Vector3 # Euler rotation in degrees scale: Vector3 type: str # bone or collision_volume + def __hash__(self): + return hash((self.name, self.type)) + @property def matrix(self): return transformations.compose_matrix( @@ -30,61 +37,64 @@ class JointNode: translate=tuple(self.translation), ) + @property + def index(self) -> int: + bone_idx = 0 + for node in self.skeleton().joint_dict.values(): + if node.type != "bone": + continue + if self is node: + return bone_idx + bone_idx += 1 + raise KeyError(f"{self.name!r} doesn't exist in skeleton") + + @property + def ancestors(self) -> Sequence[JointNode]: + joint_node = self + ancestors = [] + while joint_node.parent: + joint_node = joint_node.parent() + ancestors.append(joint_node) + return ancestors + -@dataclasses.dataclass class Skeleton: - joint_dict: Dict[str, JointNode] + def __init__(self, root_node: etree.ElementBase): + self.joint_dict: Dict[str, JointNode] = {} + self._parse_node_children(root_node, None) def __getitem__(self, item: str) -> JointNode: return self.joint_dict[item] - @classmethod - def _parse_node_children(cls, joint_dict: Dict[str, JointNode], node: etree.ElementBase, parent: MAYBE_JOINT_REF): + def _parse_node_children(self, node: etree.ElementBase, parent: MAYBE_JOINT_REF): name = node.get('name') joint = JointNode( name=name, parent=parent, + skeleton=weakref.ref(self), translation=_get_vec_attr(node, "pos", Vector3()), pivot=_get_vec_attr(node, "pivot", Vector3()), rotation=_get_vec_attr(node, "rot", Vector3()), scale=_get_vec_attr(node, "scale", Vector3(1, 1, 1)), type=node.tag, ) - joint_dict[name] = joint + self.joint_dict[name] = joint for child in node.iterchildren(): - cls._parse_node_children(joint_dict, child, weakref.ref(joint)) - - @classmethod - def from_xml(cls, node: etree.ElementBase): - joint_dict = {} - cls._parse_node_children(joint_dict, node, None) - return cls(joint_dict) - - def get_required_joints(self, joint_names: Collection[str]) -> Set[str]: - """Get all joints required to have a chain from all joints up to the root joint""" - required = set(joint_names) - for joint_name in joint_names: - joint_node = self.joint_dict.get(joint_name) - while joint_node: - required.add(joint_node.name) - if not joint_node.parent: - break - joint_node = joint_node.parent() - return required + self._parse_node_children(child, weakref.ref(joint)) -def load_avatar_skeleton() -> Skeleton: - skel_path = get_resource_filename("lib/base/data/avatar_skeleton.xml") - with open(skel_path, 'r') as f: - skel_root = etree.fromstring(f.read()) - return Skeleton.from_xml(skel_root.getchildren()[0]) - - -def _get_vec_attr(node, attr_name, default) -> Vector3: +def _get_vec_attr(node, attr_name: str, default: Vector3) -> Vector3: attr_val = node.get(attr_name, None) if not attr_val: return default return Vector3(*(float(x) for x in attr_val.split(" ") if x)) +def load_avatar_skeleton() -> Skeleton: + skel_path = get_resource_filename("lib/base/data/avatar_skeleton.xml") + with open(skel_path, 'r') as f: + skel_root = etree.fromstring(f.read()) + return Skeleton(skel_root.getchildren()[0]) + + AVATAR_SKELETON = load_avatar_skeleton() diff --git a/tests/base/test_skeleton.py b/tests/base/test_skeleton.py new file mode 100644 index 0000000..29e54b8 --- /dev/null +++ b/tests/base/test_skeleton.py @@ -0,0 +1,31 @@ +import unittest + +import numpy as np + +from hippolyzer.lib.base.mesh_skeleton import load_avatar_skeleton + + +class TestSkeleton(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.skeleton = load_avatar_skeleton() + + def test_get_joint(self): + node = self.skeleton["mNeck"] + self.assertEqual("mNeck", node.name) + self.assertEqual(self.skeleton, node.skeleton()) + + def test_get_joint_index(self): + self.assertEqual(7, self.skeleton["mNeck"].index) + + def test_get_joint_parent(self): + self.assertEqual("mChest", self.skeleton["mNeck"].parent().name) + + def test_get_joint_matrix(self): + expected_mat = np.array([ + [1., 0., 0., -0.01], + [0., 1., 0., 0.], + [0., 0., 1., 0.251], + [0., 0., 0., 1.] + ]) + np.testing.assert_equal(expected_mat, self.skeleton["mNeck"].matrix)