Copied initial tools.

scripts/utils_2d_pose.py

@@ -0,0 +1,71 @@
+import os
+
+import numpy as np
+from mmpose.apis import MMPoseInferencer
+
+# ==================================================================================================
+
+filepath = os.path.dirname(os.path.realpath(__file__)) + "/"
+
+# ==================================================================================================
+
+
+def load_model():
+    print("Loading mmpose model ...")
+
+    model = MMPoseInferencer(
+        pose2d="/mmpose/projects/rtmpose/rtmpose/body_2d_keypoint/rtmpose-l_8xb256-420e_coco-384x288.py",
+        pose2d_weights="https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/rtmpose-l_simcc-body7_pt-body7_420e-384x288-3f5a1437_20230504.pth",
+        det_model="/mmpose/projects/rtmpose/rtmdet/person/rtmdet_nano_320-8xb32_coco-person.py",
+        det_weights="https://download.openmmlab.com/mmpose/v1/projects/rtmpose/rtmdet_nano_8xb32-100e_coco-obj365-person-05d8511e.pth",
+        det_cat_ids=[0],
+    )
+
+    print("Loaded mmpose model")
+    return model
+
+
+def load_wb_model():
+    print("Loading mmpose whole body model ...")
+
+    model = MMPoseInferencer(
+        pose2d="/mmpose/projects/rtmpose/rtmpose/wholebody_2d_keypoint/rtmpose-l_8xb32-270e_coco-wholebody-384x288.py",
+        pose2d_weights="https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/rtmpose-l_simcc-coco-wholebody_pt-aic-coco_270e-384x288-eaeb96c8_20230125.pth",
+        det_model="/mmpose/projects/rtmpose/rtmdet/person/rtmdet_nano_320-8xb32_coco-person.py",
+        det_weights="https://download.openmmlab.com/mmpose/v1/projects/rtmpose/rtmdet_nano_8xb32-100e_coco-obj365-person-05d8511e.pth",
+        det_cat_ids=[0],
+    )
+
+    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"""
+
+    result_generator = model(imgs, show=False)
+    new_poses = []
+    for _ in range(len(imgs)):
+        result = next(result_generator)
+
+        poses = []
+        for i in range(len(result["predictions"][0])):
+            kpts = result["predictions"][0][i]["keypoints"]
+            scores = result["predictions"][0][i]["keypoint_scores"]
+
+            kpts = np.array(kpts)
+            scores = np.array(scores).reshape(-1, 1)
+            scores = np.clip(scores, 0, 1)
+            pose = np.concatenate((kpts, 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