# Python API: SceneSeg (Autoware VisionPilot) on Blackhole

Use this API from Python code (a pipeline, a notebook, a ROS 2 node wrapper). You do not need the HTTP server: the
API and the server share the decoders, the device graph and the post-processing, so the outputs and the speed are
the same.

## Install

Install the package on top of an environment that already has `ttnn` (a tt-metal `python_env` at `44d66500520`
with `patches/tt-metal-eth-dispatch.patch`, or the tt-model container). From the root of the model repository (the
directory that holds `pyproject.toml`, `README.md` and `code/`):

```bash
pip install -e .                        # the Python API
pip install -e ".[server,test]"         # + the HTTP server and the tests
```

The pip project is the repository's top-level `pyproject.toml`; it installs the package from `code/tt_sceneseg` (there is no
`pyproject.toml` inside `code/`, because the container build copies `code/` over the tt-metal tree).

The package carries `tt_sceneseg.ttaw`, the shared code of the Autoware ports to Blackhole (device open, trace runner,
decoders, model base class, HTTP app), vendored at the version recorded in `code/tt_sceneseg/ttaw/VENDORED.json`.

| You want to run | Extras |
|---|---|
| the Python API | none |
| the HTTP server (`tt_sceneseg.server.app`, see `SERVING.md`) | `server` |
| host tests (no device; device tests are skipped): see "Tests" | `server,test` |
| device tests: `python -m pytest -q -s code/tt_sceneseg/tests/test_pcc_device.py code/tt_sceneseg/tests/test_e2e_device.py` | `test` |

## Quickstart

Run from the repository root (`code/tt_sceneseg/samples/highway_normal_1.png` is a path relative to it):

```python
from tt_sceneseg import SceneSeg

with SceneSeg.from_pretrained(device_id=0) as model:      # weights -> your HF cache, trace captured
    out = model("code/tt_sceneseg/samples/highway_normal_1.png")

mask = out.mask                    # uint8 (H, W) at the input size: 255 = foreground (the VisionPilot topic)
print(out.class_counts())          # pixels per class of out.class_map (background / foreground / road)
print(out.to_dict()["meta"])       # the same JSON as POST /predict
```

`examples/quickstart.py` runs this snippet and writes the `/predict` JSON, the mask and a tinted picture.

## `SceneSeg.from_pretrained(...)`

```python
SceneSeg.from_pretrained(
    model_id=None,          # HF repo or a local directory with the weights files; default changh95/sceneseg-p150 (weights/)
    *,
    revision=None,          # default for the default repo: the validated weights commit a6829b246a0
    variant=None,           # load-time variant (bgr, rgb); default $SCENESEG_VARIANT or "bgr"
    device_id=None,         # chip to open; default $TT_DEVICE_ID or 0
    device=None,            # an already-opened ttnn device (tt_sceneseg.device.open_device); close() does not close it
    dispatch=None,          # "eth" (p150 target, 12x10 grid) | "worker" (A/B only, 11x10) | "auto"; default $SCENESEG_DISPATCH or "eth"
    num_command_queues=None,  # default $SCENESEG_NUM_CQS or 1
    weights_dir=None,       # explicit local weights directory; no Hub access
    warmup_variants="default",  # trace variants to capture now; see "Warm-up"
    verbose=False,
    **compile_params,       # none for this model: the network input is fixed at 640x320
) -> SceneSeg
```

What it does: resolves the weights first, so a Hub problem never claims the chip (`weights_dir` >
`$SCENESEG_WEIGHTS_DIR` > a local `model_id` directory > the HF snapshot at the pinned revision, restricted to
`weights/SceneSeg_FP32.onnx, weights/provenance.json, weights/LICENSE, weights/NOTICE`, with an offline fallback to the cache), opens the chip (ETH dispatch, 12×10, 1 CQ;
the other open parameters are `DEVICE_DEFAULTS` in `tt_sceneseg/device.py`, overridable with `SCENESEG_*`), builds the
graph of the selected channel order, then compiles and captures the metal trace. If ETH dispatch cannot open
(tt-metal without the patch), it warns and falls back to WORKER dispatch.

`variant` is the **input channel order**, a load-time choice folded into the first conv's weights on the device:

