import numpy as np

PART_NAMES = [
    "nose",
    "leftEye",
    "rightEye",
    "leftEar",
    "rightEar",
    "leftShoulder",
    "rightShoulder",
    "leftElbow",
    "rightElbow",
    "leftWrist",
    "rightWrist",
    "leftHip",
    "rightHip",
    "leftKnee",
    "rightKnee",
    "leftAnkle",
    "rightAnkle",
]

NUM_KEYPOINTS = len(PART_NAMES)  # 17


#############################################################
# argmax2d
#
# Input:
#   t (H, W, D)
# Returns:
#   matrix (D, 2), with each row as [y, x] of argmax for D
############################################################
def argmax2d(t):
    if len(t.shape) > 3:
        t = np.squeeze(t)

    if not len(t.shape) == 3:
        print("Input must be a 3D matrix, or be able to be squeezed into one.")
        return

    height, width, depth = t.shape

    reshaped_t = np.reshape(t, [height * width, depth])
    argmax_coords = np.argmax(reshaped_t, axis=0)
    y_coords = argmax_coords // width
    x_coords = argmax_coords % width

    return np.concatenate(
        [np.expand_dims(y_coords, 1), np.expand_dims(x_coords, 1)], axis=1
    )


###################################################################################
# get_offset_vectors
#
# Input:
#   heatmap_coords (NUM_KEYPOINTS, 2)
#   offsets (height, width, NUM_KEYPOINTS * 2)
# Returns:
#   matrix (NUM_KEYPOINTS, 2), with each row as [y, x] of offset for each keypoint
###################################################################################
def get_offset_vectors(heatmaps_coords, offsets):
    result = []

    for keypoint in range(NUM_KEYPOINTS):
        heatmap_y = heatmaps_coords[keypoint, 0]
        heatmap_x = heatmaps_coords[keypoint, 1]

        offset_y = offsets[heatmap_y, heatmap_x, keypoint]
        offset_x = offsets[heatmap_y, heatmap_x, keypoint + NUM_KEYPOINTS]

        result.append([offset_y, offset_x])

    return result


############################################################################################
# get_offset_points
#
# Input:
#   heatmap_coords (NUM_KEYPOINTS, 2)
#   offsets (H, W, NUM_KEYPOINTS * 2)
#   output_stride (scalar)
# Returns:
#   matrix (NUM_KEYPOINTS, 2), with each row as [y, x] location prediction for each keypoint
#############################################################################################
def get_offset_points(heatmaps_coords, offsets, output_stride):
    offset_vectors = get_offset_vectors(heatmaps_coords, offsets)
    scaled_heatmap = heatmaps_coords * output_stride
    return scaled_heatmap + offset_vectors


##############################################################
# get_points_confidence
#
# Input:
#   heatmaps (H, W, NUM_KEYPOINTS)
#   heatmap_coords (NUM_KEYPOINTS, 2)
# Returns:
#   matrix (NUM_KEYPOINTS), with confidence for each keypoint
##############################################################
def get_points_confidence(heatmaps, heatmaps_coords):
    result = []

    for keypoint in range(NUM_KEYPOINTS):
        # Get max value of heatmap for each keypoint
        result.append(
            heatmaps[
                heatmaps_coords[keypoint, 0], heatmaps_coords[keypoint, 1], keypoint
            ]
        )

    return result


#####################################################################################################
# decode_single_pose
#
# Input:
#   heatmaps (H, W, NUM_KEYPOINTS)
#   offsets (H, W, NUM_KEYPOINTS * 2)
#   output_stride (scalar)
# Returns:
#   prediction in the form outlined at https://github.com/tensorflow/tfjs-models/tree/master/posenet
#####################################################################################################
def decode_single_pose(heatmaps, offsets, output_stride, width_factor, height_factor):
    # Squeeze into 3D arrays
    poses = []
    heatmaps = np.squeeze(heatmaps)
    offsets = np.squeeze(offsets)

    heatmaps_coords = argmax2d(heatmaps)
    offset_points = get_offset_points(heatmaps_coords, offsets, output_stride)
    keypoint_confidence = get_points_confidence(heatmaps, heatmaps_coords)

    keypoints = [
        {
            "position": {
                "y": offset_points[keypoint, 0] * height_factor,
                "x": offset_points[keypoint, 1] * width_factor,
            },
            "part": PART_NAMES[keypoint],
            "score": score,
        }
        for keypoint, score in enumerate(keypoint_confidence)
    ]

    poses.append(
        {
            "keypoints": keypoints,
            "score": (sum(keypoint_confidence) / len(keypoint_confidence)),
        }
    )
    return poses

24ce62ce7692ea9811cb42e8823d9d4fb3e37018
<mask token>


