Moved image normalization into onnx models.

extras/mmdeploy/README.md

@@ -22,7 +22,7 @@ python3 ./tools/deploy.py \
     --work-dir work_dir \
     --show
 
-mv /mmdeploy/work_dir/end2end.onnx /RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet_nano_320.onnx
+mv /mmdeploy/work_dir/end2end.onnx /RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_320x320.onnx
 ```
 
 ```bash
@@ -37,5 +37,5 @@ python3 ./tools/deploy.py \
     --work-dir work_dir \
     --show
 
-mv /mmdeploy/work_dir/end2end.onnx /RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384.onnx
+mv /mmdeploy/work_dir/end2end.onnx /RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384x288.onnx
 ```

extras/mmdeploy/add_norm_step.py

@@ -0,0 +1,82 @@
+import numpy as np
+import onnx
+from onnx import helper, numpy_helper
+
+# ==================================================================================================
+
+base_path = "/RapidPoseTriangulation/extras/mmdeploy/exports/"
+pose_model_path = base_path + "rtmpose-m_384x288.onnx"
+det_model_path = base_path + "rtmdet-nano_320x320.onnx"
+
+norm_mean = -1 * np.array([103.53, 116.28, 123.675])
+norm_std = 1.0 / np.array([57.375, 57.12, 58.395])
+
+
+# ==================================================================================================
+
+
+def add_steps_to_onnx(model_path, use_bgr=False):
+
+    # Load existing model
+    model = onnx.load(model_path)
+    graph = model.graph
+
+    mean = norm_mean.astype(np.float32)
+    std = norm_std.astype(np.float32)
+
+    if use_bgr:
+        mean = mean[::-1]
+        std = std[::-1]
+
+    mean = np.reshape(mean, (1, 3, 1, 1)).astype(np.float32)
+    std = np.reshape(std, (1, 3, 1, 1)).astype(np.float32)
+
+    # Add the initializers to the graph
+    mean_initializer = numpy_helper.from_array(mean, name="norm_mean")
+    std_initializer = numpy_helper.from_array(std, name="norm_std")
+    graph.initializer.extend([mean_initializer, std_initializer])
+
+    # Define layer names, assuming the first input is the image tensor
+    input_name = graph.input[0].name
+    mean_added_output = "mean_added_output"
+    normalized_output = "normalized_output"
+
+    # Node to add mean
+    mean_add_node = helper.make_node(
+        "Add",
+        inputs=[input_name, "norm_mean"],
+        outputs=[mean_added_output],
+        name="Mean_Addition",
+    )
+
+    # Node to multiply by std
+    std_mul_node = helper.make_node(
+        "Mul",
+        inputs=[mean_added_output, "norm_std"],
+        outputs=[normalized_output],
+        name="Std_Multiplication",
+    )
+
+    # Replace original input of the model with the output of normalization
+    for node in graph.node:
+        for idx, input_name_in_node in enumerate(node.input):
+            if input_name_in_node == input_name:
+                node.input[idx] = normalized_output
+
+    # Add the new nodes to the graph
+    graph.node.insert(0, mean_add_node)
+    graph.node.insert(1, std_mul_node)
+
+    path = model_path.replace(".onnx", "_with-norm.onnx")
+    onnx.save(model, path)
+
+
+def main():
+    add_steps_to_onnx(pose_model_path, use_bgr=True)
+    add_steps_to_onnx(det_model_path, use_bgr=False)
+
+
+# ==================================================================================================
+
+if __name__ == "__main__":
+    main()