From ee8b9bafb361bb8dfe04b0fa78bb91a99a4db051 Mon Sep 17 00:00:00 2001
From: Daniel <daniel@mail.de>
Date: Fri, 6 Dec 2024 17:35:49 +0100
Subject: [PATCH] Updated easypose scripts.

---
 extras/easypose/base_model.py                    | 10 +++++++---
 extras/easypose/detection.py                     |  4 ++--
 extras/easypose/pose.py                          |  5 +----
 {scripts => extras/easypose}/utils_2d_pose_ep.py |  7 ++-----
 4 files changed, 12 insertions(+), 14 deletions(-)
 rename {scripts => extras/easypose}/utils_2d_pose_ep.py (84%)

diff --git a/extras/easypose/base_model.py b/extras/easypose/base_model.py
index b6dc7c5..4eab24d 100644
--- a/extras/easypose/base_model.py
+++ b/extras/easypose/base_model.py
@@ -29,10 +29,14 @@ class BaseModel(ABC):
         self.input_shape = self.session.get_inputs()[0].shape
 
         input_type = self.session.get_inputs()[0].type
-        if input_type == 'tensor(float16)':
-            self.input_type = np.float16
-        else:
+        if input_type == 'tensor(float32)':
             self.input_type = np.float32
+        elif input_type == 'tensor(float16)':
+            self.input_type = np.float16
+        elif input_type == 'tensor(uint8)':
+            self.input_type = np.uint8
+        else:
+            raise ValueError('Unknown input type: ', input_type)
 
         if warmup > 0:
             self.warmup(warmup)
diff --git a/extras/easypose/detection.py b/extras/easypose/detection.py
index 0b970b4..e8cf15a 100644
--- a/extras/easypose/detection.py
+++ b/extras/easypose/detection.py
@@ -20,12 +20,12 @@ class RTMDet(BaseModel):
         self.scale = 0
 
     def preprocess(self, image: np.ndarray):
-        th, tw = self.input_shape[2:]
+        th, tw = self.input_shape[1:3]
         image, self.dx, self.dy, self.scale = letterbox(
             image, (tw, th), fill_value=114
         )
         tensor = np.asarray(image).astype(self.input_type, copy=False)[..., ::-1]
-        tensor = np.expand_dims(tensor, axis=0).transpose((0, 3, 1, 2))
+        tensor = np.expand_dims(tensor, axis=0)
         return tensor
 
     def postprocess(self, tensor: List[np.ndarray]):
diff --git a/extras/easypose/pose.py b/extras/easypose/pose.py
index 070ef0c..b0e6327 100644
--- a/extras/easypose/pose.py
+++ b/extras/easypose/pose.py
@@ -39,13 +39,10 @@ class Heatmap(BaseModel):
 class SimCC(BaseModel):
     def __init__(self, model_path: str, device: str = 'CUDA', warmup: int = 30):
         super(SimCC, self).__init__(model_path, device, warmup)
-        self.dx = 0
-        self.dy = 0
-        self.scale = 0
 
     def preprocess(self, image: np.ndarray):
         tensor = np.asarray(image).astype(self.input_type, copy=False)
-        tensor = np.expand_dims(tensor, axis=0).transpose((0, 3, 1, 2))
+        tensor = np.expand_dims(tensor, axis=0)
         return tensor
 
     def postprocess(self, tensor: List[np.ndarray]):
diff --git a/scripts/utils_2d_pose_ep.py b/extras/easypose/utils_2d_pose_ep.py
similarity index 84%
rename from scripts/utils_2d_pose_ep.py
rename to extras/easypose/utils_2d_pose_ep.py
index d1ec06e..67fe3e2 100644
--- a/scripts/utils_2d_pose_ep.py
+++ b/extras/easypose/utils_2d_pose_ep.py
@@ -14,13 +14,10 @@ filepath = os.path.dirname(os.path.realpath(__file__)) + "/"
 def load_model():
     print("Loading mmpose model ...")
 
-    # model = ep.TopDown("rtmpose_m", "SimCC", "rtmdet_s")
     model = ep.TopDown(
-        "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384x288_with-norm.onnx",
-        # "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384x288_fp16_with-norm.onnx",
+        "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384x288_fp16_extra-steps.onnx",
         "SimCC",
-        "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_320x320_with-norm.onnx",
-        # "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_320x320_fp16_with-norm.onnx",
+        "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_320x320_fp16_extra-steps.onnx",
         conf_threshold=0.3,
         iou_threshold=0.3,
         warmup=10,