File size: 6,395 Bytes
ac1609a | 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 | #!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Run a public-dataset sequence through the Python API (the device graph) and score it against the CPU goldens.
bin/devrun -t 1800 -- python code/scripts/run_public_tt.py --scene ps019 [--frames 0:80] [--out DIR]
Sequences (``research/streampetr/public_data/inputs/<scene>``, built by the research agents; never shipped):
``ps019`` (PandaSet 019 frames 0-79, CC BY 4.0), ``ns0103`` / ``ns0916`` (nuScenes v1.0-mini scenes 0103 / 0916,
every CAM_FRONT image; CC BY-NC-SA 4.0: internal validation only). One stream per run, reset at the first frame:
images (paths), inline calibration (K / P of the image as sent, ``T_ego_cam``), the per-camera stamps (the node's
0.15 s sync check) and the anchor stamp + ego pose (``ego_motion.npz`` ``T_map_base``).
Writes ``<out>/<scene>/detections.json`` in the research format (``{"meta", "results": [{"frame", "dets": [...]}]}``,
the input of ``research/streampetr/public_data/scripts/compare_runs.py det`` and ``eval_nuscenes_ap.py
--det-root``) and ``<out>/<scene>/f1.json``: per-frame detection F1 against the research ORT goldens (1 m, same label;
the free-running gate of PLAN 2.9 is the mean over >= 30 frames).
"""
from __future__ import annotations
import argparse
import json
import os
import time
from pathlib import Path
import numpy as np
ROOT = Path(os.environ.get("TT_MODELS_ROOT", "/home/ubuntu/experiments/tt-models"))
PUBLIC = ROOT / "research" / "streampetr" / "public_data"
SLOTS = ("CAM_FRONT", "CAM_FRONT_LEFT", "CAM_BACK_LEFT", "CAM_FRONT_RIGHT", "CAM_BACK_RIGHT")
def requests(scene: str, frames=None):
"""Yield ``(frame, request kwargs)`` for the API (module docstring)."""
src = PUBLIC / "inputs" / scene
man = json.loads((src / "manifest.json").read_text())
em = np.load(src / "ego_motion.npz")
slots = man["slots"]
h, w = man["img_hw"]
cams = {}
for s in SLOTS:
c = man["cams"][slots[s]]
cams[s] = {"intrinsics": c["K"], "P": c["P"], "T_ref_from_camera": c["T_ego_cam"]}
calib = {"frame_id": "base_link", "image_size": [w, h], "cameras": cams}
a, b = frames or (0, len(man["frames"]))
for k in range(a, b):
fr = man["frames"][k]
images = [{"camera": s, "path": str((src / fr["imgs"][slots[s]]).resolve()),
"timestamp_s": float(fr["stamps"][s])} for s in SLOTS]
stream = {"id": f"public-{scene}", "timestamp_s": float(fr["anchor_stamp"]),
"T_world_from_ego": np.asarray(em["T_map_base"][k], np.float64).tolist()}
if k == a:
stream["reset"] = True
yield int(fr["frame"]), {"images": images, "calibration": calib, "stream": stream}
def research_dets(out) -> list:
"""API output -> the research ``dets`` rows (``index`` / ``layer`` / ``query`` from ``meta``)."""
body = out.to_dict()
meta = body.get("meta") or {}
layers, queries = meta.get("decoder_layer") or [], meta.get("query") or []
rows = []
for i, d in enumerate(body["detections"]):
layer = int(layers[i]) if i < len(layers) else -1
query = int(queries[i]) if i < len(queries) else -1
cx, cy, cz = d["center"]
sx, sy, sz = d["size"]
vx, vy = (d.get("velocity") or [0.0, 0.0])[:2]
rows.append({"index": layer * 900 + query if layer >= 0 else -1, "layer": layer, "query": query,
"label": int(d["label_id"]), "score": float(d["score"]), "x": cx, "y": cy, "z": cz,
"dim_x": sx, "dim_y": sy, "dim_z": sz, "yaw": float(d["yaw"]), "vx": vx, "vy": vy})
return rows
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--scene", default="ps019", choices=["ps019", "ns0103", "ns0916"])
ap.add_argument("--frames", default=None, help="a:b (default: all)")
ap.add_argument("--out", default=str(ROOT / "logs" / "streampetr" / "public"))
ap.add_argument("--weights-dir", default=None)
a = ap.parse_args()
from tt_streampetr import StreamPETR
from tt_streampetr.host.agreement import detection_f1
frames = tuple(int(x) for x in a.frames.split(":")) if a.frames else None
out_dir = Path(a.out) / a.scene
out_dir.mkdir(parents=True, exist_ok=True)
gold_path = PUBLIC / "golden" / a.scene / "detections.json"
gold = {int(r["frame"]): r["dets"] for r in json.loads(gold_path.read_text())["results"]} \
if gold_path.is_file() else {}
results, f1 = [], []
with StreamPETR.from_pretrained(weights_dir=a.weights_dir) as model:
info = model.info
for f, req in requests(a.scene, frames):
t0 = time.perf_counter()
out = model(**req)
dets = research_dets(out)
results.append({"frame": f, "num_final": len(dets), "timing_ms": dict(out.timing_ms),
"wall_ms": (time.perf_counter() - t0) * 1e3, "dets": dets})
if f in gold:
m = detection_f1(dets, gold[f], dist=1.0)
f1.append({"frame": f, **m})
print(f"{a.scene} frame {f}: {len(dets)} objects, F1 {m['f1']:.3f} (tp {m['tp']}/{m['n_ref']}), "
f"trace {out.timing_ms.get('trace', float('nan')):.1f} ms, total {out.timing_ms['total']:.0f} ms",
flush=True)
meta = {"scene": a.scene, "frames": frames, "model": info.get("model"), "device": info.get("device"),
"weights": info.get("weights"), "knobs": info.get("knobs"),
"pipeline": "tt_streampetr API (device graph: backbone + head + memory update in one trace)"}
(out_dir / "detections.json").write_text(json.dumps({"meta": meta, "results": results}, indent=1, default=float))
summary = {"frames": len(f1), "mean_f1": float(np.mean([x["f1"] for x in f1])) if f1 else None,
"min_f1": float(np.min([x["f1"] for x in f1])) if f1 else None,
"mean_f1_from_frame1": float(np.mean([x["f1"] for x in f1[1:]])) if len(f1) > 1 else None,
"trace_ms_p50": float(np.median([r["timing_ms"].get("trace", np.nan) for r in results])),
"total_ms_p50": float(np.median([r["timing_ms"]["total"] for r in results]))}
(out_dir / "f1.json").write_text(json.dumps({"summary": summary, "frames": f1}, indent=1, default=float))
print(json.dumps(summary, indent=1))
if __name__ == "__main__":
main()
|