51 lines
2.3 KiB
Python
51 lines
2.3 KiB
Python
"""Export Dyn-HaMR MANO joints, preserving the exact reconstruction FK convention."""
|
|
import os
|
|
os.environ.setdefault('OMP_NUM_THREADS', '4')
|
|
import sys
|
|
from pathlib import Path
|
|
import numpy as np
|
|
import torch
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(ROOT / 'third_party/Dyn-HaMR/dyn-hamr'))
|
|
from body_model.mano_wrapper import MANO
|
|
|
|
def main():
|
|
torch.set_num_threads(4)
|
|
source = ROOT / 'output/results_2047635068/world_results.npz'
|
|
out = ROOT / 'output/dex_2047635068'
|
|
out.mkdir(exist_ok=True)
|
|
data = np.load(source)
|
|
candidates = np.flatnonzero(np.all(data['is_right'] > .5, axis=1))
|
|
assert len(candidates) == 1, 'Expected exactly one consistent right-hand track'
|
|
track = int(candidates[0])
|
|
pose = data['pose_body'][track].reshape(-1, 45)
|
|
model = MANO(model_path=str(ROOT / 'third_party/Dyn-HaMR/_DATA/data/mano'),
|
|
batch_size=128, pose2rot=True)
|
|
joints = []
|
|
wrists = []
|
|
with torch.no_grad():
|
|
for start in range(0, len(pose), 128):
|
|
p = torch.from_numpy(pose[start:start+128]).float()
|
|
n = len(p)
|
|
result = model(hand_pose=p, global_orient=torch.zeros(n, 3),
|
|
transl=torch.zeros(n, 3),
|
|
betas=torch.from_numpy(data['betas'][track]).float().expand(n, -1))
|
|
j = result.joints.numpy()
|
|
joints.append(j - j[:, :1])
|
|
world = model(hand_pose=p,
|
|
global_orient=torch.from_numpy(data['root_orient'][track,start:start+n]).float(),
|
|
transl=torch.from_numpy(data['trans'][track,start:start+n]).float(),
|
|
betas=torch.from_numpy(data['betas'][track]).float().expand(n,-1))
|
|
wrists.append(world.joints[:,0].numpy())
|
|
joints = np.concatenate(joints)
|
|
assert joints.shape == (len(pose), 21, 3) and np.isfinite(joints).all()
|
|
np.savez_compressed(out / 'human_joints.npz', joints=joints, fps=30,
|
|
source_track=track, source=str(source.resolve()),
|
|
wrist_world=np.concatenate(wrists),
|
|
root_orient=data['root_orient'][track], trans=data['trans'][track])
|
|
print(f'Exported {joints.shape}, right track {track}, to {out}', flush=True)
|
|
|
|
if __name__ == '__main__':
|
|
main()
|