File size: 8,405 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 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 | #!/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()
|