| |
| """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" |
| DEFAULT_REPO = "AutowareFoundation/ptv3" |
| DEFAULT_TAG = "v4.0" |
| |
| DEFAULT_REVISION = "1673104843af3c740145ac74dc4af6b8a4502bf3" |
| ALLOW_PATTERNS = ["ptv3_encoder.onnx", "ptv3_seg3d_head.onnx", "ptv3_det3d_head.onnx", |
| "ml_package_ptv3_*.param.yaml", "deploy_metadata.yaml"] |
| |
| |
| VARIANTS = ["seg_det", "seg"] |
| DEFAULT_VARIANT = "seg_det" |
| INPUT_KIND = "lidar" |
| CAMERA_ORDER = () |
| POINT_FIELDS = tio.DEFAULT_POINT_FIELDS |
| LABELS = list(SEG_CLASSES) |
| |
| |
| |
| 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 |
|
|
| 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__ |
|
|
| |
| 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_<bucket>`` (encoder + seg head) and, |
| for the variant ``seg_det``, trace B ``det_<bucket>`` (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(): |
| 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(): |
| 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 |
| 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: |
| 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: |
| 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: |
| 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: |
| 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) |
|
|
|
|
|
|
