Spaces:
Running on Zero
Running on Zero
File size: 6,158 Bytes
8f1bcd9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | """
skeleton.py
Pure-Python armature / pose-bone system.
Design matches Blender's pose-mode semantics:
- bone.rest_matrix_local = 4×4 rest pose in parent space (edit-mode)
- bone.pose_rotation_quat = local rotation DELTA from rest (≡ bone.rotation_quaternion)
- bone.pose_location = local translation DELTA from rest (≡ bone.location)
- bone.pose_scale = local scale (≡ bone.scale)
- bone.matrix_armature = FK-computed 4×4 in armature space (≡ bone.matrix in pose mode)
Armature.world_matrix corresponds to arm.matrix_world.
"""
from __future__ import annotations
import numpy as np
from typing import Dict, List, Optional, Tuple
from .math3d import (
quat_identity, quat_normalize, quat_mul,
quat_to_matrix4, matrix4_to_quat,
translation_matrix, scale_matrix, trs_to_matrix4, matrix4_to_trs,
vec3,
)
class PoseBone:
def __init__(
self,
name: str,
rest_matrix_local: np.ndarray, # 4×4, in parent local space
parent: Optional["PoseBone"] = None,
):
self.name = name
self.parent: Optional[PoseBone] = parent
self.children: List[PoseBone] = []
self.rest_matrix_local: np.ndarray = rest_matrix_local.copy()
# Pose state — start at rest (delta = identity)
self.pose_rotation_quat: np.ndarray = quat_identity()
self.pose_location: np.ndarray = vec3()
self.pose_scale: np.ndarray = np.ones(3)
# Cached FK result — call armature.update_fk() to refresh
self._matrix_armature: np.ndarray = np.eye(4)
# -----------------------------------------------------------------------
# Properties
# -----------------------------------------------------------------------
@property
def matrix_armature(self) -> np.ndarray:
"""4×4 FK result in armature space. Refresh with armature.update_fk()."""
return self._matrix_armature
@property
def head(self) -> np.ndarray:
"""Bone head position in armature space."""
return self._matrix_armature[:3, 3].copy()
@property
def tail(self) -> np.ndarray:
"""
Approximate tail position (Y-axis in bone space, length 1).
Works for Y-along-bone convention (Blender / BVH default).
"""
y_axis = self._matrix_armature[:3, :3] @ np.array([0.0, 1.0, 0.0])
return self._matrix_armature[:3, 3] + y_axis
# -----------------------------------------------------------------------
# FK
# -----------------------------------------------------------------------
def _compute_local_matrix(self) -> np.ndarray:
"""rest_local @ T(pose_loc) @ R(pose_rot) @ S(pose_scale)."""
T = translation_matrix(self.pose_location)
R = quat_to_matrix4(self.pose_rotation_quat)
S = scale_matrix(self.pose_scale)
return self.rest_matrix_local @ T @ R @ S
def _fk(self, parent_matrix: np.ndarray) -> None:
self._matrix_armature = parent_matrix @ self._compute_local_matrix()
for child in self.children:
child._fk(self._matrix_armature)
# -----------------------------------------------------------------------
# Parent-chain helpers (Blender: bone.parent_recursive)
# -----------------------------------------------------------------------
@property
def parent_recursive(self) -> List["PoseBone"]:
chain: List[PoseBone] = []
cur = self.parent
while cur is not None:
chain.append(cur)
cur = cur.parent
return chain
class Armature:
"""
Collection of PoseBones with a world transform.
Corresponds to a Blender armature object.
"""
def __init__(self, name: str = "Armature"):
self.name = name
self.world_matrix: np.ndarray = np.eye(4) # arm.matrix_world
self._bones: Dict[str, PoseBone] = {}
self._roots: List[PoseBone] = []
# -----------------------------------------------------------------------
# Construction helpers
# -----------------------------------------------------------------------
def add_bone(self, bone: PoseBone, parent_name: Optional[str] = None) -> PoseBone:
self._bones[bone.name] = bone
if parent_name and parent_name in self._bones:
parent = self._bones[parent_name]
bone.parent = parent
parent.children.append(bone)
elif bone.parent is None:
self._roots.append(bone)
return bone
@property
def pose_bones(self) -> Dict[str, PoseBone]:
return self._bones
def get_bone(self, name: str) -> PoseBone:
if name not in self._bones:
raise KeyError(f"Bone '{name}' not found in armature '{self.name}'")
return self._bones[name]
def has_bone(self, name: str) -> bool:
return name in self._bones
# -----------------------------------------------------------------------
# FK update
# -----------------------------------------------------------------------
def update_fk(self) -> None:
"""Recompute all bone armature-space matrices via FK."""
for root in self._roots:
root._fk(np.eye(4))
# -----------------------------------------------------------------------
# Snapshot / restore (for calc-correction passes)
# -----------------------------------------------------------------------
def snapshot(self) -> Dict[str, Tuple[np.ndarray, np.ndarray, np.ndarray]]:
return {
name: (
bone.pose_rotation_quat.copy(),
bone.pose_location.copy(),
bone.pose_scale.copy(),
)
for name, bone in self._bones.items()
}
def restore(self, snap: Dict[str, Tuple[np.ndarray, np.ndarray, np.ndarray]]) -> None:
for name, (r, t, s) in snap.items():
if name in self._bones:
self._bones[name].pose_rotation_quat = r.copy()
self._bones[name].pose_location = t.copy()
self._bones[name].pose_scale = s.copy()
self.update_fk()
|