Files
spoterembedding/normalization/blazepose_mapping.py
2023-04-07 09:44:12 +00:00

88 lines
2.5 KiB
Python

_BODY_KEYPOINT_MAPPING = {
"nose": "nose",
"left_eye": "leftEye",
"right_eye": "rightEye",
"left_ear": "leftEar",
"right_ear": "rightEar",
"left_shoulder": "leftShoulder",
"right_shoulder": "rightShoulder",
"left_elbow": "leftElbow",
"right_elbow": "rightElbow",
"left_wrist": "leftWrist",
"right_wrist": "rightWrist"
}
_HAND_KEYPOINT_MAPPING = {
"wrist": "wrist",
"index_finger_tip": "indexTip",
"index_finger_dip": "indexDIP",
"index_finger_pip": "indexPIP",
"index_finger_mcp": "indexMCP",
"middle_finger_tip": "middleTip",
"middle_finger_dip": "middleDIP",
"middle_finger_pip": "middlePIP",
"middle_finger_mcp": "middleMCP",
"ring_finger_tip": "ringTip",
"ring_finger_dip": "ringDIP",
"ring_finger_pip": "ringPIP",
"ring_finger_mcp": "ringMCP",
"pinky_tip": "littleTip",
"pinky_dip": "littleDIP",
"pinky_pip": "littlePIP",
"pinky_mcp": "littleMCP",
"thumb_tip": "thumbTip",
"thumb_ip": "thumbIP",
"thumb_mcp": "thumbMP",
"thumb_cmc": "thumbCMC"
}
def map_blazepose_keypoint(column):
# Remove _x, _y suffixes
suffix = column[-2:].upper()
column = column[:-2]
if column.startswith("left_hand_"):
hand = "left"
finger_name = column[10:]
elif column.startswith("right_hand_"):
hand = "right"
finger_name = column[11:]
else:
if column not in _BODY_KEYPOINT_MAPPING:
return None
mapped = _BODY_KEYPOINT_MAPPING[column]
return mapped + suffix
if finger_name not in _HAND_KEYPOINT_MAPPING:
return None
mapped = _HAND_KEYPOINT_MAPPING[finger_name]
return f"{mapped}_{hand}{suffix}"
def map_blazepose_df(df):
for index, row in df.iterrows():
lsx = row["leftShoulder_X"]
rsx = row["rightShoulder_X"]
lsy = row["leftShoulder_Y"]
rsy = row["rightShoulder_Y"]
# convert all to list
lsx = lsx[1:-1].split(",")
rsx = rsx[1:-1].split(",")
lsy = lsy[1:-1].split(",")
rsy = rsy[1:-1].split(",")
sequence_size = len(lsx)
neck_x = []
neck_y = []
# Treat each element of the sequence (analyzed frame) individually
for sequence_index in range(sequence_size):
neck_x.append((float(lsx[sequence_index]) + float(rsx[sequence_index])) / 2)
neck_y.append((float(lsy[sequence_index]) + float(rsy[sequence_index])) / 2)
df.loc[index, "neck_X"] = str(neck_x)
df.loc[index, "neck_Y"] = str(neck_y)
return df