Testing onnx runtime with easypose.

scripts/utils_2d_pose_ep.py

@@ -0,0 +1,61 @@
+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("rtmpose_m", "SimCC", "rtmdet_s")
+
+    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