def get_offset_points(heatmaps_coords, offsets, output_stride):
    offset_vectors = get_offset_vectors(heatmaps_coords, offsets)
    scaled_heatmap = heatmaps_coords * output_stride
    return scaled_heatmap + offset_vectors


def get_points_confidence(heatmaps, heatmaps_coords):
    result = []
    for keypoint in range(NUM_KEYPOINTS):
        result.append(heatmaps[heatmaps_coords[keypoint, 0],
            heatmaps_coords[keypoint, 1], keypoint])
    return result


def decode_single_pose(heatmaps, offsets, output_stride, width_factor,
    height_factor):
    poses = []
    heatmaps = np.squeeze(heatmaps)
    offsets = np.squeeze(offsets)
    heatmaps_coords = argmax2d(heatmaps)
    offset_points = get_offset_points(heatmaps_coords, offsets, output_stride)
    keypoint_confidence = get_points_confidence(heatmaps, heatmaps_coords)
    keypoints = [{'position': {'y': offset_points[keypoint, 0] *
        height_factor, 'x': offset_points[keypoint, 1] * width_factor},
        'part': PART_NAMES[keypoint], 'score': score} for keypoint, score in
        enumerate(keypoint_confidence)]
    poses.append({'keypoints': keypoints, 'score': sum(keypoint_confidence) /
        len(keypoint_confidence)})
    return poses

<mask token>


def argmax2d(t):
    if len(t.shape) > 3:
        t = np.squeeze(t)
    if not len(t.shape) == 3:
        print('Input must be a 3D matrix, or be able to be squeezed into one.')
        return
    height, width, depth = t.shape
    reshaped_t = np.reshape(t, [height * width, depth])
    argmax_coords = np.argmax(reshaped_t, axis=0)
    y_coords = argmax_coords // width
    x_coords = argmax_coords % width
    return np.concatenate([np.expand_dims(y_coords, 1), np.expand_dims(
        x_coords, 1)], axis=1)


def get_offset_vectors(heatmaps_coords, offsets):
    result = []
    for keypoint in range(NUM_KEYPOINTS):
        heatmap_y = heatmaps_coords[keypoint, 0]
        heatmap_x = heatmaps_coords[keypoint, 1]
        offset_y = offsets[heatmap_y, heatmap_x, keypoint]
        offset_x = offsets[heatmap_y, heatmap_x, keypoint + NUM_KEYPOINTS]
        result.append([offset_y, offset_x])
    return result


def get_offset_points(heatmaps_coords, offsets, output_stride):
    offset_vectors = get_offset_vectors(heatmaps_coords, offsets)
    scaled_heatmap = heatmaps_coords * output_stride
    return scaled_heatmap + offset_vectors


def get_points_confidence(heatmaps, heatmaps_coords):
    result = []
    for keypoint in range(NUM_KEYPOINTS):
        result.append(heatmaps[heatmaps_coords[keypoint, 0],
            heatmaps_coords[keypoint, 1], keypoint])
    return result


def decode_single_pose(heatmaps, offsets, output_stride, width_factor,
    height_factor):
    poses = []
    heatmaps = np.squeeze(heatmaps)
    offsets = np.squeeze(offsets)
    heatmaps_coords = argmax2d(heatmaps)
    offset_points = get_offset_points(heatmaps_coords, offsets, output_stride)
    keypoint_confidence = get_points_confidence(heatmaps, heatmaps_coords)
    keypoints = [{'position': {'y': offset_points[keypoint, 0] *
        height_factor, 'x': offset_points[keypoint, 1] * width_factor},
        'part': PART_NAMES[keypoint], 'score': score} for keypoint, score in
        enumerate(keypoint_confidence)]
    poses.append({'keypoints': keypoints, 'score': sum(keypoint_confidence) /
        len(keypoint_confidence)})
    return poses

<mask token>
PART_NAMES = ['nose', 'leftEye', 'rightEye', 'leftEar', 'rightEar',
    'leftShoulder', 'rightShoulder', 'leftElbow', 'rightElbow', 'leftWrist',
    'rightWrist', 'leftHip', 'rightHip', 'leftKnee', 'rightKnee',
    'leftAnkle', 'rightAnkle']
NUM_KEYPOINTS = len(PART_NAMES)


def argmax2d(t):
    if len(t.shape) > 3:
        t = np.squeeze(t)
    if not len(t.shape) == 3:
        print('Input must be a 3D matrix, or be able to be squeezed into one.')
        return
    height, width, depth = t.shape
    reshaped_t = np.reshape(t, [height * width, depth])
    argmax_coords = np.argmax(reshaped_t, axis=0)
    y_coords = argmax_coords // width
    x_coords = argmax_coords % width
    return np.concatenate([np.expand_dims(y_coords, 1), np.expand_dims(
        x_coords, 1)], axis=1)


