| |
| """Dump gpt_mvsl generations in NSLP-G's dump format, so both models can be scored |
| with one FGD/MAEJ/DTW protocol. |
| |
| Two alignment details, both required for the comparison to be fair: |
| |
| * joints -- gpt_mvsl emits the 124-kpt `upper` layout; NSLP-G emits 50 joints |
| (8 body + 42 hands, no face). We slice to the same 50 here. |
| KEEP_50 = upper[0:8] + upper[82:124], matching |
| 0.NSLP-G/.../modules/data/mvsl.py. |
| * clips -- NSLP-G selects with np.random.default_rng(0); eval_vsl.py uses |
| np.random.RandomState(0). Those give DIFFERENT subsets, so nothing is |
| assumed: every clip is dumped with its name and the scorer intersects |
| on names. |
| |
| Output matches their `score(..., dump_dir)`: poses (object array of [T,50,2]), |
| labels, names. |
| """ |
| import argparse |
| import os |
|
|
| import numpy as np |
| import torch |
| from tqdm import tqdm |
|
|
| 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 = list(range(8)) + list(range(82, 124)) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--data-dir', default='./dataset/MVSL') |
| ap.add_argument('--vq', default='output_vsl/vq_mvsl/net_best.pth') |
| ap.add_argument('--gpt', default='output_vsl/gpt_mvsl/net_best.pth') |
| ap.add_argument('--split', default='test') |
| ap.add_argument('--n', type=int, default=0, help='0 = all clips in the split') |
| ap.add_argument('--dump-dir', default='dumps_fgd_t2m') |
| ap.add_argument('--sampling', default='categorial', choices=['categorial', 'greedy']) |
| ap.add_argument('--device', default='cuda') |
| args = ap.parse_args() |
|
|
| device = torch.device(args.device) |
| net, targs, _ = 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 |
| ids = list(range(len(st.index))) |
| if args.n: |
| ids = ids[:args.n] |
| unit = 2 ** targs.down_t |
|
|
| gen_p, ceil_p, gt_p, labels, names = [], [], [], [], [] |
| with torch.no_grad(): |
| for i in tqdm(ids, desc='dump gpt_mvsl'): |
| c = st.index[i] |
| motion, _ = st.get(i) |
| gt = (motion * st.std + st.mean).reshape(-1, NK, 2)[:, KEEP_50] |
|
|
| mt = torch.from_numpy(motion).unsqueeze(0).to(device) |
| T = (len(motion) // unit) * unit |
| rec = net.decode_batch(net.encode(mt[:, :T]))[0].cpu().numpy() |
| rec = (rec * st.std + st.mean).reshape(-1, NK, 2)[:, KEEP_50] |
|
|
| feat = text_enc([c[g.text_field]]) |
| idx = gpt.sample(feat, if_categorial=(args.sampling == 'categorial')) |
| if idx is None or idx.numel() == 0: |
| continue |
| idx = idx.clamp(max=g.nb_code - 1) |
| gen = net.decode_batch(idx)[0].cpu().numpy() |
| gen = (gen * st.std + st.mean).reshape(-1, NK, 2)[:, KEEP_50] |
|
|
| gen_p.append(gen.astype(np.float32)) |
| ceil_p.append(rec.astype(np.float32)) |
| gt_p.append(gt.astype(np.float32)) |
| labels.append(int(c.get('label', -1))) |
| names.append(c['name']) |
|
|
| os.makedirs(args.dump_dir, exist_ok=True) |
| for tag, arr in (('t2mgpt_gen', gen_p), ('t2mgpt_ceiling', ceil_p), ('gt', gt_p)): |
| np.savez(os.path.join(args.dump_dir, f'{tag}.npz'), |
| poses=np.array(arr, dtype=object), |
| labels=np.array(labels), names=np.array(names)) |
| print(f' {tag}: {len(arr)} clips -> {args.dump_dir}/{tag}.npz') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
