File size: 6,554 Bytes
cc0dd3c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Dict, List, Tuple


def get_eye_keypoint_ids(dataset_meta: Dict) -> Tuple[int, int]:
    """A helper function to get the keypoint indices of left and right eyes
    from the dataset meta information.

    Args:
        dataset_meta (dict): dataset meta information.

    Returns:
        tuple[int, int]: The keypoint indices of left eye and right eye.
    """
    left_eye_idx = None
    right_eye_idx = None

    # try obtaining eye point ids from dataset_meta
    keypoint_name2id = dataset_meta.get('keypoint_name2id', {})
    left_eye_idx = keypoint_name2id.get('left_eye', None)
    right_eye_idx = keypoint_name2id.get('right_eye', None)

    if left_eye_idx is None or right_eye_idx is None:
        # Fall back to hard coded keypoint id
        dataset_name = dataset_meta.get('dataset_name', 'unknown dataset')
        if dataset_name in {'coco', 'coco_wholebody'}:
            left_eye_idx = 1
            right_eye_idx = 2
        elif dataset_name in {'animalpose', 'ap10k'}:
            left_eye_idx = 0
            right_eye_idx = 1
        else:
            raise ValueError('Can not determine the eye keypoint id of '
                             f'{dataset_name}')

    return left_eye_idx, right_eye_idx


def get_face_keypoint_ids(dataset_meta: Dict) -> List:
    """A helper function to get the keypoint indices of the face from the
    dataset meta information.

    Args:
        dataset_meta (dict): dataset meta information.

    Returns:
        list[int]: face keypoint indices. The length depends on the dataset.
    """
    face_indices = []

    # try obtaining nose point ids from dataset_meta
    keypoint_name2id = dataset_meta.get('keypoint_name2id', {})
    for id in range(68):
        face_indices.append(keypoint_name2id.get(f'face-{id}', None))

    if None in face_indices:
        # Fall back to hard coded keypoint id
        dataset_name = dataset_meta.get('dataset_name', 'unknown dataset')
        if dataset_name in {'coco_wholebody'}:
            face_indices = list(range(23, 91))
        else:
            raise ValueError('Can not determine the face id of '
                             f'{dataset_name}')

    return face_indices


def get_wrist_keypoint_ids(dataset_meta: Dict) -> Tuple[int, int]:
    """A helper function to get the keypoint indices of left and right wrists
    from the dataset meta information.

    Args:
        dataset_meta (dict): dataset meta information.
    Returns:
        tuple[int, int]: The keypoint indices of left and right wrists.
    """

    # try obtaining wrist point ids from dataset_meta
    keypoint_name2id = dataset_meta.get('keypoint_name2id', {})
    left_wrist_idx = keypoint_name2id.get('left_wrist', None)
    right_wrist_idx = keypoint_name2id.get('right_wrist', None)

    if left_wrist_idx is None or right_wrist_idx is None:
        # Fall back to hard coded keypoint id
        dataset_name = dataset_meta.get('dataset_name', 'unknown dataset')
        if dataset_name in {'coco', 'coco_wholebody'}:
            left_wrist_idx = 9
            right_wrist_idx = 10
        elif dataset_name == 'animalpose':
            left_wrist_idx = 16
            right_wrist_idx = 17
        elif dataset_name == 'ap10k':
            left_wrist_idx = 7
            right_wrist_idx = 10
        else:
            raise ValueError('Can not determine the eye keypoint id of '
                             f'{dataset_name}')

    return left_wrist_idx, right_wrist_idx


def get_mouth_keypoint_ids(dataset_meta: Dict) -> int:
    """A helper function to get the mouth keypoint index from the dataset meta
    information.

    Args:
        dataset_meta (dict): dataset meta information.
    Returns:
        int: The mouth keypoint index
    """
    # try obtaining mouth point ids from dataset_info
    keypoint_name2id = dataset_meta.get('keypoint_name2id', {})
    mouth_index = keypoint_name2id.get('face-62', None)

    if mouth_index is None:
        # Fall back to hard coded keypoint id
        dataset_name = dataset_meta.get('dataset_name', 'unknown dataset')
        if dataset_name == 'coco_wholebody':
            mouth_index = 85
        else:
            raise ValueError('Can not determine the eye keypoint id of '
                             f'{dataset_name}')

    return mouth_index


def get_hand_keypoint_ids(dataset_meta: Dict) -> List[int]:
    """A helper function to get the keypoint indices of left and right hand
    from the dataset meta information.

    Args:
        dataset_meta (dict): dataset meta information.
    Returns:
        list[int]: hand keypoint indices. The length depends on the dataset.
    """
    # try obtaining hand keypoint ids from dataset_meta
    keypoint_name2id = dataset_meta.get('keypoint_name2id', {})
    hand_indices = []
    hand_indices.append(keypoint_name2id.get('left_hand_root', None))

    for id in range(1, 5):
        hand_indices.append(keypoint_name2id.get(f'left_thumb{id}', None))
    for id in range(1, 5):
        hand_indices.append(keypoint_name2id.get(f'left_forefinger{id}', None))
    for id in range(1, 5):
        hand_indices.append(
            keypoint_name2id.get(f'left_middle_finger{id}', None))
    for id in range(1, 5):
        hand_indices.append(
            keypoint_name2id.get(f'left_ring_finger{id}', None))
    for id in range(1, 5):
        hand_indices.append(
            keypoint_name2id.get(f'left_pinky_finger{id}', None))

    hand_indices.append(keypoint_name2id.get('right_hand_root', None))

    for id in range(1, 5):
        hand_indices.append(keypoint_name2id.get(f'right_thumb{id}', None))
    for id in range(1, 5):
        hand_indices.append(
            keypoint_name2id.get(f'right_forefinger{id}', None))
    for id in range(1, 5):
        hand_indices.append(
            keypoint_name2id.get(f'right_middle_finger{id}', None))
    for id in range(1, 5):
        hand_indices.append(
            keypoint_name2id.get(f'right_ring_finger{id}', None))
    for id in range(1, 5):
        hand_indices.append(
            keypoint_name2id.get(f'right_pinky_finger{id}', None))

    if None in hand_indices:
        # Fall back to hard coded keypoint id
        dataset_name = dataset_meta.get('dataset_name', 'unknown dataset')
        if dataset_name in {'coco_wholebody'}:
            hand_indices = list(range(91, 133))
        else:
            raise ValueError('Can not determine the hand id of '
                             f'{dataset_name}')

    return hand_indices