Download dump_t2m_for_render.py from Tri1/t2m-gpt-vsl-code: direct link, hf CLI and curl.
- Browser
- Download file 4.28 kB
-
https://huggingface.co/Tri1/t2m-gpt-vsl-code/resolve/main/dump_t2m_for_render.py
- Command line
-
hf download hf://Tri1/t2m-gpt-vsl-code/dump_t2m_for_render.py
-
curl -L -o dump_t2m_for_render.py https://huggingface.co/Tri1/t2m-gpt-vsl-code/resolve/main/dump_t2m_for_render.py
4.28 kB
| #!/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() | |