Add YOLOv26n RDK X5 model and quantized artifacts

x5_quantization/extract_yolo26_bpu.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+"""Extract YOLO26 raw detection heads from a standard Ultralytics ONNX export."""
+
+from __future__ import annotations
+
+import argparse
+import re
+from pathlib import Path
+
+import onnx
+from onnx import TensorProto, helper, shape_inference
+
+
+TERMINAL_HEAD = re.compile(
+    r"one2one_cv([23])\.(\d+)/one2one_cv\1\.\2\.2/Conv$"
+)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    args = parser.parse_args()
+
+    model = shape_inference.infer_shapes(onnx.load(args.input))
+    shapes = {
+        value.name: [dim.dim_value for dim in value.type.tensor_type.shape.dim]
+        for value in list(model.graph.value_info) + list(model.graph.output)
+    }
+    heads: dict[tuple[int, int], str] = {}
+    for node in model.graph.node:
+        match = TERMINAL_HEAD.search(node.name)
+        if match and node.op_type == "Conv":
+            branch, scale = int(match.group(1)), int(match.group(2))
+            heads[(branch, scale)] = node.output[0]
+
+    expected = {(branch, scale) for branch in (2, 3) for scale in range(3)}
+    if set(heads) != expected:
+        missing = sorted(expected - set(heads))
+        raise SystemExit(f"could not locate all YOLO26 one-to-one heads; missing {missing}")
+
+    del model.graph.output[:]
+    output_names = []
+    for scale, stride in enumerate((8, 16, 32)):
+        for branch, kind in ((3, "cls"), (2, "box")):
+            source = heads[(branch, scale)]
+            source_shape = shapes[source]
+            if len(source_shape) != 4:
+                raise SystemExit(f"unexpected head shape for {source}: {source_shape}")
+            output_name = f"{kind}_s{stride}"
+            model.graph.node.append(
+                helper.make_node(
+                    "Transpose",
+                    [source],
+                    [output_name],
+                    name=f"bpu_output_{kind}_s{stride}",
+                    perm=[0, 2, 3, 1],
+                )
+            )
+            model.graph.output.append(
+                helper.make_tensor_value_info(
+                    output_name,
+                    TensorProto.FLOAT,
+                    [source_shape[0], source_shape[2], source_shape[3], source_shape[1]],
+                )
+            )
+            output_names.append(output_name)
+
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    temporary = args.output.with_suffix(".unpruned.onnx")
+    onnx.save(model, temporary)
+    onnx.utils.extract_model(str(temporary), str(args.output), [model.graph.input[0].name], output_names)
+    temporary.unlink()
+    extracted = onnx.load(args.output)
+    onnx.checker.check_model(extracted)
+    print(f"wrote {args.output} with outputs:")
+    for output in extracted.graph.output:
+        dims = [dim.dim_value for dim in output.type.tensor_type.shape.dim]
+        print(f"  {output.name}: {dims}")
+
+
+if __name__ == "__main__":
+    main()