Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,936 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 |
# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Union
from mmpose.codecs import * # noqa: F401, F403
from mmpose.registry import TRANSFORMS
from .common_transforms import RandomFlip
@TRANSFORMS.register_module()
class HandRandomFlip(RandomFlip):
"""Data augmentation with random image flip. A child class of
`TopDownRandomFlip`.
Required Keys:
- img
- joints_3d
- joints_3d_visible
- center
- hand_type
- rel_root_depth
- ann_info
Modified Keys:
- img
- joints_3d
- joints_3d_visible
- center
- hand_type
- rel_root_depth
Args:
prob (float | list[float]): The flipping probability. If a list is
given, the argument `direction` should be a list with the same
length. And each element in `prob` indicates the flipping
probability of the corresponding one in ``direction``. Defaults
to 0.5
"""
def __init__(self, prob: Union[float, List[float]] = 0.5) -> None:
super().__init__(prob=prob, direction='horizontal')
def transform(self, results: dict) -> dict:
"""The transform function of :class:`HandRandomFlip`.
See ``transform()`` method of :class:`BaseTransform` for details.
Args:
results (dict): The result dict
Returns:
dict: The result dict.
"""
# base flip augmentation
results = super().transform(results)
# flip hand type and root depth
hand_type = results['hand_type']
rel_root_depth = results['rel_root_depth']
flipped = results['flip']
if flipped:
hand_type[..., [0, 1]] = hand_type[..., [1, 0]]
rel_root_depth = -rel_root_depth
results['hand_type'] = hand_type
results['rel_root_depth'] = rel_root_depth
return results
|