Spaces:
Running on Zero
Running on Zero
File size: 3,756 Bytes
9a25493 | 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 | """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
}
|