| variant | what it is |
|---|---|
| `bgr` (default) | exactly what VisionPilot's `scene_seg_model` node computes: the deployed C++ pre-processing feeds B, G, R planes (with the ImageNet constants reordered to match) to the network |
| `rgb` | the order the network was trained with (`scene_seg_trainer.py` @ `ca58cb50`): clearly better masks, most of all on two-wheelers (README "Channel order") |

Both run the same device graph and are validated against the fp32 CPU reference of their own order.

## Warm-up

`from_pretrained` returns a warm model: it builds the graph, runs it once eagerly (this first run compiles the
kernels and prepares the conv weights on the device) and captures the whole forward as ONE metal trace, so no call
compiles anything. There are no shape buckets: every image, whatever its size, is squashed on the host to the same
640x320 uint8 canvas, so one trace serves every input.

- Measured on the shared host (2026-10-13, `logs/sceneseg/republish/final/`): with the JIT cache warm, a load takes
  18.5-18.7 s (the served app: graph build 11.4 s, which prepares the S1 and K11 weights on the host, and warm-up +
  capture 2.8 s). With an empty JIT cache the first load took 275 s (measured before R8.1, which saves about 21 s of
  graph build). The first call took 28-30 ms, the later ones 22.5-24 ms on the shipped sample.
- `warmup_variants="none"` (server: `SCENESEG_WARMUP=none`) skips the capture at load; call `model.warmup()` before
  the first frame (it is idempotent).
- The JIT cache is tt-metal's (`~/.cache/tt-metal-cache`, or `TT_METAL_CACHE`); inside the container it is
  `~/.cache/tt-model/sceneseg-p150/cache` of the host.

## Call: `model(image, **params)`

| Argument | Type | Description |
|---|---|---|
| `image` (first positional) or `images=[image]` | path, `bytes`, `PIL.Image`, uint8 array / tensor, float array, `CameraImage` | one camera image of any size (see "Input types") |
| `class_map` | bool, default `True` | also encode the 3-class map in `to_dict()` (the counts and `out.class_map` are always there). The node itself has no runtime parameter: its decision is a plain argmax |

### Input types

- The image (first argument, or `images=[x]` with exactly one entry): a path to a PNG / JPEG file, encoded PNG / JPEG
  `bytes`, a `PIL.Image`, a uint8 array or torch tensor in **RGB** order (`HxWx3`, `3xHxW`; `HxW` grey is replicated,
  a 4th channel is dropped), a float array in [0, 1], or a `tt_sceneseg.io.CameraImage`. Any size from 2x2 up.
- **Channel order of the input.** Arrays are read as RGB and converted to BGR8, which is what the VisionPilot node
  receives (`bgr8` image topic). An OpenCV or ROS `bgr8` array must be flipped first: `model(bgr[:, :, ::-1])`.
  Encoded files and PIL images carry their own channel order. (This is independent of `variant`, which selects
  the plane order the *network* is fed.)
- The node squashes the image to 640x320 with OpenCV's `cv::resize` (INTER_LINEAR, aspect ratio not kept); the port
  reproduces it bit for bit on the host. SceneSeg's Python tooling also had a bicubic resize; it is not offered,
  because the deployed node never used it.
- There is no calibration, no camera model and no temporal state: one image per call, as one node instance
  subscribes to one image topic. Inline intrinsics or transforms in a `/predict` image entry are accepted by the
  shared schema and ignored.

## Output

`model(...)` returns a `tt_sceneseg.Output` (= `SceneSegMask`, a `ttaw.outputs.Mask2D`); `out.to_dict()` is exactly the
`POST /predict` body (SERVING.md section 3.2).

| field | type | meaning |
|---|---|---|
| `mask` | uint8 `[H, W]` | what the node publishes on `/autoseg/scene_seg/mask`: mono8 at the input's size, 255 where the per-pixel argmax is class 1 (foreground), else 0; resized back from 640x320 with INTER_NEAREST |
| `class_map` | uint8 `[H, W]` | the 3-class argmax on the same grid: 0 background (incl. sky), 1 foreground objects, 2 drivable road (a by-product the node does not publish) |
| `class_names` | tuple | `("background", "foreground", "road")` |
| `meta` | dict | `channel_order`, `source_hw`, `network_hw` `[320, 640]`, `foreground_class` 1, `mask_value` 255, `foreground_fraction` |
| `timing_ms` | dict | `preprocess`, `device` (H2D + trace + D2H), `postprocess`, `total` of this call |

