File size: 9,200 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 196 197 198 199 200 201 202 203 204 205 | #!/usr/bin/env python3
"""One TriVis table across model families and settings, on a single protocol.
Both projects already evaluate on `dataset/VSL_upper`, but with three different
conventions, so their published numbers are not directly comparable:
joints clips n
NSLP-G sentence-level 50 (8 body+42 hand) default_rng(0) 200
T2M-GPT eval_vsl.py 124 (incl. 68 face) RandomState(0) 300
This script re-scores T2M-GPT under NSLP-G's convention -- their exact 200 clip ids
(read from their results json), their 50-joint layout, their metric code -- and merges
the result with their rows into one table.
Units. Both report `frame` (the raw per-axis frame-normalized coordinates the pack
stores) and `shoulder` (per-clip: divide by that clip's own median shoulder width,
after undoing the per-axis anisotropy). The per-clip method is theirs; a single global
factor is an approximation that differs by ~17% and is not used here.
"""
import argparse
import csv
import json
import os
import sys
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 = list(range(8)) + list(range(82, 124)) # inside the 124-kpt upper layout
GROUPS_50 = {'all': (0, 50), 'body': (0, 8), 'hands': (8, 50)}
NECK, RSHO, LSHO = 1, 2, 5 # inside the 50-joint layout
def anchor_of(xy, valid, W, H):
"""Per-clip neck position and shoulder width, in pixels (NSLP-G's method)."""
px = xy * np.array([W, H], np.float32)
ok = valid[:, NECK].astype(bool) & valid[:, RSHO].astype(bool) & valid[:, LSHO].astype(bool)
if ok.sum() < 3:
ok = np.ones(len(xy), 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 to_shoulder(xy, neck, sw, W, H):
return (xy * np.array([W, H], np.float32) - neck[None, None, :].astype(np.float32)) / sw
def _dtw(pred, gt, valid):
"""Same DTW as eval_vsl.dtw_mje but with the 50-joint groups."""
import eval_vsl
saved = eval_vsl.GROUPS
eval_vsl.GROUPS = GROUPS_50
try:
return eval_vsl.dtw_mje(pred, gt, valid)
finally:
eval_vsl.GROUPS = saved
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--data-dir', default='./dataset/VSL_upper')
ap.add_argument('--vq', default='output_vsl/vq_vsl_upper/net_best.pth')
ap.add_argument('--gpt', default='output_vsl/gpt_vsl_upper/net_best.pth')
ap.add_argument('--nslpg-json', default=os.path.join(
REPO, '0.NSLP-G/sentence-level/results_sent_shoulder.json'))
ap.add_argument('--clipid-json', default=os.path.join(
REPO, '0.NSLP-G/sentence-level/results_sent_scratch.json'))
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('--device', default='cuda')
ap.add_argument('--out-json', default='output_vsl/trivis_unified.json')
args = ap.parse_args()
W, H = args.frame_w, args.frame_h
device = torch.device(args.device)
with open(args.clipid_json) as f:
clip_ids = json.load(f)['clip_ids']
print(f'using NSLP-G\'s exact {len(clip_ids)} test clip ids')
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, 'test')
NK = st.layout.n_kpts
pred_gloss = {}
if os.path.exists(args.pred_gloss):
with open(args.pred_gloss, encoding='utf-8') as f:
pred_gloss = {k: v['pred_gloss'] for k, v in json.load(f).items()}
conds = ['ceiling', 'gloss', 'pred_gloss', 'sentence', 'shuffled']
acc = {c: {u: {k: [] for k in GROUPS_50} for u in ('frame', 'shoulder')} for c in conds}
lens = {c: [] for c in conds}
unit_len = 2 ** targs.down_t
with torch.no_grad():
for i in tqdm(clip_ids, desc='T2M-GPT on NSLP-G protocol'):
c = st.index[i]
motion, mask = st.get(i)
valid_full = mask[:, ::2]
gt_full = (motion * st.std + st.mean).reshape(-1, NK, 2)
gt = gt_full[:, KEEP_50]
valid = valid_full[:, KEEP_50]
neck, sw = anchor_of(gt, valid, W, H)
gt_sh = to_shoulder(gt, neck, sw, W, H)
outs = {}
mt = torch.from_numpy(motion).unsqueeze(0).to(device)
T = (len(motion) // unit_len) * unit_len
rec = net.decode_batch(net.encode(mt[:, :T]))[0].cpu().numpy()
outs['ceiling'] = (rec * st.std + st.mean).reshape(-1, NK, 2)[:, KEEP_50]
other = st.index[(i + len(st.index) // 2) % len(st.index)]
texts = {'gloss': c['gloss'],
'pred_gloss': pred_gloss.get(c['name'], c['gloss']),
'sentence': c['sentence'],
'shuffled': other['gloss']}
for tag, txt in texts.items():
idx = gpt.sample(text_enc([txt]), if_categorial=True)
if idx is None or idx.numel() == 0:
continue
idx = idx.clamp(max=g.nb_code - 1)
p = net.decode_batch(idx)[0].cpu().numpy()
outs[tag] = (p * st.std + st.mean).reshape(-1, NK, 2)[:, KEEP_50]
for tag, p in outs.items():
for u, (pp, gg) in (('frame', (p, gt)),
('shoulder', (to_shoulder(p, neck, sw, W, H), gt_sh))):
d = _dtw(pp.astype(np.float64), gg.astype(np.float64),
valid.astype(np.float64))
for k, v in d.items():
acc[tag][u][k].append(v)
lens[tag].append(len(p) / len(gt))
t2m = {}
for cd in conds:
if not acc[cd]['frame']['hands']:
continue
t2m[cd] = {u: {k: float(np.mean(v)) for k, v in acc[cd][u].items()}
for u in ('frame', 'shoulder')}
t2m[cd]['len_ratio'] = float(np.mean(lens[cd]))
# ---- merge with NSLP-G's rows ----
ns = json.load(open(args.nslpg_json))['results']
print(f"\n{'system / setting':<40}{'frame_h':>10}{'shldr_h':>10}{'frame_all':>11}{'len_r':>8}")
print('-' * 79)
rows = []
order = [('NSLP-G stage-1 ceiling', ns.get('pre/ceiling'), None),
('T2M-GPT stage-1 ceiling', None, t2m.get('ceiling')),
(None, None, None),
('NSLP-G pretrained on Multi-VSL', ns.get('pre/nslpg_pred'), None),
('NSLP-G from scratch', ns.get('scratch/nslpg_pred'), None),
('NSLP-G pretrained + handshape loss', ns.get('shape/nslpg_pred'), None),
('NSLP-G oracle length', ns.get('pre/nslpg_oracle'), None),
(None, None, None),
('T2M-GPT gloss', None, t2m.get('gloss')),
('T2M-GPT predicted gloss (BARTpho)', None, t2m.get('pred_gloss')),
('T2M-GPT raw sentence (direct)', None, t2m.get('sentence')),
(None, None, None),
('NSLP-G shuffled gloss (control)', ns.get('pre/shuffled_gloss'), None),
('T2M-GPT shuffled gloss (control)', None, t2m.get('shuffled')),
('NSLP-G random init (floor)', ns.get('pre/random_init'), None),
('NSLP-G global mean pose (floor)', ns.get('pre/global_mean'), None)]
for name, a, b in order:
if name is None:
print('-' * 79); continue
if a is not None:
fr, sh, al = a['frame']['hands']['mean'], a['shoulder']['hands']['mean'], a['frame']['all']['mean']
lr = float('nan')
elif b is not None:
fr, sh, al, lr = b['frame']['hands'], b['shoulder']['hands'], b['frame']['all'], b['len_ratio']
else:
print(f'{name:<40} (missing)'); continue
lrs = '-' if lr != lr else f'{lr:.3f}'
print(f'{name:<40}{fr:>10.4f}{sh:>10.4f}{al:>11.4f}{lrs:>8}')
rows.append({'setting': name, 'frame_hands': fr, 'shoulder_hands': sh,
'frame_all': al, 'len_ratio': None if lr != lr else lr})
os.makedirs(os.path.dirname(args.out_json) or '.', exist_ok=True)
with open(args.out_json, 'w') as f:
json.dump({'n_clips': len(clip_ids), 'joints': 50, 'protocol': 'NSLP-G sentence-level',
't2mgpt': t2m, 'rows': rows}, f, indent=2)
print(f'\nwrote {args.out_json}')
if __name__ == '__main__':
main()
|