Buckets:
| """cwm3/aide/features/build_baton_pose_corpus.py — BATON -> AIDE reverse-transfer pretraining corpus INDEX (GPU-baselines experiment A; | |
| aide/docs/protocol.md item C21). | |
| Source: the route-level causal in-cabin pose of the BATON routes (results/pose_causal_routes_v2/<uid>.npz for the 74 re-extracted routes, | |
| results/pose_causal_routes/<uid>.npz for the rest; HALPE-136 RTMW keypoints at 10 Hz on the causal driver box, `kpts` f16 [F,136,3] = | |
| (x_norm, y_norm, score) normalised by the frame size, `box` f32 [F,4] xyxy in native pixels, `driver_det`, `t` = k/10 s). | |
| Devices / splits: derived from the s12 anchor manifests (s12_wm/results/targets_{manual,engaged}_{train,val,test}.npz: uid, dongle). A route | |
| enters the corpus only if it is a TRAIN-split route AND its device (dongle) appears in NO val / test split of either regime (the s12 splits are | |
| device-disjoint: 96 train devices vs 66 val/test devices, 0 overlap — asserted here). No AIDE file other than the keypoint cache statistics of | |
| the AIDE TRAIN split is read; no AIDE val / test row enters anything. | |
| Window shape = the AIDE protocol (aide/data/protocol.py): 10 frames, 5 observed -> 5 future. AIDE samples frames 0,4,9,...,44 of a 15 fps | |
| clip (steps of 4 / 5 frames = 0.267 / 0.333 s, mean 0.296 s); BATON is 10 Hz, so a window takes every 3rd 10-Hz sample (0.3 s steps, span | |
| 2.7 s vs AIDE 2.93 s) — the closest realisable stride; documented as a deviation. Anchors every ANCHOR_STRIDE (3) samples; a window is kept iff | |
| its 10 samples are contiguous in time (no span gap) and every frame is valid: driver detected, finite box, finite keypoints, and >= 80 % of the | |
| body joints (HALPE 0-25) AND >= 80 % of the face joints (26-93) have score >= 0.3 (the extractor's score_th). Hands are not filtered (AIDE | |
| hand confidences are low too). HALPE joint 17 (head-top) is structurally absent from the RTMW output (NaN in every frame; the | |
| extractor fills neck 18 / pelvis 19 but not 17): it is synthesised as nose + HEADTOP_K * (nose - neck) with score min(nose, neck) (`fill_missing_joints`, | |
| shared with the pretrainer); any other NaN joint takes the nose position with score 0 (= the AIDE joint-dropout convention). A frame needs | |
| finite nose (0), neck (18) and box. Windows are capped at MAX_WINDOWS (1,000,000) by a seeded uniform subsample; 5 % of the routes (seeded, >= 20) | |
| are flagged `route_val` = held-out routes for the pretraining early stop (never used for fitting). | |
| Coordinate convention (the per-window driver normalisation into the AIDE pixel frame): AIDE keypoints are pixels of a 1920x1080 frame and the model's | |
| channels are (dx,dy)/200, x/W, y/H, conf. For every window an ISOTROPIC similarity maps the driver into the AIDE frame using a body-intrinsic scale: | |
| scale = SW_AIDE / sw_window (AIDE TRAIN median shoulder width 309 px / the window's median shoulder width, HALPE joints 5-6, native px) and the window's | |
| median shoulder centre -> the AIDE TRAIN median shoulder centre (373, 762 px): x_aide = (x_norm * native_w - cx_sh) * scale + CX_AIDE, likewise y. | |
| (A first build scaled by the driver-box height instead; BATON person boxes cover more torso than AIDE boxes, which left the BATON shoulders at 201 px | |
| vs AIDE 309 px — a 1.5x under-scaling — so the box is now used only as a driver-tracking validity flag.) Isotropic so that motion magnitudes (the | |
| (dx,dy)/200 channels and the pixel metrics) scale consistently with body size. Confidence: RTMW scores (mean 0.6-0.9) and AIDE's AlphaPose scores | |
| (mean 0.11) differ in distribution -> per part (body, face, lhand, rhand) a quantile map BATON-score -> AIDE-TRAIN-score is stored (4,096 quantiles | |
| each) and applied by the pretrainer. Frame validity additionally requires both shoulders with score >= 0.3. | |
| Output aide/results/cache/baton_pose_corpus.npz (an INDEX, ~40 MB; the frames are re-read from the route npz by the pretrainer): | |
| route_uid [R] str, route_src [R] str (v1|v2 dir), route_native_wh [R,2] i32, route_n_frames [R] i32, route_val [R] bool, route_dongle [R] str, | |
| win_route [M] i32, win_start [M] i32 (frame index of sample 0; samples = start + 3*j), win_cx [M] f32, win_cy [M] f32 (window-median shoulder centre, native px), | |
| win_scale [M] f32, conf_q_baton [4,4096] f32, conf_q_aide [4,4096] f32 (quantile maps, part order body/face/lhand/rhand), aide_target [3] f32 (shoulder cx, cy, width, AIDE px), | |
| frame_stride 3, anchor_stride 3, n_frames_total, params (json string), stats (json string) | |
| + baton_pose_corpus_summary.json and a run json. No openpilot DMS field is read (the pose npz carry keypoints / boxes only). | |
| python aide/features/build_baton_pose_corpus.py [--max-windows 1000000] [--limit-routes K] [--out PATH] | |
| """ | |
| from __future__ import annotations | |
| import os | |
| for _v in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS", "NUMEXPR_NUM_THREADS"): | |
| os.environ.setdefault(_v, "4") | |
| import argparse # noqa: E402 | |
| import json # noqa: E402 | |
| import sys # noqa: E402 | |
| import time # noqa: E402 | |
| from pathlib import Path # noqa: E402 | |
| import numpy as np # noqa: E402 | |
| CWM3 = Path(__file__).resolve().parents[2] | |
| for p in (CWM3, CWM3 / "common"): | |
| sys.path.insert(0, str(p)) | |
| from aide.data import protocol as P # noqa: E402 | |
| from runmeta import write_run_json # noqa: E402 | |
| S12R = CWM3.parent / "s12_wm" / "results" | |
| POSE_V1, POSE_V2 = CWM3 / "results" / "pose_causal_routes", CWM3 / "results" / "pose_causal_routes_v2" | |
| OUT = CWM3 / "aide" / "results" / "cache" / "baton_pose_corpus.npz" | |
| FRAME_STRIDE, ANCHOR_STRIDE, MAX_WINDOWS = 3, 3, 1_000_000 | |
| CONF_TH, MIN_FRAC = 0.3, 0.8 | |
| N_Q = 4096 | |
| VAL_ROUTE_FRAC, VAL_ROUTE_MIN, SEED = 0.05, 20, 0 | |
| PART_SLICES = {k: (int(v[0]), int(v[-1]) + 1) for k, v in P.PARTS.items()} # body 0-26, face 26-94, lhand 94-115, rhand 115-136 | |
| NOSE, HEADTOP, NECK = 0, 17, 18 | |
| HEADTOP_K = 0.65 # head-top ~ nose + 0.65 (nose - neck): nose->vertex ~ 12 cm vs nose->neck ~ 18 cm (anthropometric proxy for the missing HALPE 17) | |
| def fill_missing_joints(k): | |
| """k f32 [F,136,3] (x, y, score; NaN = absent) -> finite copy: joint 17 = nose + HEADTOP_K (nose - neck) with score min(nose, neck) wherever it is | |
| NaN; every other NaN joint = nose position with score 0. Frames whose nose / neck are NaN stay NaN (they are invalid anyway).""" | |
| k = k.copy() | |
| nose, neck = k[:, NOSE], k[:, NECK] | |
| miss17 = np.isnan(k[:, HEADTOP, 0]) | |
| proxy = np.concatenate([nose[:, :2] + HEADTOP_K * (nose[:, :2] - neck[:, :2]), np.minimum(nose[:, 2:3], neck[:, 2:3])], 1) | |
| k[miss17, HEADTOP] = proxy[miss17] | |
| other = np.isnan(k[..., 0]) | |
| if other.any(): | |
| fx, fj = np.nonzero(other) | |
| k[fx, fj, :2] = nose[fx, :2] | |
| k[fx, fj, 2] = 0.0 | |
| return k | |
| def split_map(): | |
| """-> (safe_uids sorted, train_devices, eval_devices, per-uid dongle) from the s12 manifests; asserts device-disjoint splits.""" | |
| dev, uids, dongle_of = {}, {}, {} | |
| for r in ("manual", "engaged"): | |
| for s in ("train", "val", "test"): | |
| with np.load(S12R / f"targets_{r}_{s}.npz", allow_pickle=True) as z: | |
| u, d = z["uid"].astype(str), z["dongle"].astype(str) | |
| dev[(r, s)], uids[(r, s)] = set(np.unique(d).tolist()), set(np.unique(u).tolist()) | |
| for a, b in zip(u, d): | |
| dongle_of[a] = b | |
| train_dev = dev[("manual", "train")] | dev[("engaged", "train")] | |
| eval_dev = set().union(*[dev[(r, s)] for r in ("manual", "engaged") for s in ("val", "test")]) | |
| train_uid = uids[("manual", "train")] | uids[("engaged", "train")] | |
| eval_uid = set().union(*[uids[(r, s)] for r in ("manual", "engaged") for s in ("val", "test")]) | |
| assert not (train_dev & eval_dev), f"s12 splits share {len(train_dev & eval_dev)} devices — the corpus rule needs device-disjoint splits" | |
| safe = sorted(u for u in train_uid if dongle_of[u] not in eval_dev and u not in eval_uid) | |
| return safe, train_dev, eval_dev, dongle_of | |
| LSHO, RSHO = 5, 6 | |
| def aide_train_targets(): | |
| """AIDE TRAIN split only: median shoulder centre (cx, cy) and shoulder width (protocol frames) and per-part confidence quantiles (all 45 frames).""" | |
| c = np.load(P.CACHE_NPZ, allow_pickle=False) | |
| tr = c["split"] == "train" | |
| S = c["S10"][tr][..., :2].astype(np.float64) | |
| sc, sw = (S[:, :, LSHO] + S[:, :, RSHO]) / 2, np.linalg.norm(S[:, :, LSHO] - S[:, :, RSHO], axis=-1) | |
| box = np.array([np.median(sc[..., 0]), np.median(sc[..., 1]), np.median(sw)], np.float32) | |
| conf = c["S"][tr][..., 2] # [Ntr,45,136] | |
| qs = np.linspace(0, 1, N_Q) | |
| q_aide = np.stack([np.quantile(conf[:, :, lo:hi].ravel(), qs) for lo, hi in PART_SLICES.values()]).astype(np.float32) | |
| return box, q_aide, int(tr.sum()) | |
| def route_path(uid): | |
| p2 = POSE_V2 / f"{uid}.npz" | |
| return (p2, "v2") if p2.exists() else (POSE_V1 / f"{uid}.npz", "v1") | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--max-windows", type=int, default=MAX_WINDOWS) | |
| ap.add_argument("--limit-routes", type=int, default=0) | |
| ap.add_argument("--out", default=str(OUT)) | |
| a = ap.parse_args() | |
| t0 = time.time() | |
| safe, train_dev, eval_dev, dongle_of = split_map() | |
| if a.limit_routes: | |
| safe = safe[:a.limit_routes] | |
| box_t, q_aide, n_aide_train = aide_train_targets() | |
| print(f"routes: {len(safe)} train-only-device routes ({len({dongle_of[u] for u in safe})} devices); train devices {len(train_dev)}, " | |
| f"val/test devices {len(eval_dev)} (disjoint); AIDE train shoulder target cx {box_t[0]:.1f} cy {box_t[1]:.1f} width {box_t[2]:.1f} (n={n_aide_train})", flush=True) | |
| rng = np.random.default_rng(SEED) | |
| R = len(safe) | |
| n_val = max(VAL_ROUTE_MIN, int(round(VAL_ROUTE_FRAC * R))) | |
| val_idx = set(rng.choice(R, n_val, replace=False).tolist()) | |
| route_uid, route_src, route_wh, route_n, route_val, route_dongle = [], [], [], [], [], [] | |
| win_route, win_start, win_cx, win_cy, win_scale = [], [], [], [], [] | |
| conf_samples = [[] for _ in PART_SLICES] | |
| st = {"frames_total": 0, "frames_driver_det": 0, "frames_body_ok": 0, "frames_face_ok": 0, "frames_valid": 0, "windows_candidates": 0, | |
| "windows_valid": 0, "windows_dropped_gap": 0, "windows_dropped_invalid": 0, "routes_v2": 0} | |
| sw_after = [] # shoulder width after the AIDE-frame transform (px), sampled | |
| disp_after = [] # mean future displacement per window after transform (px), sampled (the HM motion score form) | |
| t_last = time.time() | |
| for ri, uid in enumerate(safe): | |
| path, src = route_path(uid) | |
| with np.load(path, allow_pickle=True) as z: | |
| t = z["t"].astype(np.float64) | |
| k = z["kpts"] | |
| box = z["box"].astype(np.float32) | |
| det = z["driver_det"].astype(bool) | |
| wh = z["native_wh"].astype(np.int32) | |
| F = len(t) | |
| k = fill_missing_joints(k.astype(np.float32)) | |
| conf = k[..., 2] | |
| fin = np.isfinite(k[:, [NOSE, NECK], 0]).all(1) & np.isfinite(box).all(1) | |
| body_ok = (np.nan_to_num(conf[:, PART_SLICES["body"][0]:PART_SLICES["body"][1]]) >= CONF_TH).mean(1) >= MIN_FRAC | |
| face_ok = (np.nan_to_num(conf[:, PART_SLICES["face"][0]:PART_SLICES["face"][1]]) >= CONF_TH).mean(1) >= MIN_FRAC | |
| sho_ok = (np.nan_to_num(conf[:, [LSHO, RSHO]]) >= CONF_TH).all(1) | |
| valid = det & fin & body_ok & face_ok & sho_ok | |
| st["frames_total"] += F | |
| st["frames_driver_det"] += int(det.sum()) | |
| st["frames_body_ok"] += int(body_ok.sum()) | |
| st["frames_face_ok"] += int(face_ok.sum()) | |
| st["frames_valid"] += int(valid.sum()) | |
| st["routes_v2"] += int(src == "v2") | |
| # windows: samples start + 3*j, j = 0..9 | |
| span = FRAME_STRIDE * (P.N_SAMPLE - 1) | |
| starts = np.arange(0, F - span, ANCHOR_STRIDE) | |
| st["windows_candidates"] += len(starts) | |
| if len(starts): | |
| idx = starts[:, None] + FRAME_STRIDE * np.arange(P.N_SAMPLE)[None, :] # [S,10] | |
| contiguous = np.abs((t[idx[:, -1]] - t[idx[:, 0]]) - 0.1 * span) < 0.02 | |
| allv = valid[idx].all(1) | |
| keep = contiguous & allv | |
| st["windows_dropped_gap"] += int((~contiguous).sum()) | |
| st["windows_dropped_invalid"] += int((contiguous & ~allv).sum()) | |
| st["windows_valid"] += int(keep.sum()) | |
| idx = idx[keep] | |
| if len(idx): | |
| kw = k[idx][:, :, [LSHO, RSHO], :2] * wh[None, None, None, :] # [S,10,2,2] shoulders in native px | |
| cen = np.median(kw.mean(2), axis=1) # window-median shoulder centre [S,2] | |
| swd = np.maximum(np.median(np.linalg.norm(kw[:, :, 0] - kw[:, :, 1], axis=-1), axis=1), 1.0) # window-median shoulder width [S] | |
| cx, cy = cen[:, 0], cen[:, 1] | |
| sc = box_t[2] / swd | |
| win_route.append(np.full(len(idx), ri, np.int32)) | |
| win_start.append(idx[:, 0].astype(np.int32)) | |
| win_cx.append(cx.astype(np.float32)); win_cy.append(cy.astype(np.float32)); win_scale.append(sc.astype(np.float32)) | |
| # sanity samples in the AIDE frame (<= 64 windows per route) | |
| sel = rng.choice(len(idx), min(64, len(idx)), replace=False) | |
| kk = k[idx[sel]] # [s,10,136,3] (filled, f32) | |
| xa = (kk[..., 0] * wh[0] - cx[sel, None, None]) * sc[sel, None, None] + box_t[0] | |
| ya = (kk[..., 1] * wh[1] - cy[sel, None, None]) * sc[sel, None, None] + box_t[1] | |
| sw_after.append(np.hypot(xa[:, :, 5] - xa[:, :, 6], ya[:, :, 5] - ya[:, :, 6]).ravel()) | |
| fut = np.stack([xa[:, 5:], ya[:, 5:]], -1) | |
| disp_after.append(np.sqrt(((fut[:, 1:] - fut[:, :-1]) ** 2).sum(-1)).mean((1, 2))) | |
| # confidence samples (valid frames) for the quantile maps | |
| vf = np.flatnonzero(valid) | |
| if len(vf): | |
| sf = rng.choice(vf, min(300, len(vf)), replace=False) | |
| for j, (lo, hi) in enumerate(PART_SLICES.values()): | |
| conf_samples[j].append(conf[sf, lo:hi].ravel()) | |
| route_uid.append(uid); route_src.append(src); route_wh.append(wh); route_n.append(F); route_val.append(ri in val_idx); route_dongle.append(dongle_of[uid]) | |
| if time.time() - t_last > 30: | |
| t_last = time.time() | |
| print(f" {ri + 1}/{R} routes, {st['windows_valid']:,} windows ({time.time() - t0:.0f}s)", flush=True) | |
| win_route, win_start = np.concatenate(win_route), np.concatenate(win_start) | |
| win_cx, win_cy, win_scale = np.concatenate(win_cx), np.concatenate(win_cy), np.concatenate(win_scale) | |
| M = len(win_route) | |
| st["windows_before_cap"] = int(M) | |
| if M > a.max_windows: | |
| keep = np.sort(rng.choice(M, a.max_windows, replace=False)) | |
| win_route, win_start, win_cx, win_cy, win_scale = win_route[keep], win_start[keep], win_cx[keep], win_cy[keep], win_scale[keep] | |
| M = len(win_route) | |
| route_val = np.array(route_val, bool) | |
| st.update(windows_final=int(M), windows_val_routes=int(route_val[win_route].sum()), routes=int(R), routes_val=int(route_val.sum()), | |
| devices=len(set(route_dongle)), frames_valid_frac=st["frames_valid"] / max(st["frames_total"], 1), | |
| windows_per_route_median=float(np.median(np.bincount(win_route, minlength=R)))) | |
| qs = np.linspace(0, 1, N_Q) | |
| q_baton = np.stack([np.quantile(np.nan_to_num(np.concatenate(cs)), qs) for cs in conf_samples]).astype(np.float32) | |
| sw_after, disp_after = np.concatenate(sw_after), np.concatenate(disp_after) | |
| # AIDE train references for the same sanity statistics | |
| c = np.load(P.CACHE_NPZ, allow_pickle=False) | |
| Str = c["S10"][c["split"] == "train"] | |
| sw_aide = np.hypot(Str[:, :, 5, 0] - Str[:, :, 6, 0], Str[:, :, 5, 1] - Str[:, :, 6, 1]).ravel() | |
| disp_aide = np.sqrt(((Str[:, 6:, :, :2] - Str[:, 5:-1, :, :2]) ** 2).sum(-1)).mean((1, 2)) | |
| st["sanity_aide_frame"] = { | |
| "shoulder_width_px_median": {"baton": float(np.median(sw_after)), "aide_train": float(np.median(sw_aide))}, | |
| "shoulder_width_px_p10_p90": {"baton": [float(np.percentile(sw_after, 10)), float(np.percentile(sw_after, 90))], | |
| "aide_train": [float(np.percentile(sw_aide, 10)), float(np.percentile(sw_aide, 90))]}, | |
| "future_disp_px_median": {"baton": float(np.median(disp_after)), "aide_train": float(np.median(disp_aide))}, | |
| "future_disp_px_p90": {"baton": float(np.percentile(disp_after, 90)), "aide_train": float(np.percentile(disp_aide, 90))}, | |
| "conf_mean_by_part": {k: {"baton_raw": float(np.nanmean(np.concatenate(conf_samples[j]))), "aide_train": float(q_aide[j].mean())} | |
| for j, k in enumerate(PART_SLICES)}} | |
| params = {"frame_stride": FRAME_STRIDE, "anchor_stride": ANCHOR_STRIDE, "step_s": 0.1 * FRAME_STRIDE, "window_span_s": 0.1 * FRAME_STRIDE * (P.N_SAMPLE - 1), | |
| "aide_step_s": [0.267, 0.333], "aide_span_s": 44 / 15, "conf_th": CONF_TH, "min_frac_body_face": MIN_FRAC, "headtop_proxy_k": HEADTOP_K, "max_windows": a.max_windows, | |
| "val_route_frac": VAL_ROUTE_FRAC, "val_route_min": VAL_ROUTE_MIN, "seed": SEED, "n_quantiles": N_Q, | |
| "similarity": "isotropic: scale = AIDE train median shoulder width / window-median shoulder width; window-median shoulder centre -> AIDE train median shoulder centre (373, 762 px)", "box_use": "driver-tracking validity flag only", | |
| "device_rule": "train-split routes whose dongle appears in no val/test split of either regime (s12 manifests); AIDE: train-split statistics only", | |
| "pose_dirs": {"v2": str(POSE_V2), "v1": str(POSE_V1)}, "aide_cache": str(P.CACHE_NPZ)} | |
| Path(a.out).parent.mkdir(parents=True, exist_ok=True) | |
| np.savez_compressed(a.out, route_uid=np.array(route_uid), route_src=np.array(route_src), route_native_wh=np.stack(route_wh), route_n_frames=np.array(route_n, np.int32), | |
| route_val=route_val, route_dongle=np.array(route_dongle), win_route=win_route, win_start=win_start, win_cx=win_cx, win_cy=win_cy, win_scale=win_scale, | |
| conf_q_baton=q_baton, conf_q_aide=q_aide, aide_target=box_t, frame_stride=np.array(FRAME_STRIDE), anchor_stride=np.array(ANCHOR_STRIDE), | |
| n_frames_total=np.array(st["frames_total"]), part_order=np.array(list(PART_SLICES)), params=np.array(json.dumps(params)), stats=np.array(json.dumps(st))) | |
| st["wall_s"] = round(time.time() - t0, 1) | |
| summ = {"params": params, "stats": st, "train_devices": len(train_dev), "eval_devices": len(eval_dev), "out": a.out} | |
| Path(a.out).with_name("baton_pose_corpus_summary.json").write_text(json.dumps(summ, indent=1)) | |
| write_run_json(Path(a.out).parent, "build_baton_pose_corpus", [Path(__file__), P.__file__], [P.CACHE_NPZ] + [S12R / f"targets_{r}_{s}.npz" for r in ("manual", "engaged") for s in ("train", "val", "test")], | |
| seed=SEED, t_start=t0, input_listing=False, extra={"summary": summ, "n_route_npz": R, "dms_fields_used": [], | |
| "note_inputs": "input digest covers the AIDE cache + the six s12 manifests; the R route npz are listed in the corpus (route_uid / route_src)"}) | |
| print(json.dumps(st, indent=1)) | |
| print(f"corpus -> {a.out} ({M:,} windows, {R} routes, {time.time() - t0:.0f}s)") | |
| print("BATON_POSE_CORPUS_DONE", flush=True) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 19.6 kB
- Xet hash:
- 202fc9c245e133de44ebceafe2ad4e88e96b5aa66f7a2168dca58b1975be1c19
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.