File size: 10,209 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 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 | #!/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()
|