def get_offset_vectors(heatmaps_coords, offsets):
    result = []
    for keypoint in range(NUM_KEYPOINTS):
        heatmap_y = heatmaps_coords[keypoint, 0]
        heatmap_x = heatmaps_coords[keypoint, 1]
        offset_y = offsets[heatmap_y, heatmap_x, keypoint]
        offset_x = offsets[heatmap_y, heatmap_x, keypoint + NUM_KEYPOINTS]
        result.append([offset_y, offset_x])
    return result


def get_offset_points(heatmaps_coords, offsets, output_stride):
    offset_vectors = get_offset_vectors(heatmaps_coords, offsets)
    scaled_heatmap = heatmaps_coords * output_stride
    return scaled_heatmap + offset_vectors


def get_points_confidence(heatmaps, heatmaps_coords):
    result = []
    for keypoint in range(NUM_KEYPOINTS):
        result.append(heatmaps[heatmaps_coords[keypoint, 0],
            heatmaps_coords[keypoint, 1], keypoint])
    return result


def decode_single_pose(heatmaps, offsets, output_stride, width_factor,
    height_factor):
    poses = []
    heatmaps = np.squeeze(heatmaps)
    offsets = np.squeeze(offsets)
    heatmaps_coords = argmax2d(heatmaps)
    offset_points = get_offset_points(heatmaps_coords, offsets, output_stride)
    keypoint_confidence = get_points_confidence(heatmaps, heatmaps_coords)
    keypoints = [{'position': {'y': offset_points[keypoint, 0] *
        height_factor, 'x': offset_points[keypoint, 1] * width_factor},
        'part': PART_NAMES[keypoint], 'score': score} for keypoint, score in
        enumerate(keypoint_confidence)]
    poses.append({'keypoints': keypoints, 'score': sum(keypoint_confidence) /
        len(keypoint_confidence)})
    return poses

import numpy as np
PART_NAMES = ['nose', 'leftEye', 'rightEye', 'leftEar', 'rightEar',
    'leftShoulder', 'rightShoulder', 'leftElbow', 'rightElbow', 'leftWrist',
    'rightWrist', 'leftHip', 'rightHip', 'leftKnee', 'rightKnee',
    'leftAnkle', 'rightAnkle']
NUM_KEYPOINTS = len(PART_NAMES)


def argmax2d(t):
    if len(t.shape) > 3:
        t = np.squeeze(t)
    if not len(t.shape) == 3:
        print('Input must be a 3D matrix, or be able to be squeezed into one.')
        return
    height, width, depth = t.shape
    reshaped_t = np.reshape(t, [height * width, depth])
    argmax_coords = np.argmax(reshaped_t, axis=0)
    y_coords = argmax_coords // width
    x_coords = argmax_coords % width
    return np.concatenate([np.expand_dims(y_coords, 1), np.expand_dims(
        x_coords, 1)], axis=1)


def get_offset_vectors(heatmaps_coords, offsets):
    result = []
    for keypoint in range(NUM_KEYPOINTS):
        heatmap_y = heatmaps_coords[keypoint, 0]
        heatmap_x = heatmaps_coords[keypoint, 1]
        offset_y = offsets[heatmap_y, heatmap_x, keypoint]
        offset_x = offsets[heatmap_y, heatmap_x, keypoint + NUM_KEYPOINTS]
        result.append([offset_y, offset_x])
    return result


def get_offset_points(heatmaps_coords, offsets, output_stride):
    offset_vectors = get_offset_vectors(heatmaps_coords, offsets)
    scaled_heatmap = heatmaps_coords * output_stride
    return scaled_heatmap + offset_vectors


def get_points_confidence(heatmaps, heatmaps_coords):
    result = []
    for keypoint in range(NUM_KEYPOINTS):
        result.append(heatmaps[heatmaps_coords[keypoint, 0],
            heatmaps_coords[keypoint, 1], keypoint])
    return result


def decode_single_pose(heatmaps, offsets, output_stride, width_factor,
    height_factor):
    poses = []
    heatmaps = np.squeeze(heatmaps)
    offsets = np.squeeze(offsets)
    heatmaps_coords = argmax2d(heatmaps)
    offset_points = get_offset_points(heatmaps_coords, offsets, output_stride)
    keypoint_confidence = get_points_confidence(heatmaps, heatmaps_coords)
    keypoints = [{'position': {'y': offset_points[keypoint, 0] *
        height_factor, 'x': offset_points[keypoint, 1] * width_factor},
        'part': PART_NAMES[keypoint], 'score': score} for keypoint, score in
        enumerate(keypoint_confidence)]
    poses.append({'keypoints': keypoints, 'score': sum(keypoint_confidence) /
        len(keypoint_confidence)})
    return poses

