Spaces:
Running
on
Zero
Running
on
Zero
| # Copyright (c) OpenMMLab. All rights reserved. | |
| from itertools import product | |
| from typing import Optional, Tuple, Union | |
| import numpy as np | |
| from scipy.spatial.distance import cdist | |
| def generate_oks_maps( | |
| heatmap_size: Tuple[int, int], | |
| keypoints: np.ndarray, | |
| keypoints_visible: np.ndarray, | |
| keypoints_visibility: np.ndarray, | |
| sigma: float = 0.55, | |
| increase_sigma_with_padding: bool = False, | |
| ) -> Tuple[np.ndarray, np.ndarray]: | |
| """Generate gaussian heatmaps of keypoints using `UDP`_. | |
| Args: | |
| heatmap_size (Tuple[int, int]): Heatmap size in [W, H] | |
| keypoints (np.ndarray): Keypoint coordinates in shape (N, K, D) | |
| keypoints_visible (np.ndarray): Keypoint visibilities in shape | |
| (N, K) | |
| sigma (float): The sigma value of the Gaussian heatmap | |
| keypoints_visibility (np.ndarray): The visibility bit for each keypoint (N, K) | |
| increase_sigma_with_padding (bool): Whether to increase the sigma | |
| value with padding. Default: False | |
| Returns: | |
| tuple: | |
| - heatmaps (np.ndarray): The generated heatmap in shape | |
| (K, H, W) where [W, H] is the `heatmap_size` | |
| - keypoint_weights (np.ndarray): The target weights in shape | |
| (N, K) | |
| .. _`UDP`: https://arxiv.org/abs/1911.07524 | |
| """ | |
| N, K, _ = keypoints.shape | |
| W, H = heatmap_size | |
| # The default sigmas are used for COCO dataset. | |
| sigmas = np.array( | |
| [2.6, 2.5, 2.5, 3.5, 3.5, 7.9, 7.9, 7.2, 7.2, 6.2, 6.2, 10.7, 10.7, 8.7, 8.7, 8.9, 8.9])/100 | |
| # sigmas = sigmas * 2 / sigmas.mean() | |
| # sigmas = np.round(sigmas).astype(int) | |
| # sigmas = np.clip(sigmas, 1, 10) | |
| heatmaps = np.zeros((K, H, W), dtype=np.float32) | |
| keypoint_weights = keypoints_visible.copy() | |
| # bbox_area = W/1.25 * H/1.25 | |
| # bbox_area = W * H * 0.53 | |
| bbox_area = np.sqrt(H/1.25 * W/1.25) | |
| # print(scales_arr) | |
| # print(scaled_sigmas) | |
| for n, k in product(range(N), range(K)): | |
| kpt_sigma = sigmas[k] | |
| # skip unlabled keypoints | |
| if keypoints_visible[n, k] < 0.5: | |
| continue | |
| y_idx, x_idx = np.indices((H, W)) | |
| dx = x_idx - keypoints[n, k, 0] | |
| dy = y_idx - keypoints[n, k, 1] | |
| dist = np.sqrt(dx**2 + dy**2) | |
| # e_map = (dx**2 + dy**2) / ((kpt_sigma*100)**2 * sigma) | |
| vars = (kpt_sigma*2)**2 | |
| s = vars * bbox_area * 2 | |
| s = np.clip(s, 0.55, 3.0) | |
| if sigma is not None and sigma > 0: | |
| s = sigma | |
| e_map = dist**2 / (2*s) | |
| oks_map = np.exp(-e_map) | |
| keypoint_weights[n, k] = (oks_map.max() > 0).astype(int) | |
| # Scale such that there is always 1 at the maximum | |
| if oks_map.max() > 1e-3: | |
| oks_map = oks_map / oks_map.max() | |
| # Scale OKS map such that 1 stays 1 and 0.5 becomes 0 | |
| # oks_map[oks_map < 0.5] = 0 | |
| # oks_map = 2 * oks_map - 1 | |
| # oks_map[oks_map > 0.95] = 1 | |
| # print("{:.4f}, {:7.1f}, {:9.3f}, {:9.3f}, {:4.2f}".format(vars, bbox_area, vars * bbox_area* 2, s, oks_map.max())) | |
| # if np.all(oks_map < 0.1): | |
| # print("\t{:d} --> {:.4f}".format(k, s)) | |
| heatmaps[k] = oks_map | |
| # breakpoint() | |
| return heatmaps, keypoint_weights | |