eval_mvsl_on_trivis.py

#!/usr/bin/env python3
"""Score the WORD-level Multi-VSL model on TriVis sentences, head-to-head with the
sentence-level model — the only way to ask "which is better on TriVis".

The word model emits ONE isolated sign per call, so a TriVis sentence is built by
running it once per gloss sign and concatenating. Three known handicaps, all reported
rather than hidden:

  * vocabulary  -- only ~17% of TriVis gloss items exist in Multi-VSL's 1,000 words.
                   Conditioning goes through frozen PhoBERT (not a class embedding), so
                   unseen words still produce output, just unlearned. `--known-only`
                   restricts to clips whose every sign is in-vocabulary.
  * duration    -- isolated signs run 3.2x slower than connected signing. DTW is
                   alignment-free so it largely forgives this; `len_ratio` exposes it.
                   `--time-scale` optionally compresses each sign.
  * coarticulation -- each sign starts and ends at rest, so joins have rest-transitions.

Coordinate conversion (word model -> TriVis space). The word model outputs isotropic
shoulder-width units centred on the neck; TriVis references are per-axis frame
normalized over a W x H frame:

    dx_px, dy_px = xy_shoulder * SHOULDER_PX
    x_frame = (neck_x_px + dx_px) / W ,  y_frame = (neck_y_px + dy_px) / H

The neck anchor is taken from the reference clip's own median neck. That hands the
composed model ground-truth global placement, which FAVOURS it -- global position is
not what we are testing, and without an anchor the comparison would be meaningless.
"""
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 eval_vsl import dtw_mje, prefix_mje
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__)), '..')
NECK = 1


def norm(s):
    s = unicodedata.normalize('NFC', str(s)).lower().strip()
    return ' '.join(s.split())


