File size: 4,637 Bytes
7e120dd | 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 | 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]) # frame 0
# parse unity
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 # + np.array([0.0, -0.02, 0.0], dtype=np.float64)
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)
# compute pel0
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
# math check
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)
# Kabsch verify manually
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()
|