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. | |
| from typing import List, Optional, Tuple | |
| 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). 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). Return ``None`` if the input | |
| ``keypoints_visible`` is ``None`` | |
| """ | |
| assert keypoints.shape[:-1] == keypoints_visible.shape, ( | |
| 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[..., flip_indices, :] | |
| if keypoints_visible is not None: | |
| keypoints_visible = keypoints_visible[..., flip_indices] | |
| # 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: int = 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 (int): 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': | |
| assert keypoints.shape[-2] > center_index | |
| x_c = keypoints[..., center_index, 0] | |
| 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 | |