def load_gloss_signs(csv_path):
    """clip name -> list of pipe-separated gloss signs (the index has pipes stripped)."""
    out = {}
    with open(csv_path, newline='', encoding='utf-8') as f:
        for r in csv.DictReader(f):
            name = os.path.splitext(os.path.basename(r['npz_path']))[0]
            out[name] = [t.strip() for t in str(r['Sign_sentence']).split('|') if t.strip()]
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--trivis-dir', default='./dataset/VSL')
    ap.add_argument('--mvsl-dir', default='./dataset/MVSL')
    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('--csv', default=os.path.join(REPO, 'Full_TriVis', 'split_lab_front.csv'))
    ap.add_argument('--lexicon', default=os.path.join(REPO, 'Multi-VSL_WACV_2025', 'data',
                                                      'glosses_1_1000.csv'))
    ap.add_argument('--n', type=int, default=300)
    ap.add_argument('--frame-w', type=float, default=1176.0)
    ap.add_argument('--frame-h', type=float, default=1288.0)
    ap.add_argument('--shoulder-px', type=float, default=356.0)
    ap.add_argument('--time-scale', type=float, default=1.0,
                    help='<1 compresses each isolated sign (e.g. 0.31 to match connected speed)')
    ap.add_argument('--known-only', action='store_true')
    ap.add_argument('--text-override', default=None,
                    help="BARTpho pred_test.json -- compose from PREDICTED gloss instead of "
                         "ground truth. Uses 'pred_raw', which keeps the | sign boundaries "
                         "the composition needs (pred_gloss has them stripped).")
    ap.add_argument('--device', default='cuda')
    ap.add_argument('--out-json', default='output_vsl/gpt_mvsl/on_trivis.json')
    args = ap.parse_args()

    device = torch.device(args.device)
    net, targs, _ = build_vqvae(args.mvsl_vq, device)
    tck = torch.load(args.mvsl_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)

    mv = dataset_vsl.VSLStore(args.mvsl_dir, 'test')     # for mean/std of the word model
    tri = dataset_vsl.VSLStore(args.trivis_dir, 'test')  # references + layout
    NK_T, NK_M = tri.layout.n_kpts, mv.layout.n_kpts
    assert NK_T == 128, 'expects the 128-kpt TriVis layout as reference'
    groups = tri.layout.metric_groups()

    # Multi-VSL 'upper' layout drops knees(9,12)/ankles(10,13); map its keypoints back
    # onto the TriVis-128 index space, leaving the dropped ones invalid.
    m2t = mv.layout.keep                     # new(MVSL) -> original(128)
    lex = {norm(r['word']) for r in csv.DictReader(open(args.lexicon)) if r['word']}
    signs_of = load_gloss_signs(args.csv)
    if args.text_override:
        with open(args.text_override, encoding='utf-8') as f:
            ov = json.load(f)
        n_ov = 0
        for k, v in ov.items():
            raw = v['pred_raw'] if isinstance(v, dict) else str(v)
            sg = [t.strip() for t in raw.split('|') if t.strip()]
            if sg:
                signs_of[k] = sg
                n_ov += 1
        print(f'text-override: composing from PREDICTED gloss for {n_ov} clips')

    rng = np.random.RandomState(0)
    items = sorted(rng.choice(len(tri.index), size=min(args.n, len(tri.index)),
                              replace=False).tolist())

    agg = {k: {q: [] for q in groups} for k in ('gen_dtw', 'gen_prefix')}
    ratios, n_used, n_skip, cover = [], 0, 0, []
    W, H, SP = args.frame_w, args.frame_h, args.shoulder_px

    with torch.no_grad():
        for i in tqdm(items, desc='mvsl->trivis'):
            c = tri.index[i]
            sg = signs_of.get(c['name'], [])
            if not sg:
                n_skip += 1
                continue
            known = [s for s in sg if norm(s) in lex]
            cover.append(len(known) / len(sg))
            if args.known_only and len(known) != len(sg):
                n_skip += 1
                continue

            motion, mask = tri.get(i)
            valid = mask[:, ::2]
            gt = (motion * tri.std + tri.mean).reshape(-1, NK_T, 2)
            neck_px = gt[:, NECK, :].mean(0) * np.array([W, H])   # anchor, in pixels

            # --- generate each sign, concatenate ---
            chunks = []
            for s in sg:
                feat = text_enc([s])
                idx = gpt.sample(feat, 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() * mv.std + mv.mean
                p = p.reshape(-1, NK_M, 2)
                if args.time_scale != 1.0:
                    T2 = max(2, int(round(len(p) * args.time_scale)))
                    ix = np.linspace(0, len(p) - 1, T2)
                    p = np.stack([np.stack([np.interp(ix, np.arange(len(p)), p[:, k, d])
                                            for d in range(2)], -1) for k in range(NK_M)], 1)
                chunks.append(p)
            if not chunks:
                n_skip += 1
                continue
            seq = np.concatenate(chunks, 0)                       # [T,NK_M,2] shoulder-w

            # --- shoulder-width, neck-centred  ->  TriVis per-axis frame coords ---
            px = neck_px[0] + seq[:, :, 0] * SP
            py = neck_px[1] + seq[:, :, 1] * SP
            pred = np.zeros((len(seq), NK_T, 2), np.float32)
            v_pred = np.zeros(NK_T, bool)
            for new_i, old_i in enumerate(m2t):
                pred[:, old_i, 0] = px[:, new_i] / W
                pred[:, old_i, 1] = py[:, new_i] / H
                v_pred[old_i] = True
            # keypoints the word model does not produce are excluded from scoring
            vmask = valid.copy()
            vmask[:, ~v_pred] = 0

            ratios.append(len(pred) / len(gt))
            for q, v in dtw_mje(pred, gt, vmask).items():
                agg['gen_dtw'][q].append(v)
            for q, v in prefix_mje(pred, gt, vmask).items():
                agg['gen_prefix'][q].append(v)
            n_used += 1

    res = {'n_scored': n_used, 'n_skipped': n_skip,
           'mean_gloss_coverage': float(np.mean(cover)) if cover else 0.0,
           'known_only': args.known_only, 'time_scale': args.time_scale,
           'len_ratio_mean': float(np.mean(ratios)) if ratios else 0.0,
           'units': 'TriVis frame-widths (converted from shoulder-widths)'}
    for k, d in agg.items():
        res[k] = {q: (float(np.mean(v)) if v else None) for q, v in d.items()}
    print(json.dumps(res, indent=2))
    print(f"\n{'metric':<26}{'all':>9}{'body':>9}{'hands':>9}")
    for k in ('gen_dtw', 'gen_prefix'):
        print(f"{k:<26}" + ''.join(f"{res[k][q]:>9.4f}" for q in ('all', 'body', 'hands')))
    print(f"len_ratio {res['len_ratio_mean']:.3f}   gloss coverage "
          f"{100*res['mean_gloss_coverage']:.1f}%   scored {n_used}, skipped {n_skip}")
    os.makedirs(os.path.dirname(args.out_json) or '.', exist_ok=True)
    with open(args.out_json, 'w') as f:
        json.dump(res, f, indent=2)


if __name__ == '__main__':
    main()