develop

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,

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()
+    )