Spaces:
Running on Zero
Running on Zero
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| import numpy as np | |
| from sapiens.engine import Config | |
| def parse_pose_metainfo(metainfo: dict): | |
| if "from_file" in metainfo: | |
| cfg_file = metainfo["from_file"] | |
| metainfo = Config.fromfile(cfg_file).dataset_info | |
| # check data integrity | |
| assert "dataset_name" in metainfo | |
| assert "keypoint_info" in metainfo | |
| assert "skeleton_info" in metainfo | |
| # assert 'joint_weights' in metainfo | |
| # assert 'sigmas' in metainfo | |
| # parse metainfo | |
| parsed = dict( | |
| dataset_name=None, | |
| num_keypoints=None, | |
| keypoint_id2name={}, | |
| keypoint_name2id={}, | |
| upper_body_ids=[], | |
| lower_body_ids=[], | |
| flip_indices=[], | |
| flip_pairs=[], | |
| keypoint_colors=[], | |
| num_skeleton_links=None, | |
| skeleton_links=[], | |
| skeleton_link_colors=[], | |
| dataset_keypoint_weights=None, | |
| sigmas=None, | |
| ) | |
| parsed["dataset_name"] = metainfo["dataset_name"] | |
| if "remove_teeth" in metainfo: | |
| parsed["remove_teeth"] = metainfo["remove_teeth"] | |
| if "min_visible_keypoints" in metainfo: | |
| parsed["min_visible_keypoints"] = metainfo["min_visible_keypoints"] | |
| if "teeth_keypoint_ids" in metainfo: | |
| parsed["teeth_keypoint_ids"] = metainfo["teeth_keypoint_ids"] | |
| if "coco_wholebody_to_goliath_mapping" in metainfo: | |
| parsed["coco_wholebody_to_goliath_mapping"] = metainfo[ | |
| "coco_wholebody_to_goliath_mapping" | |
| ] | |
| if "coco_wholebody_to_goliath_keypoint_info" in metainfo: | |
| parsed["coco_wholebody_to_goliath_keypoint_info"] = metainfo[ | |
| "coco_wholebody_to_goliath_keypoint_info" | |
| ] | |
| if "idx_to_original_idx_mapping" in metainfo: | |
| parsed["idx_to_original_idx_mapping"] = metainfo["idx_to_original_idx_mapping"] | |
| # parse keypoint information | |
| parsed["num_keypoints"] = len(metainfo["keypoint_info"]) | |
| for kpt_id, kpt in metainfo["keypoint_info"].items(): | |
| kpt_name = kpt["name"] | |
| parsed["keypoint_id2name"][kpt_id] = kpt_name | |
| parsed["keypoint_name2id"][kpt_name] = kpt_id | |
| parsed["keypoint_colors"].append(kpt.get("color", [255, 128, 0])) | |
| kpt_type = kpt.get("type", "") | |
| if kpt_type == "upper": | |
| parsed["upper_body_ids"].append(kpt_id) | |
| elif kpt_type == "lower": | |
| parsed["lower_body_ids"].append(kpt_id) | |
| swap_kpt = kpt.get("swap", "") | |
| if swap_kpt == kpt_name or swap_kpt == "": | |
| parsed["flip_indices"].append(kpt_name) | |
| else: | |
| parsed["flip_indices"].append(swap_kpt) | |
| pair = (swap_kpt, kpt_name) | |
| if pair not in parsed["flip_pairs"]: | |
| parsed["flip_pairs"].append(pair) | |
| # parse skeleton information | |
| parsed["num_skeleton_links"] = len(metainfo["skeleton_info"]) | |
| for _, sk in metainfo["skeleton_info"].items(): | |
| parsed["skeleton_links"].append(sk["link"]) | |
| parsed["skeleton_link_colors"].append(sk.get("color", [96, 96, 255])) | |
| # parse extra information | |
| if "joint_weights" in metainfo: | |
| parsed["dataset_keypoint_weights"] = np.array( | |
| metainfo["joint_weights"], dtype=np.float32 | |
| ) | |
| if "sigmas" in metainfo: | |
| parsed["sigmas"] = np.array(metainfo["sigmas"], dtype=np.float32) | |
| if "stats_info" in metainfo: | |
| parsed["stats_info"] = {} | |
| for name, val in metainfo["stats_info"].items(): | |
| parsed["stats_info"][name] = np.array(val, dtype=np.float32) | |
| # formatting | |
| def _map(src, mapping: dict): | |
| if isinstance(src, (list, tuple)): | |
| cls = type(src) | |
| return cls(_map(s, mapping) for s in src) | |
| else: | |
| return mapping[src] | |
| parsed["flip_pairs"] = _map( | |
| parsed["flip_pairs"], mapping=parsed["keypoint_name2id"] | |
| ) | |
| parsed["flip_indices"] = _map( | |
| parsed["flip_indices"], mapping=parsed["keypoint_name2id"] | |
| ) | |
| parsed["skeleton_links"] = _map( | |
| parsed["skeleton_links"], mapping=parsed["keypoint_name2id"] | |
| ) | |
| parsed["keypoint_colors"] = np.array(parsed["keypoint_colors"], dtype=np.uint8) | |
| parsed["skeleton_link_colors"] = np.array( | |
| parsed["skeleton_link_colors"], dtype=np.uint8 | |
| ) | |
| return parsed | |