engine/pose_estimation/pose_utils/humans.py

# Multi-HMR
# Copyright (c) 2024-present NAVER Corp.
# CC BY-NC-SA 4.0 license

import numpy as np
import torch.nn.functional as F
import torch
import roma
from smplx.joint_names import JOINT_NAMES


def rot6d_to_rotmat(x):
    """
    6D rotation representation to 3x3 rotation matrix.
    Args:
        x: (B,6) Batch of 6-D rotation representations.
    Returns:
        torch.Tensor: Batch of corresponding rotation matrices with shape (B,3,3).
    """
    x = x.reshape(-1, 2, 3).permute(0, 2, 1).contiguous()
    y = roma.special_gramschmidt(x)
    return y


def get_smplx_joint_names(*args, **kwargs):
    return JOINT_NAMES[:127]


COCO17_JOINTS_NAME = {
    0: "nose",
    1: "left_eye",
    2: "right_eye",
    3: "left_ear",
    4: "right_ear",
    5: "left_shoulder",
    6: "right_shoulder",
    7: "left_elbow",
    8: "right_elbow",
    9: "left_wrist",
    10: "right_wrist",
    11: "left_hip",
    12: "right_hip",
    13: "left_knee",
    14: "right_knee",
    15: "left_ankle",
    16: "right_ankle",
}

OPENPOSE25_JOINTS_NAME = {
    0: "nose",
    1: "neck",
    2: "right_shoulder",
    3: "right_elbow",
    4: "right_wrist",
    5: "left_shoulder",
    6: "left_elbow",
    7: "left_wrist",
    8: "MidHip",
    9: "right_hip",
    10: "right_knee",
    11: "right_ankle",
    12: "left_hip",
    13: "left_knee",
    14: "left_ankle",
    15: "right_eye",
    16: "left_eye",
    17: "right_ear",
    18: "left_ear",
    19: "LBigToe",
    20: "LSmallToe",
    21: "left_heel",
    22: "RBigToe",
    23: "RSmallToe",
    24: "right_heel",
}


def joints_smplx_to_coco():
    smplx_joints_name = get_smplx_joint_names()
    joints_idx = []
    for k, v in COCO17_JOINTS_NAME.items():
        joints_idx.append(smplx_joints_name.index(v))

    return joints_idx


def joints_openpose25_to_coco17():
    idx_list = [0] * 17
    is_found = False
    for coco_key, coco_value in COCO17_JOINTS_NAME.items():
        is_found = False
        for openpose_key, openpose_value in OPENPOSE25_JOINTS_NAME.items():
            if coco_value == openpose_value:
                idx_list[coco_key] = openpose_key
                is_found = True
                break
        assert is_found, f"{coco_key} is not found in openpose keypoints"
    return idx_list


