AsianPlayer's picture
download
raw
9.09 kB
"""Build the AIDE keypoint cache (results/cache/aide_kp.npz) from the official splits.
python aide/data/build_cache.py [--workers 8]
Per clip: parse annotation/<id>.json, select the driver per frame (protocol.select_driver,
protocol.DRIVER_RULE = first detection, locked by Gate 0), store
S [N,45,136,3] per-frame driver keypoints (long clips: uniform 45-index grid)
S10 [N,10,136,3] the 10 protocol frames = sample_idx(n_frames) on the ORIGINAL pose_list
(protocol.DRIVER_RULE / protocol.SAMPLE_GRID, locked by Gate 0)
S10_alt, bbox10_alt, frame_idx10_alt same with protocol.ALT_DRIVER_RULE / ALT_SAMPLE_GRID
bbox45/bbox10, frame_idx10 (integer frame ids from imgname -> <view>frames/<idx>.jpg),
ndet10, n_nodet, labels y_<task> (ints in protocol.CLASSES order), raw label strings,
clip_id, split, n_frames.
Counts and label histograms go to results/cache/cache_summary.json; a run json is written.
"""
from __future__ import annotations
import os
# ops rule (docs/decisions.md 2026-09-10, host CPU saturation): cap BLAS/OMP threads before numpy import
for _v in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS", "NUMEXPR_NUM_THREADS"):
os.environ.setdefault(_v, "4")
import argparse
import json
import os
import sys
import time
from collections import Counter
from multiprocessing import Pool
from pathlib import Path
import numpy as np
import pandas as pd
ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT))
from aide.data import protocol as P # noqa: E402
from aide.data.protocol import bbox_area # noqa: E402
from common.runmeta import write_run_json # noqa: E402
def read_split_ids(split: str) -> list[str]:
df = pd.read_csv(P.SPLIT_CSV[split], dtype=str)
ids = []
for a, b in zip(df.iloc[:, 0], df.iloc[:, 1]):
cid = a.strip().strip("/").split("/")[-1]
assert b.strip().endswith(f"/{cid}.json"), (a, b)
ids.append(cid)
assert len(set(ids)) == len(ids), f"duplicate ids in {split}"
return ids
def load_clip(cid: str, rule: str = P.DRIVER_RULE) -> dict:
d = json.load(open(P.AIDE_ANN_DIR / f"{cid}.json"))
pl = d["pose_list"]
names = [int(os.path.splitext(fr["imgname"])[0]) for fr in pl]
assert names == sorted(names), cid
n = len(pl)
kp, bb, ndet, filled = P.select_driver(pl, rule)
assert not np.isnan(kp).any(), f"{cid}: clip without any detection"
g45 = np.round(np.linspace(0, n - 1, P.N_FRAMES_STD)).astype(int)
g10 = P.sample_idx(n)
kp_alt, bb_alt, _, _ = P.select_driver(pl, P.ALT_DRIVER_RULE)
g10_alt = P.sample_idx(n, P.ALT_SAMPLE_GRID)
out = {
"S": kp[g45], "bbox45": bb[g45],
"S10": kp[g10], "bbox10": bb[g10],
"frame_idx10": np.asarray(names)[g10].astype(np.int32),
"S10_alt": kp_alt[g10_alt], "bbox10_alt": bb_alt[g10_alt],
"frame_idx10_alt": np.asarray(names)[g10_alt].astype(np.int32),
"ndet10": ndet[g10], "n_nodet": int(filled.sum()), "n_frames": n,
"n_multi": int((ndet > 1).sum()),
"frac_first_is_largest": float(np.mean([
(not fr["result"]) or bbox_area(fr["result"][0]["bbox"]) >= max(bbox_area(r["bbox"]) for r in fr["result"])
for fr in pl])),
}
for t in P.TASKS:
raw = d[P.TASK_KEY[t]]
out[f"raw_{t}"] = str(raw)
out[f"y_{t}"] = P.label_index(t, raw)
return out
def _work(args):
cid, split = args
o = load_clip(cid)
o["clip_id"], o["split"] = cid, split
return o
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--workers", type=int, default=8)
ap.add_argument("--out", default=str(P.CACHE_NPZ))
a = ap.parse_args()
t0 = time.time()
local = {p.name for p in P.AIDE_ROOT.iterdir() if p.is_dir() and p.name.isdigit()}
items, summary = [], {"splits": {}, "n_local_clip_dirs": len(local)}
for split in ("train", "val", "test"):
ids = read_split_ids(split)
missing = [i for i in ids if i not in local]
kept = [i for i in ids if i in local]
summary["splits"][split] = {"csv_rows": len(ids), "local": len(kept), "missing": missing}
assert len(kept) == P.EXPECTED_SPLIT_SIZES[split], (split, len(kept), missing[:5])
items += [(i, split) for i in kept]
all_ids = [i for i, _ in items]
assert len(set(all_ids)) == len(all_ids), "clip appears in more than one split"
summary["n_unused_local_clips"] = len(local - set(all_ids))
print({s: v["local"] for s, v in summary["splits"].items()}, "unused local:",
summary["n_unused_local_clips"], flush=True)
with Pool(a.workers) as pool:
recs = pool.map(_work, items, chunksize=16)
N = len(recs)
arr = {
"S": np.stack([r["S"] for r in recs]).astype(np.float32),
"S10": np.stack([r["S10"] for r in recs]).astype(np.float32),
"bbox45": np.stack([r["bbox45"] for r in recs]).astype(np.float32),
"bbox10": np.stack([r["bbox10"] for r in recs]).astype(np.float32),
"frame_idx10": np.stack([r["frame_idx10"] for r in recs]),
"S10_alt": np.stack([r["S10_alt"] for r in recs]).astype(np.float32),
"bbox10_alt": np.stack([r["bbox10_alt"] for r in recs]).astype(np.float32),
"frame_idx10_alt": np.stack([r["frame_idx10_alt"] for r in recs]),
"frac_first_is_largest": np.array([r["frac_first_is_largest"] for r in recs], dtype=np.float32),
"driver_rule": np.array(P.DRIVER_RULE), "sample_grid": np.array(P.SAMPLE_GRID),
"alt_driver_rule": np.array(P.ALT_DRIVER_RULE), "alt_sample_grid": np.array(P.ALT_SAMPLE_GRID),
"ndet10": np.stack([r["ndet10"] for r in recs]),
"n_nodet": np.array([r["n_nodet"] for r in recs], dtype=np.int32),
"n_multi": np.array([r["n_multi"] for r in recs], dtype=np.int32),
"n_frames": np.array([r["n_frames"] for r in recs], dtype=np.int32),
"clip_id": np.array([r["clip_id"] for r in recs]),
"split": np.array([r["split"] for r in recs]),
}
for t in P.TASKS:
arr[f"y_{t}"] = np.array([r[f"y_{t}"] for r in recs], dtype=np.int64)
arr[f"raw_{t}"] = np.array([r[f"raw_{t}"] for r in recs])
assert arr["S"].shape == (N, P.N_FRAMES_STD, P.J, 3) and arr["S10"].shape == (N, P.N_SAMPLE, P.J, 3)
assert np.isfinite(arr["S"]).all() and np.isfinite(arr["S10"]).all()
out = Path(a.out)
out.parent.mkdir(parents=True, exist_ok=True)
np.savez(out, **arr)
# ---- summary
summary["N"] = N
summary["n_long_clips"] = int((arr["n_frames"] != P.N_FRAMES_STD).sum())
summary["long_clip_lengths"] = {cid: int(n) for cid, n in zip(arr["clip_id"], arr["n_frames"])
if n != P.N_FRAMES_STD}
summary["frames_without_detection_total"] = int(arr["n_nodet"].sum())
summary["clips_with_missing_frames"] = int((arr["n_nodet"] > 0).sum())
summary["frac_frames_multi_detection"] = float(arr["n_multi"].sum() / arr["n_frames"].sum())
summary["frac_sampled_frames_multi_detection"] = float((arr["ndet10"] > 1).mean())
summary["conf_lt_0.05_fraction"] = float((arr["S10"][..., 2] < 0.05).mean())
summary["label_hist"] = {}
summary["raw_label_variants"] = {}
for t in P.TASKS:
summary["raw_label_variants"][t] = dict(Counter(arr[f"raw_{t}"].tolist()))
summary["label_hist"][t] = {}
for split in ("train", "val", "test", "all"):
m = np.ones(N, bool) if split == "all" else arr["split"] == split
c = np.bincount(arr[f"y_{t}"][m], minlength=len(P.CLASSES[t]))
summary["label_hist"][t][split] = {P.CLASSES[t][k]: int(c[k]) for k in range(len(c))}
summary["cache"] = str(out)
summary["driver_rule"] = P.DRIVER_RULE
summary["sample_grid"] = P.SAMPLE_GRID
summary["sample_idx_45"] = P.sample_idx(45).tolist()
summary["alt_driver_rule"], summary["alt_sample_grid"] = P.ALT_DRIVER_RULE, P.ALT_SAMPLE_GRID
summary["alt_sample_idx_45"] = P.sample_idx(45, P.ALT_SAMPLE_GRID).tolist()
summary["frac_frames_first_det_is_largest"] = float(
(arr["frac_first_is_largest"] * arr["n_frames"]).sum() / arr["n_frames"].sum())
summary["frac_S10_equal_S10_alt_xy"] = float(np.all(arr["S10"][..., :2] == arr["S10_alt"][..., :2], axis=(1, 2, 3)).mean())
summary["wall_s"] = round(time.time() - t0, 1)
(out.parent / "cache_summary.json").write_text(json.dumps(summary, indent=1))
print(json.dumps({k: v for k, v in summary.items() if k not in ("label_hist", "raw_label_variants",
"long_clip_lengths")}, indent=1))
print("label_hist test:", json.dumps({t: summary["label_hist"][t]["test"] for t in P.TASKS}))
write_run_json(out.parent, "build_cache", [__file__, P.__file__],
[P.SPLIT_CSV[s] for s in ("train", "val", "test")], seed=None, t_start=t0,
extra={"N": N, "splits": {s: v["local"] for s, v in summary["splits"].items()}},
input_listing=False)
if __name__ == "__main__":
main()

Xet Storage Details

Size:
9.09 kB
·
Xet hash:
04c69b9ca807e406ed7ca885a240b928a7313f5f2b69f92f781868c9e8dccbe4

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.