crosstyan/RapidPoseTriangulation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Optional batched pose processing.

Browse Source
This commit is contained in:
Daniel
2024-12-18 16:22:08 +01:00
parent 7b8d209601
commit 07426fac2f
8 changed files with 151 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
extras/mmdeploy/add_extra_steps.py
View File

@ -1,12 +1,15 @@
import re
import numpy as np
import onnx
from onnx import TensorProto, compose, helper, numpy_helper
from onnx import TensorProto, 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"
det_model_path = base_path + "rtmdet-nano_1x3x320x320.onnx"
pose_model_path1 = base_path + "rtmpose-m_Bx3x384x288.onnx"
pose_model_path2 = base_path + "rtmpose-m_1x3x384x288.onnx"
norm_mean = -1 * (np.array([0.485, 0.456, 0.406]) * 255)
norm_std = 1.0 / (np.array([0.229, 0.224, 0.225]) * 255)
@ -97,6 +100,11 @@ def add_steps_to_onnx(model_path):
for i, j in enumerate([0, 3, 1, 2]):
input_shape[j].dim_value = dims[i]
# Set the batch size to a defined string
input_shape = graph.input[0].type.tensor_type.shape.dim
if input_shape[0].dim_value == 0:
input_shape[0].dim_param = "batch_size"
# Rename the input tensor
main_input_image_name = model.graph.input[0].name
for node in model.graph.node:
@ -108,7 +116,8 @@ def add_steps_to_onnx(model_path):
# 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")
path = re.sub(r"(x)(\d+)x(\d+)x(\d+)", r"\1\3x\4x\2", model_path)
path = path.replace(".onnx", "_extra-steps.onnx")
onnx.save(model, path)
@ -116,10 +125,12 @@ def add_steps_to_onnx(model_path):
def main():
add_steps_to_onnx(pose_model_path)
add_steps_to_onnx(det_model_path)
add_steps_to_onnx(pose_model_path1)
add_steps_to_onnx(pose_model_path2)
add_steps_to_onnx(det_model_path.replace(".onnx", "_fp16.onnx"))
add_steps_to_onnx(pose_model_path.replace(".onnx", "_fp16.onnx"))
add_steps_to_onnx(pose_model_path1.replace(".onnx", "_fp16.onnx"))
add_steps_to_onnx(pose_model_path2.replace(".onnx", "_fp16.onnx"))
# ==================================================================================================