AI十行代码系列2.人体关键点定位-MediaPipe Python

Posted 2022-03-15 朱铭德

tags:

篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了AI十行代码系列2.人体关键点定位-MediaPipe Python相关的知识，希望对你有一定的参考价值。

【TLAIP系列简介】Ten-Lines-AI-Projects

现阶段有非常多优秀的开源的AI工程，为了有更好的可扩展性，接口往往设计得十分复杂，这需要投入一定的时间和精力来处理，对于不熟悉或者刚入门的小伙伴，这可能需要花很久。本系列设计的出发点也很纯粹，进一步降低小伙伴们挑选和使用轮子的时间，让更多小伙伴能更快地验证算法效果，算法满足效果了再扒开看看，不满意直接看下一个。有任何问题和建议欢迎随时评论，目前系列里的十行代码会按照下列方式进行：

import  关键py
1.定义好输入
2.网络初始化配置
3.开始跑并输出结果
4.可视化结果

包含空格和一些必要的说明，估计正好十行。哈哈哈哈（不能再少了，再少我的强迫症就不能忍受了）。

如果遇到什么很牛逼的AI工程或者自己不想跑的但是一眼看上去效果很好的工程，欢迎留言交流哈，有空可以一并整理到这个工程，一键可运行，无繁杂使用过程。

【正文】

本文的主角依旧是谷歌大名鼎鼎的MediaPipe，人体姿态近些年的落地场景也是非常多，同样，这个仓库入门也非常合适。惯例先放结果（完整仓库代码和测试视频、图片请直接拖到最后）：

十行关键代码

正好十行，一行都不多，如果看过上一篇手势，可以发现基本一样，哈哈，谷歌封装得真不错。

from utils.pose_mediapipe import InputData, InitPoseTracker, ShowResult

# 初始化输入源, file支持数字（相机）以及视频文件路径，图片路径或文件夹路径
input_data = InputData("test/Star-Lord.mp4")  # input_data = InputData("test/imgs")
# 初始化人体姿态追踪tracker
pose_track = InitPoseTracker(use_static_mode=input_data.use_img_list, up_body_only=False)
# 获取图像以及结果的generator
run_pose_result = pose_track.run_pose_tracking(input_data.get_next_img())
# 显示结果， ESC退出，图片模式按任意键继续
ShowResult(input_data.wait_key, pose_track.up_body_only).show_result(run_pose_result)

InputData用来处理各种输入，目前兼容下列各种情况的输入（是不是很良心）：

系统相机，输入相机编号即可，例如 InputData(file=0)
一个视频文件，输入视频路径即可，例如 InputData(file="test/hand_tracking.mp4")
一张图片，输入图片路径即可，例如 InputData(file="test/imgs/0.jpg")
一个包含很多图像的目录，输入目录路径即可，例如 InputData(file="test/imgs")

InitPoseTracker用来初始化网络的参数，主要参数如下：

use_static_mode 图片模型还是视频模式，如果是图片，会先运行检测算法，否则直接追踪
up_body_only 用来控制是半身模式还是全身模式
detect_conf 检测网络的置信度（默认0.5就行，一般不用改）
track_conf 追踪网络的置信度（默认0.5就行，一般不用改）

ShowResult用来可视化结果

waitkey_mode 默认1，会刷新显示视频；输入0的话会阻塞图像显示，按下任意键继续
up_body_only 和pose tracker里的模式对应，画关键点的方式有些区别

环境配置和使用说明

配置conda（可选）

conda create -n mediapipe python=3.6
conda activate mediapipe

安装必要的依赖包

注意这里使用的mediapipe版本，不同版本的mediapipe接口不太一样。

pip install opencv-python
pip install mediapipe==0.8.3

配置好环境后，直接运行 “运行这个.py" 即可

功能代码剖析

Mediepipe封装地非常好了，和上一篇手部追踪的代码基本一致（可以直接跳到最后下载源码及测试视频），之前讲过的一些代码就不重复阐述了，主要讲下关键的一些配置。

网络初始化模块InitPoseTracker

其他配置都和手部追踪一致，主要关注up_body_only参数，用来控制半身模式和全身模式，如果设置为True，只会显示上半身，下面这张图很直观了。

可以参照下这块的代码：

class InitPoseTracker:
    def __init__(self, use_static_mode=False, detect_conf=0.5, track_conf=0.5, up_body_only=False):
        self.use_static_mode = use_static_mode
        self.detect_conf = detect_conf
        self.track_conf = track_conf
        self.up_body_only = up_body_only
        self.pose = None
        self.init_network()

    def init_network(self):
        self.pose = mp.solutions.pose.Pose(
            static_image_mode=self.use_static_mode,
            upper_body_only=self.up_body_only,
            min_detection_confidence=self.detect_conf,
            min_tracking_confidence=self.track_conf
        )
        print("init hand tracking down")

    def run_pose_tracking(self, get_nex_img):
        while True:
            img_origin = next(get_nex_img)
            if img_origin is None:
                yield [None, None]
                break
            img = cv2.cvtColor(img_origin, cv2.COLOR_BGR2RGB)
            yield [img_origin, self.pose.process(img)]