MCPcopy Create free account
hub / github.com/MotrixLab/AiOS / search_limbs

Function search_limbs

detrsmpl/utils/keypoint_utils.py:12–61  ·  view source on GitHub ↗

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 Non

(
        data_source: str,
        mask: Optional[Union[np.ndarray, tuple, list]] = None,
        keypoints_factory: dict = KEYPOINTS_FACTORY)

Source from the content-addressed store, hash-verified

10
11
12def search_limbs(
13 data_source: str,
14 mask: Optional[Union[np.ndarray, tuple, list]] = None,
15 keypoints_factory: dict = KEYPOINTS_FACTORY) -> Tuple[dict, dict]:
16 """Search the corresponding limbs following the basis human_data limbs. The
17 mask could mask out the incorrect keypoints.
18
19 Args:
20 data_source (str): data source type.
21 mask (Optional[Union[np.ndarray, tuple, list]], optional):
22 refer to keypoints_mapping. Defaults to None.
23 keypoints_factory (dict, optional): Dict of all the conventions.
24 Defaults to KEYPOINTS_FACTORY.
25 Returns:
26 Tuple[dict, dict]: (limbs_target, limbs_palette).
27 """
28 limbs_source = HUMAN_DATA_LIMBS_INDEX
29 limbs_palette = HUMAN_DATA_PALETTE
30 keypoints_source = keypoints_factory['human_data']
31 keypoints_target = keypoints_factory[data_source]
32 limbs_target = {}
33 for k, part_limbs in limbs_source.items():
34 limbs_target[k] = []
35 for limb in part_limbs:
36 flag = False
37 if (keypoints_source[limb[0]]
38 in keypoints_target) and (keypoints_source[limb[1]]
39 in keypoints_target):
40 if mask is not None:
41 if mask[keypoints_target.index(keypoints_source[
42 limb[0]])] != 0 and mask[keypoints_target.index(
43 keypoints_source[limb[1]])] != 0:
44 flag = True
45 else:
46 flag = True
47 if flag:
48 limbs_target.setdefault(k, []).append([
49 keypoints_target.index(keypoints_source[limb[0]]),
50 keypoints_target.index(keypoints_source[limb[1]])
51 ])
52 if k in limbs_target:
53 if k == 'body':
54 np.random.seed(0)
55 limbs_palette[k] = np.random.randint(0,
56 high=255,
57 size=(len(
58 limbs_target[k]), 3))
59 else:
60 limbs_palette[k] = np.array(limbs_palette[k])
61 return limbs_target, limbs_palette

Callers 2

visualize_kp3dFunction · 0.90
__init__Method · 0.90

Calls 1

itemsMethod · 0.45

Tested by

no test coverage detected