`out.class_counts()` gives the pixels per class of `class_map`; `out.to_dicts()` the same as
`[{"label", "label_id", "pixels"}]`. Ties between classes resolve to the lowest index, as the node's strict `>` scan does.

## Lifetime and information

- `model.close()` releases the trace and the device tensors and closes the chip if the model opened it; idempotent.
  `with` calls it for you; an unclosed model is closed when Python exits.
- `model.info`: model name, weights (repo, revision, path, license), device (`dispatch`, `grid`, `cores`,
  `num_command_queues`, `fallback`), variant and `channel_order`, warm variants, warm-up times, runtime parameter
  defaults, plus `graph` (the device graph description), `trace` and `knobs` (the build-time `SCENESEG_*` knobs).
- Calls from several threads are safe: the device calls are serialised by a lock. One model per process per chip.

## Speed

Warm calls, batch 1, ETH dispatch, 1 CQ, 12x10, `bgr`, the published numerics (`code/scripts/bench.py`, 100
iterations, 2026-10-13 on a shared host; `OPT_BASELINE.md` "After optimization"):

| stage | shipped sample 727x414 PNG | PandaSet 1920x1080 JPEG |
|---|---:|---:|
| `model(path)` end to end, p50 (p99) | 22.54 (23.38) ms | 30.00 (32.01) ms |
| host decode + squash to 640x320 (PIL decode; exact `cv::resize` as a numba twin on 4 threads) | 8.53 ms | 13.75 ms |
| H2D of the uint8 640x320 canvas (320 rows of 1,920 B) | 0.21 ms | 0.34 ms |
| device trace, one forward | 12.63 ms | 12.64 ms |
| D2H (the uint8 class map) | 0.23 ms | 0.30 ms |
| host post-processing (INTER_NEAREST to the source size, numba twin) | 0.62 ms | 3.06 ms |

- Back-to-back replays take 12.58 ms per frame (79.5 frames/s of device throughput); a synchronous call gives
  44.4 frames/s on the sample. The `rgb` profile replays in the same time (e2e p50 22.66 ms on the sample).
- The first release measured 90.85 ms (sample) and 116.12 ms (PandaSet) end to end and 61.06 ms per replay;
  `OPT_REPORT.md` lists every step in between.
- A raw RGB array skips the decode. The device path does not depend on the image size; the host stages grow with it.

## Tests

From the repository root, with `code/` on `PYTHONPATH` (or the package installed):

- Host (no device, no ttnn call): `TT_VISIBLE_DEVICES=none python -m pytest -q
  code/tt_sceneseg/tests/test_bundle_host.py code/tt_sceneseg/tests/test_io_host.py
  code/tt_sceneseg/tests/test_server_host.py code/tt_sceneseg/tests/test_host_preprocess.py
  code/tt_sceneseg/tests/test_host_postprocess.py code/tt_sceneseg/tests/test_reference_host.py`. Inside the
  tt-models workspace the suite runs on its fake ttnn (`-p fake_ttnn_plugin`), which also enables
  `test_tt_graph_host.py` and `test_tt_layers_host.py`; `test_reference_cpu.py` (the CPU reference against ONNX
  Runtime) needs `onnxruntime`.
- Device (one p150): `python -m pytest -q -s code/tt_sceneseg/tests/test_pcc_device.py
  code/tt_sceneseg/tests/test_e2e_device.py` (per-module PCC against the fp32 CPU reference, end-to-end agreement on
  the shipped tutorial and comma10k samples, replay == eager, API == `/predict`, the stored-reference smoke gates;
  `SCENESEG_VARIANT=rgb` for the other order). The public-frame gates (23 frames with logits, 415 held-out frames)
  run only inside the tt-models workspace, where their data live: PandaSet, nuScenes and KITTI frames, of which
  nuScenes and KITTI may not be redistributed. There, a missing input fails the test instead of skipping it.
- The gates are frozen in `tests/*.gates.json` at the first green run and are never loosened.

## Limits

- Batch 1 on the chip; one model per process; one image per call.
- The network input is fixed at 640x320 (the SceneContext reshape hard-codes the 10x20 stride-32 grid, batch 1); any
  source image size is accepted, squashed on the host without keeping the aspect ratio, and the mask is resized back
  with nearest neighbour.
- Only `SceneSeg_FP32.onnx`, the network the node deploys. The INT8 QAT ONNX (never deployed; it disagrees with FP32)
  and SceneSegLite (a different 19-class network) are not supported.