#!/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()