poser.py

#############################################################################
#
#   Source from:
#   https://www.tensorflow.org/hub/tutorials/movenet
#
#
#############################################################################
import PIL.Image
import PIL.ImageOps
import numpy as np
import tensorflow as tf
from PIL import ImageDraw
from huggingface_hub import snapshot_download


# Dictionary that maps from joint names to keypoint indices.
KEYPOINT_DICT = {
    'nose': 0,
    'left_eye': 1, 'right_eye': 2,
    'left_ear': 3, 'right_ear': 4,
    'left_shoulder': 5, 'right_shoulder': 6,
    'left_elbow': 7, 'right_elbow': 8,
    'left_wrist': 9, 'right_wrist': 10,
    'left_hip': 11, 'right_hip': 12,
    'left_knee': 13, 'right_knee': 14,
    'left_ankle': 15, 'right_ankle': 16
}

COLOR_DICT = {
    (0, 1): 'Magenta',
    (0, 2): 'Cyan',
    (1, 3): 'Magenta',
    (2, 4): 'Cyan',
    (0, 5): 'Magenta',
    (0, 6): 'Cyan',
    (5, 7): 'Magenta',
    (7, 9): 'Magenta',
    (6, 8): 'Cyan',
    (8, 10): 'Cyan',
    (5, 6): 'Yellow',
    (5, 11): 'Magenta',
    (6, 12): 'Cyan',
    (11, 12): 'Yellow',
    (11, 13): 'Magenta',
    (13, 15): 'Magenta',
    (12, 14): 'Cyan',
    (14, 16): 'Cyan'
}


def process_keypoints(keypoints, height, width, threshold=0.22):
    """Returns high confidence keypoints and edges for visualization.

      Args:
        keypoints: A numpy array with shape [1, 1, 17, 3] representing
          the keypoint coordinates and scores returned from the MoveNet model.
        height: height of the image in pixels.
        width: width of the image in pixels.
        threshold: minimum confidence score for a keypoint to be
          visualized.

      Returns:
        A (joints, bones, colors) containing:
          * the coordinates of all keypoints of all detected entities;
          * the coordinates of all skeleton edges of all detected entities;
          * the colors in which the edges should be plotted.
      """
    keypoints_all = []
    keypoint_edges_all = []
    colors = []
    num_instances, _, _, _ = keypoints.shape
    for idx in range(num_instances):
        kpts_x = keypoints[0, idx, :, 1]
        kpts_y = keypoints[0, idx, :, 0]
        kpts_scores = keypoints[0, idx, :, 2]
        kpts_absolute_xy = np.stack(
            [width * np.array(kpts_x), height * np.array(kpts_y)], axis=-1)
        kpts_above_thresh_absolute = kpts_absolute_xy[
                                     kpts_scores > threshold, :]
        keypoints_all.append(kpts_above_thresh_absolute)

        for edge_pair, color in COLOR_DICT.items():
            if (kpts_scores[edge_pair[0]] > threshold and
                    kpts_scores[edge_pair[1]] > threshold):
                x_start = kpts_absolute_xy[edge_pair[0], 0]
                y_start = kpts_absolute_xy[edge_pair[0], 1]
                x_end = kpts_absolute_xy[edge_pair[1], 0]
                y_end = kpts_absolute_xy[edge_pair[1], 1]
                line_seg = np.array([[x_start, y_start], [x_end, y_end]])
                keypoint_edges_all.append(line_seg)
                colors.append(color)
    if keypoints_all:
        joints = np.concatenate(keypoints_all, axis=0)
    else:
        joints = np.zeros((0, 17, 2))

    if keypoint_edges_all:
        bones = np.stack(keypoint_edges_all, axis=0)
    else:
        bones = np.zeros((0, 2, 2))
    return joints, bones, colors


def draw_bones(pixmap: PIL.Image, keypoints):
    draw = ImageDraw.Draw(pixmap)
    joints, bones, colors = process_keypoints(keypoints, pixmap.height, pixmap.width)

    for bone, color in zip(bones.tolist(), colors):
        draw.line((*bone[0], *bone[1]), fill=color, width=4)

    radio = 3

    for c_x, c_y in joints:
        shape = [(c_x - radio, c_y - radio), (c_x + radio, c_y + radio)]
        draw.ellipse(shape, fill="red", outline="red")


def movenet(image):
    """Runs detection on an input image.

        Args:
          image: A [1, height, width, 3] tensor represents the input image
            pixels. Note that the height/width should already be resized and match the
            expected input resolution of the model before passing into this function.

        Returns:
          A [1, 1, 17, 3] float numpy array representing the predicted keypoint
          coordinates and scores.
    """
    model_path = snapshot_download("leonelhs/movenet")
    module = tf.saved_model.load(model_path)
    model = module.signatures['serving_default']
    # SavedModel format expects tensor type of int32.
    image = tf.cast(image, dtype=tf.int32)
    # Run model inference.
    outputs = model(image)
    # Output is a [1, 1, 17, 3] tensor.
    return outputs['output_0'].numpy()