File size: 9,291 Bytes
aec3f16 e0adbab aec3f16 e0adbab aec3f16 e0adbab aec3f16 e0adbab aec3f16 | 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 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | # SPDX-License-Identifier: Apache-2.0
"""Python API: SceneSeg (Autoware VisionPilot) on one Tenstorrent Blackhole p150.
from tt_sceneseg import SceneSeg
with SceneSeg.from_pretrained(device_id=0) as model: # weights -> HF cache, device open, traces captured
out = model("code/tt_sceneseg/samples/highway_normal_1.png") # see PYTHON.md for every input form
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_sceneseg.server.app``) calls this class, so ``/predict`` and ``model(...)`` agree bit for bit.
Input: ONE camera image per call, as VisionPilot's ``scene_seg_model`` node takes one image topic: the first
positional argument (path, bytes, PIL image, RGB uint8 HxWx3 array or tensor) or ``images=[...]`` like the server
(exactly one entry; its camera name is free); no calibration. ``io.single_camera`` decodes it for the API and the
server alike; ``tt_sceneseg.host`` reproduces the node's pre-processing bit for bit (BGR8, OpenCV INTER_LINEAR squash
to 640x320) and post-processing (argmax with lowest-index ties, class 1 -> 255, INTER_NEAREST back to the input size).
Variants (load-time, PLAN.md 2.2 / D10): the input channel order, folded into the stem conv on the device. ``bgr``
(default) is the deployed VisionPilot C++ order; ``rgb`` is the training order, better on two-wheelers (model card).
Importing this module has no side effects: ttnn / torch are imported inside the hooks, never at module level.
"""
from __future__ import annotations
from typing import Any, Dict, Mapping, Optional
from . import io as tio
from .device import DEVICE_DEFAULTS
from .host.output import SceneSegMask
from .reference import config
from .ttaw.api_base import ModelBase
__all__ = ["SceneSeg", "Output"]
# The result class of this model: a ttaw.outputs.Mask2D (the published mono8 mask) with the 3-class map.
Output = SceneSegMask
class SceneSeg(ModelBase):
"""SceneSeg (Autoware VisionPilot) on one Blackhole p150. Create it with :meth:`from_pretrained`."""
MODEL_NAME = "sceneseg-p150"
ENV_PREFIX = "SCENESEG" # prefix of the environment knobs (SERVING.md section 3.4)
# The weights are re-hosted in this bundle's own repo under weights/ (PLAN.md D10: the upstream
# AutowareFoundation/SceneSeg card holds no weights; the files come from VisionPilot's Google Drive links).
DEFAULT_REPO = config.WEIGHTS_REPO
DEFAULT_TAG = None # the weights are a commit of this repo, not a tag
# the commit of the weights upload (deferred until it exists: PLAN.md 4.1 gates the publish, not the port)
DEFAULT_REVISION = "a6829b246a05dba55c356323bbf75f7b1d148499"
ALLOW_PATTERNS = ["weights/SceneSeg_FP32.onnx", "weights/provenance.json", "weights/LICENSE", "weights/NOTICE"]
WEIGHTS_LICENSE = config.WEIGHTS_LICENSE
VARIANTS = ["bgr", "rgb"] # load-time: the input channel order (folded into the stem weights on the device)
DEFAULT_VARIANT = "bgr"
INPUT_KIND = "camera" # lidar | camera | multicam | lidar+multicam | planner
CAMERA_ORDER = () # one image per request (one node instance per camera)
REQUIRE_CALIBRATION = False # the node uses no CameraInfo / TF
POINT_FIELDS = tio.DEFAULT_POINT_FIELDS # unused (camera model); kept for the shared /info schema
# class names in network output order (load_data_scene_seg.py:136-141; class 1 is the published mask)
LABELS = list(config.CLASS_NAMES)
# Per-request knobs: name -> (type, min, max, default). The node has none (pure argmax, no threshold); this port
# adds one output toggle. Anything that changes a device shape or the weights is a load-time option
# (BUNDLE_CONVENTIONS.md section 9).
RUNTIME_PARAMS = {
"class_map": (bool, None, None, True), # also return the 3-class map (a by-product the node does not publish)
}
EXTRA_INPUTS = ()
DEVICE_DEFAULTS = DEVICE_DEFAULTS # validated open parameters (device.py)
# ---- port-specific hooks (called by ModelBase; keep host work out of _forward) ---------------------------
def _build(self) -> None:
"""Load the weights (``reference.weights.load_weights``: the deployed ONNX's initializers by consuming node,
checked against the 108-layer architecture and the pinned sha256) and build the ttnn graph of
``tt_sceneseg.tt`` for the selected channel order with its ``TraceRunner`` (persistent uint8 input, uint32
class-map readback; build-time knobs ``SCENESEG_*`` of ``tt/knobs.py``). No capture here (``_warm_one``)."""
from .reference import assets
from .reference.weights import load_weights
from .tt.knobs import KNOBS
from .tt.runner import SceneSegRunner
onnx = assets.find_in(self.weights_path, config.ONNX_FILE)
if onnx is None:
raise FileNotFoundError(f"{config.ONNX_FILE} not found under {self.weights_path}")
self.knobs = KNOBS.read()
if "numba" in (self.knobs.HOST_SQUASH, self.knobs.HOST_POST):
from .host import fast
fast.warm_up() # compile the host twins before serving
self.runner = SceneSegRunner(self.device, load_weights(onnx), self.variant,
weight_terms=self.knobs.WEIGHT_TERMS, act_terms=self.knobs.ACT_TERMS,
name=self.MODEL_NAME, input_layout=self.knobs.INPUT_LAYOUT)
def _warm_one(self, variant: Dict[str, Any]) -> None:
"""Warm up and capture the traces (``TraceRunner.capture``: every pending variant eagerly first, then the
captures with program-cache misses forbidden; idempotent)."""
self.runner.capture()
def _prepare(self, points: Any = None, *, images: Any = None, calibration: Optional[Mapping[str, Any]] = None,
sweeps: Any = None, stream: Any = None, inputs: Any = None, **_: Any) -> Dict[str, Any]:
"""The node's pre-processing (``tt_sceneseg.host.preprocess``): the one image (``io.single_camera``, RGB) ->
BGR8 -> the exact OpenCV INTER_LINEAR squash to 640x320 uint8 (the device input). ``tio.InputError`` for
client mistakes (no image, two images, an undecodable image)."""
from .host.preprocess import prepare_rgb
other = sorted(k for k, v in {"sweeps": sweeps, "stream": stream, "inputs": inputs}.items() if v is not None)
if other:
raise tio.InputError(f"this model takes one camera image, not {other}")
camera = tio.single_camera(points, images, calibration=calibration,
require_calibration=self.requires_calibration())
return {"prepared": prepare_rgb(camera.image, channel_order=self.variant), "camera": camera.name}
def _forward(self, prepared: Dict[str, Any]) -> Any:
"""Upload ``prepared["prepared"].device_input()`` (the uint8 squash) into the persistent device input, replay
the trace and read the class decision back: ``{"class_map": uint8 [320, 640]}``. Runs under the model
lock."""
return self.runner(prepared["prepared"].device_input())
def _postprocess(self, raw: Any, prepared: Dict[str, Any], params: Dict[str, Any]) -> Output:
"""The node's post-processing (``tt_sceneseg.host.postprocess``) of the device's class map (or of logits:
CPU reference / debug readback): class 1 -> 255, INTER_NEAREST back to the source size, + the class map."""
from .host.output import to_output
from .host.postprocess import postprocess
prep = prepared["prepared"]
res = postprocess(prep.src_hw, logits=raw.get("logits"), class_map=raw.get("class_map"))
return to_output(res, model=self.MODEL_NAME, channel_order=self.variant, include_class_map=params["class_map"])
def _release(self) -> None:
"""Release the traces and persistent device tensors; also called when ``from_pretrained`` fails half-way."""
runner = getattr(self, "runner", None)
if runner is not None:
runner.release()
def extra_info(self) -> Dict[str, Any]:
"""Additions to ``model.info`` and ``/info``: the channel order, the network size and the classes, plus the
trace variants, CQs and persistent tensors of the runner."""
runner = getattr(self, "runner", None)
info: Dict[str, Any] = {"channel_order": getattr(self, "variant", None), "input_hw": list(config.INPUT_HW),
"classes": list(config.CLASS_NAMES), "mask_class": config.FOREGROUND_CLASS}
if runner is not None:
d = runner.describe()
info.update({"graph": d["model"], "trace": d["trace"], "knobs": self.knobs.as_dict()})
return info
|