File size: 4,278 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 | #!/usr/bin/env python3
"""Generate T2M-GPT poses for named clips and save them in NSLP-G's 50-joint format.
Purpose: a visual comparison of the two families that is actually about quality rather
than about rendering. T2M-GPT emits 128 DWPose keypoints (including 68 face points),
NSLP-G emits 50 (8 body + 21 + 21 hand). Rendered with their own renderers the two look
different for reasons that have nothing to do with how good the signing is, so this maps
T2M-GPT's output onto exactly NSLP-G's 50 joints and writes the `poses`/`names` npz that
`0.NSLP-G/Word-level/NSLP-G/render_vsl.py` reads. Both families are then drawn by the
same code, in the same coordinate space, for the same clips.
Coordinate space: the VSL3 pack (T2M-GPT) and the VSL3_upper pack (NSLP-G) are both
per-axis frame-normalized [0,1] over the same clips, so no rescaling is needed -- only a
keypoint-index gather:
original DWPose index = upper_layout.keep[j] for j in NSLP-G's KEEP_50
"""
import argparse
import json
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
# NSLP-G's modules/data/mvsl.py: 8 body + lhand + rhand inside the 124-kpt upper layout
KEEP_50 = list(range(8)) + list(range(82, 124))
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--data-dir', default='./dataset/VSL3')
ap.add_argument('--upper-dir', default='./dataset/VSL3_upper',
help='only its layout.json is read, for the index mapping')
ap.add_argument('--vq', default='output_vsl/vq_vsl_3view/net_best.pth')
ap.add_argument('--gpt', default='output_vsl/gpt_vsl_3view/net_best.pth')
ap.add_argument('--split', default='test')
ap.add_argument('--names', nargs='+', required=True, help='clip names to generate')
ap.add_argument('--sampling', default='categorial', choices=['greedy', 'categorial'])
ap.add_argument('--seed', type=int, default=0)
ap.add_argument('--device', default='cuda')
ap.add_argument('--out', required=True, help='output .npz')
args = ap.parse_args()
device = torch.device(args.device)
torch.manual_seed(args.seed)
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
with open(os.path.join(args.upper_dir, 'layout.json')) as f:
upper_keep = json.load(f)['keep']
take = [upper_keep[j] for j in KEEP_50] # -> original DWPose indices
assert len(take) == 50
by_name = {c['name']: k for k, c in enumerate(st.index)}
poses, names, texts = [], [], []
with torch.no_grad():
for nm in args.names:
if nm not in by_name:
print(f' [skip] {nm} not in {args.split}')
continue
c = st.index[by_name[nm]]
txt = c[g.text_field]
idx = gpt.sample(text_enc([txt]), if_categorial=(args.sampling == 'categorial'))
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)[:, take] # [T,50,2]
poses.append(p.astype(np.float32))
names.append(nm)
texts.append(txt)
print(f' {nm}: T={len(p)} gloss="{txt}"')
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, frame-normalized coords')
if __name__ == '__main__':
main()
|