Spaces:
Running on Zero
Running on Zero
| """Streaming SMPL pose emission base (generic machinery only). | |
| Each mesh-capable representation package provides a StreamPoseRecovery | |
| subclass implementing `_reset_state()` and `_step(frame) -> (root_rotmat, | |
| body_rotmat, transl)` in ITS OWN layout; this base contributes the | |
| representation-agnostic parts: | |
| - first-frame SE(2) canonicalization (frame-0 pelvis faces +Z at the XZ | |
| origin — the same convention as the offline renderers / rep2smpl), applied | |
| causally to every streamed frame; | |
| - optional SLERP/linear smoothing (smoothing_alpha < 1); | |
| - neutral-body FK to 22 joints. | |
| reset() cooperatively clears integration state (via _reset_state), the | |
| canonical anchor and the smoothing state. | |
| """ | |
| import numpy as np | |
| from .rotations import matrix_to_quaternion, quat_slerp, quaternion_to_matrix | |
| from .smpl_canonical import pelvis_yaw | |
| from .smpl_skeleton import fk_neutral_22, yaw_rotation_matrix | |
| class PoseStreamBase: | |
| def __init__(self, smoothing_alpha: float = 1.0, fps: float = 30.0): | |
| self.smoothing_alpha = float(np.clip(smoothing_alpha, 0.0, 1.0)) | |
| self.fps = fps | |
| self.reset() | |
| # -- subclass interface ------------------------------------------------- | |
| def _reset_state(self): | |
| raise NotImplementedError | |
| def _step(self, frame): | |
| """One feature frame -> (root_rotmat (3,3), body_rotmat (21,3,3), | |
| transl (3,) pelvis world position), in the representation's own | |
| (non-canonicalized) decode frame.""" | |
| raise NotImplementedError | |
| # -- generic machinery -------------------------------------------------- | |
| def reset(self): | |
| self._canon_R = None | |
| self._canon_xz = None | |
| self.prev_root_quat = None | |
| self.prev_body_quat = None | |
| self.prev_transl = None | |
| self._reset_state() | |
| def process_frame(self, frame_data: np.ndarray) -> dict: | |
| root_rotmat, body_rotmat, transl = self._step( | |
| np.asarray(frame_data, dtype=np.float64) | |
| ) | |
| # first-frame SE(2) canonicalization (same convention as offline) | |
| if self._canon_R is None: | |
| self._canon_R = yaw_rotation_matrix(np.array(-pelvis_yaw(root_rotmat))) | |
| t0 = self._canon_R @ transl | |
| self._canon_xz = np.array([t0[0], 0.0, t0[2]]) | |
| root_rotmat = self._canon_R @ root_rotmat | |
| transl = self._canon_R @ transl - self._canon_xz | |
| if self.smoothing_alpha < 1.0: | |
| root_quat = matrix_to_quaternion(root_rotmat[None])[0] | |
| body_quat = matrix_to_quaternion(body_rotmat) | |
| if self.prev_root_quat is None: | |
| self.prev_root_quat, self.prev_body_quat = root_quat.copy(), body_quat.copy() | |
| self.prev_transl = transl.copy() | |
| else: | |
| root_quat = quat_slerp(self.prev_root_quat, root_quat, self.smoothing_alpha) | |
| body_quat = quat_slerp(self.prev_body_quat, body_quat, self.smoothing_alpha) | |
| transl = self.smoothing_alpha * transl + (1.0 - self.smoothing_alpha) * self.prev_transl | |
| self.prev_root_quat, self.prev_body_quat = root_quat.copy(), body_quat.copy() | |
| self.prev_transl = transl.copy() | |
| root_rotmat = quaternion_to_matrix(root_quat[None])[0] | |
| body_rotmat = quaternion_to_matrix(body_quat) | |
| joints = fk_neutral_22(root_rotmat[None], body_rotmat[None], transl[None])[0] | |
| return { | |
| "joints": joints.astype(np.float32), | |
| "root_rotmat": root_rotmat.astype(np.float32), | |
| "body_rotmat": body_rotmat.astype(np.float32), | |
| "transl": transl.astype(np.float64), # pelvis WORLD position | |
| } | |