| |
| """Dump TriVis poses for the T2M-GPT family in NSLP-G's dump format. |
| |
| Produces, on NSLP-G's exact 200 test clip ids and their 50-joint layout: |
| |
| gt ground truth (the shared reference for every scorer) |
| t2m_ceiling stage-1 encode->decode of GT |
| t2m_gloss gloss -> pose |
| t2m_predgloss BARTpho-predicted gloss -> pose |
| t2m_sentence raw Vietnamese sentence -> pose |
| t2m_shuffled another clip's gloss -> pose (control) |
| composed_gloss Multi-VSL WORD model, one sign per gloss item, concatenated |
| (recipe C: train on Multi-VSL only + a sentence->gloss stage) |
| composed_predgloss same, from BARTpho's predicted gloss |
| |
| The composed rows convert the word model's shoulder-width output into TriVis frame |
| coordinates using the reference clip's own median neck and shoulder width, which hands |
| them ground-truth global placement -- deliberately generous, since global position is |
| not what is being tested. |
| """ |
| import argparse |
| import csv |
| import json |
| import os |
| import unicodedata |
|
|
| 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 |
|
|
| REPO = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..') |
| KEEP_50_UPPER = list(range(8)) + list(range(82, 124)) |
| KEEP_50_MVSL = list(range(8)) + list(range(82, 124)) |
| NECK, RSHO, LSHO = 1, 2, 5 |
|
|
|
|
| def load_gpt(vq, gpt_ckpt, device): |
| net, targs, _ = build_vqvae(vq, device) |
| tck = torch.load(gpt_ckpt, map_location='cpu') |
| g = argparse.Namespace(**tck['args']) |
| te = ViTextEncoder(g.text_model, device=str(device)) |
| m = trans.Text2Motion_Transformer( |
| num_vq=g.nb_code, embed_dim=g.embed_dim_gpt, clip_dim=te.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) |
| m.load_state_dict(tck['trans'], strict=True) |
| m.eval().to(device) |
| return net, targs, m, te, g |
|
|
|
|
| def anchor_of(xy50, valid50, W, H): |
| px = xy50 * np.array([W, H], np.float32) |
| ok = (valid50[:, NECK] > 0) & (valid50[:, RSHO] > 0) & (valid50[:, LSHO] > 0) |
| if ok.sum() < 3: |
| ok = np.ones(len(xy50), bool) |
| neck = np.median(px[ok, NECK, :], axis=0) |
| sw = float(np.median(np.linalg.norm(px[ok, RSHO, :] - px[ok, LSHO, :], axis=-1))) |
| return neck, max(sw, 1e-3) |
|
|
|
|
| def norm_txt(s): |
| return ' '.join(unicodedata.normalize('NFC', str(s)).lower().split()) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--trivis-dir', default='./dataset/VSL_upper') |
| ap.add_argument('--mvsl-dir', default='./dataset/MVSL') |
| ap.add_argument('--tri-vq', default='output_vsl/vq_vsl_upper/net_best.pth') |
| ap.add_argument('--tri-gpt', default='output_vsl/gpt_vsl_upper/net_best.pth') |
| ap.add_argument('--mvsl-vq', default='output_vsl/vq_mvsl/net_best.pth') |
| ap.add_argument('--mvsl-gpt', default='output_vsl/gpt_mvsl/net_best.pth') |
| ap.add_argument('--clipid-json', default=os.path.join( |
| REPO, '0.NSLP-G/sentence-level/results_sent_scratch.json')) |
| ap.add_argument('--csv', default=os.path.join(REPO, 'Full_TriVis/split_lab_front.csv')) |
| ap.add_argument('--pred-gloss', default='output_vsl/text2gloss/pred_test.json') |
| ap.add_argument('--frame-w', type=float, default=1176.0) |
| ap.add_argument('--frame-h', type=float, default=1288.0) |
| ap.add_argument('--out-dir', default='dumps_trivis') |
| ap.add_argument('--skip-composed', action='store_true', |
| help='omit the Multi-VSL-composed rows (recipe C); they cost ~8 extra ' |
| 'word generations per clip and are model-independent here') |
| ap.add_argument('--device', default='cuda') |
| args = ap.parse_args() |
|
|
| W, H = args.frame_w, args.frame_h |
| device = torch.device(args.device) |
| clip_ids = json.load(open(args.clipid_json))['clip_ids'] |
| print(f'{len(clip_ids)} clip ids') |
|
|
| tri_net, tri_targs, tri_gpt, te, tg = load_gpt(args.tri_vq, args.tri_gpt, device) |
| mv_net, mv_targs, mv_gpt, _, mg = load_gpt(args.mvsl_vq, args.mvsl_gpt, device) |
| st = dataset_vsl.VSLStore(args.trivis_dir, 'test') |
| mv = dataset_vsl.VSLStore(args.mvsl_dir, 'test') |
| NK = st.layout.n_kpts |
|
|
| pred = {} |
| if os.path.exists(args.pred_gloss): |
| pred = json.load(open(args.pred_gloss, encoding='utf-8')) |
| signs_of = {} |
| with open(args.csv, newline='', encoding='utf-8') as f: |
| for r in csv.DictReader(f): |
| nm = os.path.splitext(os.path.basename(r['npz_path']))[0] |
| signs_of[nm] = [t.strip() for t in str(r['Sign_sentence']).split('|') if t.strip()] |
|
|
| tags = ['gt', 't2m_ceiling', 't2m_gloss', 't2m_predgloss', 't2m_sentence', |
| 't2m_shuffled'] |
| if not args.skip_composed: |
| tags += ['composed_gloss', 'composed_predgloss'] |
| out = {t: [] for t in tags} |
| names = [] |
| unit = 2 ** tri_targs.down_t |
| mv_unit = 2 ** mv_targs.down_t |
|
|
| def gen(model, netv, txt, store, keep): |
| idx = model.sample(te([txt]), if_categorial=True) |
| if idx is None or idx.numel() == 0: |
| return None |
| idx = idx.clamp(max=(mg.nb_code if model is mv_gpt else tg.nb_code) - 1) |
| p = netv.decode_batch(idx)[0].cpu().numpy() |
| return (p * store.std + store.mean).reshape(len(p), -1, 2)[:, keep] |
|
|
| with torch.no_grad(): |
| for i in tqdm(clip_ids, desc='dump TriVis'): |
| c = st.index[i] |
| motion, mask = st.get(i) |
| v_full = mask[:, ::2] |
| gt_full = (motion * st.std + st.mean).reshape(-1, NK, 2) |
| gt = gt_full[:, KEEP_50_UPPER] |
| v50 = v_full[:, KEEP_50_UPPER] |
| neck, sw = anchor_of(gt, v50, W, H) |
|
|
| mt = torch.from_numpy(motion).unsqueeze(0).to(device) |
| T = (len(motion) // unit) * unit |
| rec = tri_net.decode_batch(tri_net.encode(mt[:, :T]))[0].cpu().numpy() |
| rec = (rec * st.std + st.mean).reshape(-1, NK, 2)[:, KEEP_50_UPPER] |
|
|
| other = st.index[(i + len(st.index) // 2) % len(st.index)] |
| pg = pred.get(c['name'], {}) |
| texts = {'t2m_gloss': c['gloss'], |
| 't2m_predgloss': pg.get('pred_gloss', c['gloss']), |
| 't2m_sentence': c['sentence'], |
| 't2m_shuffled': other['gloss']} |
| row = {'gt': gt, 't2m_ceiling': rec} |
| ok = True |
| for tag, txt in texts.items(): |
| p = gen(tri_gpt, tri_net, txt, st, KEEP_50_UPPER) |
| if p is None: |
| ok = False; break |
| row[tag] = p |
| if not ok: |
| continue |
|
|
| |
| for tag, src in () if args.skip_composed else (('composed_gloss', signs_of.get(c['name'], [])), |
| ('composed_predgloss', |
| [t.strip() for t in str(pg.get('pred_raw', '')).split('|') |
| if t.strip()] or signs_of.get(c['name'], []))): |
| chunks = [] |
| for s_ in src: |
| q = gen(mv_gpt, mv_net, s_, mv, KEEP_50_MVSL) |
| if q is not None: |
| chunks.append(q) |
| if not chunks: |
| row[tag] = None |
| continue |
| seq = np.concatenate(chunks, 0) |
| px = neck[0] + seq[:, :, 0] * sw |
| py = neck[1] + seq[:, :, 1] * sw |
| row[tag] = np.stack([px / W, py / H], -1).astype(np.float32) |
|
|
| if any(row.get(t) is None for t in tags): |
| continue |
| for t in tags: |
| out[t].append(row[t].astype(np.float32)) |
| names.append(c['name']) |
|
|
| os.makedirs(args.out_dir, exist_ok=True) |
| for t in tags: |
| np.savez(os.path.join(args.out_dir, f'{t}.npz'), |
| poses=np.array(out[t], dtype=object), |
| labels=np.full(len(names), -1), names=np.array(names)) |
| print(f' {t:<20} {len(out[t])} clips, mean len ' |
| f'{np.mean([len(p) for p in out[t]]):.0f}') |
| print(f'-> {args.out_dir}') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
