基于 pytorch-openpose 实现 “多目标” 人体姿态估计

前言

还记得上次通过 MediaPipe 估计人体姿态关键点驱动 3D 角色模型，虽然节省了动作 K 帧时间，但是网上还有一种似乎更方便的方法。MagicAnimate 就是其一，说是只要提供一张人物图片和一段动作视频 (舞蹈武术等)，就可以完成图片人物转视频。

于是我就去官网体验了一下，发现动作的视频长度不能超过 5 秒，当然，如果说要整长视频可以切多段处理再合成解决。主要的还是视频需要那种背景相对较纯的，不然提交表单一直报错，还有他也不能处理画面内多人物的姿态估计。

多目标人体姿态估计

为什么我要弄多目标，其实是我有次拿了一舞团的视频用 MediaPipe 检测，发现一个画面中只能采集到一个人的动作数据。虽然齐舞可能就一套动作，其他的角色模型可以复制粘贴，但是有些编舞为了好看，伴舞也会根据节奏作不同的变化。所以说对于我用来采集舞蹈数据，这个很重要了，当然他也可以用在多人互动的 AR 游戏，或用在同时培训多人的动作规范检测等等场景。

要从单一人体检测到多人体姿态估计，开始我是打算用 YOLO 对画面中的多 Person 区块读出来，然后再将这些方块遍历交给 MediaPipe 对指定区域作人物动作节点识别。但是最后发现有现成的算法，就是 pytorch-openpose，所以果断先用这个来体验了一下。

pytorch-openpose 简介

PyTorch-OpenPose 是一个基于 PyTorch 的开源库，它实现了 OpenPose 的功能，可以进行人的面部表情、躯干和四肢甚至手指的跟踪。它不仅适用于单人也适用于多人，同时具有较好的鲁棒性。要运行 PyTorch-OpenPose，需要安装支持 CUDA 的 PyTorch，以下例子有使用作者提供的预训练模型，通过拆分视频帧，绘制多人物动作线条保存图片，最后将图片合成为视频。

环境

• scikit-image
• opencv-python
• scipy
• matplotlib
• numpy

编码

帧拆分绘制

import cv2
import matplotlib.pyplot as plt
import copy
import numpy as np
import torch
from src import model
from src import util
from src.body import Body
from src.hand import Hand

body_estimation = Body('model/body_pose_model.pth')
hand_estimation = Hand('model/hand_pose_model.pth')

print(f"Torch device: {torch.cuda.get_device_name()}")

cap = cv2.VideoCapture("D:/3code/6pytorch/opencv_demo/12_open_pose/11.mp4")
cap.set(3, 640)
cap.set(4, 480)

indices = 1
while True:
    ret, oriImg = cap.read()
    if not ret:
        break

    candidate, subset = body_estimation(oriImg)

    canvas = copy.deepcopy(oriImg)
    canvas = util.draw_bodypose(canvas, candidate, subset)

    # detect hand
    hands_list = util.handDetect(candidate, subset, oriImg)

    all_hand_peaks = []
    for x, y, w, is_left in hands_list:
        peaks = hand_estimation(oriImg[y:y+w, x:x+w, :])
        peaks[:, 0] = np.where(peaks[:, 0]==0, peaks[:, 0], peaks[:, 0]+x)
        peaks[:, 1] = np.where(peaks[:, 1]==0, peaks[:, 1], peaks[:, 1]+y)
        all_hand_peaks.append(peaks)

    canvas = util.draw_handpose(canvas, all_hand_peaks)

    cv2.imwrite('image_out/img_{}.jpg'.format(indices), canvas)
    indices += 1

    print("images:", indices)

    # cv2.imshow('demo', canvas)#一个窗口用以显示原视频
    # if cv2.waitKey(1) & 0xFF == ord('q'):
    #     break

cap.release()
cv2.destroyAllWindows()

视频合成

from pathlib import Path
import cv2
import os

# 将视频video_path分割成图片和音频文件，保存到save_path文件夹中
def video2mp3_img(video_path, save_path, audio_path):
    def video_split(video_path, save_path):
        if not os.path.exists(save_path):
            os.makedirs(save_path)
        cap = cv2.VideoCapture(video_path)
        i = 0
        while True:
            ret, frame = cap.read()
            if ret:
                cv2.imwrite(save_path + '/' + str(i) + '.jpg', frame)
                i += 1
            else:
                break
        cap.release()

    if not os.path.exists(save_path):
        os.makedirs(save_path)

    # 视频分割
    video_split(video_path, save_path)

    # 视频转音频
    # os.system("ffmpeg -i {} -vn -acodec copy {}/audio.mp3".format(video_path, audio_path))
    os.system("ffmpeg -i {} -q:a 0 -map a {}/audio.mp3".format(video_path, audio_path))
    # 音频转wav
    # os.system("ffmpeg -i {}/audio.mp3 {}/audio.wav".format(save_path, save_path))

# 将video_imgout文件夹中的图片合成视频并且添加音频文件video_img/audio.mp3
def img2mp4(image_out, save_name):

    BASE_PATH = os.path.dirname(__file__)
    # 读取img size
    img = cv2.imread("{}/img_1.jpg".format(image_out))

    imgInfo = img.shape
    size = (imgInfo[1], imgInfo[0])

    files = []
    for dirpath, dirnames, filenames in os.walk(image_out):
        for filename in filenames:
            fileName = Path(os.path.join(dirpath, filename))
            files.append(os.path.join(dirpath, filename))

    files = [file.replace('\', '/') for file in files]
    files.sort(key=lambda x: int(x.split('/')[-1].split('.')[0].split('_')[-1]))

    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    videoWrite = cv2.VideoWriter(f'videos/{save_name}.mp4', fourcc, 25, size)  # 写入对象 1 file name  3: 视频帧率

    for i in files:
        img = cv2.imread(str(i))
        videoWrite.write(img)

    print(f'videos/{save_name}.mp4')

    # 将video_img中的音频文件添加到视频中
    # os.system("ffmpeg -i {}/videos/{}.mp4 -c:v copy -c:a aac -strict experimental {}/videos/{}.mp4".format(BASE_PATH, save_name, BASE_PATH, save_name))

if __name__ == '__main__':
    BASE = os.path.dirname(__file__)
    video_path = os.path.join(BASE, "videos/yangguo.mp4")  # 视频路径
    save_path = os.path.join(BASE, "video_img")            # 拆解视频保存路径
    audio_path = os.path.join(BASE, "audio")               # 分离音频保存路径

    # 视频  ==> imgs
    # video2mp3_img(video_path, save_path, audio_path)

    # # imgs ==> 视频
    img2mp4("image_out", save_name='ldh')