310 lines
12 KiB
Python
310 lines
12 KiB
Python
"""
|
|
Phase 5 — Adaptive PCHIP per-joint tokenizer.
|
|
|
|
For each window from Phase 4 (WINDOW_FRAMES frames -- 8 originally, 24 as of
|
|
2026-07-22, see REPORT.md #38), each of the 17 joints gets an independent
|
|
PCHIP compression with adaptive control-point count based on per-joint
|
|
curvature over the WHOLE window (all WINDOW_FRAMES frames considered, not a
|
|
fixed sub-grid):
|
|
|
|
Tier 2 (2 CPs: start + end) — max curvature < tau_low
|
|
Tier 4 (4 CPs: start + end + top-2 interior) — tau_low <= curvature < tau_high
|
|
Tier MAX_CPS (MAX_CPS CPs: start+end+top-(MAX_CPS-2) interior, chosen by
|
|
curvature out of ALL WINDOW_FRAMES candidates) — curvature >= tau_high
|
|
|
|
MAX_CPS is fixed at 8 regardless of WINDOW_FRAMES -- widening the window
|
|
gives the top tier more candidate positions to pick its best 8 from, not more
|
|
tokens/joint. See MAX_CPS's own docstring below for why this matters.
|
|
|
|
Token stream per window (WINDOW_FRAMES=24 example; t values now range over
|
|
however many of the window's real frame indices got chosen as CPs, e.g.
|
|
t_0 and t_23 for tier 2, or t_0/t_5/t_14/t_23 for tier 4 if frames 5 and 14
|
|
had the highest curvature):
|
|
<fps_30>
|
|
<pelvis> <pelvis_t_0> <pelvis_x_N> <pelvis_y_N> <pelvis_z_N>
|
|
<pelvis_t_23> <pelvis_x_N> <pelvis_y_N> <pelvis_z_N> </pelvis>
|
|
<r_hip> <r_hip_t_0> <r_hip_x_N> ... </r_hip>
|
|
...
|
|
|
|
Quantization: [-2.0 m, +2.0 m] -> [0, 255] (precision ~15.7 mm)
|
|
Time tokens: real frame index within the window, 0 to WINDOW_FRAMES-1 --
|
|
requires <{joint}_t_N> tokens up to WINDOW_FRAMES-1 in the
|
|
tokenizer vocab (only 0-7 existed before 2026-07-22).
|
|
|
|
Input: outputs/yolo_cleaned_30fps/{video_id}_cleaned.jsonl
|
|
Output: outputs/agent_tokens_adaptive/{video_id}_tokens.jsonl
|
|
Each line: {"video_id", "window_id", "fps", "token_str", "cp_counts"}
|
|
"""
|
|
|
|
import argparse
|
|
import glob
|
|
import json
|
|
import os
|
|
|
|
import numpy as np
|
|
|
|
# ── Constants ─────────────────────────────────────────────────────────────────
|
|
|
|
TARGET_FPS = 30
|
|
WINDOW_FRAMES = 8
|
|
N_JOINTS = 17
|
|
COORD_RANGE = 2.0
|
|
STRIDE = 8
|
|
# 2026-07-22 (REPORT.md #38): cap on control points per joint, independent of
|
|
# WINDOW_FRAMES. Before this change WINDOW_FRAMES==MAX_CPS==8 always (the top
|
|
# tier was literally "use every frame"), so widening the window to 24 frames
|
|
# would have silently tripled worst-case tokens/joint (24 CPs instead of 8)
|
|
# with no code change needed to trigger it. Keeping MAX_CPS fixed at 8 means
|
|
# the top tier now means "pick the best 8 of WINDOW_FRAMES candidates by
|
|
# curvature" instead of "use all of them" -- same worst-case token cost as
|
|
# before, but the 8 chosen points can be anywhere in the (now wider) window
|
|
# instead of forced onto a fixed 8-slot grid. This is the whole point of
|
|
# "Option 2" (dense pose, no subsampling before curve-fitting) agreed with
|
|
# the user: Phase 3 keeps every real frame, Phase 5 decides freely which
|
|
# frames matter most.
|
|
MAX_CPS = 8
|
|
|
|
TAU_LOW = 0.005
|
|
TAU_HIGH = 0.05
|
|
|
|
JOINT_NAMES = [
|
|
"pelvis", "r_hip", "r_knee", "r_ankle",
|
|
"l_hip", "l_knee", "l_ankle",
|
|
"spine", "thorax", "nose", "head_top",
|
|
"l_shoulder", "l_elbow", "l_wrist",
|
|
"r_shoulder", "r_elbow", "r_wrist",
|
|
]
|
|
|
|
# ── Quantization ──────────────────────────────────────────────────────────────
|
|
|
|
def quantize(v: float) -> int:
|
|
return int(np.clip(round((v + COORD_RANGE) / (2.0 * COORD_RANGE) * 255), 0, 255))
|
|
|
|
|
|
def dequantize(n: int) -> float:
|
|
return n / 255.0 * (2.0 * COORD_RANGE) - COORD_RANGE
|
|
|
|
|
|
# ── Per-joint adaptive CP selection ───────────────────────────────────────────
|
|
|
|
def joint_curvature(trajectory: np.ndarray) -> float:
|
|
"""Max curvature (acceleration norm) for a single joint trajectory (8, 3)."""
|
|
if trajectory.shape[0] < 3:
|
|
return 0.0
|
|
vel = np.diff(trajectory, axis=0)
|
|
acc = np.diff(vel, axis=0)
|
|
return float(np.max(np.linalg.norm(acc, axis=1)))
|
|
|
|
|
|
def select_cp_indices(trajectory: np.ndarray, tau_low: float, tau_high: float) -> np.ndarray:
|
|
"""Choose which frame indices become control points for one joint.
|
|
|
|
Tier sizes are fixed at 2 / 4 / MAX_CPS regardless of how many frames are
|
|
in the window (see MAX_CPS's docstring) -- the top tier picks the
|
|
MAX_CPS-2 highest-curvature *interior* frames out of every candidate in
|
|
the window, not a fixed grid position. With WINDOW_FRAMES==MAX_CPS==8
|
|
(the original config) this is exactly equivalent to the old
|
|
`np.arange(WINDOW_FRAMES)` behavior, since "top 6 of 6 interior
|
|
candidates" is all of them.
|
|
"""
|
|
curv = joint_curvature(trajectory)
|
|
n_frames = trajectory.shape[0]
|
|
|
|
if curv < tau_low:
|
|
return np.array([0, n_frames - 1])
|
|
|
|
n_interior = 2 if curv < tau_high else (MAX_CPS - 2)
|
|
n_interior = min(n_interior, max(n_frames - 2, 0))
|
|
|
|
vel = np.diff(trajectory, axis=0)
|
|
acc = np.diff(vel, axis=0)
|
|
acc_norms = np.linalg.norm(acc, axis=1) # (n_frames-2,)
|
|
|
|
# acc[i] corresponds to frame i+1 (second derivative offset)
|
|
interior_curv = np.zeros(n_frames)
|
|
for i in range(len(acc_norms)):
|
|
interior_curv[i + 1] = acc_norms[i]
|
|
|
|
# Exclude endpoints (already included), pick top-n_interior interior frames
|
|
interior_curv[0] = -1.0
|
|
interior_curv[-1] = -1.0
|
|
top_n = np.argsort(interior_curv)[-n_interior:] if n_interior > 0 else np.array([], dtype=int)
|
|
|
|
indices = np.unique(np.sort(np.concatenate(([0], top_n, [n_frames - 1]))))
|
|
return indices.astype(int)
|
|
|
|
|
|
# ── Token builder ─────────────────────────────────────────────────────────────
|
|
|
|
def build_token_str(
|
|
states: np.ndarray,
|
|
fps: int = TARGET_FPS,
|
|
tau_low: float = TAU_LOW,
|
|
tau_high: float = TAU_HIGH,
|
|
) -> tuple:
|
|
"""
|
|
states : (8, 17, 3) float32, root-centred metric coordinates
|
|
Returns (token_str, cp_counts_dict)
|
|
"""
|
|
parts = [f"<fps_{fps}>"]
|
|
cp_counts = {}
|
|
|
|
for j in range(N_JOINTS):
|
|
name = JOINT_NAMES[j]
|
|
trajectory = states[:, j, :] # (8, 3)
|
|
cp_idx = select_cp_indices(trajectory, tau_low, tau_high)
|
|
cp_counts[name] = len(cp_idx)
|
|
|
|
parts.append(f"<{name}>")
|
|
for fi in cp_idx:
|
|
x, y, z = trajectory[fi]
|
|
parts.append(f"<{name}_t_{fi}>")
|
|
parts.append(f"<{name}_x_{quantize(x)}>")
|
|
parts.append(f"<{name}_y_{quantize(y)}>")
|
|
parts.append(f"<{name}_z_{quantize(z)}>")
|
|
parts.append(f"</{name}>")
|
|
|
|
return " ".join(parts), cp_counts
|
|
|
|
|
|
# ── Per-file processing ──────────────────────────────────────────────────────
|
|
|
|
def process_file(
|
|
input_path: str,
|
|
output_jsonl: str,
|
|
video_id: str,
|
|
stride: int = STRIDE,
|
|
tau_low: float = TAU_LOW,
|
|
tau_high: float = TAU_HIGH,
|
|
) -> int:
|
|
records = []
|
|
|
|
with open(input_path, "r", encoding="utf-8") as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
|
|
data = json.loads(line)
|
|
window_id = int(data["window_id"])
|
|
|
|
if window_id % stride != 0:
|
|
continue
|
|
|
|
states = np.array(data["states"], dtype=np.float32)
|
|
|
|
if states.shape != (WINDOW_FRAMES, N_JOINTS, 3):
|
|
continue
|
|
if np.isnan(states).any():
|
|
continue
|
|
|
|
token_str, cp_counts = build_token_str(states, TARGET_FPS, tau_low, tau_high)
|
|
|
|
records.append({
|
|
"video_id": video_id,
|
|
"window_id": window_id,
|
|
"fps": TARGET_FPS,
|
|
"token_str": token_str,
|
|
"cp_counts": cp_counts,
|
|
})
|
|
|
|
if not records:
|
|
return 0
|
|
|
|
tmp = output_jsonl + ".tmp"
|
|
os.makedirs(os.path.dirname(output_jsonl), exist_ok=True)
|
|
with open(tmp, "w", encoding="utf-8") as f:
|
|
for rec in records:
|
|
f.write(json.dumps(rec, ensure_ascii=False) + "\n")
|
|
os.replace(tmp, output_jsonl)
|
|
|
|
return len(records)
|
|
|
|
|
|
# ── Entry point ───────────────────────────────────────────────────────────────
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
p = argparse.ArgumentParser(
|
|
description="Phase 5 — Adaptive PCHIP per-joint tokenizer."
|
|
)
|
|
p.add_argument("--input-dir", required=True,
|
|
help="Directory with *_cleaned.jsonl from Phase 4.")
|
|
p.add_argument("--output-dir", required=True,
|
|
help="Directory to write *_tokens.jsonl files.")
|
|
p.add_argument("--stride", type=int, default=STRIDE,
|
|
help=f"Keep windows where window_id %% stride == 0. Default: {STRIDE}")
|
|
p.add_argument("--tau-low", type=float, default=TAU_LOW,
|
|
help=f"Curvature threshold for 2-CP tier. Default: {TAU_LOW}")
|
|
p.add_argument("--tau-high", type=float, default=TAU_HIGH,
|
|
help=f"Curvature threshold for 8-CP tier. Default: {TAU_HIGH}")
|
|
p.add_argument("--file-list", default=None,
|
|
help="Optional text file listing specific *_cleaned.jsonl paths.")
|
|
p.add_argument("--window-frames", type=int, default=WINDOW_FRAMES,
|
|
help=f"Must match Phase 3/4's --window-size. Default: {WINDOW_FRAMES}. "
|
|
f"2026-07-22: use 24 to match the wider cosmos chunk window -- "
|
|
f"see REPORT.md #38. MAX_CPS (top-tier control-point count) is "
|
|
f"NOT tied to this and stays 8 either way.")
|
|
return p.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
global WINDOW_FRAMES
|
|
args = parse_args()
|
|
WINDOW_FRAMES = args.window_frames
|
|
os.makedirs(args.output_dir, exist_ok=True)
|
|
|
|
if args.file_list:
|
|
with open(args.file_list) as f:
|
|
all_files = [l.strip() for l in f if l.strip()]
|
|
else:
|
|
all_files = sorted(glob.glob(os.path.join(args.input_dir, "*_cleaned.jsonl")))
|
|
|
|
task_id = int(os.environ.get("SLURM_ARRAY_TASK_ID", 0))
|
|
num_tasks = int(os.environ.get("SLURM_ARRAY_TASK_COUNT", 1))
|
|
|
|
my_files = [f for i, f in enumerate(all_files) if i % num_tasks == task_id]
|
|
total = len(my_files)
|
|
|
|
print(f"\n[Worker {task_id}/{num_tasks}] {total} files to process.")
|
|
print("=" * 60)
|
|
|
|
processed = skipped = empty = 0
|
|
tier_counts = {2: 0, 4: 0, 8: 0}
|
|
|
|
for idx, input_path in enumerate(my_files, start=1):
|
|
base = os.path.basename(input_path)
|
|
video_id = base[: -len("_cleaned.jsonl")]
|
|
|
|
out_jsonl = os.path.join(args.output_dir, f"{video_id}_tokens.jsonl")
|
|
|
|
if os.path.exists(out_jsonl):
|
|
skipped += 1
|
|
print(f"[{idx}/{total}] {video_id} — already done", end="\r")
|
|
continue
|
|
|
|
try:
|
|
n = process_file(
|
|
input_path, out_jsonl, video_id,
|
|
stride=args.stride,
|
|
tau_low=args.tau_low,
|
|
tau_high=args.tau_high,
|
|
)
|
|
if n > 0:
|
|
processed += 1
|
|
pct = (processed + skipped + empty) / total * 100
|
|
print(f"[{idx}/{total}] {pct:.1f}% | {video_id} — {n} windows")
|
|
else:
|
|
empty += 1
|
|
except Exception as e:
|
|
print(f"[{idx}/{total}] ERROR {video_id} — {e}")
|
|
for p in (out_jsonl + ".tmp",):
|
|
if os.path.exists(p):
|
|
os.remove(p)
|
|
|
|
print("\n" + "=" * 60)
|
|
print(f"[Worker {task_id}] done — processed: {processed}, skipped: {skipped}, empty: {empty}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|