COCO_WHOLEBODY_KEYPOINTS = [
    "nose",
    "left_eye",
    "right_eye",
    "left_ear",
    "right_ear",
    "left_shoulder",
    "right_shoulder",
    "left_elbow",
    "right_elbow",
    "left_wrist",
    "right_wrist",
    "left_hip",
    "right_hip",
    "left_knee",
    "right_knee",
    "left_ankle",
    "right_ankle",
    "left_bigtoe",
    "left_smalltoe",
    "left_heel",
    "right_bigtoe",
    "right_smalltoe",
    "right_heel",
    "right_contour_1",  # original name: face_contour_1
    "right_contour_2",  # original name: face_contour_2
    "right_contour_3",  # original name: face_contour_3
    "right_contour_4",  # original name: face_contour_4
    "right_contour_5",  # original name: face_contour_5
    "right_contour_6",  # original name: face_contour_6
    "right_contour_7",  # original name: face_contour_7
    "right_contour_8",  # original name: face_contour_8
    "contour_middle",  # original name: face_contour_9
    "left_contour_8",  # original name: face_contour_10
    "left_contour_7",  # original name: face_contour_11
    "left_contour_6",  # original name: face_contour_12
    "left_contour_5",  # original name: face_contour_13
    "left_contour_4",  # original name: face_contour_14
    "left_contour_3",  # original name: face_contour_15
    "left_contour_2",  # original name: face_contour_16
    "left_contour_1",  # original name: face_contour_17
    "right_eyebrow_1",
    "right_eyebrow_2",
    "right_eyebrow_3",
    "right_eyebrow_4",
    "right_eyebrow_5",
    "left_eyebrow_5",
    "left_eyebrow_4",
    "left_eyebrow_3",
    "left_eyebrow_2",
    "left_eyebrow_1",
    "nosebridge_1",
    "nosebridge_2",
    "nosebridge_3",
    "nosebridge_4",
    "right_nose_2",  # original name: nose_1
    "right_nose_1",  # original name: nose_2
    "nose_middle",  # original name: nose_3
    "left_nose_1",  # original name: nose_4
    "left_nose_2",  # original name: nose_5
    "right_eye_1",
    "right_eye_2",
    "right_eye_3",
    "right_eye_4",
    "right_eye_5",
    "right_eye_6",
    "left_eye_4",
    "left_eye_3",
    "left_eye_2",
    "left_eye_1",
    "left_eye_6",
    "left_eye_5",
    "right_mouth_1",  # original name: mouth_1
    "right_mouth_2",  # original name: mouth_2
    "right_mouth_3",  # original name: mouth_3
    "mouth_top",  # original name: mouth_4
    "left_mouth_3",  # original name: mouth_5
    "left_mouth_2",  # original name: mouth_6
    "left_mouth_1",  # original name: mouth_7
    "left_mouth_5",  # original name: mouth_8
    "left_mouth_4",  # original name: mouth_9
    "mouth_bottom",  # original name: mouth_10
    "right_mouth_4",  # original name: mouth_11
    "right_mouth_5",  # original name: mouth_12
    "right_lip_1",  # original name: lip_1
    "right_lip_2",  # original name: lip_2
    "lip_top",  # original name: lip_3
    "left_lip_2",  # original name: lip_4
    "left_lip_1",  # original name: lip_5
    "left_lip_3",  # original name: lip_6
    "lip_bottom",  # original name: lip_7
    "right_lip_3",  # original name: lip_8
    "left_hand_root",
    "left_thumb_1",
    "left_thumb_2",
    "left_thumb_3",
    "left_thumb",
    "left_index_1",
    "left_index_2",
    "left_index_3",
    "left_index",
    "left_middle_1",
    "left_middle_2",
    "left_middle_3",
    "left_middle",
    "left_ring_1",
    "left_ring_2",
    "left_ring_3",
    "left_ring",
    "left_pinky_1",
    "left_pinky_2",
    "left_pinky_3",
    "left_pinky",
    "right_hand_root",
    "right_thumb_1",
    "right_thumb_2",
    "right_thumb_3",
    "right_thumb",
    "right_index_1",
    "right_index_2",
    "right_index_3",
    "right_index",
    "right_middle_1",
    "right_middle_2",
    "right_middle_3",
    "right_middle",
    "right_ring_1",
    "right_ring_2",
    "right_ring_3",
    "right_ring",
    "right_pinky_1",
    "right_pinky_2",
    "right_pinky_3",
    "right_pinky",
]

