Spaces:
Runtime error
Runtime error
File size: 4,482 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 |
# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Tuple, Union
import numpy as np
from mmcv.transforms import BaseTransform
from mmpose.registry import TRANSFORMS
@TRANSFORMS.register_module()
class KeypointConverter(BaseTransform):
"""Change the order of keypoints according to the given mapping.
Required Keys:
- keypoints
- keypoints_visible
Modified Keys:
- keypoints
- keypoints_visible
Args:
num_keypoints (int): The number of keypoints in target dataset.
mapping (list): A list containing mapping indexes. Each element has
format (source_index, target_index)
Example:
>>> import numpy as np
>>> # case 1: 1-to-1 mapping
>>> # (0, 0) means target[0] = source[0]
>>> self = KeypointConverter(
>>> num_keypoints=3,
>>> mapping=[
>>> (0, 0), (1, 1), (2, 2), (3, 3)
>>> ])
>>> results = dict(
>>> keypoints=np.arange(34).reshape(2, 3, 2),
>>> keypoints_visible=np.arange(34).reshape(2, 3, 2) % 2)
>>> results = self(results)
>>> assert np.equal(results['keypoints'],
>>> np.arange(34).reshape(2, 3, 2)).all()
>>> assert np.equal(results['keypoints_visible'],
>>> np.arange(34).reshape(2, 3, 2) % 2).all()
>>>
>>> # case 2: 2-to-1 mapping
>>> # ((1, 2), 0) means target[0] = (source[1] + source[2]) / 2
>>> self = KeypointConverter(
>>> num_keypoints=3,
>>> mapping=[
>>> ((1, 2), 0), (1, 1), (2, 2)
>>> ])
>>> results = dict(
>>> keypoints=np.arange(34).reshape(2, 3, 2),
>>> keypoints_visible=np.arange(34).reshape(2, 3, 2) % 2)
>>> results = self(results)
"""
def __init__(self, num_keypoints: int,
mapping: Union[List[Tuple[int, int]], List[Tuple[Tuple,
int]]]):
self.num_keypoints = num_keypoints
self.mapping = mapping
source_index, target_index = zip(*mapping)
src1, src2 = [], []
interpolation = False
for x in source_index:
if isinstance(x, (list, tuple)):
assert len(x) == 2, 'source_index should be a list/tuple of ' \
'length 2'
src1.append(x[0])
src2.append(x[1])
interpolation = True
else:
src1.append(x)
src2.append(x)
# When paired source_indexes are input,
# keep a self.source_index2 for interpolation
if interpolation:
self.source_index2 = src2
self.source_index = src1
self.target_index = target_index
self.interpolation = interpolation
def transform(self, results: dict) -> dict:
num_instances = results['keypoints'].shape[0]
keypoints = np.zeros((num_instances, self.num_keypoints, 2))
keypoints_visible = np.zeros((num_instances, self.num_keypoints))
# When paired source_indexes are input,
# perform interpolation with self.source_index and self.source_index2
if self.interpolation:
keypoints[:, self.target_index] = 0.5 * (
results['keypoints'][:, self.source_index] +
results['keypoints'][:, self.source_index2])
keypoints_visible[:, self.target_index] = results[
'keypoints_visible'][:, self.source_index] * \
results['keypoints_visible'][:, self.source_index2]
else:
keypoints[:,
self.target_index] = results['keypoints'][:, self.
source_index]
keypoints_visible[:, self.target_index] = results[
'keypoints_visible'][:, self.source_index]
results['keypoints'] = keypoints
results['keypoints_visible'] = keypoints_visible
return results
def __repr__(self) -> str:
"""print the basic information of the transform.
Returns:
str: Formatted string.
"""
repr_str = self.__class__.__name__
repr_str += f'(num_keypoints={self.num_keypoints}, '\
f'mapping={self.mapping})'
return repr_str
|