Updated easypose scripts.

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)

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]):

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]):

extras/easypose/utils_2d_pose_ep.py

@@ -0,0 +1,68 @@
+import os
+
+import cv2
+import numpy as np
+import easypose as ep
+
+# ==================================================================================================
+
+filepath = os.path.dirname(os.path.realpath(__file__)) + "/"
+
+# ==================================================================================================
+
+
+def load_model():
+    print("Loading mmpose model ...")
+
+    model = ep.TopDown(
+        "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384x288_fp16_extra-steps.onnx",
+        "SimCC",
+        "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_320x320_fp16_extra-steps.onnx",
+        conf_threshold=0.3,
+        iou_threshold=0.3,
+        warmup=10,
+    )
+
+    print("Loaded mmpose model")
+    return model
+
+
+def load_wb_model():
+    print("Loading mmpose whole body model ...")
+
+    model = None
+
+    print("Loaded mmpose model")
+    return model
+
+
+# ==================================================================================================
+
+
+def get_2d_pose(model, imgs, num_joints=17):
+    """See: https://mmpose.readthedocs.io/en/latest/user_guides/inference.html#basic-usage"""
+
+    new_poses = []
+    for i in range(len(imgs)):
+        img = imgs[i]
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+        poses = []
+        dets = model.predict(img)
+        for pose in dets:
+            pose = pose.keypoints
+            pose = np.asarray(pose)
+
+            scores = pose[:, 2].reshape(-1, 1)
+            scores = np.clip(scores, 0, 1)
+            pose = np.concatenate((pose[:, :2], scores), axis=-1)
+
+            poses.append(pose)
+
+        if len(poses) == 0:
+            poses.append(np.zeros([num_joints, 3]))
+
+        poses = np.array(poses)
+        new_poses.append(poses)
+
+    return new_poses