Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,366 Bytes
a249588 |
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 |
# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Tuple, Union
import numpy as np
def transform_sigmas(sigmas: Union[List, np.ndarray], num_keypoints: int,
mapping: Union[List[Tuple[int, int]], List[Tuple[Tuple,
int]]]):
"""Transforms the sigmas based on the mapping."""
if len(mapping):
source_index, target_index = map(list, zip(*mapping))
else:
source_index, target_index = [], []
list_input = False
if isinstance(sigmas, list):
sigmas = np.array(sigmas)
list_input = True
new_sigmas = np.ones(num_keypoints, dtype=sigmas.dtype)
new_sigmas[target_index] = sigmas[source_index]
if list_input:
new_sigmas = new_sigmas.tolist()
return new_sigmas
def transform_ann(ann_info: Union[dict, list], num_keypoints: int,
mapping: Union[List[Tuple[int, int]], List[Tuple[Tuple,
int]]]):
"""Transforms COCO-format annotations based on the mapping."""
if len(mapping):
source_index, target_index = map(list, zip(*mapping))
else:
source_index, target_index = [], []
list_input = True
if not isinstance(ann_info, list):
ann_info = [ann_info]
list_input = False
for each in ann_info:
if 'keypoints' in each:
keypoints = np.array(each['keypoints'])
C = 3 # COCO-format: x, y, score
keypoints = keypoints.reshape(-1, C)
new_keypoints = np.zeros((num_keypoints, C), dtype=keypoints.dtype)
new_keypoints[target_index] = keypoints[source_index]
each['keypoints'] = new_keypoints.reshape(-1).tolist()
if 'num_keypoints' in each:
each['num_keypoints'] = num_keypoints
if not list_input:
ann_info = ann_info[0]
return ann_info
def transform_pred(pred_info: Union[dict, list], num_keypoints: int,
mapping: Union[List[Tuple[int, int]], List[Tuple[Tuple,
int]]]):
"""Transforms predictions based on the mapping."""
if len(mapping):
source_index, target_index = map(list, zip(*mapping))
else:
source_index, target_index = [], []
list_input = True
if not isinstance(pred_info, list):
pred_info = [pred_info]
list_input = False
for each in pred_info:
if 'keypoints' in each:
keypoints = np.array(each['keypoints'])
N, _, C = keypoints.shape
new_keypoints = np.zeros((N, num_keypoints, C),
dtype=keypoints.dtype)
new_keypoints[:, target_index] = keypoints[:, source_index]
each['keypoints'] = new_keypoints
keypoint_scores = np.array(each['keypoint_scores'])
new_scores = np.zeros((N, num_keypoints),
dtype=keypoint_scores.dtype)
new_scores[:, target_index] = keypoint_scores[:, source_index]
each['keypoint_scores'] = new_scores
if 'num_keypoints' in each:
each['num_keypoints'] = num_keypoints
if not list_input:
pred_info = pred_info[0]
return pred_info
|