#!/usr/bin/env python3 """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)) # inside the 124-kpt upper layout KEEP_50_MVSL = list(range(8)) + list(range(82, 124)) # MVSL pack is also `upper` 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 # ---- recipe C: Multi-VSL word model, one sign per gloss item ---- 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) # shoulder-width, neck-centred 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()