feat(axelera): drive decoder geometry from the compiled model manifest

Read the padded/original input shapes, the exact content offset inside the
padded input (n_padded_ch_inputs), per-output padded + real channel counts
and dequantize (scale, zero-point) from 1/manifest.json (fallback
quantized/manifest.json) instead of hardcoding YOLO-NAS constants, so the
decoder works for any compiled YOLO-family model shipping a postprocess
graph. Map each postprocess input to its model output by (C, H, W) shape
(the compile order of heads is not guaranteed to match the graph's input
order). Accept two postprocess output shapes: two outputs (scores + xyxy
boxes, YOLO-NAS) or one fused output [cx, cy, w, h, (objectness), C class
scores] in CENTER format (YOLOX/SDK DecodeYolo), converted to xyxy and
scaled by objectness. Verified on hardware against both YOLO-NAS M and
YOLOX-S: identical boxes on the same frames, YOLOX-S yields the same
bus/persons as the SDK stream reference.
This commit is contained in:
patrykorwat 2026-08-23 23:09:00 +00:00
parent 86637985ea
commit 4e33dda6da

View File

@ -56,26 +56,13 @@ DETECTOR_KEY = "axelera"
AX_DEFAULT_NETWORK = "yolox-s-coco-onnx"
AX_DEFAULT_BUILD_ROOT = "/opt/voyager-sdk/build"
MAX_ROWS = 20
# YOLO-NAS M compiled model input (padded) from the manifest:
# input_shapes=[1, 642, 656, 4], input_shapes_original=[1, 3, 640, 640]
# Padding: H (1,1), W (1,15), channels (0,1); quantize scale ~= 1/255, zp -128.
MODEL_INPUT_H, MODEL_INPUT_W = 640, 640
PADDED_H, PADDED_W = 642, 656 # includes top/left 1px border introduced by compile
INPUT_Q_SCALE = 0.003921568859368563
# int8 zero-point the compiled inputs are offset by (px - 128); also the
# filler for the padded border. Everything else about a specific compiled
# model (padded input size, content offset, per-output padded/real channels
# and dequant scale/zero-point) is read at init from 1/manifest.json, so the
# decoder works for any compiled YOLO-family model that ships a baked
# postprocess_graph.onnx (verified: YOLO-NAS M and YOLOX-S).
INPUT_Q_ZERO = -128
# dequantize_params from the compiled manifest, one per head output, in the
# AIPU output order: (cls, reg, cls, reg, cls, reg) per scale (80, 40, 20).
# AIPU tensors are padded NHWC channels (128); the real channel counts come from
# output_shapes_original: [80, 68, 80, 68, 80, 68].
HEAD_REAL_CHANNELS = (80, 68, 80, 68, 80, 68)
HEAD_DEQUANT = (
(0.07594172656536102, 84),
(0.08785015344619751, -63),
(0.06776463985443115, 66),
(0.07863091677427292, -67),
(0.06013955920934677, 44),
(0.06746677309274673, -61),
)
# confidence / NMS defaults mirror the SDK DecodeYolo operator defaults
CONF_THRESHOLD = 0.25
NMS_IOU_THRESHOLD = 0.45
@ -213,6 +200,52 @@ def _nms_xyxy(boxes, scores, iou_threshold: float) -> list:
return keep
def _read_model_manifest(model_dir: str) -> dict:
"""Read the compiled model's manifest (geometry + dequantize params).
The compiler writes two JSONs next to the kernel: ``1/model.json`` (padded
kernel I/O only) and ``1/manifest.json`` (original/padded shapes, the exact
content offset inside the padded input, per-output padded channel counts
and dequantize (scale, zero-point) for every head). ``1/manifest.json`` is
the cross-model source of truth (verified for YOLO-NAS M and YOLOX-S);
fall back to ``<parent>/quantized/manifest.json`` for older compilations.
"""
candidates = [
os.path.join(model_dir, "manifest.json"),
os.path.join(
os.path.dirname(model_dir), "quantized", "manifest.json"
),
]
for candidate in candidates:
if not os.path.isfile(candidate):
continue
try:
with open(candidate) as f:
manifest = json.load(f)
except Exception as exc:
raise ValueError(
f"could not parse compiled manifest {candidate}: {exc}"
) from exc
required = (
"input_shapes",
"input_shapes_original",
"output_shapes",
"output_shapes_original",
"dequantize_params",
)
missing = [k for k in required if k not in manifest]
if missing:
raise ValueError(
f"compiled manifest {candidate} missing required keys: {missing}"
)
return manifest
raise FileNotFoundError(
"compiled manifest not found next to the model kernel. Point "
"model.path at a directory compiled by the Voyager SDK that contains "
"1/model.json + 1/manifest.json + 1/postprocess_graph.onnx."
)
# ----------------- AxeleraRuntimeInference (engine) ----------------- #
@ -276,6 +309,58 @@ class _AxeleraRuntimeInference:
model, num_sub_devices=1, aipu_cores=1
)
manifest = _read_model_manifest(model_dir)
# input: padded NHWC (1, H, W, 4) exactly as the runtime expects it,
# plus the original NCHW (1, C, H, W) workspace the letterbox targets
in_padded = [int(x) for x in manifest["input_shapes"][0]]
in_orig = [int(x) for x in manifest["input_shapes_original"][0]]
if len(in_padded) != 4 or in_padded[3] != 4 or in_orig[1] != 3:
raise ValueError(
f"unsupported compiled input: padded {in_padded}, "
f"original {in_orig}"
)
self._padded_h, self._padded_w = in_padded[1], in_padded[2]
self._model_h, self._model_w = in_orig[2], in_orig[3]
# where the compiler placed the unpadded content inside the padded
# buffer: n_padded_ch_inputs = [[N(lo,hi), H(lo,hi), W(lo,hi),
# C(lo,hi)]] (one entry per input)
npad_data = manifest["n_padded_ch_inputs"]
if (
isinstance(npad_data, list)
and npad_data
and isinstance(npad_data[0], list)
):
npad_data = npad_data[0]
npad = [int(x) for x in npad_data]
if len(npad) == 8:
self._content_top, self._content_left = npad[2], npad[4]
else:
self._content_top = (self._padded_h - self._model_h) // 2
self._content_left = (self._padded_w - self._model_w) // 2
out_padded = [
tuple(int(x) for x in o) for o in manifest["output_shapes"]
]
out_real = [
tuple(int(x) for x in o)
for o in manifest["output_shapes_original"]
]
dequant = manifest["dequantize_params"]
if not (len(out_padded) == len(out_real) == len(dequant)):
raise ValueError(
"compiled manifest output counts mismatch: padded "
f"{len(out_padded)}, original {len(out_real)}, "
f"dequant {len(dequant)}"
)
# per head: real channel count (crop target) + dequant (scale, zp)
self._head_real = [o[1] for o in out_real]
self._head_dequant = [
(np.float32(d[0]), np.float32(d[1])) for d in dequant
]
self._out_shapes = out_padded
pp_path = os.path.join(model_dir, "postprocess_graph.onnx")
if not os.path.isfile(pp_path):
raise FileNotFoundError(
@ -284,12 +369,38 @@ class _AxeleraRuntimeInference:
self._pp = onnx_factory.InferenceSession(
pp_path, providers=["CPUExecutionProvider"]
)
self._pp_input_names = [i.name for i in self._pp.get_inputs()]
if len(self._pp_input_names) != len(HEAD_REAL_CHANNELS):
pp_inputs = self._pp.get_inputs()
self._pp_input_names = [i.name for i in pp_inputs]
if len(pp_inputs) != len(self._head_real):
raise ValueError(
"unexpected postprocess graph: expected "
f"{len(HEAD_REAL_CHANNELS)} inputs, got {len(self._pp_input_names)}"
"postprocess graph input count mismatch: "
f"{len(pp_inputs)} graph inputs vs "
f"{len(self._head_real)} model outputs"
)
# Map each postprocess input to the model output whose dequantised NCHW
# crop has exactly that (1, C, H, W) shape: the compile order of the
# raw heads is not guaranteed to match the postprocess graph's input
# order (they differ for YOLOX vs YOLO-NAS), shapes are unique per head.
self._feed_index: list[int] = []
for pin in pp_inputs:
want = tuple(int(x) for x in pin.shape)
matches = [
i
for i in range(len(self._head_real))
if (
1,
self._head_real[i],
out_padded[i][1],
out_padded[i][2],
)
== want
]
if len(matches) != 1:
raise ValueError(
f"cannot map postprocess input {pin.name} {want} to a "
"unique model output"
)
self._feed_index.append(matches[0])
# pooled letterbox-input + raw-int8-output sets, ping-ponged between
# the producer (runs the model) and the decoder (postprocesses the
@ -297,30 +408,23 @@ class _AxeleraRuntimeInference:
self._pool: queue.Queue = queue.Queue(maxsize=BUFFER_SETS)
for _ in range(BUFFER_SETS):
in_buf = np.full(
(1, PADDED_H, PADDED_W, 4), INPUT_Q_ZERO, np.int8
(1, self._padded_h, self._padded_w, 4),
INPUT_Q_ZERO,
np.int8,
)
out_bufs = [
np.zeros(
(1, MODEL_INPUT_H // s, MODEL_INPUT_W // s, 128), np.int8
)
for s in (8, 8, 16, 16, 32, 32)
]
out_bufs = [np.zeros(shape, np.int8) for shape in out_padded]
self._pool.put((in_buf, out_bufs))
self._decode_q: queue.Queue = queue.Queue(maxsize=BUFFER_SETS)
# letterbox geometry from the frame coords (width x height) into 640x640
# letterbox geometry from the frame coords (width x height) into the
# model's own workspace (640x640 for the YOLO presets)
self._scale = min(
MODEL_INPUT_W / self._width, MODEL_INPUT_H / self._height
self._model_w / self._width, self._model_h / self._height
)
new_w = round(self._width * self._scale)
new_h = round(self._height * self._scale)
self._new_w, self._new_h = max(1, new_w), max(1, new_h)
self._pad_t = (MODEL_INPUT_H - self._new_h) // 2
self._pad_l = (MODEL_INPUT_W - self._new_w) // 2
# dequantization per head: (int8 - zp) * scale (tuple stores scale, zp)
self._head_dequant = [
(np.float32(scale), np.float32(zp)) for scale, zp in HEAD_DEQUANT
]
self._pad_t = (self._model_h - self._new_h) // 2
self._pad_l = (self._model_w - self._new_w) // 2
# -- public contract used by AxeleraDetector --------------------------
@ -446,29 +550,51 @@ class _AxeleraRuntimeInference:
buf.fill(INPUT_Q_ZERO)
y0, y1 = self._pad_t, self._pad_t + self._new_h
x0, x1 = self._pad_l, self._pad_l + self._new_w
ct, cl = self._content_top, self._content_left
# channels 0..2 = quantized RGB (BGR -> RGB flip as in the SDK examples;
# the compiled input expects RGB semantics), channel 3 = alpha-255
buf[0, y0 + 1 : y1 + 1, x0 + 1 : x1 + 1, 0:3] = q[..., ::-1]
buf[0, y0 + 1 : y1 + 1, x0 + 1 : x1 + 1, 3] = 127
buf[0, ct + y0 : ct + y1, cl + x0 : cl + x1, 0:3] = q[..., ::-1]
buf[0, ct + y0 : ct + y1, cl + x0 : cl + x1, 3] = 127
def _decode(self, out_bufs) -> np.ndarray:
"""Dequantise heads + postprocess graph + conf threshold + NMS -> rows."""
# dequantize + crop to the real channel count + transpose NHWC -> NCHW
feeds = {}
for i, (name, obuf) in enumerate(
zip(self._pp_input_names, out_bufs)
):
scale, zp = self._head_dequant[i]
real = HEAD_REAL_CHANNELS[i]
for feed_i, name in enumerate(self._pp_input_names):
idx = self._feed_index[feed_i]
obuf = out_bufs[idx]
scale, zp = self._head_dequant[idx]
real = self._head_real[idx]
arr = obuf[..., :real].astype(np.float32)
arr = (arr - zp) * scale
feeds[name] = np.ascontiguousarray(arr.transpose(0, 3, 1, 2))
class_scores, boxes = self._pp.run(None, feeds)
scores = class_scores[0] # (8400, 80), sigmoid baked into the graph
boxes = boxes[0] # (8400, 4), xyxy in 640x640 input space
conf = scores.max(axis=1)
cls_idx = scores.argmax(axis=1).astype(np.int32)
pp_out = self._pp.run(None, feeds)
if len(pp_out) == 2:
# two outputs: class scores (sigmoid baked into the graph) + xyxy
class_scores, boxes = pp_out
scores = class_scores[0] # (8400, 80)
boxes = boxes[0]
conf = scores.max(axis=1)
cls_idx = scores.argmax(axis=1).astype(np.int32)
else:
# one fused output: [cx, cy, w, h, (objectness), C class scores].
# The SDK DecodeYolo emits CENTER-FORMAT boxes (confirmed by the
# vendor's axruntime_yolo11_onnxruntime.py example), absolute in
# the model input space.
raw = pp_out[0][0]
cx, cy, w, h = raw[:, 0], raw[:, 1], raw[:, 2], raw[:, 3]
boxes = np.stack(
[cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2], axis=1
)
total = raw.shape[-1]
cls = raw[:, 4:]
if total == 4 + 1 + len(self._labels):
cls = raw[:, 5:] * raw[:, 4:5] # YOLOX: score = obj * cls
if cls.shape[1] == 0:
return np.zeros((MAX_ROWS, 6), np.float32)
conf = cls.max(axis=1)
cls_idx = cls.argmax(axis=1).astype(np.int32)
keep = conf > CONF_THRESHOLD
boxes, conf, cls_idx = boxes[keep], conf[keep], cls_idx[keep]
if len(conf) == 0:
@ -560,8 +686,8 @@ class AxeleraDetector(DetectionApi):
logger.info(
"axelera: runtime model=%s (%dx%d input, frame %dx%d) ready",
self.working_model,
MODEL_INPUT_W,
MODEL_INPUT_H,
self._inference._model_w,
self._inference._model_h,
self.width,
self.height,
)