From 1bb7dd966e836afed09fe03361686eac967811a2 Mon Sep 17 00:00:00 2001
From: Daniel <daniel@mail.de>
Date: Wed, 18 Dec 2024 16:53:32 +0100
Subject: [PATCH] Converted wb model as well.

---
 extras/mmdeploy/README.md          | 20 ++++++++++++++++++--
 extras/mmdeploy/add_extra_steps.py |  6 ++++++
 scripts/test_skelda_dataset.py     |  1 +
 scripts/utils_2d_pose_ort.py       | 14 ++++++++++----
 4 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/extras/mmdeploy/README.md b/extras/mmdeploy/README.md
index fb30c9f..c382915 100644
--- a/extras/mmdeploy/README.md
+++ b/extras/mmdeploy/README.md
@@ -39,7 +39,6 @@ python3 ./tools/deploy.py \
     /mmpose/projects/rtmpose/examples/onnxruntime/human-pose.jpeg \
     --work-dir work_dir \
     --show
-
 mv /mmdeploy/work_dir/end2end.onnx /RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_1x3x384x288"$withFP16".onnx
 
 python3 ./tools/deploy.py \
@@ -49,8 +48,25 @@ python3 ./tools/deploy.py \
     /mmpose/projects/rtmpose/examples/onnxruntime/human-pose.jpeg \
     --work-dir work_dir \
     --show
-
 mv /mmdeploy/work_dir/end2end.onnx /RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_Bx3x384x288"$withFP16".onnx
+
+python3 ./tools/deploy.py \
+    configs/mmpose/pose-detection_simcc_onnxruntime_static-384x288"$withFP16".py \
+    /mmpose/projects/rtmpose/rtmpose/wholebody_2d_keypoint/rtmpose-l_8xb32-270e_coco-wholebody-384x288.py \
+    https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/rtmpose-l_simcc-coco-wholebody_pt-aic-coco_270e-384x288-eaeb96c8_20230125.pth \
+    /mmpose/projects/rtmpose/examples/onnxruntime/human-pose.jpeg \
+    --work-dir work_dir \
+    --show
+mv /mmdeploy/work_dir/end2end.onnx /RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-l_wb_1x3x384x288"$withFP16".onnx
+
+python3 ./tools/deploy.py \
+    configs/mmpose/pose-detection_simcc_onnxruntime_dynamic-384x288"$withFP16".py \
+    /mmpose/projects/rtmpose/rtmpose/wholebody_2d_keypoint/rtmpose-l_8xb32-270e_coco-wholebody-384x288.py \
+    https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/rtmpose-l_simcc-coco-wholebody_pt-aic-coco_270e-384x288-eaeb96c8_20230125.pth \
+    /mmpose/projects/rtmpose/examples/onnxruntime/human-pose.jpeg \
+    --work-dir work_dir \
+    --show
+mv /mmdeploy/work_dir/end2end.onnx /RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-l_wb_Bx3x384x288"$withFP16".onnx
 ```
 
 ```bash
diff --git a/extras/mmdeploy/add_extra_steps.py b/extras/mmdeploy/add_extra_steps.py
index 6ab35fa..4eb822b 100644
--- a/extras/mmdeploy/add_extra_steps.py
+++ b/extras/mmdeploy/add_extra_steps.py
@@ -10,6 +10,8 @@ base_path = "/RapidPoseTriangulation/extras/mmdeploy/exports/"
 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"
+pose_model_path3 = base_path + "rtmpose-l_wb_Bx3x384x288.onnx"
+pose_model_path4 = base_path + "rtmpose-l_wb_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)
@@ -128,9 +130,13 @@ def main():
     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(pose_model_path3)
+    add_steps_to_onnx(pose_model_path4)
     add_steps_to_onnx(det_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"))
+    add_steps_to_onnx(pose_model_path3.replace(".onnx", "_fp16.onnx"))
+    add_steps_to_onnx(pose_model_path4.replace(".onnx", "_fp16.onnx"))
 
 
 # ==================================================================================================
diff --git a/scripts/test_skelda_dataset.py b/scripts/test_skelda_dataset.py
index 1fb65bb..a5bb04e 100644
--- a/scripts/test_skelda_dataset.py
+++ b/scripts/test_skelda_dataset.py
@@ -116,6 +116,7 @@ datasets = {
         "path": "/datasets/human36m/skelda/wb/test.json",
         "take_interval": 100,
         "min_bbox_score": 0.4,
+        "batch_poses": False,
     },
     "egohumans_tagging": {
         "path": "/datasets/egohumans/skelda/all.json",
diff --git a/scripts/utils_2d_pose_ort.py b/scripts/utils_2d_pose_ort.py
index fe29cca..944d5a4 100644
--- a/scripts/utils_2d_pose_ort.py
+++ b/scripts/utils_2d_pose_ort.py
@@ -479,12 +479,18 @@ def load_model(min_bbox_score=0.3, min_bbox_area=0.1 * 0.1, batch_poses=False):
     return model
 
 
-def load_wb_model():
-    print("Loading mmpose whole body model ...")
+def load_wb_model(min_bbox_score=0.3, min_bbox_area=0.1 * 0.1, batch_poses=False):
+    print("Loading 2D-WB model ...")
 
-    model = None
+    model = TopDown(
+        "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_1x320x320x3_fp16_extra-steps.onnx",
+        f"/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-l_wb_{'B' if batch_poses else '1'}x384x288x3_fp16_extra-steps.onnx",
+        box_conf_threshold=min_bbox_score,
+        box_min_area=min_bbox_area,
+        warmup=30,
+    )
 
-    print("Loaded mmpose model")
+    print("Loaded 2D-WB model")
     return model