70 lines
1.8 KiB
Python
70 lines
1.8 KiB
Python
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")
|
|
model = ep.TopDown(
|
|
"/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384.onnx",
|
|
"SimCC",
|
|
"/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet_nano_320.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
|