prepare_multivsl_data.py

#!/usr/bin/env python3
"""Pack Multi-VSL (WACV 2025) front-view DWPose skeletons for T2M-GPT training.

Isolated-sign data: one gloss -> one clip, so there is no gloss/pose alignment
problem. That makes this the clean diagnostic for "can T2M-GPT learn gloss->pose
at all", separate from the sentence-level alignment question.

TWO NORMALIZATIONS ARE REQUIRED HERE, unlike Full_TriVis:

1. Geometry. The dataset ships a *different YOLO crop per clip* (resolutions seen:
   560x712, 698x986, 1396x2106, ...). DWPose coordinates are normalized per axis to
   [0,1], so they are (a) anisotropically scaled, because W != H, and (b) not
   comparable across clips, because the crop scale differs. So:
     xy_px  = xy_norm * [W, H]                    -> back to isotropic pixels
     xy_out = (xy_px - neck) / shoulder_width     -> crop-invariant, body-relative
   The reference (neck position, shoulder width) is computed ONCE PER CLIP from the
   median over valid frames -- deliberately not per frame. A per-clip affine removes
   only the crop artifact; a per-frame one would also erase genuine body sway and
   hand displacement relative to the torso, which is exactly the signal.
   Full_TriVis is left un-normalized; only this dataset needs it.

2. Frame rate. Clips are 29.97 / 30 / 59.94 fps. Anything above --fps-thresh is
   decimated by 2 so token duration means the same thing everywhere.

Splits are the dataset's official signer-disjoint ones (20,161 / 3,713 / 4,538, all
1,000 classes present in each), so no split is invented here.

Outputs mirror prepare_vsl_data.py: {split}_xy.npy / {split}_valid.npy /
{split}_index.json / mean.npy / std.npy / layout.json.
"""
import argparse
import json
import os

import numpy as np
from tqdm import tqdm

from dataset.layout import Layout, PRESETS

REPO = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..')
NECK, RSHO, LSHO = 1, 2, 5          # OpenPose-18 body indices
MIN_SHOULDER_PX = 8.0               # below this the reference is untrustworthy


def load_npz(path):
    with np.load(path) as d:
        xy = d['all_xy'].astype(np.float32)          # [T,128,2] in [0,1] per axis
        sc = d['all_score'].astype(np.float32)       # [T,128]
        wh = d['frame_size'].astype(np.float32) if 'frame_size' in d else None
        fps = float(d['fps']) if 'fps' in d else 30.0
    return xy, sc, wh, fps


