diff --git a/frigate/detectors/plugins/axelera.py b/frigate/detectors/plugins/axelera.py index 0ba5ca41da..b8cd91636e 100644 --- a/frigate/detectors/plugins/axelera.py +++ b/frigate/detectors/plugins/axelera.py @@ -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 ``/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, )