实时多人姿态检测| Tensorflow real time multipose detection

系列文章

单人的姿态检测-图片

多人的姿态检测-图片

利用OpenCV对读取视频每一帧或保存视频

有了对图片的多人和单人姿态检测基础，我们就可以利用opencv的相关知识实时的检测视频中多人的姿态了。

在视频中每一帧画人的每个姿态点的方法和模型的获取方法都和之前一样，你可以参考我的系列网站。不同的地方是下面代码。大概的步骤是：1.打开摄像头，2.读每一帧，3.在每一帧里画姿态点，4.显示在视频上面。

cap = cv2.VideoCapture(1)
while cap.isOpened():
    ret, frame = cap.read()
    # Resize image
    img = frame.copy()
    # print(img.shape)
    img = tf.image.resize_with_pad(tf.expand_dims(img, axis=0), 480, 640)
    keypoints_with_scores = multipose_model(img)

    # Render keypoints
    loop_through_people(frame, keypoints_with_scores, EDGES, 0.1)  # 0.1 is our confidence_threshold
    print(frame.shape)
    cv2.imshow('Video', frame)

    if cv2.waitKey(10) & 0xFF == ord('q'):
        break
cap.release()
cv2.destroyAllWindows()

如果要对保存视频，代码如下：

cap = cv2.VideoCapture(1)

# 设置需要保存视频的格式“xvid”
# 该参数是MPEG-4编码类型，文件名后缀为.avi
fourcc = cv2.VideoWriter_fourcc(*'XVID')
# 设置视频帧频
fps = cap.get(cv2.CAP_PROP_FPS)
# 设置视频大小
size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
# VideoWriter方法是cv2库提供的保存视频方法
# 按照设置的格式来out输出
out = cv2.VideoWriter('C:/Users/hl77319/Desktop/out.avi',fourcc ,fps, size)

while cap.isOpened():
    ret, frame = cap.read()
    # Resize image
    img = frame.copy()
    # # print(img.shape)
    img = tf.image.resize_with_pad(tf.expand_dims(img, axis=0), 480, 640)
    keypoints_with_scores = multipose_model(img)
    #
    # # Render keypoints
    loop_through_people(frame, keypoints_with_scores, EDGES, 0.1)  # 0.1 is our confidence_threshold
    if ret == True:
        # 垂直翻转矩阵
        # frame = cv2.flip(frame, 0)
        out.write(frame)
        cv2.imshow('frame', frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    else:
        break

cap.release()
out.release()
cv2.destroyAllWindows()

完整代码

#pip install google-colab -i https://pypi.tuna.tsinghua.edu.cn/simple

import tensorflow as tf
import tensorflow_hub as hub
import cv2
from matplotlib import pyplot as plt
import numpy as np


gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

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_keypoints(frame, keypoints, confidence_threshold):
    y, x, c = frame.shape
    shaped = np.squeeze(np.multiply(keypoints, [y, x, 1]))

    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)


def draw_connections(frame, keypoints, edges, confidence_threshold):
    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)

#model = hub.load('https://tfhub.dev/google/movenet/multipose/lightning/1')
model = hub.load('D:/ML/Model_weight/multipose')
movenet = model.signatures['serving_default']
def multipose_model(img):
        input_img = tf.cast(img, dtype=tf.int32)
        # Detection section
        results = movenet(input_img)
        keypoints_with_scores = results['output_0'].numpy()[:, :, :51].reshape((6, 17, 3))
        return keypoints_with_scores


cap = cv2.VideoCapture(0)
while cap.isOpened():
        ret, frame = cap.read()
        # Resize image
        img = frame.copy()
        print(img.shape)
        img = tf.image.resize_with_pad(tf.expand_dims(img, axis=0), 480, 640)
        keypoints_with_scores =  multipose_model(img)

        # Render keypoints
        loop_through_people(frame, keypoints_with_scores, EDGES, 0.1)  # 0.1 is our confidence_threshold

        # cv2_imshow(frame)
        cv2.imshow('Video',frame)

        if cv2.waitKey(10) & 0xFF==ord('q'):
            break
cap.release()
cv2.destroyAllWindows()