diff --git a/privacy_blur/runtime.py b/privacy_blur/runtime.py index 84045a0..8087dc0 100644 --- a/privacy_blur/runtime.py +++ b/privacy_blur/runtime.py @@ -115,7 +115,15 @@ def _dynamic_onnx_model(path: Path, kind: str, identity_size: int) -> bytes: dimensions[0].dim_param = "batch" dimensions[2].dim_value = identity_size dimensions[3].dim_value = identity_size - for value in list(model.graph.output) + list(model.graph.value_info): + + # The bundled YuNet model contains intermediate value_info shapes inferred + # for its original 640x640 input. ONNX Runtime uses those stale values to + # restore static 6400/1600/400 output dimensions even after the graph + # outputs are marked dynamic. CUDA I/O binding then rejects the correctly + # sized buffers produced for a 320x320 input. Intermediate value_info is + # optional, so discard it and let ONNX Runtime infer it from the new input. + del model.graph.value_info[:] + for value in model.graph.output: output_dimensions = value.type.tensor_type.shape.dim if len(output_dimensions) > 1 and output_dimensions[1].dim_value in { 400, diff --git a/tests/test_runtime_onnx.py b/tests/test_runtime_onnx.py new file mode 100644 index 0000000..23eae10 --- /dev/null +++ b/tests/test_runtime_onnx.py @@ -0,0 +1,22 @@ +from pathlib import Path + +import onnxruntime as ort + +from privacy_blur.runtime import _dynamic_onnx_model + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] + + +def test_yunet_output_anchor_dimensions_remain_dynamic(): + model = _dynamic_onnx_model( + PROJECT_ROOT / "models" / "face_detection_yunet_2023mar.onnx", + "yunet", + 320, + ) + + session = ort.InferenceSession(model, providers=["CPUExecutionProvider"]) + + assert all( + output.shape[1] == "batched_anchors" for output in session.get_outputs() + )