#!/usr/bin/env python3 """Generate output-only ONNX variants for the real yolo11n depth_35 model.""" from __future__ import annotations import argparse from pathlib import Path import onnx from onnx import TensorProto, helper, shape_inference ORIGINAL_DEBUG_OUTPUTS = [ "/model.20/act/Mul_output_0", "/model.9/cv1/act/Mul_output_0", "/model.8/m.0/m/m.0/cv2/conv/Conv_output_0", "/model.23/cv3.1/cv3.1.0/cv3.1.0.0/act/Mul_output_0", "/model.8/m.0/m/m.1/cv1/act/Mul_output_0", "/model.23/dfl/Transpose_output_0", "output0", "/model.23/cv2.1/cv2.1.0/act/Mul_output_0", ] LOCALIZATION_NODE_NAMES = [ "/model.23/dfl/Transpose", "/model.23/dfl/Softmax", "/model.23/dfl/conv/Conv", "/model.23/dfl/Reshape_1", "/model.23/Slice", "/model.23/Slice_1", "/model.23/Sub", "/model.23/Add_1", "/model.23/Sub_1", "/model.23/Concat_4", "/model.23/Mul_2", "/model.23/Sigmoid", "/model.23/Concat_5", ] def collect_value_infos(model: onnx.ModelProto) -> dict[str, onnx.ValueInfoProto]: infos: dict[str, onnx.ValueInfoProto] = {} for value in list(model.graph.input) + list(model.graph.output) + list(model.graph.value_info): infos[value.name] = value return infos def clone_value_info(value: onnx.ValueInfoProto, name: str | None = None) -> onnx.ValueInfoProto: cloned = onnx.ValueInfoProto() cloned.CopyFrom(value) if name is not None: cloned.name = name return cloned def make_tensor_value_info_from_type(name: str, tensor_type: onnx.TypeProto.Tensor) -> onnx.ValueInfoProto: shape = [] for dim in tensor_type.shape.dim: if dim.HasField("dim_value"): shape.append(dim.dim_value) elif dim.HasField("dim_param"): shape.append(dim.dim_param) else: shape.append(None) return helper.make_tensor_value_info(name, tensor_type.elem_type, shape) def lookup_output_value_info(model: onnx.ModelProto, value_infos: dict[str, onnx.ValueInfoProto], output_name: str) -> onnx.ValueInfoProto: value = value_infos.get(output_name) if value is not None: return clone_value_info(value) for initializer in model.graph.initializer: if initializer.name != output_name: continue dims = list(initializer.dims) return helper.make_tensor_value_info(output_name, initializer.data_type, dims) raise KeyError(f"missing value info for output {output_name}") def build_model_with_outputs( base_model: onnx.ModelProto, inferred_model: onnx.ModelProto, output_names: list[str], extra_nodes: list[onnx.NodeProto] | None = None, extra_value_infos: list[onnx.ValueInfoProto] | None = None, ) -> onnx.ModelProto: value_infos = collect_value_infos(inferred_model) if extra_value_infos: for value in extra_value_infos: value_infos[value.name] = value model = onnx.ModelProto() model.CopyFrom(base_model) del model.graph.output[:] for output_name in output_names: model.graph.output.append(lookup_output_value_info(inferred_model, value_infos, output_name)) if extra_nodes: model.graph.node.extend(extra_nodes) if extra_value_infos: model.graph.value_info.extend(extra_value_infos) onnx.checker.check_model(model) return model def find_node_output(model: onnx.ModelProto, node_name: str) -> str: for node in model.graph.node: if node.name == node_name: return node.output[0] matching_names = sorted(node.name for node in model.graph.node if "model.23" in node.name and "dfl" in node.name) suffix = "" if matching_names: suffix = "\nmatching /model.23 dfl nodes:\n " + "\n ".join(matching_names) raise KeyError(f"could not find node named {node_name}{suffix}") def save_variant(model: onnx.ModelProto, out_dir: Path, variant_name: str) -> None: variant_dir = out_dir / variant_name variant_dir.mkdir(parents=True, exist_ok=True) onnx.save(model, variant_dir / f"{variant_name}.onnx") def unique_preserving_order(names: list[str]) -> list[str]: seen: set[str] = set() unique_names: list[str] = [] for name in names: if name in seen: continue seen.add(name) unique_names.append(name) return unique_names def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--input", required=True, help="Path to yolo11n depth_35 ONNX model.") parser.add_argument("--out-dir", required=True, help="Directory where variants will be generated.") args = parser.parse_args() input_path = Path(args.input).resolve() out_dir = Path(args.out_dir).resolve() out_dir.mkdir(parents=True, exist_ok=True) base_model = onnx.load(input_path) inferred_model = shape_inference.infer_shapes(base_model) output0_only = build_model_with_outputs(base_model, inferred_model, ["output0"]) save_variant(output0_only, out_dir, "output0_only") output0_first = build_model_with_outputs( base_model, inferred_model, ["output0"] + [name for name in ORIGINAL_DEBUG_OUTPUTS if name != "output0"], ) save_variant(output0_first, out_dir, "output0_first_with_original_debug_outputs") output0_last = build_model_with_outputs( base_model, inferred_model, [name for name in ORIGINAL_DEBUG_OUTPUTS if name != "output0"] + ["output0"], ) save_variant(output0_last, out_dir, "output0_last_with_original_debug_outputs") identity_name = "output0_identity" identity_node = helper.make_node("Identity", ["output0"], [identity_name], name="output0_identity_node") output0_value = lookup_output_value_info(inferred_model, collect_value_infos(inferred_model), "output0") duplicated = build_model_with_outputs( base_model, inferred_model, ["output0", identity_name], extra_nodes=[identity_node], extra_value_infos=[make_tensor_value_info_from_type(identity_name, output0_value.type.tensor_type)], ) save_variant(duplicated, out_dir, "output0_duplicated") localization_outputs = [ "/model.23/dfl/Transpose_output_0", find_node_output(base_model, "/model.23/dfl/Softmax"), find_node_output(base_model, "/model.23/dfl/conv/Conv"), find_node_output(base_model, "/model.23/dfl/Reshape_1"), find_node_output(base_model, "/model.23/Slice"), find_node_output(base_model, "/model.23/Slice_1"), find_node_output(base_model, "/model.23/Sub"), find_node_output(base_model, "/model.23/Add_1"), find_node_output(base_model, "/model.23/Sub_1"), find_node_output(base_model, "/model.23/Concat_4"), find_node_output(base_model, "/model.23/Mul_2"), find_node_output(base_model, "/model.23/Sigmoid"), find_node_output(base_model, "/model.23/Concat_5"), "output0", ] localization = build_model_with_outputs( base_model, inferred_model, unique_preserving_order(localization_outputs), ) save_variant(localization, out_dir, "yolo_tail_localization_outputs") return 0 if __name__ == "__main__": raise SystemExit(main())