def normalize_clip(xy, sc, wh, score_thr):
    """-> (xy_norm [T,128,2], valid [T,128] bool, info) or (None, None, reason)."""
    T = len(xy)
    valid = np.isfinite(xy).all(-1) & (sc > score_thr)
    xy = np.nan_to_num(xy, nan=0.0, posinf=0.0, neginf=0.0)

    # (a) per-axis [0,1] -> isotropic pixels
    xy = xy * wh[None, None, :]

    # (b) per-clip reference from the median over frames where it is observed
    ok_ref = valid[:, NECK] & valid[:, RSHO] & valid[:, LSHO]
    if ok_ref.sum() < max(3, 0.1 * T):
        return None, None, 'no reliable neck/shoulder reference'
    neck = np.median(xy[ok_ref, NECK, :], axis=0)                     # [2]
    sw = np.median(np.linalg.norm(xy[ok_ref, RSHO, :] - xy[ok_ref, LSHO, :], axis=-1))
    if not np.isfinite(sw) or sw < MIN_SHOULDER_PX:
        return None, None, f'shoulder width {sw:.1f}px too small'

    xy = (xy - neck[None, None, :]) / sw
    return xy, valid, {'shoulder_px': float(sw), 'neck_px': neck.tolist()}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--root', default=os.path.join(REPO, 'Multi-VSL_front'))
    ap.add_argument('--out-dir', default='./dataset/MVSL')
    ap.add_argument('--score-thr', type=float, default=0.3)
    ap.add_argument('--layout', default='upper', choices=list(PRESETS))
    ap.add_argument('--min-frames', type=int, default=16)
    ap.add_argument('--fps-thresh', type=float, default=45.0,
                    help='clips above this fps are decimated by 2')
    ap.add_argument('--limit', type=int, default=0)
    args = ap.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)
    layout = Layout.preset(args.layout)
    layout.save(args.out_dir)
    print(layout)
    NKP, DIM = layout.n_kpts, layout.dim

    with open(os.path.join(args.root, 'clips.json'), encoding='utf-8') as f:
        clips = json.load(f)
    by_split = {}
    for c in clips:
        by_split.setdefault(c['split'], []).append(c)
    print({k: len(v) for k, v in by_split.items()})

    skel = os.path.join(args.root, 'skeleton')
    stats_all = {}

    for split, rows in sorted(by_split.items()):
        if args.limit:
            rows = rows[:args.limit]

        # ---- pass 1: normalize into RAM (these clips are short; ~2.4M frames total)
        kept, drop = [], {'missing': 0, 'short': 0, 'noref': 0}
        for c in tqdm(rows, desc=f'{split}: load'):
            p = os.path.join(skel, os.path.splitext(c['name'])[0] + '.npz')
            if not os.path.exists(p):
                drop['missing'] += 1
                continue
            try:
                xy, sc, wh, fps = load_npz(p)
            except Exception:
                drop['missing'] += 1
                continue
            if wh is None or len(xy) < args.min_frames:
                drop['short'] += 1
                continue
            if fps > args.fps_thresh:                      # 59.94 -> ~30
                xy, sc = xy[::2], sc[::2]
            nxy, valid, info = normalize_clip(xy, sc, wh, args.score_thr)
            if nxy is None:
                drop['noref'] += 1
                continue
            if len(nxy) < args.min_frames:
                drop['short'] += 1
                continue
            kept.append((c, nxy[:, layout.keep, :].reshape(len(nxy), DIM),
                         valid[:, layout.keep].astype(np.uint8), info))
        print(f'{split}: kept {len(kept)}, dropped {drop}')

        total = sum(len(k[1]) for k in kept)
        xy_mm = np.lib.format.open_memmap(os.path.join(args.out_dir, f'{split}_xy.npy'),
                                          mode='w+', dtype=np.float16, shape=(total, DIM))
        vd_mm = np.lib.format.open_memmap(os.path.join(args.out_dir, f'{split}_valid.npy'),
                                          mode='w+', dtype=np.uint8, shape=(total, NKP))
        index, off = [], 0
        s1 = np.zeros(DIM, np.float64); s2 = np.zeros(DIM, np.float64)
        cnt = np.zeros(DIM, np.float64)

        for c, x, v, info in tqdm(kept, desc=f'{split}: pack'):
            T = len(x)
            xy_mm[off:off + T] = x.astype(np.float16)
            vd_mm[off:off + T] = v
            if split == 'train':
                m = np.repeat(v, 2, axis=1).astype(np.float64)
                x64 = x.astype(np.float64)
                s1 += (x64 * m).sum(0); s2 += (x64 * x64 * m).sum(0); cnt += m.sum(0)
            index.append({'name': os.path.splitext(c['name'])[0], 'start': off, 'length': T,
                          'label': c['label'], 'gloss': c['word'], 'sentence': c['word'],
                          'shoulder_px': info['shoulder_px']})
            off += T

        xy_mm.flush(); vd_mm.flush(); del xy_mm, vd_mm
        with open(os.path.join(args.out_dir, f'{split}_index.json'), 'w', encoding='utf-8') as f:
            json.dump(index, f, ensure_ascii=False)

        if split == 'train':
            cnt = np.maximum(cnt, 1.0)
            mean = s1 / cnt
            std = np.sqrt(np.maximum(s2 / cnt - mean ** 2, 0.0))
            std[cnt < 100] = 1.0
            std = np.maximum(std, 1e-2)
            np.save(os.path.join(args.out_dir, 'mean.npy'), mean.astype(np.float32))
            np.save(os.path.join(args.out_dir, 'std.npy'), std.astype(np.float32))
            print(f'mean/std saved; std range {std.min():.4f}..{std.max():.4f}')

        L = np.array([len(k[1]) for k in kept])
        u = 4
        print(f'{split}: frames min {L.min()} median {int(np.median(L))} max {L.max()} '
              f'| tokens median {int(np.median(L)//u)} max {L.max()//u}')
        stats_all[split] = {'clips': len(kept), 'frames': int(total), 'dropped': drop,
                            'tokens_median': int(np.median(L) // u),
                            'tokens_max': int(L.max() // u)}

    with open(os.path.join(args.out_dir, 'prep_stats.json'), 'w') as f:
        json.dump({'layout': layout.name, 'splits': stats_all}, f, indent=2)
    print('\n' + json.dumps(stats_all, indent=2))


if __name__ == '__main__':
    main()