SMPLX_KEYPOINTS = [
    "pelvis",
    "left_hip",
    "right_hip",
    "spine_1",
    "left_knee",
    "right_knee",
    "spine_2",
    "left_ankle",
    "right_ankle",
    "spine_3",
    "left_foot",
    "right_foot",
    "neck",
    "left_collar",
    "right_collar",
    "head",
    "left_shoulder",
    "right_shoulder",
    "left_elbow",
    "right_elbow",
    "left_wrist",
    "right_wrist",
    "jaw",
    "left_eyeball",
    "right_eyeball",
    "left_index_1",
    "left_index_2",
    "left_index_3",
    "left_middle_1",
    "left_middle_2",
    "left_middle_3",
    "left_pinky_1",
    "left_pinky_2",
    "left_pinky_3",
    "left_ring_1",
    "left_ring_2",
    "left_ring_3",
    "left_thumb_1",
    "left_thumb_2",
    "left_thumb_3",
    "right_index_1",
    "right_index_2",
    "right_index_3",
    "right_middle_1",
    "right_middle_2",
    "right_middle_3",
    "right_pinky_1",
    "right_pinky_2",
    "right_pinky_3",
    "right_ring_1",
    "right_ring_2",
    "right_ring_3",
    "right_thumb_1",
    "right_thumb_2",
    "right_thumb_3",
    "nose",
    "right_eye",
    "left_eye",
    "right_ear",
    "left_ear",
    "left_bigtoe",
    "left_smalltoe",
    "left_heel",
    "right_bigtoe",
    "right_smalltoe",
    "right_heel",
    "left_thumb",
    "left_index",
    "left_middle",
    "left_ring",
    "left_pinky",
    "right_thumb",
    "right_index",
    "right_middle",
    "right_ring",
    "right_pinky",
    "right_eyebrow_1",
    "right_eyebrow_2",
    "right_eyebrow_3",
    "right_eyebrow_4",
    "right_eyebrow_5",
    "left_eyebrow_5",
    "left_eyebrow_4",
    "left_eyebrow_3",
    "left_eyebrow_2",
    "left_eyebrow_1",
    "nosebridge_1",
    "nosebridge_2",
    "nosebridge_3",
    "nosebridge_4",
    "right_nose_2",  # original name: nose_1
    "right_nose_1",  # original name: nose_2
    "nose_middle",  # original name: nose_3
    "left_nose_1",  # original name: nose_4
    "left_nose_2",  # original name: nose_5
    "right_eye_1",
    "right_eye_2",
    "right_eye_3",
    "right_eye_4",
    "right_eye_5",
    "right_eye_6",
    "left_eye_4",
    "left_eye_3",
    "left_eye_2",
    "left_eye_1",
    "left_eye_6",
    "left_eye_5",
    "right_mouth_1",  # original name: mouth_1
    "right_mouth_2",  # original name: mouth_2
    "right_mouth_3",  # original name: mouth_3
    "mouth_top",  # original name: mouth_4
    "left_mouth_3",  # original name: mouth_5
    "left_mouth_2",  # original name: mouth_6
    "left_mouth_1",  # original name: mouth_7
    "left_mouth_5",  # original name: mouth_8
    "left_mouth_4",  # original name: mouth_9
    "mouth_bottom",  # original name: mouth_10
    "right_mouth_4",  # original name: mouth_11
    "right_mouth_5",  # original name: mouth_12
    "right_lip_1",  # original name: lip_1
    "right_lip_2",  # original name: lip_2
    "lip_top",  # original name: lip_3
    "left_lip_2",  # original name: lip_4
    "left_lip_1",  # original name: lip_5
    "left_lip_3",  # original name: lip_6
    "lip_bottom",  # original name: lip_7
    "right_lip_3",  # original name: lip_8
    "right_contour_1",  # original name: face_contour_1
    "right_contour_2",  # original name: face_contour_2
    "right_contour_3",  # original name: face_contour_3
    "right_contour_4",  # original name: face_contour_4
    "right_contour_5",  # original name: face_contour_5
    "right_contour_6",  # original name: face_contour_6
    "right_contour_7",  # original name: face_contour_7
    "right_contour_8",  # original name: face_contour_8
    "contour_middle",  # original name: face_contour_9
    "left_contour_8",  # original name: face_contour_10
    "left_contour_7",  # original name: face_contour_11
    "left_contour_6",  # original name: face_contour_12
    "left_contour_5",  # original name: face_contour_13
    "left_contour_4",  # original name: face_contour_14
    "left_contour_3",  # original name: face_contour_15
    "left_contour_2",  # original name: face_contour_16
    "left_contour_1",  # original name: face_contour_17
]

