Allow updating skeleton definitions with attributes from mesh

hippolyzer/lib/base/gltftools.py

@@ -429,8 +429,8 @@ class GLTFBuilder:
 
         # Add each joint to the child list of their respective parent
         for joint_name, joint_ctx in built_joints.items():
-            if parent := AVATAR_SKELETON[joint_name].parent:
-                built_joints[parent().name].node.children.append(self.model.nodes.index(joint_ctx.node))
+            if parent_name := AVATAR_SKELETON[joint_name].parent_name:
+                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]:

hippolyzer/lib/base/mesh_skeleton.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import copy
 import dataclasses
 import weakref
 from typing import *
@@ -9,16 +10,16 @@ from lxml import etree
 
 from hippolyzer.lib.base.datatypes import Vector3, RAD_TO_DEG
 from hippolyzer.lib.base.helpers import get_resource_filename
+from hippolyzer.lib.base.mesh import MeshAsset, SkinSegmentDict, llsd_to_mat4
 
-
-MAYBE_JOINT_REF = Optional[Callable[[], "JointNode"]]
+MAYBE_JOINT_REF = Optional[str]
 SKELETON_REF = Optional[Callable[[], "Skeleton"]]
 
 
 @dataclasses.dataclass
 class JointNode:
     name: str
-    parent: MAYBE_JOINT_REF
+    parent_name: MAYBE_JOINT_REF
     skeleton: SKELETON_REF
     translation: Vector3
     pivot: Vector3  # pivot point for the joint, generally the same as translation
@@ -38,6 +39,12 @@ class JointNode:
             translate=tuple(self.translation),
         )
 
+    @property
+    def parent(self) -> Optional[JointNode]:
+        if self.parent_name:
+            return self.skeleton()[self.parent_name]
+        return None
+
     @property
     def index(self) -> int:
         bone_idx = 0
@@ -52,47 +59,56 @@ class JointNode:
     @property
     def ancestors(self) -> Sequence[JointNode]:
         joint_node = self
-        ancestors = []
-        while joint_node.parent:
-            joint_node = joint_node.parent()
+        skeleton = self.skeleton()
+        ancestors: List[JointNode] = []
+        while joint_node.parent_name:
+            joint_node = skeleton.joint_dict.get(joint_node.parent_name)
             ancestors.append(joint_node)
         return ancestors
 
     @property
     def children(self) -> Sequence[JointNode]:
-        children = []
+        children: List[JointNode] = []
         for node in self.skeleton().joint_dict.values():
-            if node.parent and node.parent() == self:
+            if node.parent_name and node.parent_name == self.name:
                 children.append(node)
         return children
 
     @property
     def descendents(self) -> Set[JointNode]:
-        descendents = set()
-        ancestors = {self}
-        last_ancestors = set()
+        descendents: Set[JointNode] = set()
+        ancestors: Set[str] = {self.name}
+        last_ancestors: Set[str] = set()
         while last_ancestors != ancestors:
-            last_ancestors = ancestors
+            last_ancestors = ancestors.copy()
             for node in self.skeleton().joint_dict.values():
-                if node.parent and node.parent() in ancestors:
-                    ancestors.add(node)
+                if node.parent_name and node.parent_name in ancestors:
+                    ancestors.add(node.name)
                     descendents.add(node)
         return descendents
 
 
 class Skeleton:
-    def __init__(self, root_node: etree.ElementBase):
+    def __init__(self, root_node: Optional[etree.ElementBase] = None):
         self.joint_dict: Dict[str, JointNode] = {}
-        self._parse_node_children(root_node, None)
+        if root_node is not None:
+            self._parse_node_children(root_node, None)
 
     def __getitem__(self, item: str) -> JointNode:
         return self.joint_dict[item]
 
-    def _parse_node_children(self, node: etree.ElementBase, parent: MAYBE_JOINT_REF):
+    def clone(self) -> Self:
+        val = copy.deepcopy(self)
+        skel_ref = weakref.ref(val)
+        for joint in val.joint_dict.values():
+            joint.skeleton = skel_ref
+        return val
+
+    def _parse_node_children(self, node: etree.ElementBase, parent_name: MAYBE_JOINT_REF):
         name = node.get('name')
         joint = JointNode(
             name=name,
-            parent=parent,
+            parent_name=parent_name,
             skeleton=weakref.ref(self),
             translation=_get_vec_attr(node, "pos", Vector3()),
             pivot=_get_vec_attr(node, "pivot", Vector3()),
@@ -103,7 +119,26 @@ class Skeleton:
         )
         self.joint_dict[name] = joint
         for child in node.iterchildren():
-            self._parse_node_children(child, weakref.ref(joint))
+            self._parse_node_children(child, joint.name)
+
+    def merge_mesh_skeleton(self, mesh: MeshAsset) -> None:
+        """Update this skeleton with a skeleton definition from a mesh asset"""
+        skin_seg: Optional[SkinSegmentDict] = mesh.segments.get('skin')
+        if not skin_seg:
+            return
+
+        for joint_name, matrix in zip(skin_seg['joint_names'], skin_seg.get('alt_inverse_bind_matrix', [])):
+            # We're only meant to use the translation component from the alt inverse bind matrix.
+            joint_decomp = transformations.decompose_matrix(llsd_to_mat4(matrix))
+            joint_node = self.joint_dict.get(joint_name)
+            if not joint_node:
+                continue
+            joint_node.translation = Vector3(*joint_decomp[3])
+
+        if pelvis_offset := skin_seg.get('pelvis_offset'):
+            # TODO: Should we even do this?
+            pelvis_node = self["mPelvis"]
+            pelvis_node.translation += Vector3(0, 0, pelvis_offset)
 
 
 def _get_vec_attr(node, attr_name: str, default: Vector3) -> Vector3: