omegaconf req

feature_extractor.py

@@ -0,0 +1,76 @@
+# - "body_transl_delta_pelv"
+# - "body_orient_xy"
+# - "z_orient_delta"
+# - "body_pose"
+# - "body_joints_local_wo_z_rot"
+from transform3d import transform_body_pose, change_for, remove_z_rot, get_z_rot, rot_diff
+from einops import rearrange
+import torch
+
+def to_tensor(array):
+    if torch.is_tensor(array):
+        return array
+    else:
+        return torch.tensor(array)
+
+def _get_body_transl_delta_pelv(data):
+    """
+    get body pelvis tranlation delta relative to pelvis coord.frame
+    v_i = t_i - t_{i-1} relative to R_{i-1}
+    """
+    trans = to_tensor(data['trans'])
+    trans_vel = trans - trans.roll(1, 0)  # shift one right and subtract
+    pelvis_orient = transform_body_pose(to_tensor(data['rots'][..., :3]), "aa->rot")
+    trans_vel_pelv = change_for(trans_vel, pelvis_orient.roll(1, 0))
+    trans_vel_pelv[0] = 0  # zero out velocity of first frame
+    return trans_vel_pelv
+
+def _get_body_orient_xy(data):
+    """get body global orientation"""
+    # default is axis-angle representation
+    pelvis_orient = to_tensor(data['rots'][..., :3])
+    # if rot_repr == "6d":
+    # axis-angle to rotation matrix & drop last row
+    pelvis_orient_xy = remove_z_rot(pelvis_orient, in_format="aa")
+    return pelvis_orient_xy
+
+def _get_body_pose(data):
+    """get body pose"""
+    # default is axis-angle representation: Frames x (Jx3) (J=21)
+    pose = to_tensor(data['rots'][..., 3:3 + 21*3])  # drop pelvis orientation
+    pose = transform_body_pose(pose, f"aa->6d")
+    return pose
+
+def _get_body_joints_local_wo_z_rot(data):
+    """get body joint coordinates relative to the pelvis"""
+    joints = to_tensor(data['joint_positions'][:, :22, :])
+    pelvis_transl = to_tensor(joints[:, 0, :])
+    joints_glob = to_tensor(joints[:, :22, :])
+    pelvis_orient = to_tensor(data['rots'][..., :3])
+
+    pelvis_orient_z = get_z_rot(pelvis_orient, in_format="aa")
+    # pelvis_orient_z = transform_body_pose(pelvis_orient_z, "aa->rot").float()
+    # relative_joints = R.T @ (p_global - pelvis_translation)
+    rel_joints = torch.einsum('fdi,fjd->fji',
+                                pelvis_orient_z,
+                                joints_glob - pelvis_transl[:, None, :])
+
+    return rearrange(rel_joints, '... j c -> ... (j c)')
+
+def _get_z_orient_delta(data):
+    """get global body orientation delta"""
+    # default is axis-angle representation
+    pelvis_orient = to_tensor(data['rots'][..., :3])
+    pelvis_orient_z = get_z_rot(pelvis_orient, in_format="aa")
+    pelvis_orient_z = transform_body_pose(pelvis_orient_z, "rot->aa")
+    z_orient_delta = rot_diff(pelvis_orient_z, in_format="aa",
+                                    out_format='6d')
+    return z_orient_delta
+
+FEAT_GET_METHODS = {
+            "body_transl_delta_pelv": _get_body_transl_delta_pelv,
+            "body_orient_xy": _get_body_orient_xy,
+            "z_orient_delta": _get_z_orient_delta,
+            "body_pose": _get_body_pose,
+            "body_joints_local_wo_z_rot": _get_body_joints_local_wo_z_rot,
+        }