import numpy as np

PART_NAMES = [
    "nose",
    "leftEye",
    "rightEye",
    "leftEar",
    "rightEar",
    "leftShoulder",
    "rightShoulder",
    "leftElbow",
    "rightElbow",
    "leftWrist",
    "rightWrist",
    "leftHip",
    "rightHip",
    "leftKnee",
    "rightKnee",
    "leftAnkle",
    "rightAnkle",
]

NUM_KEYPOINTS = len(PART_NAMES)  # 17


#############################################################
# argmax2d
#
# Input:
#   t (H, W, D)
# Returns:
#   matrix (D, 2), with each row as [y, x] of argmax for D
############################################################
def argmax2d(t):
    if len(t.shape) > 3:
        t = np.squeeze(t)

    if not len(t.shape) == 3:
        print("Input must be a 3D matrix, or be able to be squeezed into one.")
        return

    height, width, depth = t.shape

    reshaped_t = np.reshape(t, [height * width, depth])
    argmax_coords = np.argmax(reshaped_t, axis=0)
    y_coords = argmax_coords // width
    x_coords = argmax_coords % width

    return np.concatenate(
        [np.expand_dims(y_coords, 1), np.expand_dims(x_coords, 1)], axis=1
    )


###################################################################################
# get_offset_vectors
#
# Input:
#   heatmap_coords (NUM_KEYPOINTS, 2)
#   offsets (height, width, NUM_KEYPOINTS * 2)
# Returns:
#   matrix (NUM_KEYPOINTS, 2), with each row as [y, x] of offset for each keypoint
###################################################################################
def get_offset_vectors(heatmaps_coords, offsets):
    result = []

    for keypoint in range(NUM_KEYPOINTS):
        heatmap_y = heatmaps_coords[keypoint, 0]
        heatmap_x = heatmaps_coords[keypoint, 1]

        offset_y = offsets[heatmap_y, heatmap_x, keypoint]
        offset_x = offsets[heatmap_y, heatmap_x, keypoint + NUM_KEYPOINTS]

        result.append([offset_y, offset_x])

    return result


############################################################################################
# get_offset_points
#
# Input:
#   heatmap_coords (NUM_KEYPOINTS, 2)
#   offsets (H, W, NUM_KEYPOINTS * 2)
#   output_stride (scalar)
# Returns:
#   matrix (NUM_KEYPOINTS, 2), with each row as [y, x] location prediction for each keypoint
#############################################################################################
def get_offset_points(heatmaps_coords, offsets, output_stride):
    offset_vectors = get_offset_vectors(heatmaps_coords, offsets)
    scaled_heatmap = heatmaps_coords * output_stride
    return scaled_heatmap + offset_vectors


##############################################################
# get_points_confidence
#
# Input:
#   heatmaps (H, W, NUM_KEYPOINTS)
#   heatmap_coords (NUM_KEYPOINTS, 2)
# Returns:
#   matrix (NUM_KEYPOINTS), with confidence for each keypoint
##############################################################
def get_points_confidence(heatmaps, heatmaps_coords):
    result = []

    for keypoint in range(NUM_KEYPOINTS):
        # Get max value of heatmap for each keypoint
        result.append(
            heatmaps[
                heatmaps_coords[keypoint, 0], heatmaps_coords[keypoint, 1], keypoint
            ]
        )

    return result


#####################################################################################################
# decode_single_pose
#
# Input:
#   heatmaps (H, W, NUM_KEYPOINTS)
#   offsets (H, W, NUM_KEYPOINTS * 2)
#   output_stride (scalar)
# Returns:
#   prediction in the form outlined at https://github.com/tensorflow/tfjs-models/tree/master/posenet
#####################################################################################################
def decode_single_pose(heatmaps, offsets, output_stride, width_factor, height_factor):
    # Squeeze into 3D arrays
    poses = []
    heatmaps = np.squeeze(heatmaps)
    offsets = np.squeeze(offsets)

    heatmaps_coords = argmax2d(heatmaps)
    offset_points = get_offset_points(heatmaps_coords, offsets, output_stride)
    keypoint_confidence = get_points_confidence(heatmaps, heatmaps_coords)

    keypoints = [
        {
            "position": {
                "y": offset_points[keypoint, 0] * height_factor,
                "x": offset_points[keypoint, 1] * width_factor,
            },
            "part": PART_NAMES[keypoint],
            "score": score,
        }
        for keypoint, score in enumerate(keypoint_confidence)
    ]

    poses.append(
        {
            "keypoints": keypoints,
            "score": (sum(keypoint_confidence) / len(keypoint_confidence)),
        }
    )
    return poses
