check_trivis_floors.py

#!/usr/bin/env python3
"""What does a TriVis sentence score with NO model? The floor rows nobody measured.

`eval_vsl.py` reports the TriVis sentence-level T2M-GPT against its tokenizer ceiling but
against no floor, so 'hands 0.1537' has never been read against what a model-free
baseline gets on the same 300 clips. Three baselines, no network involved:

  global-mean        the mean TriVis *train* pose, held for the reference length.
  random-train-clip  a real train clip's skeleton, as-is: generic connected signing of
                     the wrong sentence. This is the strong floor -- it has correct
                     dynamics and correct statistics, and zero information about the text.
  len-matched-clip   the same, but drawn from train clips within +/-10% of the reference
                     length, so a length mismatch cannot be what it is being penalised for.

Reported in both units: TriVis frame-widths (what eval_vsl.py prints) and shoulder widths
(isotropic, per-clip -- the unit the NSLP-G tables use). Same clip set as everything else:
`RandomState(0).choice(...)` sorted.
"""
import argparse
import json

import numpy as np
from tqdm import tqdm

import eval_vsl
from dataset import dataset_vsl

NECK, RSHO, LSHO = 1, 2, 5
# `prepare_multivsl_data.normalize_clip` drops a clip whose median shoulder width is under
# this, because the scale reference is then untrustworthy. Scoring has to honour the same
# rule: on the 3-view pack one `right`-view clip has both shoulders on the same pixel, so
# sw = 0 and dividing by it turns the whole shoulder-unit column into nan.
MIN_SHOULDER_PX = 8.0


def anchor_of(gt, valid, W, H):
    px = gt * np.array([W, H], np.float32)
    ok = valid[:, NECK] & valid[:, RSHO] & valid[:, LSHO]
    if ok.sum() < 3:
        ok = np.ones(len(gt), 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, sw


def to_shoulder(gt, neck, sw, W, H):
    return (gt * np.array([W, H], np.float32) - neck[None, None, :].astype(np.float32)) / sw


class Agg:
    def __init__(self, groups):
        self.vals = {g: [] for g in groups}

    def add(self, d):
        for g, v in d.items():
            self.vals[g].append(v)

    def summary(self):
        out = {}
        for g, v in self.vals.items():
            a = np.asarray(v, np.float64)
            out[g] = {"mean": float(a.mean()), "n": len(a),
                      "sem": float(a.std(ddof=1) / np.sqrt(len(a))) if len(a) > 1 else 0.0}
        return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--data-dir', default='./dataset/VSL')
    ap.add_argument('--split', default='test')
    ap.add_argument('--n', type=int, default=300)
    ap.add_argument('--seed', type=int, default=0)
    ap.add_argument('--frame-w', type=float, default=1176.0)
    ap.add_argument('--frame-h', type=float, default=1288.0)
    ap.add_argument('--out-json', default='output_vsl/trivis_floors.json')
    args = ap.parse_args()

    W, H = args.frame_w, args.frame_h
    te = dataset_vsl.VSLStore(args.data_dir, args.split)
    tr = dataset_vsl.VSLStore(args.data_dir, 'train')
    NK = te.layout.n_kpts
    groups = te.layout.metric_groups()
    eval_vsl.GROUPS = groups

    rs = np.random.RandomState(args.seed)
    items = sorted(rs.choice(len(te.index), size=min(args.n, len(te.index)),
                             replace=False).tolist())
    print(f'{len(items)} clips of {args.split}, {NK} keypoints')

    def raw(store, i):
        c = store.index[i]
        s, T = c['start'], c['length']
        xy = np.asarray(store.xy[s:s + T], np.float32).reshape(T, NK, 2)
        vd = np.asarray(store.valid[s:s + T], bool)
        return xy, vd

    refs, dropped = [], 0
    for i in items:
        gt, vd = raw(te, i)
        neck, sw = anchor_of(gt, vd, W, H)
        if not np.isfinite(sw) or sw < MIN_SHOULDER_PX:
            dropped += 1
            continue
        refs.append((gt, vd, neck, sw, to_shoulder(gt, neck, sw, W, H)))
    if dropped:
        print(f'[drop] {dropped} clip(s) with median shoulder width < {MIN_SHOULDER_PX} px '
              f'-- no usable scale reference; scored on {len(refs)}')

    # global mean training pose, in frame coords
    step = max(1, len(tr.index) // 2000)
    acc = np.zeros((NK, 2), np.float64)
    wsum = np.zeros((NK, 1), np.float64)
    for i in range(0, len(tr.index), step):
        xy, vd = raw(tr, i)
        w = vd[..., None].astype(np.float64)
        acc += (xy * w).sum(0)
        wsum += w.sum(0)
    gmean = acc / np.maximum(wsum, 1e-6)

    tr_len = np.array([c['length'] for c in tr.index])
    rs2 = np.random.RandomState(args.seed + 1)

    results = {}

    def run(tag, make):
        """make(k, ref) -> (pred_frame [T,NK,2], pred_shoulder [T,NK,2])"""
        af, as_ = Agg(groups), Agg(groups)
        ratios = []
        for k, ref in enumerate(tqdm(refs, desc=tag, leave=False)):
            gt, vd, neck, sw, gt_sw = ref
            pf, ps = make(k, ref)
            af.add(eval_vsl.dtw_mje(pf.astype(np.float64), gt.astype(np.float64), vd))
            as_.add(eval_vsl.dtw_mje(ps.astype(np.float64), gt_sw.astype(np.float64), vd))
            ratios.append(len(pf) / len(gt))
        f, s = af.summary(), as_.summary()
        print(f"{tag:<20} frame: " + " ".join(f"{q} {f[q]['mean']:.4f}" for q in ('all', 'body', 'hands'))
              + f" | shoulder: " + " ".join(f"{q} {s[q]['mean']:.4f}" for q in ('all', 'body', 'hands'))
              + f"  len_ratio {np.mean(ratios):.3f}", flush=True)
        results[tag] = {'frame': f, 'shoulder': s, 'len_ratio': float(np.mean(ratios))}

    def global_mean(k, ref):
        gt, vd, neck, sw, _ = ref
        pf = np.repeat(gmean[None], len(gt), 0).astype(np.float32)
        return pf, to_shoulder(pf, neck, sw, W, H)

    def donor(k, ref, pool=None):
        gt, vd, neck, sw, _ = ref
        j = int(rs2.choice(pool)) if pool is not None else int(rs2.randint(len(tr.index)))
        dxy, dvd = raw(tr, j)
        dneck, dsw = anchor_of(dxy, dvd, W, H)
        # a real clip carries its own body scale and position; in frame coords it is used
        # as-is, in shoulder coords it is normalized by its OWN anchor, exactly as the
        # reference is by its own
        return dxy, to_shoulder(dxy, dneck, dsw, W, H)

    run('global-mean', global_mean)
    run('random-train-clip', lambda k, ref: donor(k, ref))
    run('len-matched-clip', lambda k, ref: donor(
        k, ref, pool=np.flatnonzero(np.abs(tr_len - len(ref[0])) <= 0.1 * len(ref[0]))))

    with open(args.out_json, 'w') as f:
        json.dump({'args': vars(args), 'clips': [te.index[i]['name'] for i in items],
                   'results': results}, f, indent=2)
    print(f'wrote {args.out_json}')


if __name__ == '__main__':
    main()