File size: 2,568 Bytes
68a3286 | 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 | #!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Quickstart: the model card's Python snippet on the shipped sample, on one Blackhole p150.
pip install -e . # once, from the repo root, on top of an environment that has ttnn (tt-metal)
python examples/quickstart.py [input] [--out-dir examples/output]
Writes <out-dir>/quickstart.json (the same JSON as POST /predict) and <out-dir>/quickstart_bev.png (a bird's-eye view
of the input points and the detections, tt_transfusion.viz). The default input is the shipped sample
code/tt_transfusion/samples/test_pcd.npz (Autoware's test.pcd, Apache-2.0), found relative to this file, so the script
runs from any directory; an input given on the command line is relative to the current directory.
"""
import argparse
import json
from pathlib import Path
REPO = Path(__file__).resolve().parents[1]
ap = argparse.ArgumentParser()
ap.add_argument("input", nargs="?", default=str(REPO / "code" / "tt_transfusion" / "samples" / "test_pcd.npz"))
ap.add_argument("--out-dir", default=str(REPO / "examples" / "output"))
ap.add_argument("--device-id", type=int, default=0)
args = ap.parse_args()
out_dir = Path(args.out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
# --- the model card snippet --------------------------------------------------------------------------------------
from tt_transfusion import TransFusion
with TransFusion.from_pretrained(device_id=args.device_id) as model: # weights -> HF cache, trace captured
out = model(args.input) # path, (N, C) array, .npy/.npz/.pcd or PointCloud; a time_lag column = pre-densified
for d in out.to_dicts()[:10]:
print(d["label"], round(d["score"], 3), d["center"], d["size"], round(d["yaw"], 3))
# ------------------------------------------------------------------------------------------------------------------
(out_dir / "quickstart.json").write_text(json.dumps(out.to_dict(), indent=1))
from tt_transfusion.viz import load_xyz, render_bev # noqa: E402
timing = {k: round(v, 1) for k, v in out.timing_ms.items()}
img = render_bev(load_xyz(args.input), out.to_dicts(),
title=f"TransFusion-L on one Blackhole p150: {Path(args.input).name}",
lines=[f"{len(out)} detections; timing_ms {timing}",
"Bird's-eye view, forward up, rings every 10 m; boxes are Autoware DetectedObjects in base_link."])
img.save(out_dir / "quickstart_bev.png")
print(f"{len(out)} results -> {out_dir / 'quickstart.json'}, {out_dir / 'quickstart_bev.png'} timing_ms={timing}")
|