Spaces:
Running
on
Zero
Running
on
Zero
| # Copyright (c) OpenMMLab. All rights reserved. | |
| from itertools import product | |
| from typing import Tuple | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import Tensor | |
| from scipy.signal import convolve2d | |
| def get_simcc_normalized(batch_pred_simcc, sigma=None): | |
| """Normalize the predicted SimCC. | |
| Args: | |
| batch_pred_simcc (torch.Tensor): The predicted SimCC. | |
| sigma (float): The sigma of the Gaussian distribution. | |
| Returns: | |
| torch.Tensor: The normalized SimCC. | |
| """ | |
| B, K, _ = batch_pred_simcc.shape | |
| # Scale and clamp the tensor | |
| if sigma is not None: | |
| batch_pred_simcc = batch_pred_simcc / (sigma * np.sqrt(np.pi * 2)) | |
| batch_pred_simcc = batch_pred_simcc.clamp(min=0) | |
| # Compute the binary mask | |
| mask = (batch_pred_simcc.amax(dim=-1) > 1).reshape(B, K, 1) | |
| # Normalize the tensor using the maximum value | |
| norm = (batch_pred_simcc / batch_pred_simcc.amax(dim=-1).reshape(B, K, 1)) | |
| # Apply normalization | |
| batch_pred_simcc = torch.where(mask, norm, batch_pred_simcc) | |
| return batch_pred_simcc | |
| def get_simcc_maximum(simcc_x: np.ndarray, | |
| simcc_y: np.ndarray, | |
| apply_softmax: bool = False | |
| ) -> Tuple[np.ndarray, np.ndarray]: | |
| """Get maximum response location and value from simcc representations. | |
| Note: | |
| instance number: N | |
| num_keypoints: K | |
| heatmap height: H | |
| heatmap width: W | |
| Args: | |
| simcc_x (np.ndarray): x-axis SimCC in shape (K, Wx) or (N, K, Wx) | |
| simcc_y (np.ndarray): y-axis SimCC in shape (K, Wy) or (N, K, Wy) | |
| apply_softmax (bool): whether to apply softmax on the heatmap. | |
| Defaults to False. | |
| Returns: | |
| tuple: | |
| - locs (np.ndarray): locations of maximum heatmap responses in shape | |
| (K, 2) or (N, K, 2) | |
| - vals (np.ndarray): values of maximum heatmap responses in shape | |
| (K,) or (N, K) | |
| """ | |
| assert isinstance(simcc_x, np.ndarray), ('simcc_x should be numpy.ndarray') | |
| assert isinstance(simcc_y, np.ndarray), ('simcc_y should be numpy.ndarray') | |
| assert simcc_x.ndim == 2 or simcc_x.ndim == 3, ( | |
| f'Invalid shape {simcc_x.shape}') | |
| assert simcc_y.ndim == 2 or simcc_y.ndim == 3, ( | |
| f'Invalid shape {simcc_y.shape}') | |
| assert simcc_x.ndim == simcc_y.ndim, ( | |
| f'{simcc_x.shape} != {simcc_y.shape}') | |
| if simcc_x.ndim == 3: | |
| N, K, Wx = simcc_x.shape | |
| simcc_x = simcc_x.reshape(N * K, -1) | |
| simcc_y = simcc_y.reshape(N * K, -1) | |
| else: | |
| N = None | |
| if apply_softmax: | |
| simcc_x = simcc_x - np.max(simcc_x, axis=1, keepdims=True) | |
| simcc_y = simcc_y - np.max(simcc_y, axis=1, keepdims=True) | |
| ex, ey = np.exp(simcc_x), np.exp(simcc_y) | |
| simcc_x = ex / np.sum(ex, axis=1, keepdims=True) | |
| simcc_y = ey / np.sum(ey, axis=1, keepdims=True) | |
| x_locs = np.argmax(simcc_x, axis=1) | |
| y_locs = np.argmax(simcc_y, axis=1) | |
| locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32) | |
| max_val_x = np.amax(simcc_x, axis=1) | |
| max_val_y = np.amax(simcc_y, axis=1) | |
| mask = max_val_x > max_val_y | |
| max_val_x[mask] = max_val_y[mask] | |
| vals = max_val_x | |
| locs[vals <= 0.] = -1 | |
| if N: | |
| locs = locs.reshape(N, K, 2) | |
| vals = vals.reshape(N, K) | |
| return locs, vals | |
| def get_heatmap_3d_maximum(heatmaps: np.ndarray | |
| ) -> Tuple[np.ndarray, np.ndarray]: | |
| """Get maximum response location and value from heatmaps. | |
| Note: | |
| batch_size: B | |
| num_keypoints: K | |
| heatmap dimension: D | |
| heatmap height: H | |
| heatmap width: W | |
| Args: | |
| heatmaps (np.ndarray): Heatmaps in shape (K, D, H, W) or | |
| (B, K, D, H, W) | |
| Returns: | |
| tuple: | |
| - locs (np.ndarray): locations of maximum heatmap responses in shape | |
| (K, 3) or (B, K, 3) | |
| - vals (np.ndarray): values of maximum heatmap responses in shape | |
| (K,) or (B, K) | |
| """ | |
| assert isinstance(heatmaps, | |
| np.ndarray), ('heatmaps should be numpy.ndarray') | |
| assert heatmaps.ndim == 4 or heatmaps.ndim == 5, ( | |
| f'Invalid shape {heatmaps.shape}') | |
| if heatmaps.ndim == 4: | |
| K, D, H, W = heatmaps.shape | |
| B = None | |
| heatmaps_flatten = heatmaps.reshape(K, -1) | |
| else: | |
| B, K, D, H, W = heatmaps.shape | |
| heatmaps_flatten = heatmaps.reshape(B * K, -1) | |
| z_locs, y_locs, x_locs = np.unravel_index( | |
| np.argmax(heatmaps_flatten, axis=1), shape=(D, H, W)) | |
| locs = np.stack((x_locs, y_locs, z_locs), axis=-1).astype(np.float32) | |
| vals = np.amax(heatmaps_flatten, axis=1) | |
| locs[vals <= 0.] = -1 | |
| if B: | |
| locs = locs.reshape(B, K, 3) | |
| vals = vals.reshape(B, K) | |
| return locs, vals | |
| def get_heatmap_maximum(heatmaps: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: | |
| """Get maximum response location and value from heatmaps. | |
| Note: | |
| batch_size: B | |
| num_keypoints: K | |
| heatmap height: H | |
| heatmap width: W | |
| Args: | |
| heatmaps (np.ndarray): Heatmaps in shape (K, H, W) or (B, K, H, W) | |
| Returns: | |
| tuple: | |
| - locs (np.ndarray): locations of maximum heatmap responses in shape | |
| (K, 2) or (B, K, 2) | |
| - vals (np.ndarray): values of maximum heatmap responses in shape | |
| (K,) or (B, K) | |
| """ | |
| assert isinstance(heatmaps, | |
| np.ndarray), ('heatmaps should be numpy.ndarray') | |
| assert heatmaps.ndim == 3 or heatmaps.ndim == 4, ( | |
| f'Invalid shape {heatmaps.shape}') | |
| if heatmaps.ndim == 3: | |
| K, H, W = heatmaps.shape | |
| B = None | |
| heatmaps_flatten = heatmaps.reshape(K, -1) | |
| else: | |
| B, K, H, W = heatmaps.shape | |
| heatmaps_flatten = heatmaps.reshape(B * K, -1) | |
| y_locs, x_locs = np.unravel_index( | |
| np.argmax(heatmaps_flatten, axis=1), shape=(H, W)) | |
| locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32) | |
| vals = np.amax(heatmaps_flatten, axis=1) | |
| locs[vals <= 0.] = -1 | |
| if B: | |
| locs = locs.reshape(B, K, 2) | |
| vals = vals.reshape(B, K) | |
| return locs, vals | |
| def gaussian_blur(heatmaps: np.ndarray, kernel: int = 11) -> np.ndarray: | |
| """Modulate heatmap distribution with Gaussian. | |
| Note: | |
| - num_keypoints: K | |
| - heatmap height: H | |
| - heatmap width: W | |
| Args: | |
| heatmaps (np.ndarray[K, H, W]): model predicted heatmaps. | |
| kernel (int): Gaussian kernel size (K) for modulation, which should | |
| match the heatmap gaussian sigma when training. | |
| K=17 for sigma=3 and k=11 for sigma=2. | |
| Returns: | |
| np.ndarray ([K, H, W]): Modulated heatmap distribution. | |
| """ | |
| assert kernel % 2 == 1 | |
| border = (kernel - 1) // 2 | |
| K, H, W = heatmaps.shape | |
| for k in range(K): | |
| origin_max = np.max(heatmaps[k]) | |
| dr = np.zeros((H + 2 * border, W + 2 * border), dtype=np.float32) | |
| dr[border:-border, border:-border] = heatmaps[k].copy() | |
| dr = cv2.GaussianBlur(dr, (kernel, kernel), 0) | |
| heatmaps[k] = dr[border:-border, border:-border].copy() | |
| heatmaps[k] *= origin_max / (np.max(heatmaps[k])+1e-12) | |
| return heatmaps | |
| def gaussian_blur1d(simcc: np.ndarray, kernel: int = 11) -> np.ndarray: | |
| """Modulate simcc distribution with Gaussian. | |
| Note: | |
| - num_keypoints: K | |
| - simcc length: Wx | |
| Args: | |
| simcc (np.ndarray[K, Wx]): model predicted simcc. | |
| kernel (int): Gaussian kernel size (K) for modulation, which should | |
| match the simcc gaussian sigma when training. | |
| K=17 for sigma=3 and k=11 for sigma=2. | |
| Returns: | |
| np.ndarray ([K, Wx]): Modulated simcc distribution. | |
| """ | |
| assert kernel % 2 == 1 | |
| border = (kernel - 1) // 2 | |
| N, K, Wx = simcc.shape | |
| for n, k in product(range(N), range(K)): | |
| origin_max = np.max(simcc[n, k]) | |
| dr = np.zeros((1, Wx + 2 * border), dtype=np.float32) | |
| dr[0, border:-border] = simcc[n, k].copy() | |
| dr = cv2.GaussianBlur(dr, (kernel, 1), 0) | |
| simcc[n, k] = dr[0, border:-border].copy() | |
| simcc[n, k] *= origin_max / np.max(simcc[n, k]) | |
| return simcc | |
| def batch_heatmap_nms(batch_heatmaps: Tensor, kernel_size: int = 5): | |
| """Apply NMS on a batch of heatmaps. | |
| Args: | |
| batch_heatmaps (Tensor): batch heatmaps in shape (B, K, H, W) | |
| kernel_size (int): The kernel size of the NMS which should be | |
| a odd integer. Defaults to 5 | |
| Returns: | |
| Tensor: The batch heatmaps after NMS. | |
| """ | |
| assert isinstance(kernel_size, int) and kernel_size % 2 == 1, \ | |
| f'The kernel_size should be an odd integer, got {kernel_size}' | |
| padding = (kernel_size - 1) // 2 | |
| maximum = F.max_pool2d( | |
| batch_heatmaps, kernel_size, stride=1, padding=padding) | |
| maximum_indicator = torch.eq(batch_heatmaps, maximum) | |
| batch_heatmaps = batch_heatmaps * maximum_indicator.float() | |
| return batch_heatmaps | |
| def get_heatmap_expected_value(heatmaps: np.ndarray, parzen_size: float = 0.1, return_heatmap: bool = False) -> Tuple[np.ndarray, np.ndarray]: | |
| """Get maximum response location and value from heatmaps. | |
| Note: | |
| batch_size: B | |
| num_keypoints: K | |
| heatmap height: H | |
| heatmap width: W | |
| Args: | |
| heatmaps (np.ndarray): Heatmaps in shape (K, H, W) or (B, K, H, W) | |
| Returns: | |
| tuple: | |
| - locs (np.ndarray): locations of maximum heatmap responses in shape | |
| (K, 2) or (B, K, 2) | |
| - vals (np.ndarray): values of maximum heatmap responses in shape | |
| (K,) or (B, K) | |
| """ | |
| assert isinstance(heatmaps, | |
| np.ndarray), ('heatmaps should be numpy.ndarray') | |
| assert heatmaps.ndim == 3 or heatmaps.ndim == 4, ( | |
| f'Invalid shape {heatmaps.shape}') | |
| assert parzen_size >= 0.0 and parzen_size <= 1.0, ( | |
| f'Invalid parzen_size {parzen_size}') | |
| if heatmaps.ndim == 3: | |
| K, H, W = heatmaps.shape | |
| B = 1 | |
| FIRST_DIM = K | |
| heatmaps_flatten = heatmaps.reshape(1, K, H, W) | |
| else: | |
| B, K, H, W = heatmaps.shape | |
| FIRST_DIM = K*B | |
| heatmaps_flatten = heatmaps.reshape(B, K, H, W) | |
| # Blur heatmaps with Gaussian | |
| # heatmaps_flatten = gaussian_blur(heatmaps_flatten, kernel=9) | |
| # Zero out pixels far from the maximum for each heatmap | |
| # heatmaps_tmp = heatmaps_flatten.copy().reshape(B*K, H*W) | |
| # y_locs, x_locs = np.unravel_index( | |
| # np.argmax(heatmaps_tmp, axis=1), shape=(H, W)) | |
| # locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32) | |
| # heatmaps_flatten = heatmaps_flatten.reshape(B*K, H, W) | |
| # for i, x in enumerate(x_locs): | |
| # y = y_locs[i] | |
| # start_x = int(max(0, x - 0.2*W)) | |
| # end_x = int(min(W, x + 0.2*W)) | |
| # start_y = int(max(0, y - 0.2*H)) | |
| # end_y = int(min(H, y + 0.2*H)) | |
| # mask = np.zeros((H, W)) | |
| # mask[start_y:end_y, start_x:end_x] = 1 | |
| # heatmaps_flatten[i] = heatmaps_flatten[i] * mask | |
| # heatmaps_flatten = heatmaps_flatten.reshape(B, K, H, W) | |
| bbox_area = np.sqrt(H/1.25 * W/1.25) | |
| kpt_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 | |
| heatmaps_covolved = np.zeros_like(heatmaps_flatten) | |
| for k in range(K): | |
| vars = (kpt_sigmas[k]*2)**2 | |
| s = vars * bbox_area * 2 | |
| s = np.clip(s, 0.55, 3.0) | |
| radius = np.ceil(s * 3).astype(int) | |
| diameter = 2*radius + 1 | |
| diameter = np.ceil(diameter).astype(int) | |
| # kernel_sizes[kernel_sizes % 2 == 0] += 1 | |
| center = diameter // 2 | |
| dist_x = np.arange(diameter) - center | |
| dist_y = np.arange(diameter) - center | |
| dist_x, dist_y = np.meshgrid(dist_x, dist_y) | |
| dist = np.sqrt(dist_x**2 + dist_y**2) | |
| oks_kernel = np.exp(-dist**2 / (2 * s)) | |
| oks_kernel = oks_kernel / oks_kernel.sum() | |
| htm = heatmaps_flatten[:, k, :, :].reshape(-1, H, W) | |
| # htm = np.pad(htm, ((0, 0), (radius, radius), (radius, radius)), mode='symmetric') | |
| # htm = torch.from_numpy(htm).float() | |
| # oks_kernel = torch.from_numpy(oks_kernel).float().to(htm.device).reshape(1, diameter, diameter) | |
| oks_kernel = oks_kernel.reshape(1, diameter, diameter) | |
| htm_conv = np.zeros_like(htm) | |
| for b in range(B): | |
| htm_conv[b, :, :] = convolve2d(htm[b, :, :], oks_kernel[b, :, :], mode='same', boundary='symm') | |
| # htm_conv = F.conv2d(htm.unsqueeze(1), oks_kernel.unsqueeze(1), padding='same') | |
| # htm_conv = htm_conv[:, :, radius:-radius, radius:-radius] | |
| htm_conv = htm_conv.reshape(-1, 1, H, W) | |
| heatmaps_covolved[:, k, :, :] = htm_conv | |
| heatmaps_covolved = heatmaps_covolved.reshape(B*K, H*W) | |
| y_locs, x_locs = np.unravel_index( | |
| np.argmax(heatmaps_covolved, axis=1), shape=(H, W)) | |
| locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32) | |
| # Apply mean-shift to get sub-pixel locations | |
| locs = _get_subpixel_maximums(heatmaps_covolved.reshape(B*K, H, W), locs) | |
| # breakpoint() | |
| # heatmaps_sums = heatmaps_flatten.sum(axis=(1, 2)) | |
| # norm_heatmaps = heatmaps_flatten.copy() | |
| # norm_heatmaps[heatmaps_sums > 0] = heatmaps_flatten[heatmaps_sums > 0] / heatmaps_sums[heatmaps_sums > 0, None, None] | |
| # # Compute Parzen window with Gaussian blur along the edge instead of simple mirroring | |
| # x_pad = int(parzen_size * W + 0.5) | |
| # y_pad = int(parzen_size * H + 0.5) | |
| # # x_pad = 0 | |
| # # y_pad = 0 | |
| # kernel_size = int(min(H, W)*parzen_size + 0.5) | |
| # if kernel_size % 2 == 0: | |
| # kernel_size += 1 | |
| # # norm_heatmaps_pad_blur = np.pad(norm_heatmaps, ((0, 0), (x_pad, x_pad), (y_pad, y_pad)), mode='symmetric') | |
| # norm_heatmaps_pad = np.pad(norm_heatmaps, ((0, 0), (y_pad, y_pad), (x_pad, x_pad)), mode='constant', constant_values=0) | |
| # norm_heatmaps_pad_blur = gaussian_blur(norm_heatmaps_pad, kernel=kernel_size) | |
| # # norm_heatmaps_pad_blur[:, x_pad:-x_pad, y_pad:-y_pad] = norm_heatmaps | |
| # norm_heatmaps_pad_sum = norm_heatmaps_pad_blur.sum(axis=(1, 2)) | |
| # norm_heatmaps_pad_blur[norm_heatmaps_pad_sum>0] = norm_heatmaps_pad_blur[norm_heatmaps_pad_sum>0] / norm_heatmaps_pad_sum[norm_heatmaps_pad_sum>0, None, None] | |
| # # # Save the blurred heatmaps | |
| # # for i in range(heatmaps.shape[0]): | |
| # # tmp_htm = norm_heatmaps_pad_blur[i].copy() | |
| # # tmp_htm = (tmp_htm - tmp_htm.min()) / (tmp_htm.max() - tmp_htm.min()) | |
| # # tmp_htm = (tmp_htm*255).astype(np.uint8) | |
| # # tmp_htm = cv2.cvtColor(tmp_htm, cv2.COLOR_GRAY2BGR) | |
| # # tmp_htm = cv2.applyColorMap(tmp_htm, cv2.COLORMAP_JET) | |
| # # tmp_htm2 = norm_heatmaps_pad[i].copy() | |
| # # tmp_htm2 = (tmp_htm2 - tmp_htm2.min()) / (tmp_htm2.max() - tmp_htm2.min()) | |
| # # tmp_htm2 = (tmp_htm2*255).astype(np.uint8) | |
| # # tmp_htm2 = cv2.cvtColor(tmp_htm2, cv2.COLOR_GRAY2BGR) | |
| # # tmp_htm2 = cv2.applyColorMap(tmp_htm2, cv2.COLORMAP_JET) | |
| # # tmp_htm = cv2.addWeighted(tmp_htm, 0.5, tmp_htm2, 0.5, 0) | |
| # # cv2.imwrite(f'heatmaps_blurred_{i}.png', tmp_htm) | |
| # # norm_heatmaps_pad = np.pad(norm_heatmaps, ((0, 0), (x_pad, x_pad), (y_pad, y_pad)), mode='edge') | |
| # y_idx, x_idx = np.indices(norm_heatmaps_pad_blur.shape[1:]) | |
| # # breakpoint() | |
| # x_locs = np.sum(norm_heatmaps_pad_blur * x_idx, axis=(1, 2)) - x_pad | |
| # y_locs = np.sum(norm_heatmaps_pad_blur * y_idx, axis=(1, 2)) - y_pad | |
| # # mean_idx = np.argmax(heatmaps_flatten, axis=1) | |
| # # x_locs, y_locs = np.unravel_index(mean_idx, shape=(H, W)) | |
| # # locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32) | |
| # # breakpoint() | |
| # # vals = heatmaps_flatten[np.arange(heatmaps_flatten.shape[0]), mean_idx] | |
| # # locs[vals <= 0.] = -1 | |
| # # mean_idx = np.argmax(norm_heatmaps, axis=1) | |
| # # y_locs, x_locs = np.unravel_index( | |
| # # mean_idx, shape=(H, W)) | |
| # locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32) | |
| # # vals = np.amax(heatmaps_flatten, axis=1) | |
| x_locs_int = np.round(x_locs).astype(int) | |
| x_locs_int = np.clip(x_locs_int, 0, W-1) | |
| y_locs_int = np.round(y_locs).astype(int) | |
| y_locs_int = np.clip(y_locs_int, 0, H-1) | |
| vals = heatmaps_flatten[np.arange(B), np.arange(K), y_locs_int, x_locs_int] | |
| # breakpoint() | |
| # locs[vals <= 0.] = -1 | |
| # print(mean_idx) | |
| # print(x_locs) | |
| # print(y_locs) | |
| # print(locs) | |
| heatmaps_covolved = heatmaps_covolved.reshape(B, K, H, W) | |
| if B > 1: | |
| locs = locs.reshape(B, K, 2) | |
| vals = vals.reshape(B, K) | |
| heatmaps_covolved = heatmaps_covolved.reshape(B, K, H, W) | |
| else: | |
| locs = locs.reshape(K, 2) | |
| vals = vals.reshape(K) | |
| heatmaps_covolved = heatmaps_covolved.reshape(K, H, W) | |
| if return_heatmap: | |
| return locs, vals, heatmaps_covolved | |
| else: | |
| return locs, vals | |
| def _get_subpixel_maximums(heatmaps, locs): | |
| # Extract integer peak locations | |
| x_locs = locs[:, 0].astype(np.int32) | |
| y_locs = locs[:, 1].astype(np.int32) | |
| # Ensure we are not near the boundaries (avoid boundary issues) | |
| valid_mask = (x_locs > 0) & (x_locs < heatmaps.shape[2] - 1) & \ | |
| (y_locs > 0) & (y_locs < heatmaps.shape[1] - 1) | |
| # Initialize the output array with the integer locations | |
| subpixel_locs = locs.copy() | |
| if np.any(valid_mask): | |
| # Extract valid locations | |
| x_locs_valid = x_locs[valid_mask] | |
| y_locs_valid = y_locs[valid_mask] | |
| # Compute gradients (dx, dy) and second derivatives (dxx, dyy) | |
| dx = (heatmaps[valid_mask, y_locs_valid, x_locs_valid + 1] - | |
| heatmaps[valid_mask, y_locs_valid, x_locs_valid - 1]) / 2.0 | |
| dy = (heatmaps[valid_mask, y_locs_valid + 1, x_locs_valid] - | |
| heatmaps[valid_mask, y_locs_valid - 1, x_locs_valid]) / 2.0 | |
| dxx = heatmaps[valid_mask, y_locs_valid, x_locs_valid + 1] + \ | |
| heatmaps[valid_mask, y_locs_valid, x_locs_valid - 1] - \ | |
| 2 * heatmaps[valid_mask, y_locs_valid, x_locs_valid] | |
| dyy = heatmaps[valid_mask, y_locs_valid + 1, x_locs_valid] + \ | |
| heatmaps[valid_mask, y_locs_valid - 1, x_locs_valid] - \ | |
| 2 * heatmaps[valid_mask, y_locs_valid, x_locs_valid] | |
| # Avoid division by zero by setting a minimum threshold for the second derivatives | |
| dxx = np.where(dxx != 0, dxx, 1e-6) | |
| dyy = np.where(dyy != 0, dyy, 1e-6) | |
| # Calculate the sub-pixel shift | |
| subpixel_x_shift = -dx / dxx | |
| subpixel_y_shift = -dy / dyy | |
| # Update subpixel locations for valid indices | |
| subpixel_locs[valid_mask, 0] += subpixel_x_shift | |
| subpixel_locs[valid_mask, 1] += subpixel_y_shift | |
| return subpixel_locs | |