diff --git a/scripts/utils_2d_pose.py b/scripts/utils_2d_pose.py deleted file mode 100644 index b94e416..0000000 --- a/scripts/utils_2d_pose.py +++ /dev/null @@ -1,71 +0,0 @@ -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-m_8xb256-420e_coco-384x288.py", - pose2d_weights="https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/rtmpose-m_simcc-body7_pt-body7_420e-384x288-65e718c4_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, bbox_thr=0.3, nms_thr=0.3) - 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.asarray(kpts) - scores = np.asarray(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