#!/usr/bin/env python3 """Pack the processed_vnhn sentence-level dataset into T2M-GPT memmaps. Mirrors `prepare_vsl_data.py` (same outputs, same normalization, same layout machinery) but reads processed_vnhn's per-clip **pickles** instead of Full_TriVis's .npz. THREE THINGS DIFFER FROM Full_TriVis AND EACH ONE MATTERS --------------------------------------------------------- 1. **Keypoints are raw COCO-WholeBody 133, not the project's 128.** Full_TriVis was extracted through `easy_dwpose`, which already converts COCO-17 -> OpenPose-18 (it synthesises a `neck` from the shoulder midpoint) and emits body18+face68+Lhand21+Rhand21. These pickles are the detector's raw output, so the conversion is done here: COCO-WholeBody 133 = 0..16 body(COCO-17), 17..22 feet, 23..90 face, 91..111 Lhand, 112..132 Rhand Verified against the data: hand root 91 sits 0.021 frame-widths from COCO Lwrist 9, and 112 sits 0.032 from Rwrist 10 -- i.e. the hand blocks really do start at the wrist, which is what makes this index map safe. 2. **`scores` are NOT [0,1] confidences.** They run 0.4..11.3 (median 8.2), roughly 10x a confidence. Full_TriVis's `--score-thr 0.3` would therefore mark *everything* valid. Calibrated by matching Full_TriVis's per-group validity (body 82.7 / face 100 / hands ~100): **thr=3.0** gives face 99.7 / hands 98.1 / body 76.3 and correctly kills feet (1.0% valid). thr=4.0 is too harsh (hands collapse to 78%), thr=2.0 leaks feet (22%). 3. **Legs are out of frame.** This is a broadcast bust shot: knees score ~1.2 and ankles ~0.85, both <0.1% valid at thr=3.0, while hips stay 99.7% valid at median y=0.97 (the very bottom edge). That is precisely the case the `upper` preset exists for, so the default here is `upper` (124 kpts: body14 + face68 + hands42) rather than `full`. Using `full` would hand the decoder four dimensions of pure noise -- the documented "phantom leg" failure. TEXT. `text` is the transcript line (Vietnamese, lowercased, unpunctuated). There is no gloss annotation, so `gloss` and `sentence` are both set to it: stage 2 conditions on `--text-field`, and with only one text source the two must agree or the field silently selects nothing. CAVEAT NOT FIXED HERE. Every clip's end timestamp was padded by +3s before clamping (see the dataset's own README), so the tail of each clip contains signing the text does not describe -- ~23% of a 331-frame sample. That is left intact: trimming it changes the dataset's semantics and should be a deliberate, separately-evaluated choice. """ import argparse import csv import json import os import pickle import numpy as np from dataset.layout import Layout, PRESETS # COCO-17 -> OpenPose-18. Index 1 (neck) is synthesised from the shoulder midpoint, so it # is encoded as None and handled separately. # 0 nose 1 neck 2 Rsho 3 Relb 4 Rwri 5 Lsho 6 Lelb 7 Lwri # 8 Rhip 9 Rknee 10 Rank 11 Lhip 12 Lknee 13 Lank 14 Reye 15 Leye 16 Rear 17 Lear COCO17_TO_OP18 = [0, None, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3] FACE = (23, 91) LHAND = (91, 112) RHAND = (112, 133) def wholebody133_to_128(xy, sc): """[T,133,2]/[T,133] COCO-WholeBody -> [T,128,2]/[T,128] body18+face68+Lhand21+Rhand21.""" T = xy.shape[0] b_xy = np.zeros((T, 18, 2), np.float32) b_sc = np.zeros((T, 18), np.float32) for j, src in enumerate(COCO17_TO_OP18): if src is None: # neck = midpoint of the two shoulders b_xy[:, j] = 0.5 * (xy[:, 5] + xy[:, 6]) b_sc[:, j] = np.minimum(sc[:, 5], sc[:, 6]) # only as good as its worse parent else: b_xy[:, j] = xy[:, src] b_sc[:, j] = sc[:, src] out_xy = np.concatenate([b_xy, xy[:, FACE[0]:FACE[1]], xy[:, LHAND[0]:LHAND[1]], xy[:, RHAND[0]:RHAND[1]]], axis=1) out_sc = np.concatenate([b_sc, sc[:, FACE[0]:FACE[1]], sc[:, LHAND[0]:LHAND[1]], sc[:, RHAND[0]:RHAND[1]]], axis=1) return out_xy, out_sc def load_clip(path, score_thr, layout): with open(path, "rb") as f: d = pickle.load(f) xy = np.asarray(d["keypoints"], np.float32) sc = np.asarray(d["scores"], np.float32) if xy.ndim == 4: # [T,1,133,2] -> [T,133,2] (one person per clip) xy, sc = xy[:, 0], sc[:, 0] if xy.shape[1] != 133: raise ValueError(f"expected 133 COCO-WholeBody keypoints, got {xy.shape[1]}") xy, sc = wholebody133_to_128(xy, sc) valid = np.isfinite(xy).all(-1) & (sc > score_thr) keep = layout.keep xy, valid = xy[:, keep, :], valid[:, keep] return xy.reshape(len(xy), layout.dim), valid.astype(np.uint8) def main(): ap = argparse.ArgumentParser() ap.add_argument("--csv-dir", default="../processed_vnhn/Cut_video/output/csv", help="dir holding {train,val,test}_fixed.csv") ap.add_argument("--root", default="../processed_vnhn/Cut_video/output", help="what pkl_path in the CSV is relative to") ap.add_argument("--out-dir", default="./dataset/VNHN") ap.add_argument("--score-thr", type=float, default=3.0, help="scores here are ~10x a confidence; 3.0 matches TriVis validity") ap.add_argument("--min-frames", type=int, default=64) # Stage 2 uses max_tokens=128 everywhere in this repo, i.e. 512 frames at down_t=2. # 3.95% of vnhn clips are longer (max 1715 frames = 428 tokens). Dropping them is an # explicit, countable choice; leaving them in means silent truncation inside the loader, # which would quietly break the text<->pose correspondence on the longest clips. ap.add_argument("--max-frames", type=int, default=512, help="drop clips longer than this (0 = keep all)") ap.add_argument("--limit", type=int, default=0) ap.add_argument("--layout", default="upper", choices=list(PRESETS), help="'upper' by default: knees/ankles are off-frame in this corpus") args = ap.parse_args() os.makedirs(args.out_dir, exist_ok=True) layout = Layout.preset(args.layout) layout.save(args.out_dir) print(layout) for split in ("train", "val", "test"): path = os.path.join(args.csv_dir, f"{split}_fixed.csv") rows = list(csv.DictReader(open(path, encoding="utf-8"))) if args.limit: rows = rows[:args.limit] # pass 1: lengths, so the memmap can be allocated exactly keep_rows, lengths = [], [] n_short = n_long = 0 for r in rows: p = os.path.join(args.root, r["pkl_path"]) try: with open(p, "rb") as f: T = len(pickle.load(f)["keypoints"]) except Exception as e: # noqa: BLE001 print(f" skip {r['pkl_path']}: {e}") continue if T < args.min_frames: n_short += 1 continue if args.max_frames and T > args.max_frames: n_long += 1 continue keep_rows.append(r) lengths.append(T) total = int(sum(lengths)) print(f"[{split}] kept {len(keep_rows)}/{len(rows)} clips ({total} frames); " f"dropped {n_short} shorter than {args.min_frames} and " f"{n_long} longer than {args.max_frames} frames " f"({100*n_long/max(len(rows),1):.2f}% would have been truncated at " f"max_tokens=128)") xy_mm = np.lib.format.open_memmap( os.path.join(args.out_dir, f"{split}_xy.npy"), mode="w+", dtype=np.float16, shape=(total, layout.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, layout.n_kpts)) s1 = np.zeros(layout.dim, np.float64) s2 = np.zeros(layout.dim, np.float64) cnt = np.zeros(layout.dim, np.float64) index, off = [], 0 for r in keep_rows: xy, vd = load_clip(os.path.join(args.root, r["pkl_path"]), args.score_thr, layout) T = len(xy) 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) x64 = xy.astype(np.float64) s1 += (x64 * m).sum(0) s2 += (x64 * x64 * m).sum(0) cnt += m.sum(0) txt = r["text"].strip() index.append({ "name": r["sample_id"], "npz_path": r["pkl_path"], # key name kept for loader compatibility "start": off, "length": T, "gloss": txt, # no gloss annotation exists; both = transcript "sentence": txt, "category": r.get("source_id", ""), "fps": float(r.get("fps", 25) or 25), "end_clamped": r.get("end_clamped", ""), }) 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 # never-observed dim: do not amplify it 100x 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 {float(std.min()):.4f}..{float(std.max()):.4f}") L = np.array(lengths) if len(L): print(f" frames per clip: min {L.min()} median {int(np.median(L))} " f"max {L.max()} mean {L.mean():.0f}") if __name__ == "__main__": main()