File size: 4,118 Bytes
8e5456b | 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 | #!/usr/bin/env python3
"""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)) # upper-layout indices
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()
|