Spaces:
Running
on
Zero
Running
on
Zero
File size: 6,158 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 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 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 |
# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Optional, Tuple, Union
import numpy as np
def flip_keypoints(keypoints: np.ndarray,
keypoints_visible: Optional[np.ndarray],
image_size: Tuple[int, int],
flip_indices: List[int],
direction: str = 'horizontal'
) -> Tuple[np.ndarray, Optional[np.ndarray]]:
"""Flip keypoints in the given direction.
Note:
- keypoint number: K
- keypoint dimension: D
Args:
keypoints (np.ndarray): Keypoints in shape (..., K, D)
keypoints_visible (np.ndarray, optional): The visibility of keypoints
in shape (..., K, 1) or (..., K, 2). Set ``None`` if the keypoint
visibility is unavailable
image_size (tuple): The image shape in [w, h]
flip_indices (List[int]): The indices of each keypoint's symmetric
keypoint
direction (str): The flip direction. Options are ``'horizontal'``,
``'vertical'`` and ``'diagonal'``. Defaults to ``'horizontal'``
Returns:
tuple:
- keypoints_flipped (np.ndarray): Flipped keypoints in shape
(..., K, D)
- keypoints_visible_flipped (np.ndarray, optional): Flipped keypoints'
visibility in shape (..., K, 1) or (..., K, 2). Return ``None`` if
the input ``keypoints_visible`` is ``None``
"""
ndim = keypoints.ndim
assert keypoints.shape[:-1] == keypoints_visible.shape[:ndim - 1], (
f'Mismatched shapes of keypoints {keypoints.shape} and '
f'keypoints_visible {keypoints_visible.shape}')
direction_options = {'horizontal', 'vertical', 'diagonal'}
assert direction in direction_options, (
f'Invalid flipping direction "{direction}". '
f'Options are {direction_options}')
# swap the symmetric keypoint pairs
if direction == 'horizontal' or direction == 'vertical':
keypoints = keypoints.take(flip_indices, axis=ndim - 2)
if keypoints_visible is not None:
keypoints_visible = keypoints_visible.take(
flip_indices, axis=ndim - 2)
# flip the keypoints
w, h = image_size
if direction == 'horizontal':
keypoints[..., 0] = w - 1 - keypoints[..., 0]
elif direction == 'vertical':
keypoints[..., 1] = h - 1 - keypoints[..., 1]
else:
keypoints = [w, h] - keypoints - 1
return keypoints, keypoints_visible
def flip_keypoints_custom_center(keypoints: np.ndarray,
keypoints_visible: np.ndarray,
flip_indices: List[int],
center_mode: str = 'static',
center_x: float = 0.5,
center_index: Union[int, List] = 0):
"""Flip human joints horizontally.
Note:
- num_keypoint: K
- dimension: D
Args:
keypoints (np.ndarray([..., K, D])): Coordinates of keypoints.
keypoints_visible (np.ndarray([..., K])): Visibility item of keypoints.
flip_indices (list[int]): The indices to flip the keypoints.
center_mode (str): The mode to set the center location on the x-axis
to flip around. Options are:
- static: use a static x value (see center_x also)
- root: use a root joint (see center_index also)
Defaults: ``'static'``.
center_x (float): Set the x-axis location of the flip center. Only used
when ``center_mode`` is ``'static'``. Defaults: 0.5.
center_index (Union[int, List]): Set the index of the root joint, whose
x location will be used as the flip center. Only used when
``center_mode`` is ``'root'``. Defaults: 0.
Returns:
np.ndarray([..., K, C]): Flipped joints.
"""
assert keypoints.ndim >= 2, f'Invalid pose shape {keypoints.shape}'
allowed_center_mode = {'static', 'root'}
assert center_mode in allowed_center_mode, 'Get invalid center_mode ' \
f'{center_mode}, allowed choices are {allowed_center_mode}'
if center_mode == 'static':
x_c = center_x
elif center_mode == 'root':
center_index = [center_index] if isinstance(center_index, int) else \
center_index
assert keypoints.shape[-2] > max(center_index)
x_c = keypoints[..., center_index, 0].mean(axis=-1)
keypoints_flipped = keypoints.copy()
keypoints_visible_flipped = keypoints_visible.copy()
# Swap left-right parts
for left, right in enumerate(flip_indices):
keypoints_flipped[..., left, :] = keypoints[..., right, :]
keypoints_visible_flipped[..., left] = keypoints_visible[..., right]
# Flip horizontally
keypoints_flipped[..., 0] = x_c * 2 - keypoints_flipped[..., 0]
return keypoints_flipped, keypoints_visible_flipped
def keypoint_clip_border(keypoints: np.ndarray, keypoints_visible: np.ndarray,
shape: Tuple[int,
int]) -> Tuple[np.ndarray, np.ndarray]:
"""Set the visibility values for keypoints outside the image border.
Args:
keypoints (np.ndarray): Input keypoints coordinates.
keypoints_visible (np.ndarray): Visibility values of keypoints.
shape (Tuple[int, int]): Shape of the image to which keypoints are
being clipped in the format of (w, h).
Note:
This function sets the visibility values of keypoints that fall outside
the specified frame border to zero (0.0).
"""
width, height = shape[:2]
# Create a mask for keypoints outside the frame
outside_mask = ((keypoints[..., 0] > width) | (keypoints[..., 0] < 0) |
(keypoints[..., 1] > height) | (keypoints[..., 1] < 0))
# Update visibility values for keypoints outside the frame
if keypoints_visible.ndim == 2:
keypoints_visible[outside_mask] = 0.0
elif keypoints_visible.ndim == 3:
keypoints_visible[outside_mask, 0] = 0.0
return keypoints, keypoints_visible
|