Buckets:
| """cwm3/baselines_dms/extractors/base.py — Extractor ABC + the resumable per-route runner that writes the feature cache (plan §3, 2026-09-14). | |
| Extractor contract | |
| setup(registry_entry) load weights (registry.resolve() output: {id, local_path, sha256, version, ...}); may be a no-op | |
| infer(crops) -> f32[n, C] one row per input crop (u8 BGR arrays; None entries allowed = no crop) — a row is all-NaN when the model | |
| found nothing (fail 3); any other non-finite value marks the row fail 4; rows of None inputs must be NaN | |
| channels: list[str] C channel names (registry vocabulary); model_ids: list[str]; weight_sha256: {name: sha}; lib_versions: {lib: ver} | |
| needs_crop (True) True = infer receives the driver crop (video_io.driver_crop, +10 %); False = the resized full frame | |
| decode_wh ((640, 400)) decode size handed to video_io.iter_frames (None = native) | |
| batch (64) crops per infer() call | |
| Runner (run_route) decodes the route's spans ONCE (video_io.FrameStream, seek+grab workers), crops, infers in batches and writes | |
| <out_dir>/<uid>.npz in the cache_schema.RouteCache contract (fail codes 0 ok / 1 no box / 2 no frame / 3 no face / | |
| 4 non-finite; feat NaN when invalid; box native px; moving from the pose cache; meta incl. code_sha256 of the | |
| extractor module + base + schema, native_fps (header fps) + seek_corrected + decode_stats of the index-exact decoder | |
| (video_io.FrameStream.decode_meta, 2026-09-14); runtime p50 / p95 per frame). | |
| Resumability (pattern of data/run_extract_pose_causal.sh / extract_pose_causal_routes.py) | |
| * a cache that exists and whose spans cover the request is skipped (code 0, status "exists"); | |
| * a cache that covers only part of the request (e.g. a previous chunk stopped on the time budget -> meta.partial) is EXTENDED: only the | |
| uncovered spans are decoded and merged (RouteCache.merged_with); | |
| * the time budget (--max-minutes, measured from the runner's t_start) is checked after every decode block; when it is hit the frames done | |
| so far are written as a partial cache (spans truncated at the last decoded frame, meta.partial = True) and the runner returns code 3; | |
| * run_routes / the CLI exit 3 = budget reached with work remaining (loop again), 0 = everything requested is done, 4 = a route failed. | |
| No concrete model is implemented here (DummyExtractor = deterministic gray-level statistics for tests / smoke runs); no openpilot field is read. | |
| python baselines_dms/extractors/base.py --self-test [--uid UID] # DummyExtractor on 1 local TRAIN route: write, skip, resume, budget exit | |
| python baselines_dms/extractors/base.py --baseline dummy --uids <file> --out-dir <dir> [--max-minutes 20 --workers 6] | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import sys | |
| import time | |
| from abc import ABC, abstractmethod | |
| from pathlib import Path | |
| import numpy as np | |
| HERE = Path(__file__).resolve().parent | |
| PKG = HERE.parent | |
| CWM3 = PKG.parent | |
| for p in (str(PKG), str(HERE)): | |
| if p not in sys.path: | |
| sys.path.insert(0, p) | |
| import cache_schema as cs # noqa: E402 | |
| import video_io as vio # noqa: E402 | |
| EXIT_DONE, EXIT_BUDGET, EXIT_FAIL = 0, 3, 4 | |
| class Extractor(ABC): | |
| baseline: str = "abstract" | |
| channels: list = [] | |
| model_ids: list = [] | |
| weight_sha256: dict = {} | |
| lib_versions: dict = {} | |
| needs_crop: bool = True | |
| decode_wh = vio.DEFAULT_DECODE_WH | |
| batch: int = 64 | |
| def setup(self, registry_entry=None): | |
| """default: remember the registry entry; concrete extractors load their weights here (after registry.resolve verified the sha256).""" | |
| self.registry_entry = registry_entry | |
| return self | |
| def infer(self, crops) -> np.ndarray: | |
| """list of u8 BGR crops (None = no crop) -> f32[n, C]; NaN rows for failures.""" | |
| def teardown(self): | |
| pass | |
| def C(self): | |
| return len(self.channels) | |
| def code_files(self): | |
| return [Path(sys.modules[type(self).__module__].__file__), HERE / "base.py", PKG / "cache_schema.py", PKG / "video_io.py"] | |
| def describe(self): | |
| return {"baseline": self.baseline, "channels": list(self.channels), "model_ids": list(self.model_ids), "weight_sha256": dict(self.weight_sha256), | |
| "lib_versions": dict(self.lib_versions), "needs_crop": bool(self.needs_crop), "decode_wh": (None if self.decode_wh is None else list(self.decode_wh)), "batch": int(self.batch)} | |
| class DummyExtractor(Extractor): | |
| """tests / smoke only: deterministic gray-level statistics of the driver crop (no model, no weights). NaN row when the crop is None / empty.""" | |
| baseline = "dummy" | |
| channels = ["dummy_mean_gray", "dummy_std_gray", "dummy_aspect"] | |
| model_ids = ["dummy gray statistics (no model)"] | |
| lib_versions = {"numpy": np.__version__} | |
| def infer(self, crops): | |
| out = np.full((len(crops), self.C), np.nan, np.float32) | |
| for i, c in enumerate(crops): | |
| if c is None or c.size == 0: | |
| continue | |
| g = c.astype(np.float32).mean(-1) | |
| out[i] = [g.mean() / 255.0, g.std() / 255.0, c.shape[1] / max(c.shape[0], 1)] | |
| return out | |
| EXTRACTORS = {"dummy": DummyExtractor} | |
| def fail_codes_for(rows: np.ndarray, had_frame: np.ndarray, had_box: np.ndarray, needs_crop: bool) -> np.ndarray: | |
| """per-row fail code from the infer output and the availability flags (precedence: 2 no frame > 1 no box > 3 no face > 4 non-finite > 0).""" | |
| fin = np.isfinite(rows) | |
| code = np.zeros(len(rows), np.uint8) | |
| code[~fin.all(1)] = 4 | |
| code[~fin.any(1)] = 3 | |
| if needs_crop: | |
| code[~had_box] = 1 | |
| code[~had_frame] = 2 | |
| return code | |
| def run_route(extractor: Extractor, uid: str, spans, out_dir, max_minutes=20.0, workers=6, hz=cs.HZ, block=300, force=False, log=print, t_start=None, vpath=None): | |
| """-> (code, path, info). code 0 = cache complete for `spans` (written now or already existing), 3 = time budget hit (partial cache written), 4 = failed.""" | |
| t_start = t_start or time.time() | |
| t_route = time.time() | |
| out_dir = Path(out_dir) | |
| path = out_dir / f"{uid}.npz" | |
| spans = cs.merge_spans(spans) | |
| info = {"uid": uid, "status": "", "n_frames": 0, "wall_s": 0.0} | |
| have = None | |
| if path.exists() and not force: | |
| try: | |
| have = cs.RouteCache.load(path) | |
| if have.channels != list(extractor.channels): | |
| log(f"{uid[:50]}: existing cache has channels {have.channels} != {list(extractor.channels)}; recomputing") | |
| have = None | |
| elif have.covers(spans): | |
| info.update(status="exists", n_frames=have.n, wall_s=round(time.time() - t_route, 2)) | |
| return EXIT_DONE, path, info | |
| except Exception as e: # unreadable / foreign file: rebuild | |
| log(f"{uid[:50]}: cannot read existing cache ({repr(e)[:80]}); recomputing") | |
| have = None | |
| todo = spans if have is None else cs.spans_subtract(spans, have.spans) | |
| if len(todo) == 0: | |
| info.update(status="exists", n_frames=have.n) | |
| return EXIT_DONE, path, info | |
| try: | |
| fs = vio.FrameStream(uid, todo, hz=hz, workers=workers, decode_wh=extractor.decode_wh, block=block, vpath=vpath, box_coords="frame") | |
| except FileNotFoundError as e: | |
| info.update(status="no_video", error=str(e)) | |
| return EXIT_FAIL, path, info | |
| N = len(fs.frame_idx) | |
| feat = np.full((N, extractor.C), np.nan, np.float32) | |
| had_frame = np.zeros(N, bool) | |
| had_box = np.isfinite(fs.box).all(1) | |
| infer_ms = [] | |
| pend_idx, pend_crops = [], [] | |
| n_done = 0 | |
| budget_hit = False | |
| def flush(): | |
| if not pend_idx: | |
| return | |
| t0 = time.time() | |
| rows = np.asarray(extractor.infer(pend_crops), np.float32).reshape(len(pend_crops), -1) | |
| assert rows.shape[1] == extractor.C, (rows.shape, extractor.C) | |
| feat[np.asarray(pend_idx)] = rows | |
| infer_ms.append(1000.0 * (time.time() - t0) / len(pend_idx)) | |
| pend_idx.clear() | |
| pend_crops.clear() | |
| it = iter(fs) | |
| i = 0 | |
| block_end = 0 | |
| for t_, fi_, fr_, bx_ in it: | |
| if fr_ is not None: | |
| had_frame[i] = True | |
| crop = vio.driver_crop(fr_, bx_) if extractor.needs_crop else fr_ | |
| if crop is not None: | |
| pend_idx.append(i) | |
| pend_crops.append(crop) | |
| if len(pend_idx) >= extractor.batch: | |
| flush() | |
| i += 1 | |
| n_done = i | |
| if i >= block_end: # block boundary: budget check (the stream itself decodes ahead by <= 2 x workers blocks) | |
| block_end = i + block | |
| if time.time() - t_start > max_minutes * 60 and i < N: | |
| budget_hit = True | |
| break | |
| flush() | |
| if budget_hit: | |
| it.close() # stop the generator (shuts the worker pool down) | |
| n = n_done | |
| fail = fail_codes_for(feat[:n], had_frame[:n], had_box[:n], extractor.needs_crop) | |
| valid = fail == 0 | |
| feat_n = feat[:n].copy() | |
| feat_n[~valid] = np.nan | |
| if budget_hit: | |
| t_last = float(fs.t[n - 1]) if n else float(todo[0, 0]) | |
| done_spans = np.array([[a, min(b, t_last)] for a, b in todo if a <= t_last], np.float32).reshape(-1, 2) | |
| else: | |
| done_spans = todo | |
| code_sha = cs.code_sha256(extractor.code_files()) | |
| meta = cs.make_meta(extractor.baseline, extractor.model_ids, extractor.weight_sha256, extractor.lib_versions, fs.decode_wh, fs.native_wh, code_sha, "local", | |
| pose_source=(fs.boxes or {}).get("source", "none"), needs_crop=bool(extractor.needs_crop), crop_expand=vio.MARGIN, hz=hz, partial=bool(budget_hit), | |
| video_path=str(fs.vpath), n_video_frames=int(fs.n_frames), registry_entry=getattr(extractor, "registry_entry", None), **fs.decode_meta()) | |
| runtime = cs.make_runtime(fs.stats["decode_ms_per_frame"], infer_ms, n, wall_s=time.time() - t_route) | |
| new = cs.RouteCache(uid, fs.t[:n].astype(np.float32), fs.frame_idx[:n].astype(np.int64), fs.box[:n].astype(np.float32), valid, fail, feat_n.astype(np.float16), | |
| list(extractor.channels), cs.merge_spans(done_spans), meta, runtime, moving=fs.moving[:n].astype(bool)) | |
| if have is not None and have.n: | |
| new = have.merged_with(new) | |
| new.meta["partial"] = bool(budget_hit) or not new.covers(spans) | |
| new.meta["resumed_from_partial"] = True | |
| new.save(path) | |
| fail_hist = {cs.FAIL_CODES[int(k)]: int(v) for k, v in zip(*np.unique(new.fail, return_counts=True))} | |
| info.update(status=("budget" if budget_hit else "done"), n_frames=int(new.n), n_new_frames=int(n), valid_rate=float(new.valid.mean()) if new.n else None, fail_hist=fail_hist, | |
| decode_worker_fps=round(fs.stats["n_grab"] / max(fs.stats["decode_worker_s"], 1e-9), 1), runtime=runtime, wall_s=round(time.time() - t_route, 2), spans_done=new.spans.tolist()) | |
| return (EXIT_BUDGET if budget_hit else EXIT_DONE), path, info | |
| def run_routes(extractor: Extractor, uids, spans_by_uid, out_dir, max_minutes=20.0, workers=6, log=print, **kw): | |
| """sequential route loop with one time budget; -> (exit_code, summary). exit 3 = budget hit / work remaining, 4 = some route failed (others done), 0 = all done.""" | |
| t_start = time.time() | |
| recs, fails = [], [] | |
| budget_hit = False | |
| for u in uids: | |
| if time.time() - t_start > max_minutes * 60: | |
| budget_hit = True | |
| break | |
| code, path, info = run_route(extractor, u, spans_by_uid[u], out_dir, max_minutes, workers, log=log, t_start=t_start, **kw) | |
| log(f"{u[:60]} {info['status']} frames {info.get('n_frames')} valid {info.get('valid_rate')} wall {info.get('wall_s')}s") | |
| if code == EXIT_FAIL: | |
| fails.append(info) | |
| else: | |
| recs.append(info) | |
| if code == EXIT_BUDGET: | |
| budget_hit = True | |
| break | |
| n_left = len(uids) - len(recs) - len(fails) | |
| summ = {"n_requested": len(uids), "n_done": len(recs), "n_failed": len(fails), "n_remaining": max(0, n_left) + int(bool(budget_hit and recs and recs[-1]["status"] == "budget")), | |
| "budget_hit": budget_hit, "wall_s": round(time.time() - t_start, 1), "fails": fails[:20], "baseline": extractor.baseline} | |
| code = EXIT_BUDGET if (budget_hit or summ["n_remaining"] > 0) else (EXIT_FAIL if fails else EXIT_DONE) | |
| log(("BUDGET_EXIT " if code == EXIT_BUDGET else "ALL_DONE ") + json.dumps({k: summ[k] for k in ("n_done", "n_failed", "n_remaining", "wall_s")})) | |
| return code, summ | |
| def _self_test(uid=None, workers=2): | |
| import tempfile | |
| ex = DummyExtractor().setup(None) | |
| if uid is None: | |
| cands = [(vio.video_path(u).stat().st_size, u) for u in vio.anchors_by_uid(("manual_train", "engaged_train")) if vio.has_video(u) and vio.pose_path(u) is not None] | |
| if not cands: | |
| print(json.dumps({"self_test": "extractors.base", "route": "none local; skipped"})) | |
| return | |
| uid = min(cands)[1] | |
| ts = vio.anchors_of(uid) | |
| W, H, n = vio.video_info(vio.video_path(uid)) | |
| sp2 = vio.spans_from_anchors(ts[:2], dur=n / vio.FPS) | |
| sp4 = vio.spans_from_anchors(ts[:6], dur=n / vio.FPS) | |
| rep = {"self_test": "extractors.base", "route": uid, "steps": []} | |
| with tempfile.TemporaryDirectory() as d: | |
| c1, p1, i1 = run_route(ex, uid, sp2, d, max_minutes=10, workers=workers, block=20) | |
| rc = cs.RouteCache.load(p1) | |
| assert c1 == 0 and i1["status"] == "done" and rc.covers(sp2) and rc.C == 3 and rc.moving is not None and rc.meta["partial"] is False | |
| rep["steps"].append({"write": i1["status"], "n_frames": rc.n, "fail_hist": i1["fail_hist"], "runtime": rc.runtime}) | |
| c2, _, i2 = run_route(ex, uid, sp2, d, max_minutes=10, workers=workers) | |
| assert c2 == 0 and i2["status"] == "exists" | |
| rep["steps"].append({"rerun": i2["status"]}) | |
| # budget exit: a tiny budget -> partial cache, code 3 | |
| c3, p3, i3 = run_route(ex, uid, sp4, d, max_minutes=1e-9, workers=workers, block=10) | |
| rc3 = cs.RouteCache.load(p3) | |
| assert c3 == 3 and i3["status"] == "budget" and rc3.meta["partial"] is True and not rc3.covers(sp4) and rc3.covers(sp2), (c3, i3["status"], rc3.spans, sp4) | |
| rep["steps"].append({"budget": i3["status"], "code": c3, "n_frames": rc3.n, "spans_done": rc3.spans.tolist()}) | |
| # resume: only the uncovered part is decoded; the merged cache covers everything and is contiguous on the grid | |
| c4, p4, i4 = run_route(ex, uid, sp4, d, max_minutes=10, workers=workers, block=20) | |
| rc4 = cs.RouteCache.load(p4) | |
| k_exp = vio.frames_in_spans(sp4, cs.HZ, n)[0] | |
| assert c4 == 0 and rc4.covers(sp4) and rc4.meta["partial"] is False and np.array_equal(np.round(rc4.t.astype(np.float64) * 10).astype(np.int64), k_exp), (c4, rc4.n, len(k_exp)) | |
| assert i4["n_new_frames"] < rc4.n | |
| rep["steps"].append({"resume": i4["status"], "n_new_frames": i4["n_new_frames"], "n_total": rc4.n}) | |
| # run_routes exit codes | |
| code, summ = run_routes(ex, [uid], {uid: sp4}, d, max_minutes=10, workers=workers, log=lambda m: None) | |
| assert code == 0 and summ["n_done"] == 1 | |
| code_b, summ_b = run_routes(ex, [uid, uid], {uid: sp4}, d, max_minutes=-1, workers=workers, log=lambda m: None) | |
| assert code_b == 3 and summ_b["budget_hit"] | |
| rep["steps"].append({"run_routes": [code, code_b]}) | |
| rep["describe"] = ex.describe() | |
| print(json.dumps(rep, indent=1, default=str)) | |
| def main(): | |
| import argparse | |
| ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) | |
| ap.add_argument("--self-test", action="store_true") | |
| ap.add_argument("--uid", default="") | |
| ap.add_argument("--uids", default="", help="file with route uids (one per line)") | |
| ap.add_argument("--baseline", default="dummy", choices=sorted(EXTRACTORS)) | |
| ap.add_argument("--out-dir", default="") | |
| ap.add_argument("--max-minutes", type=float, default=20.0) | |
| ap.add_argument("--workers", type=int, default=6) | |
| ap.add_argument("--hist-s", type=float, default=vio.HIST_S) | |
| ap.add_argument("--post-s", type=float, default=vio.POST_S) | |
| ap.add_argument("--packs", default=",".join(vio.ALL_PACKS), help="s12 packs whose anchors define the spans") | |
| a = ap.parse_args() | |
| if a.self_test or not (a.uids or a.uid): | |
| _self_test(a.uid or None, min(a.workers, 2)) | |
| return 0 | |
| ex = EXTRACTORS[a.baseline]().setup(None) | |
| uids = [a.uid] if a.uid else [l.strip() for l in open(a.uids) if l.strip()] | |
| out_dir = Path(a.out_dir) if a.out_dir else cs.cache_dir(ex.baseline) | |
| packs = tuple(a.packs.split(",")) | |
| spans = {u: vio.anchor_spans(u, packs, a.hist_s, a.post_s) for u in uids} | |
| code, summ = run_routes(ex, uids, spans, out_dir, a.max_minutes, a.workers) | |
| (out_dir / "logs").mkdir(parents=True, exist_ok=True) | |
| (out_dir / "logs" / f"run_{time.strftime('%Y%m%dT%H%M%SZ', time.gmtime())}.json").write_text(json.dumps(summ, indent=1, default=str)) | |
| return code | |
| if __name__ == "__main__": | |
| sys.exit(main()) | |
Xet Storage Details
- Size:
- 17.5 kB
- Xet hash:
- aae01d8eac2d44447cb3974846ea622c313dcd8ab39c5b0a3ac0582cddbc0fcf
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.