File size: 2,939 Bytes
5662e7f | 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 | #!/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] [--variant seg_det|seg]
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 coloured by the predicted class, with the 3D boxes; tt_ptv3.viz). The default input is the shipped
sample code/tt_ptv3/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_ptv3" / "samples" / "test_pcd.npz"))
ap.add_argument("--out-dir", default=str(REPO / "examples" / "output"))
ap.add_argument("--device-id", type=int, default=0)
ap.add_argument("--variant", default="seg_det", choices=("seg_det", "seg"))
args = ap.parse_args()
out_dir = Path(args.out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
# --- the model card snippet --------------------------------------------------------------------------------------
from tt_ptv3 import PTv3
with PTv3.from_pretrained(device_id=args.device_id, variant=args.variant) as model: # weights -> HF cache, traces captured
out = model(args.input) # base_link cloud: path (.npz/.npy/.pcd/.bin), (N, 4) array, or a PointCloud
print({k: v for k, v in out.class_counts().items() if v}) # input points per class (255 = outside the crop)
for d in (out.detections.to_dicts() if out.detections is not None else [])[: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_ptv3.viz import load_xyz, render_bev # noqa: E402
dets = out.detections.to_dicts() if out.detections is not None else []
timing = {k: round(v, 1) for k, v in out.timing_ms.items()}
img = render_bev(load_xyz(args.input), out.label_ids, dets,
title=f"PTv3 (Autoware ptv3 v4.0) on one Blackhole p150: {Path(args.input).name}",
lines=[f"{len(out.label_ids)} points, {len(dets)} boxes; timing_ms {timing}",
"Bird's-eye view, forward up, rings every 10 m; points coloured by the predicted class."])
img.save(out_dir / "quickstart_bev.png")
print(f"-> {out_dir / 'quickstart.json'}, {out_dir / 'quickstart_bev.png'} timing_ms={timing}")
|