Spaces:
Sleeping
Sleeping
File size: 2,545 Bytes
d7e58f0 |
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 |
from typing import Optional, Tuple, Union
import numpy as np
from detrsmpl.core.conventions.keypoints_mapping import KEYPOINTS_FACTORY
from detrsmpl.core.conventions.keypoints_mapping.human_data import (
HUMAN_DATA_LIMBS_INDEX,
HUMAN_DATA_PALETTE,
)
def search_limbs(
data_source: str,
mask: Optional[Union[np.ndarray, tuple, list]] = None,
keypoints_factory: dict = KEYPOINTS_FACTORY) -> Tuple[dict, dict]:
"""Search the corresponding limbs following the basis human_data limbs. The
mask could mask out the incorrect keypoints.
Args:
data_source (str): data source type.
mask (Optional[Union[np.ndarray, tuple, list]], optional):
refer to keypoints_mapping. Defaults to None.
keypoints_factory (dict, optional): Dict of all the conventions.
Defaults to KEYPOINTS_FACTORY.
Returns:
Tuple[dict, dict]: (limbs_target, limbs_palette).
"""
limbs_source = HUMAN_DATA_LIMBS_INDEX
limbs_palette = HUMAN_DATA_PALETTE
keypoints_source = keypoints_factory['human_data']
keypoints_target = keypoints_factory[data_source]
limbs_target = {}
for k, part_limbs in limbs_source.items():
limbs_target[k] = []
for limb in part_limbs:
flag = False
if (keypoints_source[limb[0]]
in keypoints_target) and (keypoints_source[limb[1]]
in keypoints_target):
if mask is not None:
if mask[keypoints_target.index(keypoints_source[
limb[0]])] != 0 and mask[keypoints_target.index(
keypoints_source[limb[1]])] != 0:
flag = True
else:
flag = True
if flag:
limbs_target.setdefault(k, []).append([
keypoints_target.index(keypoints_source[limb[0]]),
keypoints_target.index(keypoints_source[limb[1]])
])
if k in limbs_target:
if k == 'body':
np.random.seed(0)
limbs_palette[k] = np.random.randint(0,
high=255,
size=(len(
limbs_target[k]), 3))
else:
limbs_palette[k] = np.array(limbs_palette[k])
return limbs_target, limbs_palette
|