Moved pose pre/post-processing into onnx graph.

2024-12-03 17:59:09 +01:00
parent 2e5c3399ed
commit 97ff32b9ce
4 changed files with 411 additions and 284 deletions
@@ -121,14 +121,55 @@ def add_steps_to_onnx(model_path):

        # Update nodes from postprocess model to use the input of the main network
        pp2_input_image_name = pp2_model.graph.input[0].name
-        main_input_name = model.graph.input[0].name
+        main_input_image_name = model.graph.input[0].name
        for node in model.graph.node:
            for idx, name in enumerate(node.input):
                if name == pp2_input_image_name:
-                    node.input[idx] = main_input_name
+                    node.input[idx] = main_input_image_name
        model.graph.input.pop(1)

-    # Set input type to int8
+    if "pose" in model_path:
+        # Add preprocess model to main network
+        pp1_model = onnx.load(base_path + "pose_preprocess.onnx")
+        model = compose.add_prefix(model, prefix="main_")
+        pp1_model = compose.add_prefix(pp1_model, prefix="preprocess_")
+        model = compose.merge_models(
+            pp1_model,
+            model,
+            io_map=[
+                (pp1_model.graph.output[0].name, model.graph.input[0].name),
+            ],
+        )
+
+        # Add postprocess model
+        pp2_model = onnx.load(base_path + "pose_postprocess.onnx")
+        pp2_model = compose.add_prefix(pp2_model, prefix="postprocess_")
+        model = compose.merge_models(
+            model,
+            pp2_model,
+            io_map=[
+                (model.graph.output[0].name, pp2_model.graph.input[2].name),
+            ],
+        )
+
+        # Update nodes from postprocess model to use the input of the main network
+        pp2_input_image_name = pp2_model.graph.input[0].name
+        pp2_input_bbox_name = pp2_model.graph.input[1].name
+        main_input_image_name = model.graph.input[0].name
+        main_input_bbox_name = model.graph.input[1].name
+        for node in model.graph.node:
+            for idx, name in enumerate(node.input):
+                if name == pp2_input_image_name:
+                    node.input[idx] = main_input_image_name
+                if name == pp2_input_bbox_name:
+                    node.input[idx] = main_input_bbox_name
+        model.graph.input.pop(2)
+        model.graph.input.pop(2)
+
+        # Set input box type to int32
+        model.graph.input[1].type.tensor_type.elem_type = TensorProto.INT32
+
+    # Set input image type to int8
    model.graph.input[0].type.tensor_type.elem_type = TensorProto.UINT8

    path = model_path.replace(".onnx", "_extra-steps.onnx")