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_3d_gaussian_heatmaps( | |
| heatmap_size: Tuple[int, int, int], | |
| keypoints: np.ndarray, | |
| keypoints_visible: np.ndarray, | |
| sigma: Union[float, Tuple[float], np.ndarray], | |
| image_size: Tuple[int, int], | |
| heatmap3d_depth_bound: float = 400.0, | |
| joint_indices: Optional[list] = None, | |
| max_bound: float = 1.0, | |
| use_different_joint_weights: bool = False, | |
| dataset_keypoint_weights: Optional[np.ndarray] = None | |
| ) -> Tuple[np.ndarray, np.ndarray]: | |
| """Generate 3d gaussian heatmaps of keypoints. | |
| Args: | |
| heatmap_size (Tuple[int, int]): Heatmap size in [W, H, D] | |
| keypoints (np.ndarray): Keypoint coordinates in shape (N, K, C) | |
| keypoints_visible (np.ndarray): Keypoint visibilities in shape | |
| (N, K) | |
| sigma (float or List[float]): A list of sigma values of the Gaussian | |
| heatmap for each instance. If sigma is given as a single float | |
| value, it will be expanded into a tuple | |
| image_size (Tuple[int, int]): Size of input image. | |
| heatmap3d_depth_bound (float): Boundary for 3d heatmap depth. | |
| Default: 400.0. | |
| joint_indices (List[int], optional): Indices of joints used for heatmap | |
| generation. If None (default) is given, all joints will be used. | |
| Default: ``None``. | |
| max_bound (float): The maximal value of heatmap. Default: 1.0. | |
| use_different_joint_weights (bool): Whether to use different joint | |
| weights. Default: ``False``. | |
| dataset_keypoint_weights (np.ndarray, optional): Keypoints weight in | |
| shape (K, ). | |
| Returns: | |
| tuple: | |
| - heatmaps (np.ndarray): The generated heatmap in shape | |
| (K * D, H, W) where [W, H, D] is the `heatmap_size` | |
| - keypoint_weights (np.ndarray): The target weights in shape | |
| (N, K) | |
| """ | |
| W, H, D = heatmap_size | |
| # select the joints used for target generation | |
| if joint_indices is not None: | |
| keypoints = keypoints[:, joint_indices, ...] | |
| keypoints_visible = keypoints_visible[:, joint_indices, ...] | |
| N, K, _ = keypoints.shape | |
| heatmaps = np.zeros([K, D, H, W], dtype=np.float32) | |
| keypoint_weights = keypoints_visible.copy() | |
| if isinstance(sigma, (int, float)): | |
| sigma = (sigma, ) * N | |
| for n in range(N): | |
| # 3-sigma rule | |
| radius = sigma[n] * 3 | |
| # joint location in heatmap coordinates | |
| mu_x = keypoints[n, :, 0] * W / image_size[0] # (K, ) | |
| mu_y = keypoints[n, :, 1] * H / image_size[1] | |
| mu_z = (keypoints[n, :, 2] / heatmap3d_depth_bound + 0.5) * D | |
| keypoint_weights[n, ...] = keypoint_weights[n, ...] * (mu_z >= 0) * ( | |
| mu_z < D) | |
| if use_different_joint_weights: | |
| keypoint_weights[ | |
| n] = keypoint_weights[n] * dataset_keypoint_weights | |
| # xy grid | |
| gaussian_size = 2 * radius + 1 | |
| # get neighboring voxels coordinates | |
| x = y = z = np.arange(gaussian_size, dtype=np.float32) - radius | |
| zz, yy, xx = np.meshgrid(z, y, x) | |
| xx = np.expand_dims(xx, axis=0) | |
| yy = np.expand_dims(yy, axis=0) | |
| zz = np.expand_dims(zz, axis=0) | |
| mu_x = np.expand_dims(mu_x, axis=(-1, -2, -3)) | |
| mu_y = np.expand_dims(mu_y, axis=(-1, -2, -3)) | |
| mu_z = np.expand_dims(mu_z, axis=(-1, -2, -3)) | |
| xx, yy, zz = xx + mu_x, yy + mu_y, zz + mu_z | |
| local_size = xx.shape[1] | |
| # round the coordinates | |
| xx = xx.round().clip(0, W - 1) | |
| yy = yy.round().clip(0, H - 1) | |
| zz = zz.round().clip(0, D - 1) | |
| # compute the target value near joints | |
| gaussian = np.exp(-((xx - mu_x)**2 + (yy - mu_y)**2 + (zz - mu_z)**2) / | |
| (2 * sigma[n]**2)) | |
| # put the local target value to the full target heatmap | |
| idx_joints = np.tile( | |
| np.expand_dims(np.arange(K), axis=(-1, -2, -3)), | |
| [1, local_size, local_size, local_size]) | |
| idx = np.stack([idx_joints, zz, yy, xx], | |
| axis=-1).astype(int).reshape(-1, 4) | |
| heatmaps[idx[:, 0], idx[:, 1], idx[:, 2], idx[:, 3]] = np.maximum( | |
| heatmaps[idx[:, 0], idx[:, 1], idx[:, 2], idx[:, 3]], | |
| gaussian.reshape(-1)) | |
| heatmaps = (heatmaps * max_bound).reshape(-1, H, W) | |
| return heatmaps, keypoint_weights | |
| def generate_gaussian_heatmaps( | |
| heatmap_size: Tuple[int, int], | |
| keypoints: np.ndarray, | |
| keypoints_visible: np.ndarray, | |
| sigma: Union[float, Tuple[float], np.ndarray], | |
| ) -> Tuple[np.ndarray, np.ndarray]: | |
| """Generate gaussian heatmaps of keypoints. | |
| 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 or List[float]): A list of sigma values of the Gaussian | |
| heatmap for each instance. If sigma is given as a single float | |
| value, it will be expanded into a tuple | |
| 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) | |
| """ | |
| N, K, _ = keypoints.shape | |
| W, H = heatmap_size | |
| heatmaps = np.zeros((K, H, W), dtype=np.float32) | |
| keypoint_weights = keypoints_visible.copy() | |
| if isinstance(sigma, (int, float)): | |
| sigma = (sigma, ) * N | |
| for n in range(N): | |
| # 3-sigma rule | |
| radius = sigma[n] * 3 | |
| # xy grid | |
| gaussian_size = 2 * radius + 1 | |
| x = np.arange(0, gaussian_size, 1, dtype=np.float32) | |
| y = x[:, None] | |
| x0 = y0 = gaussian_size // 2 | |
| for k in range(K): | |
| # skip unlabled keypoints | |
| if keypoints_visible[n, k] < 0.5: | |
| continue | |
| # get gaussian center coordinates | |
| mu = (keypoints[n, k] + 0.5).astype(np.int64) | |
| # check that the gaussian has in-bounds part | |
| left, top = (mu - radius).astype(np.int64) | |
| right, bottom = (mu + radius + 1).astype(np.int64) | |
| if left >= W or top >= H or right < 0 or bottom < 0: | |
| keypoint_weights[n, k] = 0 | |
| continue | |
| # The gaussian is not normalized, | |
| # we want the center value to equal 1 | |
| gaussian = np.exp(-((x - x0)**2 + (y - y0)**2) / (2 * sigma[n]**2)) | |
| # valid range in gaussian | |
| g_x1 = max(0, -left) | |
| g_x2 = min(W, right) - left | |
| g_y1 = max(0, -top) | |
| g_y2 = min(H, bottom) - top | |
| # valid range in heatmap | |
| h_x1 = max(0, left) | |
| h_x2 = min(W, right) | |
| h_y1 = max(0, top) | |
| h_y2 = min(H, bottom) | |
| heatmap_region = heatmaps[k, h_y1:h_y2, h_x1:h_x2] | |
| gaussian_regsion = gaussian[g_y1:g_y2, g_x1:g_x2] | |
| _ = np.maximum( | |
| heatmap_region, gaussian_regsion, out=heatmap_region) | |
| return heatmaps, keypoint_weights | |
| def generate_unbiased_gaussian_heatmaps( | |
| heatmap_size: Tuple[int, int], | |
| keypoints: np.ndarray, | |
| keypoints_visible: np.ndarray, | |
| sigma: float, | |
| ) -> Tuple[np.ndarray, np.ndarray]: | |
| """Generate gaussian heatmaps of keypoints using `Dark Pose`_. | |
| 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) | |
| 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) | |
| .. _`Dark Pose`: https://arxiv.org/abs/1910.06278 | |
| """ | |
| N, K, _ = keypoints.shape | |
| W, H = heatmap_size | |
| heatmaps = np.zeros((K, H, W), dtype=np.float32) | |
| keypoint_weights = keypoints_visible.copy() | |
| # 3-sigma rule | |
| radius = sigma * 3 | |
| # xy grid | |
| x = np.arange(0, W, 1, dtype=np.float32) | |
| y = np.arange(0, H, 1, dtype=np.float32)[:, None] | |
| for n, k in product(range(N), range(K)): | |
| # skip unlabled keypoints | |
| if keypoints_visible[n, k] < 0.5: | |
| continue | |
| mu = keypoints[n, k] | |
| # check that the gaussian has in-bounds part | |
| left, top = mu - radius | |
| right, bottom = mu + radius + 1 | |
| if left >= W or top >= H or right < 0 or bottom < 0: | |
| keypoint_weights[n, k] = 0 | |
| continue | |
| gaussian = np.exp(-((x - mu[0])**2 + (y - mu[1])**2) / (2 * sigma**2)) | |
| _ = np.maximum(gaussian, heatmaps[k], out=heatmaps[k]) | |
| return heatmaps, keypoint_weights | |
| def generate_udp_gaussian_heatmaps( | |
| heatmap_size: Tuple[int, int], | |
| keypoints: np.ndarray, | |
| keypoints_visible: np.ndarray, | |
| sigma, | |
| keypoints_visibility: np.ndarray, | |
| 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 | |
| heatmaps = np.zeros((K, H, W), dtype=np.float32) | |
| keypoint_weights = keypoints_visible.copy() | |
| if isinstance(sigma, (int, float)): | |
| scaled_sigmas = sigma * np.ones((N, K), dtype=np.float32) | |
| sigmas = np.array([sigma] * K).reshape(1, -1).repeat(N, axis=0) | |
| else: | |
| scaled_sigmas = np.array(sigma).reshape(1, -1).repeat(N, axis=0) | |
| sigmas = np.array(sigma).reshape(1, -1).repeat(N, axis=0) | |
| scales_arr = np.ones((N, K), dtype=np.float32) | |
| if increase_sigma_with_padding: | |
| diag = np.sqrt(W**2 + H**2) | |
| for n in range(N): | |
| image_kpts = keypoints[n, :].squeeze() | |
| vis_kpts = image_kpts[keypoints_visibility[n, :] > 0.5] | |
| # Compute the distance between img_kpts and visible_kpts | |
| if vis_kpts.size == 0: | |
| min_dists = np.ones(image_kpts.shape[0]) * diag | |
| else: | |
| dists = cdist(image_kpts, vis_kpts, metric='euclidean') | |
| min_dists = np.min(dists, axis=1) | |
| scales = min_dists / diag * 2.0 # Maximum distance (diagonal) results in .0*sigma | |
| scales_arr[n, :] = scales | |
| scaled_sigmas[n, :] = sigma * (1+scales) | |
| # print(scales_arr) | |
| # print(scaled_sigmas) | |
| for n, k in product(range(N), range(K)): | |
| scaled_sigma = scaled_sigmas[n, k] | |
| # skip unlabled keypoints | |
| if keypoints_visible[n, k] < 0.5: | |
| continue | |
| # 3-sigma rule | |
| radius = scaled_sigma * 3 | |
| # xy grid | |
| gaussian_size = 2 * radius + 1 | |
| x = np.arange(0, gaussian_size, 1, dtype=np.float32) | |
| y = x[:, None] | |
| mu = (keypoints[n, k] + 0.5).astype(np.int64) | |
| # check that the gaussian has in-bounds part | |
| left, top = (mu - radius).round().astype(np.int64) | |
| right, bottom = (mu + radius + 1).round().astype(np.int64) | |
| # left, top = (mu - radius).astype(np.int64) | |
| # right, bottom = (mu + radius + 1).astype(np.int64) | |
| if left >= W or top >= H or right < 0 or bottom < 0: | |
| keypoint_weights[n, k] = 0 | |
| continue | |
| mu_ac = keypoints[n, k] | |
| x0 = y0 = gaussian_size // 2 | |
| x0 += mu_ac[0] - mu[0] | |
| y0 += mu_ac[1] - mu[1] | |
| gaussian = np.exp(-((x - x0)**2 + (y - y0)**2) / (2 * scaled_sigma**2)) | |
| # Normalize Gaussian such that scaled_sigma = sigma is the norm | |
| gaussian = gaussian / (scaled_sigma / sigmas[n, k]) | |
| # valid range in gaussian | |
| g_x1 = max(0, -left) | |
| g_x2 = min(W, right) - left | |
| g_y1 = max(0, -top) | |
| g_y2 = min(H, bottom) - top | |
| # valid range in heatmap | |
| h_x1 = max(0, left) | |
| h_x2 = min(W, right) | |
| h_y1 = max(0, top) | |
| h_y2 = min(H, bottom) | |
| # breakpoint() | |
| heatmap_region = heatmaps[k, h_y1:h_y2, h_x1:h_x2] | |
| gaussian_regsion = gaussian[g_y1:g_y2, g_x1:g_x2] | |
| _ = np.maximum(heatmap_region, gaussian_regsion, out=heatmap_region) | |
| return heatmaps, keypoint_weights | |
| def generate_onehot_heatmaps( | |
| heatmap_size: Tuple[int, int], | |
| keypoints: np.ndarray, | |
| keypoints_visible: np.ndarray, | |
| sigma, | |
| keypoints_visibility: np.ndarray, | |
| 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 | |
| heatmaps = np.zeros((K, H, W), dtype=np.float32) | |
| keypoint_weights = keypoints_visible.copy() | |
| for n, k in product(range(N), range(K)): | |
| # skip unlabled keypoints | |
| if keypoints_visible[n, k] < 0.5: | |
| continue | |
| mu = (keypoints[n, k] + 0.5).astype(np.int64) | |
| if mu[0] < 0 or mu[0] >= W or mu[1] < 0 or mu[1] >= H: | |
| keypoint_weights[n, k] = 0 | |
| continue | |
| heatmaps[k, mu[1], mu[0]] = 1 | |
| return heatmaps, keypoint_weights | |