import torch


def get_keypoints(heatmaps, thr=0):
    """
    Gets the keypoints from the heatmaps
    :param heatmaps: The heatmaps of shape (n_keypoints, height, width)
    :param thr: Threshold for the location of the keypoints.
                Default =0. So if the keypoint is located a (0,0), it set as non-existing.
    :return: The keypoints
    """
    n, h, w = heatmaps.size()
    flat = heatmaps.view(n, -1)

    max_val, max_idx = flat.max(dim=1)

    xx = (max_idx % w).view(-1, 1)  # (-1,1) for column vector
    yy = (max_idx // w).view(-1, 1)  # (-1,1) for column vector

    xx[max_val <= thr] = -1
    yy[max_val <= thr] = -1

    keypoints = torch.cat((xx, yy), dim=1)

    return keypoints