From 7b9505ca024181e0cd1ae732cbeaa283e73d8b14 Mon Sep 17 00:00:00 2001 From: Daniel Date: Fri, 29 Nov 2024 17:05:43 +0100 Subject: [PATCH] Moved image normalization into onnx models. --- extras/easypose/README.md | 9 ++++ extras/easypose/detection.py | 7 --- extras/easypose/pose.py | 7 --- extras/mmdeploy/README.md | 4 +- extras/mmdeploy/add_norm_step.py | 82 ++++++++++++++++++++++++++++++++ scripts/utils_2d_pose_ep.py | 4 +- 6 files changed, 95 insertions(+), 18 deletions(-) create mode 100644 extras/mmdeploy/add_norm_step.py diff --git a/extras/easypose/README.md b/extras/easypose/README.md index af15c02..b9c769d 100644 --- a/extras/easypose/README.md +++ b/extras/easypose/README.md @@ -2,8 +2,17 @@ Code files originally from: https://github.com/Dominic23331/EasyPose.git +
+ ```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 +``` diff --git a/extras/easypose/detection.py b/extras/easypose/detection.py index 0c29a2c..cc5f458 100644 --- a/extras/easypose/detection.py +++ b/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 diff --git a/extras/easypose/pose.py b/extras/easypose/pose.py index 677493b..fe5aa72 100644 --- a/extras/easypose/pose.py +++ b/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 diff --git a/extras/mmdeploy/README.md b/extras/mmdeploy/README.md index e8c4eb5..cf5c094 100644 --- a/extras/mmdeploy/README.md +++ b/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 ``` diff --git a/extras/mmdeploy/add_norm_step.py b/extras/mmdeploy/add_norm_step.py new file mode 100644 index 0000000..a8450fb --- /dev/null +++ b/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() diff --git a/scripts/utils_2d_pose_ep.py b/scripts/utils_2d_pose_ep.py index 4947b0b..76da721 100644 --- a/scripts/utils_2d_pose_ep.py +++ b/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,