# SPDX-License-Identifier: Apache-2.0 """Python API: Point Transformer V3 (Autoware autoware_ptv3) on one Tenstorrent Blackhole p150. from tt_ptv3 import PTv3 with PTv3.from_pretrained(device_id=0) as model: # weights -> HF cache, device open, traces captured out = model("code/tt_ptv3/samples/test_pcd.npz") # see PYTHON.md for every input form print(out.class_counts(), len(out.detections)) print(out.to_dict()) # the same JSON as POST /predict The contract shared by every bundle of the Autoware collection is the vendored ``ttaw.api_base.ModelBase`` (BUNDLE_CONVENTIONS.md section 8): ``from_pretrained`` resolves the pinned weights before it claims the chip, opens it (ETH dispatch, 12x10), builds the graph and captures every trace variant in ``warmup_variants``, so the first call is as fast as the later ones; calls are serialised by a lock (one chip, batch 1) and fill ``timing_ms``; ``close()`` is idempotent, also runs at interpreter exit, and closes the chip only if the model opened it. The HTTP server (``tt_ptv3.server.app``) calls this class, so ``/predict`` and ``model(...)`` agree bit for bit. What the model reproduces (research/ptv3/SPEC.md; PLAN.md 2.7 / D3): the ``AutowareFoundation/ptv3@v4.0`` weights with the pre- / post-processing of autoware_universe node @5eb9286. Outputs: the segmentation cloud of the node (``~/output/pointcloud/segmentation``: per input point the class, its probability and the normalised entropy; Autoware's PointCloudClassification values in ``classification``), the obstacle-cloud membership (``~/output/pointcloud/filtered``) and, with ``run_detection`` (the node runs the det head only when ``~/output/objects`` has a subscriber), the 3D boxes. Hooks: ``_prepare`` / ``_postprocess`` are the host code of ``tt_ptv3.host`` (shared with the CPU reference); ``_build`` / ``_forward`` run the traced device graph of ``tt_ptv3.tt`` (trace A encoder + seg head, trace B det head; one pair per capacity bucket of ``PTV3_BUCKETS``, ``tt.model.bucket_set``; ``_prepare`` picks the smallest bucket that holds the frame's five stage counts). Importing it has no side effects: ttnn / torch are imported inside the hooks. """ from __future__ import annotations import threading from dataclasses import dataclass from typing import Any, Dict, Optional import numpy as np from . import io as tio from .device import DEVICE_DEFAULTS from .reference.config import (DEFAULT_NEIGHBOUR_MODE, DET_CLASSES, NEIGHBOUR_MODES, OBJECT_LABELS, RECONSTRUCTION_MODES, SEG_CLASSES, PreprocessConfig, PTv3Config) from .ttaw.api_base import ModelBase from .ttaw.io import encode_array from .ttaw.outputs import Detections3D, Segmentation3D __all__ = ["PTv3", "Output", "PTv3Output", "DETECTION_LABELS", "detections_from_boxes"] DETECTION_LABELS = list(DET_CLASSES) _CALL = threading.local() @dataclass class PTv3Output(Segmentation3D): """The node's outputs for one cloud. ``label_ids`` (uint8, model class, 255 = outside the crop / clipped) and ``scores`` (probability of that class, 0 for invalid points) per output point (input order in ``full`` reconstruction); ``classification`` (PointCloudClassification, INVALID 255), ``entropy`` (NaN for invalid), ``filter_keep`` (obstacle-cloud membership); ``detections`` (``Detections3D`` with the 7 det classes; ``meta`` holds the ObjectClassification label of each box after the remapper) or None.""" classification: Optional[np.ndarray] = None entropy: Optional[np.ndarray] = None filter_keep: Optional[np.ndarray] = None detections: Optional[Detections3D] = None def to_dict(self, output_format: str = "json") -> Dict[str, Any]: body = super().to_dict(output_format) for name in ("classification", "entropy", "filter_keep"): v = getattr(self, name) if v is not None: body[name] = encode_array(np.asarray(v), key=name) if self.detections is not None: det = self.detections.to_dict(output_format) obj = self.detections.meta.get("autoware_labels") if obj is not None: for d, o in zip(det["detections"], obj): d["autoware_label"] = o body["num_detections"] = det["num_detections"] body["detections"] = det["detections"] if "arrays" in det: body["detection_arrays"] = det["arrays"] return body Output = PTv3Output def detections_from_boxes(boxes, *, model: str = "", timing_ms: Optional[dict] = None) -> Detections3D: """``host.postprocess.Boxes`` (after NMS / remapper, score order) -> ``Detections3D`` (yaw: the node's).""" arr = np.stack([boxes.x, boxes.y, boxes.z, boxes.length, boxes.width, boxes.height, boxes.yaw], axis=1) \ if len(boxes) else np.zeros((0, 7), np.float32) vel = np.stack([boxes.vel_x, boxes.vel_y], axis=1) if len(boxes) else np.zeros((0, 2), np.float32) obj = [OBJECT_LABELS[int(v)] for v in (boxes.object_label if boxes.object_label is not None else [])] return Detections3D(arr, boxes.score, boxes.label, vel, labels=DETECTION_LABELS, model=model, timing_ms=dict(timing_ms or {}), meta={"autoware_labels": obj}) class PTv3(ModelBase): """Point Transformer V3 (Autoware autoware_ptv3) on one Blackhole p150. Create it with :meth:`from_pretrained`.""" MODEL_NAME = "ptv3-p150" ENV_PREFIX = "PTV3" # prefix of the environment knobs (SERVING.md section 3.4) DEFAULT_REPO = "AutowareFoundation/ptv3" DEFAULT_TAG = "v4.0" # the Autoware ansible artifacts pin # the commit DEFAULT_TAG points to (pinned: tags can move) DEFAULT_REVISION = "1673104843af3c740145ac74dc4af6b8a4502bf3" ALLOW_PATTERNS = ["ptv3_encoder.onnx", "ptv3_seg3d_head.onnx", "ptv3_det3d_head.onnx", "ml_package_ptv3_*.param.yaml", "deploy_metadata.yaml"] # LOAD: which traces are built -- "seg_det" (encoder + seg trace, det trace; the node's default heads) or "seg" # (encoder + seg only: the autoware_launch@main wiring, where the det head is idle) VARIANTS = ["seg_det", "seg"] DEFAULT_VARIANT = "seg_det" INPUT_KIND = "lidar" CAMERA_ORDER = () POINT_FIELDS = tio.DEFAULT_POINT_FIELDS # x, y, z, intensity (0-255; divided by 255 as the node's uint8 path) LABELS = list(SEG_CLASSES) # segmentation classes (ml_package_ptv3_seg3d_head.param.yaml) # RT-host knobs (PLAN.md 2.7 parameter table). Thresholds / NMS / reconstruction are host post-processing; the # neighbour mode only changes the host-built index tables (no device shape); min_voxels is the node's 16k floor # (0 disables it, e.g. for single-sweep nuScenes clouds, which all have < 16k voxels). RUNTIME_PARAMS = { "run_detection": (bool, None, None, True), "score_threshold": (float, 0.0, 1.0, 0.1), "reconstruction": (str, None, None, "full"), "neighbour_mode": (str, None, None, DEFAULT_NEIGHBOUR_MODE), "min_voxels": (int, 0, 256000, 16000), } EXTRA_INPUTS = () DEVICE_DEFAULTS = DEVICE_DEFAULTS # validated open parameters (device.py) def __call__(self, points: Any = None, **kwargs: Any): """``ModelBase.__call__``; the neighbour mode (an RT-host param that shapes the host tables), ``min_voxels`` and ``run_detection`` (which decide whether the traces run at all) are handed to ``_prepare`` through a thread-local (each call runs prepare / forward / post on its thread).""" mode = kwargs.get("neighbour_mode", DEFAULT_NEIGHBOUR_MODE) if mode not in NEIGHBOUR_MODES: raise tio.InputError(f"neighbour_mode must be one of {list(NEIGHBOUR_MODES)}") params = self.validate_params({k: v for k, v in kwargs.items() if k in self.RUNTIME_PARAMS}) _CALL.neighbour_mode = mode _CALL.params = params try: return super().__call__(points, **kwargs) finally: _CALL.neighbour_mode = None _CALL.params = None predict = __call__ # ---- hooks -------------------------------------------------------------------------------------------- def _build(self) -> None: """Weights (``reference.weights.load_weights``: OnnxWeights by consuming node, BN folded) -> the device graph ``tt.model.TtPTv3`` and its ``TraceRunner``: trace A ``encseg_`` (encoder + seg head) and, for the variant ``seg_det``, trace B ``det_`` (det head; skipped per call by ``run_detection``). Knobs (``PTV3_*``) and the precision policy (``PTV3_PRECISION``) are read once here; nothing is captured (``_warm_one``).""" from .reference.weights import load_weights from .tt.model import DEFAULT_BUCKET, KNOBS, TtPTv3, bucket_set, build_runner self.with_det = self.variant == "seg_det" self.knobs = KNOBS.read() self.host_maps = str(self.knobs.HOST_MAPS) if self.host_maps == "numba": from .host import neighbours_nb if not neighbours_nb.warm_up(): # numba missing: the numpy oracles (same tables) self.host_maps = "numpy" self.host_pre = str(self.knobs.HOST_PRE) if self.host_pre == "numba": from .host import pre_nb if not pre_nb.warm_up(): # numba missing: the numpy oracles (same tables) self.host_pre = "numpy" self.host_nms = str(self.knobs.HOST_NMS) self.host_threads = max(1, int(self.knobs.HOST_THREADS)) self.buckets = bucket_set(self.knobs.BUCKETS, self.knobs.POOL_PROFILES) self.bucket = DEFAULT_BUCKET # the largest bucket (n256k); a frame runs in bucket_for(counts) self.caps = tuple(self.buckets[self.bucket]) weights = load_weights(self.weights_path, det=self.with_det) self.graph = TtPTv3(self.device, weights, knobs=self.knobs, det=self.with_det) variants = ("encseg", "det") if self.with_det else ("encseg",) self.runner = build_runner(self.graph, self.buckets, variants=variants, name=f"{self.MODEL_NAME}/{self.variant}") def bucket_for(self, counts) -> tuple: """``(name, caps)`` of the smallest built bucket holding the five stage ``counts`` (RT-host choice).""" from .tt.model import select_bucket b = select_bucket(self.buckets, counts) return b, tuple(self.buckets[b]) def _warm_one(self, variant: Dict[str, Any]) -> None: """Warm up (eager, every variant) and capture the traces (idempotent).""" self.runner.capture() def _prepare(self, points: Any = None, **inputs: Any) -> Dict[str, Any]: """Decode the cloud and build the host tables (``host.preprocess``) in the call's neighbour mode.""" from .host import preprocess as pre if points is None: raise tio.InputError("this model needs 'points'") extra = sorted(k for k, v in inputs.items() if v is not None) if extra: # the v4.0 node takes one concatenated cloud: no densification, no cameras, no state raise tio.InputError(f"this model takes only 'points' (one concatenated cloud in base_link); " f"refused: {extra}") cloud = tio.load_points(points) pts = np.asarray(cloud.points, np.float32) mode = getattr(_CALL, "neighbour_mode", None) or DEFAULT_NEIGHBOUR_MODE frame = pre.prepare_frame(pts, PreprocessConfig(), neighbour_mode=mode, neighbour_impl=getattr(self, "host_maps", "numpy"), pre_impl=getattr(self, "host_pre", "numpy"), threads=getattr(self, "host_threads", 1)) params = getattr(_CALL, "params", None) or {} min_vox = int(params.get("min_voxels", PreprocessConfig().min_num_voxels)) skip = frame.voxels.num_unique_voxels < min_vox or frame.voxels.num_voxels == 0 if frame.overflow and not skip: raise tio.InputError(f"stage voxel counts {frame.counts} exceed the device capacities") bucket = None if (skip or frame.overflow) else self.bucket_for(frame.counts) return {"frame": frame, "frame_id": cloud.frame_id, "skip": skip, "bucket": bucket, "run_detection": bool(params.get("run_detection", True))} def _forward(self, prepared: Dict[str, Any]) -> Any: """Upload the frame's tables into the persistent inputs of the bucket, replay trace A (one D2H: label, probability and entropy per voxel) and, when detection runs, trace B (one packed D2H). A frame the node rejects (``min_voxels``) runs nothing. Runs under the model lock.""" from .tt import inputs as tin from .tt.model import unpack_seg if prepared["skip"]: return None frame = prepared["frame"] b, caps = prepared.get("bucket") or self.bucket_for(frame.counts) inputs = {f"{b}.{k}": v for k, v in tin.frame_inputs(frame, caps, det=self.with_det, impl=self.host_pre, threads=self.host_threads).items()} seg = unpack_seg(self.runner(f"encseg_{b}", inputs=inputs)["seg"], frame.counts[0], caps[0]) raw: Dict[str, Any] = {"pred_labels": seg["labels"], "voxel_prob": seg["prob"], "voxel_entropy": seg["entropy"], "det": None} if self.with_det and prepared["run_detection"]: raw["det"] = self.graph.assemble(self.runner(f"det_{b}")) return raw def _postprocess(self, raw: Any, prepared: Dict[str, Any], params: Dict[str, Any]) -> Output: """``raw`` = ``{"pred_labels" (N0,), "voxel_prob" (N0,), "voxel_entropy" (N0,), "det": {det outputs} | None}`` (device readback; the CPU reference gives ``"pred_probs" (N0, 15)`` instead) or None for a frame the node rejects -> :class:`PTv3Output`.""" return assemble_output(raw, prepared["frame"], params, model=self.MODEL_NAME, frame_id=prepared.get("frame_id", "base_link"), nms=getattr(self, "host_nms", "ref")) def _release(self) -> None: runner = getattr(self, "runner", None) if runner is not None: runner.release() self.runner = None graph = getattr(self, "graph", None) if graph is not None: graph.release() self.graph = None def extra_info(self) -> Dict[str, Any]: runner = getattr(self, "runner", None) info: Dict[str, Any] = {"detection_labels": DETECTION_LABELS, "config": PTv3Config().describe()} if runner is not None: info["trace"] = runner.describe() graph = getattr(self, "graph", None) if graph is not None: info["graph"] = graph.describe() info["bucket"] = {"name": self.bucket, "capacities": list(self.caps)} info["buckets"] = {b: list(c) for b, c in getattr(self, "buckets", {}).items()} info["host_maps"] = getattr(self, "host_maps", None) info["host_pre"] = getattr(self, "host_pre", None) info["host_nms"] = getattr(self, "host_nms", None) info["host_threads"] = getattr(self, "host_threads", None) return info def assemble_output(raw: Dict[str, Any], frame, params: Dict[str, Any], *, model: str = "", frame_id: str = "base_link", nms: str = "ref") -> PTv3Output: """Host post-processing shared by the API and the CPU reference pipeline (``tt_ptv3.host.postprocess``); ``nms``: ``"ref"`` or ``"fast"`` (same kept boxes; ``PTV3_HOST_NMS``).""" from .host import postprocess as post cfg = PTv3Config() mode = params.get("reconstruction", "full") if mode not in RECONSTRUCTION_MODES: raise tio.InputError(f"reconstruction must be one of {list(RECONSTRUCTION_MODES)}") meta: Dict[str, Any] = {"frame": frame.describe(), "neighbour_mode": frame.neighbour_mode, "reconstruction": mode} min_vox = int(params.get("min_voxels", cfg.pre.min_num_voxels)) if raw is None or frame.voxels.num_unique_voxels < min_vox: # the node publishes nothing (ptv3_trt.cpp:940-946) meta["rejected"] = f"too few voxels ({frame.voxels.num_unique_voxels} < {min_vox})" return PTv3Output(np.zeros(0, np.uint8), SEG_CLASSES, np.zeros(0, np.float32), model=model, frame_id=frame_id, meta=meta) v = frame.voxels if "pred_probs" in raw: # the CPU reference: the 15 probabilities per voxel lab, prob, src = post.reconstruct(np.asarray(raw["pred_labels"]), np.asarray(raw["pred_probs"], np.float32), crop_mask=v.crop_mask, inverse_map=v.inverse_map, mode=mode, representative=v.representative, cropped_index=v.cropped_index) seg = post.segmentation_outputs(lab, prob, src, cfg.seg) else: # the device: label, probability and entropy per voxel seg = post.segmentation_from_voxels(raw["pred_labels"], raw["voxel_prob"], raw["voxel_entropy"], crop_mask=v.crop_mask, inverse_map=v.inverse_map, mode=mode, representative=v.representative, cropped_index=v.cropped_index, cfg=cfg.seg, num_classes=len(SEG_CLASSES)) dets = None if params.get("run_detection", True) and raw.get("det") is not None: dcfg = cfg.det.with_score_threshold(float(params.get("score_threshold", 0.1))) dets = detections_from_boxes(post.postprocess_detections(raw["det"], dcfg, nms=nms), model=model) return PTv3Output(seg.labels, SEG_CLASSES, seg.probability, model=model, frame_id=frame_id, meta=meta, classification=seg.classification, entropy=seg.entropy, filter_keep=seg.filter_keep, detections=dets)