| |
| """Dump generated + GT poses for the hand-weight sweep in the shared 50-joint space, |
| so `eval_fgd_maej.py` can score FGD/MAEJ on exactly the DTW table's 300 clips. |
| |
| Joint mapping. NSLP-G's KEEP_50 is expressed in the 124-kpt UPPER layout |
| (`range(8) + range(82,124)`), but the sweep models are trained on the 128-kpt FULL |
| layout. The upper layout only drops body 9/10/12/13, so the same physical joints are |
| `range(8) + range(86,128)` in full indices -- body 0-7, then both 21-joint hands. |
| Verified equal: upper body 0-7 == full body 0-7 (all drops are >= 9), and the hand |
| blocks shift by exactly the 4 dropped body joints (82->86, 103->107). |
| |
| MIND THE AE. `eval_fgd_maej.py` applies KEEP_50 directly to its `--data-dir` store, so |
| it must be pointed at `dataset/VSL_upper` (124 kpt), NOT `dataset/VSL`. On a 128-kpt |
| store those indices silently select face-tail + the wrong hand joints and the FGD comes |
| out plausible but wrong. Same clips either way, so the AE is unaffected. |
| |
| Clip subset is `RandomState(--seed).choice(...)` sorted -- byte-identical to |
| `eval_trivis_sent_sw.py`, so FGD and DTW describe the same 300 clips. |
| """ |
| import argparse |
| import os |
|
|
| import numpy as np |
| import torch |
|
|
| import models.t2m_trans as trans |
| from dataset import dataset_vsl |
| from models.text_encoder_vi import ViTextEncoder |
| from train_t2m_trans_vsl import build_vqvae |
|
|
| KEEP_50_FULL = list(range(8)) + list(range(86, 128)) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--data-dir', default='./dataset/VSL') |
| ap.add_argument('--vq', required=True) |
| ap.add_argument('--gpt', required=True) |
| ap.add_argument('--split', default='test') |
| ap.add_argument('--n', type=int, default=300) |
| ap.add_argument('--seed', type=int, default=0, help='clip-subset seed; keep at 0') |
| ap.add_argument('--sample-seed', type=int, default=0, help='generation seed only') |
| ap.add_argument('--out', required=True) |
| ap.add_argument('--gt-out', default=None, help='also write the GT dump here') |
| ap.add_argument('--device', default='cuda') |
| args = ap.parse_args() |
|
|
| device = torch.device(args.device) |
| torch.manual_seed(args.sample_seed) |
|
|
| net, _, _ = build_vqvae(args.vq, device) |
| tck = torch.load(args.gpt, map_location='cpu') |
| g = argparse.Namespace(**tck['args']) |
| text_enc = ViTextEncoder(g.text_model, device=args.device) |
| gpt = trans.Text2Motion_Transformer( |
| num_vq=g.nb_code, embed_dim=g.embed_dim_gpt, clip_dim=text_enc.dim, |
| block_size=g.max_tokens + 1, num_layers=g.num_layers, n_head=g.n_head_gpt, |
| drop_out_rate=g.drop_out_rate, fc_rate=g.ff_rate) |
| gpt.load_state_dict(tck['trans'], strict=True) |
| gpt.eval().to(device) |
|
|
| st = dataset_vsl.VSLStore(args.data_dir, args.split) |
| NK = st.layout.n_kpts |
| if NK != 128: |
| raise SystemExit(f'expected the 128-kpt full layout, got {NK} -- ' |
| f'KEEP_50_FULL would select the wrong joints') |
| rs = np.random.RandomState(args.seed) |
| items = sorted(rs.choice(len(st.index), size=min(args.n, len(st.index)), |
| replace=False).tolist()) |
|
|
| poses, names, texts = [], [], [] |
| gt_poses, gt_names = [], [] |
| with torch.no_grad(): |
| for k in items: |
| c = st.index[k] |
| nm, txt = c['name'], c[g.text_field] |
| idx = gpt.sample(text_enc([txt]), if_categorial=True) |
| if idx is None or idx.numel() == 0: |
| print(f' [skip] {nm}: empty sample') |
| continue |
| p = net.decode_batch(idx.clamp(max=g.nb_code - 1))[0].cpu().numpy() |
| p = (p * st.std + st.mean).reshape(-1, NK, 2)[:, KEEP_50_FULL] |
| poses.append(p.astype(np.float32)); names.append(nm); texts.append(txt) |
| if args.gt_out: |
| m, _ = st.get(k) |
| gq = (m * st.std + st.mean).reshape(-1, NK, 2)[:, KEEP_50_FULL] |
| gt_poses.append(gq.astype(np.float32)); gt_names.append(nm) |
|
|
| os.makedirs(os.path.dirname(args.out) or '.', exist_ok=True) |
| np.savez(args.out, poses=np.array(poses, dtype=object), |
| names=np.array(names), texts=np.array(texts)) |
| print(f'wrote {args.out}: {len(poses)} clips, 50 joints') |
| if args.gt_out: |
| np.savez(args.gt_out, poses=np.array(gt_poses, dtype=object), |
| names=np.array(gt_names), texts=np.array(gt_names)) |
| print(f'wrote {args.gt_out}: {len(gt_poses)} GT clips') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
