单人的姿态检测|tensorflow singlepose

import tensorflow as tf

import tensorflow_hub as hub

import cv2

from matplotlib import pyplot as plt

import numpy as np

def movenet(input_image):

    """Runs detection on an input image.

    Args:

      input_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.

    """

    # Download the model from TF Hub.

    model = hub.load("https://tfhub.dev/google/movenet/singlepose/lightning/4")

    model = model.signatures['serving_default']

    # SavedModel format expects tensor type of int32.

    input_image = tf.cast(input_image, dtype=tf.int32)

    # Run model inference.

    outputs = model(input_image)

    # Output is a [1, 1, 17, 3] tensor.

    keypoints_with_scores = outputs['output_0'].numpy()

    keypoints_with_scores = keypoints_with_scores.reshape((1, 17, 3))

    return keypoints_with_scores

def draw_keypoints(frame, keypoints, confidence_threshold):

    y, x, c = frame.shape

    shaped = np.squeeze(np.multiply(keypoints, [y, x, 1]))

    print("shaped in draw_keypoints:", shaped)

    for kp in shaped:

        ky, kx, kp_conf = kp

        if kp_conf > confidence_threshold:

            cv2.circle(frame, (int(kx), int(ky)), 6, (0, 255, 0), -1)

EDGES = {

    (0, 1): 'm',

    (0, 2): 'c',

    (1, 3): 'm',

    (2, 4): 'c',

    (0, 5): 'm',

    (0, 6): 'c',

    (5, 7): 'm',

    (7, 9): 'm',

    (6, 8): 'c',

    (8, 10): 'c',

    (5, 6): 'y',

    (5, 11): 'm',

    (6, 12): 'c',

    (11, 12): 'y',

    (11, 13): 'm',

    (13, 15): 'm',

    (12, 14): 'c',

    (14, 16): 'c'

}

def draw_connections(frame, keypoints, edges, confidence_threshold):

    print('frame', frame)

    y, x, c = frame.shape

    shaped = np.squeeze(np.multiply(keypoints, [y, x, 1]))

    for edge, color in edges.items():

        p1, p2 = edge

        y1, x1, c1 = shaped[p1]

        y2, x2, c2 = shaped[p2]

        if (c1 > confidence_threshold) & (c2 > confidence_threshold):

            cv2.line(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 0, 255), 4)

def loop_through_people(frame, keypoints_with_scores, edges, confidence_threshold):

    for person in keypoints_with_scores:

        draw_connections(frame, person, edges, confidence_threshold)

        draw_keypoints(frame, person, confidence_threshold)

image_path =外汇跟单gendan5.com 'C:/Users/Harry/Desktop/fitness.jpeg'

image = tf.io.read_file(image_path)

# image = tf.compat.v1.image.decode_image(image)

image = tf.compat.v1.image.decode_jpeg(image)

# Resize and pad the image to keep the aspect ratio and fit the expected size.

input_size = 192

input_image = tf.expand_dims(image, axis=0)

input_image = tf.image.resize_with_pad(input_image, input_size, input_size)

# Run model inference.

keypoints_with_scores = movenet(input_image)

display_image = tf.cast(tf.image.resize_with_pad(image, 1280, 1280), dtype = tf.int32)

display_image = np.array(display_image)

origin_image = np.copy(display_image)

loop_through_people(display_image, keypoints_with_scores, EDGES, 0.1)

plt.subplot(1, 2, 1)

plt.imshow(origin_image)

plt.subplot(1, 2, 2)

plt.imshow(display_image)

plt.show()