| import os |
| import json |
| import numpy as np |
| import torch |
| from scipy.spatial.transform import Rotation as R |
| from hmr4d.utils.smplx_utils import make_smplx |
|
|
| def test_kabsch(): |
| input_dir = "/mnt/c/Temp/SyntheticDataset" |
| seqs = ["sequence_100_biboo_birthday_speech_explosion_1.jsonl", "sequence_101_biboo_birthday_speech_explosion_2.jsonl"] |
| |
| device = "cuda" |
| model = make_smplx("supermotion").to(device).eval() |
|
|
| fix_rot = R.from_euler("z", 180, degrees=True).as_matrix() |
| C = np.diag([1.0, -1.0, 1.0]) |
| C4 = np.diag([1.0, -1.0, 1.0, 1.0]) |
| fix_mat = np.eye(4); fix_mat[:3, :3] = fix_rot |
|
|
| for seq in seqs: |
| p = os.path.join(input_dir, seq) |
| with open(p, "r") as f: lines = f.readlines() |
| |
| row = json.loads(lines[1]) |
| |
| |
| R_cam_w = R.from_quat(row["cam_rot_world"]).as_matrix() |
| R_pel_w = R.from_quat(row["pelvis_rot_world"]).as_matrix() |
| R_rel_unity = R_cam_w.T @ R_pel_w |
| R_cv = C @ R_rel_unity @ C |
| R_final = R_cv @ R.from_euler("z", 180, degrees=True).as_matrix() |
| go_c = R.from_matrix(R_final).as_rotvec().astype(np.float32) |
|
|
| pelvis_cam_cv = np.asarray(row["smpl_incam_transl"], dtype=np.float64).reshape(3) |
| root_cam_cv = np.asarray(row.get("smpl_root_incam_transl", [0.0, 0.0, 0.0]), dtype=np.float64).reshape(3) |
| pelvis_cam_cv = pelvis_cam_cv |
| target_cam_cv = pelvis_cam_cv |
|
|
| R_cam_w_cv = fix_rot @ (C @ R_cam_w @ C) |
| R_pelvis_w_cv = R_cam_w_cv @ R_final |
| go_w = R.from_matrix(R_pelvis_w_cv).as_rotvec().astype(np.float32) |
| |
| pos_cv_raw = (C @ row['pelvis_pos_world']) |
| pelvis_pos_w_cv = fix_rot @ pos_cv_raw |
|
|
| pose = np.asarray(row["smplx_pose"], dtype=np.float32) |
| body_pose = pose[3:66].astype(np.float32) |
| betas = np.zeros(10, dtype=np.float32) |
| |
| |
| with torch.no_grad(): |
| bt = torch.from_numpy(betas[None]).to(device) |
| bp = torch.from_numpy(body_pose[None]).to(device) |
| |
| go_c_t = torch.from_numpy(go_c[None]).to(device) |
| pel0_c = model(betas=bt, global_orient=go_c_t, body_pose=bp).joints[0, 0].cpu().numpy() |
| |
| go_w_t = torch.from_numpy(go_w[None]).to(device) |
| pel0_w = model(betas=bt, global_orient=go_w_t, body_pose=bp).joints[0, 0].cpu().numpy() |
|
|
| tr_c = target_cam_cv - pel0_c |
| tr_w = pelvis_pos_w_cv - pel0_w |
| |
| |
| R_c_v = R.from_rotvec(go_c).as_matrix() |
| R_w_v = R.from_rotvec(go_w).as_matrix() |
| R_w2c_v = R_c_v @ R_w_v.T |
| |
| j0_c = pel0_c + tr_c |
| j0_w = pel0_w + tr_w |
| t_w2c_derived = j0_c - R_w2c_v @ j0_w |
|
|
| cam_T_wc = np.eye(4) |
| cam_T_wc[:3, :3] = R_cam_w |
| cam_T_wc[:3, 3] = row["cam_pos_world"] |
| cam_T_wc_cv = fix_mat @ (C4 @ cam_T_wc @ C4) |
| T_w2c_exported = np.linalg.inv(cam_T_wc_cv) |
| |
| |
| T_c2w_exported = np.linalg.inv(T_w2c_exported) |
| |
| with torch.no_grad(): |
| tr_c_t = torch.from_numpy(tr_c[None]).to(device) |
| j_c = model(betas=bt, global_orient=go_c_t, body_pose=bp, transl=tr_c_t).joints[0, :22].cpu().numpy() |
| |
| tr_w_t = torch.from_numpy(tr_w[None]).to(device) |
| j_w = model(betas=bt, global_orient=go_w_t, body_pose=bp, transl=tr_w_t).joints[0, :22].cpu().numpy() |
| |
| src_mean = j_c.mean(axis=0) |
| dst_mean = j_w.mean(axis=0) |
| X = j_c - src_mean |
| Y = j_w - dst_mean |
| H = X.T @ Y |
| U, _, Vt = np.linalg.svd(H) |
| R_align = Vt.T @ U.T |
| if np.linalg.det(R_align) < 0: |
| Vt[-1, :] *= -1 |
| R_align = Vt.T @ U.T |
| t_align = dst_mean - src_mean @ R_align |
| |
| T_c2w_kabsch = np.eye(4) |
| T_c2w_kabsch[:3, :3] = R_align |
| T_c2w_kabsch[:3, 3] = t_align |
| |
| t_err = np.linalg.norm(T_c2w_kabsch[:3, 3] - T_c2w_exported[:3, 3]) |
| |
| print(f"[{seq}]") |
| print(f" t_err: {t_err:.3f}m") |
| print(f" t_c2w_kabsch: {T_c2w_kabsch[:3, 3]}") |
| print(f" T_c2w_exported: {T_c2w_exported[:3, 3]}") |
| print(f" diff(kabsch - exported): {T_c2w_kabsch[:3, 3] - T_c2w_exported[:3, 3]}") |
| print(f" t_w2c_derived: {t_w2c_derived}") |
| print(f" T_w2c_exported: {T_w2c_exported[:3, 3]}") |
| print(f" delta_T_w2c (derived - exported): {t_w2c_derived - T_w2c_exported[:3, 3]}") |
|
|
| if __name__ == "__main__": |
| test_kabsch() |
|
|
