Moved image normalization into onnx models.
This commit is contained in:
@ -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
|
||||
```
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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
|
||||
```
|
||||
|
||||
82
extras/mmdeploy/add_norm_step.py
Normal file
82
extras/mmdeploy/add_norm_step.py
Normal file
@ -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()
|
||||
@ -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,
|
||||
|
||||
Reference in New Issue
Block a user