單人的姿態檢測|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()