File size: 6,615 Bytes
8e5456b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | #!/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()
|