prepare_vsl_data.py

#!/usr/bin/env python3
"""Pack Full_TriVis DWPose skeletons into flat memmaps for T2M-GPT training.

Reads `Full_TriVis/split_lab_front.csv` (columns: split, npz_path, Sentence,
Sign_sentence, ...), loads each clip's DWPose `.npz` (all_xy [T,128,2],
all_score [T,128]) and concatenates every clip of a split into one contiguous
float16 array so that training can random-access frames without touching 24k
individual files on the shared volume.

Per split it writes into --out-dir:
  {split}_xy.npy     float16 [total_frames, 256]  xy flattened, NaN -> 0
  {split}_valid.npy  uint8   [total_frames, 128]  1 = keypoint usable
  {split}_index.json           per-clip offsets + text fields
And from the *train* split only: mean.npy / std.npy [256] (valid-only stats).

No body normalization is applied -- coordinates stay DWPose frame-normalized
[0,1] so absolute hand motion is preserved; the only transform is the global
per-dim standardization, which is invertible from mean/std.
"""
import argparse
import csv
import json
import os

import numpy as np
from tqdm import tqdm

from dataset.layout import Layout, PRESETS

FULL_NKP = 128


def read_rows(csv_path):
    with open(csv_path, newline="", encoding="utf-8") as f:
        return list(csv.DictReader(f))


def clip_len(path):
    with np.load(path) as d:
        return int(d["all_xy"].shape[0])


def load_clip(path, score_thr, layout):
    with np.load(path) as d:
        xy = d["all_xy"].astype(np.float32)  # [T,128,2]
        sc = d["all_score"].astype(np.float32)  # [T,128]
    valid = np.isfinite(xy).all(-1) & (sc > score_thr)  # [T,128]
    xy = np.nan_to_num(xy, nan=0.0, posinf=0.0, neginf=0.0)
    # select + reorder to the layout's keypoint set
    keep = layout.keep
    xy, valid = xy[:, keep, :], valid[:, keep]
    return xy.reshape(len(xy), layout.dim), valid.astype(np.uint8)


def gloss_to_text(s):
    """'1 | nam | mua | may ?' -> '1 nam mua may ?' (drop the gloss separator)."""
    return " ".join(t.strip() for t in str(s).split("|") if t.strip())


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--csv", default="../Full_TriVis/split_lab_front.csv")
    ap.add_argument("--root", default="..", help="repo root that npz_path is relative to")
    ap.add_argument("--out-dir", default="./dataset/VSL")
    ap.add_argument("--score-thr", type=float, default=0.3)
    ap.add_argument("--min-frames", type=int, default=64, help="drop clips shorter than this")
    ap.add_argument("--limit", type=int, default=0, help="debug: only N clips per split")
    ap.add_argument("--layout", default="full", choices=list(PRESETS),
                    help="keypoint set: 'upper' drops knees+ankles (never detected here)")
    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
    rows = read_rows(args.csv)
    by_split = {}
    for r in rows:
        by_split.setdefault(r["split"], []).append(r)
    print({k: len(v) for k, v in by_split.items()})

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

        # pass 1: lengths (so we can allocate the memmap exactly)
        keep, lengths = [], []
        for r in tqdm(srows, desc=f"{split}: scan"):
            p = os.path.join(args.root, r["npz_path"])
            try:
                T = clip_len(p)
            except Exception as e:  # noqa: BLE001 - a corrupt npz should not kill the run
                print(f"  skip {r['npz_path']}: {e}")
                continue
            if T < args.min_frames:
                continue
            keep.append(r)
            lengths.append(T)
        total = int(sum(lengths))
        print(f"{split}: {len(keep)} clips, {total} frames")

        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
        # running valid-only moments, for the train split's mean/std
        s1 = np.zeros(DIM, np.float64)
        s2 = np.zeros(DIM, np.float64)
        cnt = np.zeros(DIM, np.float64)

        for r, T in zip(tqdm(keep, desc=f"{split}: pack"), lengths):
            p = os.path.join(args.root, r["npz_path"])
            xy, vd = load_clip(p, args.score_thr, layout)
            assert len(xy) == T, (len(xy), T)
            xy_mm[off:off + T] = xy.astype(np.float16)
            vd_mm[off:off + T] = vd
            if split == "train":
                m = np.repeat(vd, 2, axis=1).astype(np.float64)  # [T,256]
                x64 = xy.astype(np.float64)
                s1 += (x64 * m).sum(0)
                s2 += (x64 * x64 * m).sum(0)
                cnt += m.sum(0)
            index.append({
                "name": os.path.splitext(os.path.basename(r["npz_path"]))[0],
                "npz_path": r["npz_path"],
                "start": off,
                "length": T,
                "gloss": gloss_to_text(r["Sign_sentence"]),
                "sentence": r["Sentence"],
                "category": r.get("Category", ""),
            })
            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
            var = np.maximum(s2 / cnt - mean ** 2, 0.0)
            std = np.sqrt(var)
            # A keypoint that is essentially never observed would give std~0.
            # Coords are frame-normalized to [0,1], so floor the scale at 1% of
            # the frame: without this, a near-static dim gets amplified ~100x by
            # z-normalization and then dominates the reconstruction loss.
            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("mean/std saved; std range", float(std.min()), float(std.max()))

        L = np.array(lengths)
        print(f"{split}: len min {L.min()} median {int(np.median(L))} p95 "
              f"{int(np.percentile(L, 95))} max {L.max()}")


if __name__ == "__main__":
    main()