def METHOD_NAME(self,
              heatmaps,
              keypoint_indices,
              tag_per_keypoint: bool,
              with_randomness: bool = True):
    K, H, W = heatmaps.shape
    N = keypoint_indices.shape[0]
    if tag_per_keypoint:
        tags = np.zeros((K, H, W), dtype=np.float32)
    else:
        tags = np.zeros((1, H, W), dtype=np.float32)
    for n, k in product(range(N), range(K)):
        y, x = np.unravel_index(keypoint_indices[n, k, 0], (H, W))
        randomness = np.random.rand() if with_randomness else 0
        if tag_per_keypoint:
            tags[k, y, x] = n + randomness
        else:
            tags[0, y, x] = n + randomness
    return tags