84 lines
2.9 KiB
Python
84 lines
2.9 KiB
Python
#!/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()
|