Moved image normalization into onnx models.

extras/easypose/README.md

@@ -2,8 +2,17 @@
 
 Code files originally from: https://github.com/Dominic23331/EasyPose.git
 
+<br>
+
 ```bash
 docker build --progress=plain -f extras/easypose/dockerfile -t rpt_easypose .
 
 ./extras/easypose/run_container.sh
 ```
+
+```bash
+export CUDA_VISIBLE_DEVICES=0
+
+python3 /RapidPoseTriangulation/scripts/test_triangulate.py
+python3 /RapidPoseTriangulation/scripts/test_skelda_dataset.py
+```

extras/easypose/detection.py

@@ -19,19 +19,12 @@ class RTMDet(BaseModel):
         self.dy = 0
         self.scale = 0
 
-        norm_mean = -1 * np.array([123.675, 116.28, 103.53])
-        norm_std = 1.0 / np.array([58.395, 57.12, 57.375])
-        self.norm_mean = np.reshape(norm_mean, (1, 1, 3)).astype(np.float32)
-        self.norm_std = np.reshape(norm_std, (1, 1, 3)).astype(np.float32)
-
     def preprocess(self, image: np.ndarray):
         th, tw = self.input_shape[2:]
         tensor, self.dx, self.dy, self.scale = letterbox(
             image, (tw, th), fill_value=114
         )
         tensor = tensor.astype(np.float32, copy=False)
-        tensor += self.norm_mean
-        tensor *= self.norm_std
         tensor = tensor[..., ::-1]
         tensor = np.expand_dims(tensor, axis=0).transpose((0, 3, 1, 2))
         return tensor

extras/easypose/pose.py

@@ -43,16 +43,9 @@ class SimCC(BaseModel):
         self.dy = 0
         self.scale = 0
 
-        norm_mean = -1 * np.array([123.675, 116.28, 103.53])
-        norm_std = 1.0 / np.array([58.395, 57.12, 57.375])
-        self.norm_mean = np.reshape(norm_mean, (1, 1, 3)).astype(np.float32)
-        self.norm_std = np.reshape(norm_std, (1, 1, 3)).astype(np.float32)
-
     def preprocess(self, image: np.ndarray):
         tensor, self.dx, self.dy, self.scale = image, 0, 0, 1
         tensor = tensor.astype(np.float32, copy=False)
-        tensor += self.norm_mean
-        tensor *= self.norm_std
         tensor = np.expand_dims(tensor, axis=0).transpose((0, 3, 1, 2))
         return tensor
 

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

scripts/utils_2d_pose_ep.py

@@ -16,9 +16,9 @@ def load_model():
 
     # model = ep.TopDown("rtmpose_m", "SimCC", "rtmdet_s")
     model = ep.TopDown(
-        "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384.onnx",
+        "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384x288_with-norm.onnx",
         "SimCC",
-        "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet_nano_320.onnx",
+        "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_320x320_with-norm.onnx",
         conf_threshold=0.3,
         iou_threshold=0.3,
         warmup=10,