Download code/scripts/make_demo.py from changh95/transfusion-p150: direct link, hf CLI and curl.
- Browser
- Download file 8.78 kB
-
https://huggingface.co/changh95/transfusion-p150/resolve/main/code/scripts/make_demo.py
- Command line
-
hf download hf://changh95/transfusion-p150/code/scripts/make_demo.py
-
curl -L -o make_demo.py https://huggingface.co/changh95/transfusion-p150/resolve/main/code/scripts/make_demo.py
8.78 kB
| #!/usr/bin/env python3 | |
| # SPDX-License-Identifier: Apache-2.0 | |
| """Demo media of the card (``media/``): the p150 output, drawn over the stored fp32 CPU reference. | |
| # 1. on the device: run the p150 on inputs and keep the /predict bodies (one model load for all inputs) | |
| bin/devrun -t 1200 -- python code/scripts/make_demo.py run --json-dir <dir> INPUT [INPUT ...] | |
| # 2. no device: the p150 boxes over the CPU reference's boxes, bird's-eye view | |
| python code/scripts/make_demo.py compare --input code/tt_transfusion/samples/test_pcd.npz \\ | |
| --tt <dir>/test_pcd.tt.json --out media/transfusion_test_pcd_tt_vs_cpu.png | |
| # 3. no device: an animated bird's-eye view of a frame sequence (p150 boxes only) | |
| python code/scripts/make_demo.py gif --tt-dir <dir> --out media/<name>.gif INPUT [INPUT ...] | |
| ``run``: ``TransFusion.from_pretrained`` (ETH dispatch, 12x10, the traced graph; ``--weights-dir`` or the HF cache) | |
| runs every input through the Python API, which returns exactly what ``POST /predict`` returns, and writes | |
| ``<json-dir>/<stem>.tt.json`` (the body plus the device description). Inputs are the API's point inputs (``.npz`` / | |
| ``.npy`` / ``.pcd``); a plain ``(M, 5)`` ``.npy`` array is read as x, y, z, intensity, time_lag (the research and | |
| public-data layout). | |
| ``compare``: the reference is ``--reference``, else ``<stem>.reference.json`` next to the input. Matching is the rule | |
| of the accuracy gates (same label, BEV centre distance <= 0.5 m, ``ttaw.metrics.match_detections``); the caption | |
| states the result. Drawing: ``tt_transfusion.viz.render_bev`` (p150 boxes in colour, reference boxes as light outlines, | |
| unmatched boxes ringed and tagged). | |
| numpy + Pillow only (no device for ``compare`` / ``gif``). | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import sys | |
| from pathlib import Path | |
| from typing import Any, Dict, List, Optional | |
| import numpy as np | |
| HERE = Path(__file__).resolve().parent | |
| CODE = HERE.parent | |
| if str(CODE) not in sys.path: | |
| sys.path.insert(0, str(CODE)) | |
| FIELDS5 = ("x", "y", "z", "intensity", "time_lag") | |
| TITLE = "TransFusion-L (Autoware lidar_transfusion v2.1)" | |
| def load_cloud(path: Path): | |
| """The API input of a file; a plain (M, 5) float array is x, y, z, intensity, time_lag.""" | |
| from tt_transfusion import io as tio | |
| if path.suffix == ".npy": | |
| a = np.load(path, allow_pickle=False) | |
| if a.dtype.names is None and a.ndim == 2 and a.shape[1] == 5: | |
| return tio.PointCloud(np.ascontiguousarray(a, np.float32), FIELDS5) | |
| return str(path) | |
| def xyzt(path: Path) -> np.ndarray: | |
| """(N, 4) x, y, z, time_lag for drawing.""" | |
| from tt_transfusion import io as tio | |
| cloud = tio.load_points(load_cloud(path)) | |
| names = list(cloud.fields) | |
| xyz = cloud.select(("x", "y", "z")) | |
| lag = cloud.select(("time_lag",)) if "time_lag" in names else np.zeros((len(xyz), 1), np.float32) | |
| return np.concatenate([np.asarray(xyz, np.float32), np.asarray(lag, np.float32).reshape(-1, 1)], 1) | |
| def rows(dets: List[Dict[str, Any]]): | |
| boxes = np.array([d["center"] + d["size"] + [d["yaw"]] for d in dets], np.float32).reshape(-1, 7) | |
| return boxes, np.array([d["label_id"] for d in dets], np.int64), np.array([d["score"] for d in dets], np.float32) | |
| def agreement(tt: Dict[str, Any], ref: Dict[str, Any]) -> Dict[str, Any]: | |
| from tt_transfusion.ttaw.metrics import match_detections | |
| m = match_detections(*rows(tt["detections"]), *rows(ref["detections"]), max_dist=0.5) | |
| out = m.to_dict() | |
| out.update(dev=len(tt["detections"]), ref=len(ref["detections"]), | |
| max_score_err=float(max(m.score_errors, default=0.0))) | |
| return out | |
| def cmd_run(a) -> int: | |
| from tt_transfusion import TransFusion | |
| a.json_dir.mkdir(parents=True, exist_ok=True) | |
| kw = {"weights_dir": str(a.weights_dir)} if a.weights_dir else {} | |
| with TransFusion.from_pretrained(device_id=a.device_id, **kw) as model: | |
| device = model.info.get("device") | |
| for path in a.inputs: | |
| out = model(load_cloud(path)) | |
| body = out.to_dict() | |
| body["device"] = device | |
| (a.json_dir / f"{path.stem}.tt.json").write_text(json.dumps(body, indent=1, default=float)) | |
| print(f"{path.name}: {len(out)} detections, pillars {body['meta']['num_pillars']}, timing_ms " | |
| f"{ {k: round(v, 1) for k, v in out.timing_ms.items()} }", flush=True) | |
| return 0 | |
| def cmd_compare(a) -> int: | |
| from tt_transfusion.viz import render_bev | |
| body = json.loads(a.tt.read_text()) | |
| ref_path = a.reference or a.input.with_name(f"{a.input.stem}.reference.json") | |
| ref = json.loads(Path(ref_path).read_text()) | |
| ag = agreement(body, ref) | |
| dev = body.get("device") or {} | |
| where = f"dispatch {dev.get('dispatch', 'eth')}, grid {dev.get('grid', '12x10')}" | |
| meta = body["meta"] | |
| lines = [f"p150 ({where}) vs the fp32 CPU reference of the same network: {ag['matched']} of {ag['ref']} " | |
| f"reference detections matched (same label, centres <= 0.5 m), {ag['dev']} on the p150; max |score " | |
| f"difference| of matched pairs {ag['max_score_err']:.4f}.", | |
| f"Input: {a.describe or a.input.name}; {meta.get('num_points', 0):,} points, " | |
| f"{meta['num_pillars']:,} pillars. Forward up, rings every 10 m."] | |
| if a.lines: | |
| lines += a.lines | |
| pts = xyzt(a.input) | |
| img = render_bev(pts, body["detections"], reference=ref["detections"], | |
| title=f"{TITLE} on one Blackhole p150 vs the fp32 CPU reference", lines=lines, | |
| x_range=tuple(a.x_range) if a.x_range else None, y_range=tuple(a.y_range) if a.y_range else None, | |
| width=a.width) | |
| a.out.parent.mkdir(parents=True, exist_ok=True) | |
| img.save(a.out, optimize=True) | |
| print(f"{a.out}: {json.dumps(ag)}") | |
| return 0 | |
| def cmd_gif(a) -> int: | |
| from PIL import Image | |
| from tt_transfusion.viz import render_bev | |
| frames = [] | |
| for i, path in enumerate(a.inputs): | |
| body = json.loads((a.tt_dir / f"{path.stem}.tt.json").read_text()) | |
| label = (a.labels[i] if a.labels and i < len(a.labels) else path.stem).replace("_", " ") | |
| img = render_bev(xyzt(path), body["detections"], title=f"{TITLE} on one Blackhole p150: {label}", | |
| lines=[f"{len(body['detections'])} detections (p150 output); {a.caption}"], | |
| x_range=tuple(a.x_range), y_range=tuple(a.y_range), width=a.width) | |
| frames.append(img.convert("P", palette=Image.ADAPTIVE, colors=64)) | |
| a.out.parent.mkdir(parents=True, exist_ok=True) | |
| frames[0].save(a.out, save_all=True, append_images=frames[1:], duration=a.ms, loop=0, optimize=True) | |
| print(f"{a.out}: {len(frames)} frames, {a.out.stat().st_size / 1e6:.2f} MB") | |
| return 0 | |
| def main(argv: Optional[List[str]] = None) -> int: | |
| ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) | |
| sub = ap.add_subparsers(dest="cmd", required=True) | |
| r = sub.add_parser("run", help="run the p150 on inputs (device)") | |
| r.add_argument("inputs", nargs="+", type=Path) | |
| r.add_argument("--json-dir", type=Path, required=True) | |
| r.add_argument("--weights-dir", type=Path) | |
| r.add_argument("--device-id", type=int, default=0) | |
| c = sub.add_parser("compare", help="p150 over the CPU reference (no device)") | |
| c.add_argument("--input", type=Path, required=True) | |
| c.add_argument("--tt", type=Path, required=True, help="the p150 body (make_demo.py run)") | |
| c.add_argument("--reference", type=Path, help="default: <stem>.reference.json next to the input") | |
| c.add_argument("--out", type=Path, required=True) | |
| c.add_argument("--describe", help="input description for the caption") | |
| c.add_argument("--lines", nargs="*", help="extra caption lines") | |
| c.add_argument("--x-range", type=float, nargs=2) | |
| c.add_argument("--y-range", type=float, nargs=2) | |
| c.add_argument("--width", type=int, default=640) | |
| g = sub.add_parser("gif", help="animated bird's-eye view of p150 outputs (no device)") | |
| g.add_argument("inputs", nargs="+", type=Path) | |
| g.add_argument("--tt-dir", type=Path, required=True) | |
| g.add_argument("--out", type=Path, required=True) | |
| g.add_argument("--labels", nargs="*", help="one per input; \"_\" prints as a space") | |
| g.add_argument("--caption", default="") | |
| g.add_argument("--x-range", type=float, nargs=2, default=(-40.0, 60.0)) | |
| g.add_argument("--y-range", type=float, nargs=2, default=(-45.0, 45.0)) | |
| g.add_argument("--width", type=int, default=560) | |
| g.add_argument("--ms", type=int, default=100, help="frame duration") | |
| a = ap.parse_args(argv) | |
| return {"run": cmd_run, "compare": cmd_compare, "gif": cmd_gif}[a.cmd](a) | |
| if __name__ == "__main__": | |
| sys.exit(main()) | |