LEFT_HAND_KEYPOINTS = [
    "left_wrist",
    "left_index_1",
    "left_index_2",
    "left_index_3",
    "left_middle_1",
    "left_middle_2",
    "left_middle_3",
    "left_pinky_1",
    "left_pinky_2",
    "left_pinky_3",
    "left_ring_1",
    "left_ring_2",
    "left_ring_3",
    "left_thumb_1",
    "left_thumb_2",
    "left_thumb_3",
]

RIGHT_HAND_KEYPOINTS = [
    "right_wrist",
    "right_index_1",
    "right_index_2",
    "right_index_3",
    "right_middle_1",
    "right_middle_2",
    "right_middle_3",
    "right_pinky_1",
    "right_pinky_2",
    "right_pinky_3",
    "right_ring_1",
    "right_ring_2",
    "right_ring_3",
    "right_thumb_1",
    "right_thumb_2",
    "right_thumb_3",
]

COCO_PLUS_KEYPOINTS = [
    "nose",
    "left_eye",
    "right_eye",
    "left_ear",
    "right_ear",
    "left_shoulder",
    "right_shoulder",
    "left_elbow",
    "right_elbow",
    "left_wrist",
    "right_wrist",
    "left_hip",
    "right_hip",
    "left_knee",
    "right_knee",
    "left_ankle",
    "right_ankle",
    "left_bigtoe",
    "left_smalltoe",
    "left_heel",
    "right_bigtoe",
    "right_smalltoe",
    "right_heel",
]

KEYPOINTS_FACTORY = {
    "smplx": SMPLX_KEYPOINTS,
    "coco_wholebody": COCO_WHOLEBODY_KEYPOINTS,
    "left_hand": LEFT_HAND_KEYPOINTS,
    "right_hand": RIGHT_HAND_KEYPOINTS,
    "coco_plus": COCO_PLUS_KEYPOINTS,
}

MAPPING_CACHE = {}


def get_mapping(
    src: str,
    dst: str,
    keypoints_factory: dict = KEYPOINTS_FACTORY,
):
    """Get mapping list from src to dst.

    Args:
        src (str): source data type from keypoints_factory.
        dst (str): destination data type from keypoints_factory.
        approximate (bool): control whether approximate mapping is allowed.
        keypoints_factory (dict, optional): A class to store the attributes.
            Defaults to keypoints_factory.

    Returns:
        list:
            [src_to_intersection_idx, dst_to_intersection_index,
             intersection_names]
    """
    if (
        src.lower() in MAPPING_CACHE.keys()
        and dst.lower() in MAPPING_CACHE[src.lower()].keys()
    ):
        return MAPPING_CACHE[src.lower()][dst.lower()]

    src_names = keypoints_factory[src.lower()]
    dst_names = keypoints_factory[dst.lower()]

    dst_idxs, src_idxs, intersection = [], [], []
    full_mapping_idx = []
    unmapped_names, approximate_names = [], []
    for dst_idx, dst_name in enumerate(dst_names):
        try:
            src_idx = src_names.index(dst_name)
        except ValueError:
            src_idx = -1
        if src_idx >= 0:
            dst_idxs.append(dst_idx)
            src_idxs.append(src_idx)
            intersection.append(dst_name)
        full_mapping_idx.append(src_idx)
        # approximate mapping

    mapping_list = (dst_idxs, src_idxs, intersection, full_mapping_idx)
    if not src.lower() in MAPPING_CACHE.keys():
        MAPPING_CACHE[src.lower()] = {}
    MAPPING_CACHE[src.lower()][dst.lower()] = mapping_list
    return mapping_list


if __name__ == "__main__":
    print(joints_smplx_to_coco())