91 lines
2.8 KiB
Python
91 lines
2.8 KiB
Python
import dataclasses
|
|
import weakref
|
|
from typing import *
|
|
|
|
import transformations
|
|
from lxml import etree
|
|
|
|
from hippolyzer.lib.base.datatypes import Vector3, RAD_TO_DEG
|
|
from hippolyzer.lib.base.helpers import get_resource_filename
|
|
|
|
|
|
MAYBE_JOINT_REF = Optional[Callable[[], "JointNode"]]
|
|
|
|
|
|
@dataclasses.dataclass(unsafe_hash=True)
|
|
class JointNode:
|
|
name: str
|
|
parent: MAYBE_JOINT_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
|
|
|
|
@property
|
|
def matrix(self):
|
|
return transformations.compose_matrix(
|
|
scale=tuple(self.scale),
|
|
angles=tuple(self.rotation / RAD_TO_DEG),
|
|
translate=tuple(self.translation),
|
|
)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Skeleton:
|
|
joint_dict: Dict[str, JointNode]
|
|
|
|
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):
|
|
name = node.get('name')
|
|
joint = JointNode(
|
|
name=name,
|
|
parent=parent,
|
|
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
|
|
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
|
|
|
|
|
|
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:
|
|
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))
|
|
|
|
|
|
AVATAR_SKELETON = load_avatar_skeleton()
|