AI十行代码系列4.全身关键点定位-MediaPipe Python

Posted 朱铭德

tags:

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

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

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

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

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

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


【正文】

        本文的主角依旧是谷歌大名鼎鼎的MediaPipe,基于人体相关的这是最后一篇了,相当于官方帮忙把所有人体相关的都整合到一个仓库了(很良心了)。惯例先放结果(完整仓库代码和测试视频、图片请直接拖到最后):

十行代码

依旧是熟悉的配方,三大块。

from utils.holistic_mediapipe import InputData, InitHolisticTracker, ShowResult

# 初始化输入源, file支持数字(相机)以及视频文件路径,图片路径或文件夹路径
input_data = InputData("test/hand_tracking.mp4")  # input_data = InputData("test/imgs")
# 初始化全身关键点追踪tracker
pose_track = InitHolisticTracker(use_static_mode=input_data.use_img_list, up_body_only=True)
# 获取图像以及结果的generator
run_pose_result = pose_track.run_face_tracking(input_data.get_next_img())
# 显示结果, ESC退出,图片模式按任意键继续
ShowResult(input_data.wait_key, up_body_only= pose_track.up_body_only).show_result(run_pose_result)

下面配置和第二篇pose的完全一致,这里不予赘述。

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" 即可 

功能代码剖析

        功能模块和pose的基本没有区别,干脆都粘上来给大家看看了(推荐直接去github下载完整的代码和测试视频)。功能代码:

import os

import cv2
import mediapipe as mp


class InputData:
    def __init__(self, file=0):
        self.cap = None
        self.img_list = []
        self.img_id = 0
        self.img_type_list = 'jpg', 'bmp', 'png', 'jpeg', 'rgb', 'tif', 'webp'
        self.deal_with_input(file)
        self.use_img_list = len(self.img_list) > 0
        self.wait_key = 0 if self.use_img_list else 1
        self.use_static_mode = self.use_img_list

    def gen_img_list(self, path):
        for item in os.listdir(path):
            if item.split(".")[-1] in self.img_type_list:
                self.img_list.append(os.path.join(path, item))

    def deal_with_input(self, file):
        path_valid = False
        if isinstance(file, int):
            # use camera
            self.cap = cv2.VideoCapture(file)
            path_valid = True
        elif isinstance(file, str):
            if os.path.isdir(file):
                # use img list
                self.gen_img_list(file)
                if len(self.img_list) > 0:
                    path_valid = True
                else:
                    print("no images in", file)
            elif os.path.isfile(file):
                if file.split(".")[-1] in self.img_type_list:
                    # only one image
                    self.img_list.append(file)
                    path_valid = True
                else:
                    self.cap = cv2.VideoCapture(file)
                    if self.cap.isOpened():
                        print("video path is", file)
                        path_valid = True
                    else:
                        print("video path is not valid, path is:", file)
        if not path_valid:
            print("Invalid input! Use camera 0 instead!")
            self.cap = cv2.VideoCapture(0)

    def get_next_img(self):
        if self.use_img_list:
            while self.img_id < len(self.img_list):
                img = cv2.imread(self.img_list[self.img_id])
                self.img_id += 1
                if img is not None:
                    yield img
            yield None
        else:
            while True:
                _, img = self.cap.read()
                if img is None:
                    yield None
                    break
                yield img

    def release_cap(self):
        if self.cap:
            self.cap.release()


class InitHolisticTracker:
    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.holistic = None
        self.init_network()

    def init_network(self):
        self.holistic = mp.solutions.holistic.Holistic(
            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
        )

    def run_face_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.holistic.process(img)]


class ShowResult:
    def __init__(self, waitkey_mode=1, up_body_only=False):
        self.waitkey = waitkey_mode
        self.up_body_only = up_body_only

    def show_result(self, run_hand_tracking):
        landmark_face = mp.solutions.drawing_utils.DrawingSpec(thickness=2, circle_radius=1, color=(255, 0, 120))
        connect_face = mp.solutions.drawing_utils.DrawingSpec(thickness=2, circle_radius=2, color=(0, 120, 255))
        landmark_hand_l = mp.solutions.drawing_utils.DrawingSpec(thickness=3, circle_radius=5, color=(122, 122, 0))
        landmark_hand_r = mp.solutions.drawing_utils.DrawingSpec(thickness=3, circle_radius=5, color=(0, 122, 122))
        connect_hand_l = mp.solutions.drawing_utils.DrawingSpec(thickness=8, circle_radius=10, color=(0, 0, 200))
        connect_hand_r = mp.solutions.drawing_utils.DrawingSpec(thickness=8, circle_radius=10, color=(200, 0, 0))
        landmark_pose = mp.solutions.drawing_utils.DrawingSpec(thickness=3, circle_radius=5, color=(0, 122, 250))
        connect_pose = mp.solutions.drawing_utils.DrawingSpec(thickness=3, circle_radius=5, color=(255, 122, 0))

        while True:
            img, results = next(run_hand_tracking)
            body_connections = mp.solutions.holistic.UPPER_BODY_POSE_CONNECTIONS \\
                if self.up_body_only else mp.solutions.holistic.POSE_CONNECTIONS
            if img is None:
                break
            mp.solutions.drawing_utils.draw_landmarks(
                img,
                results.face_landmarks,
                mp.solutions.holistic.FACE_CONNECTIONS,
                connection_drawing_spec=connect_face,
                landmark_drawing_spec=landmark_face
            )
            mp.solutions.drawing_utils.draw_landmarks(
                img,
                results.pose_landmarks,
                body_connections,
                connection_drawing_spec=connect_pose,
                landmark_drawing_spec=landmark_pose
            )
            mp.solutions.drawing_utils.draw_landmarks(
                img,
                results.left_hand_landmarks,
                mp.solutions.hands.HAND_CONNECTIONS,
                connection_drawing_spec=connect_hand_l,
                landmark_drawing_spec=landmark_hand_l)
            mp.solutions.drawing_utils.draw_landmarks(
                img,
                results.right_hand_landmarks,
                mp.solutions.hands.HAND_CONNECTIONS,
                connection_drawing_spec=connect_hand_r,
                landmark_drawing_spec=landmark_hand_r)
            cv2.imshow('MediaPipe-Pose', img)
            if cv2.waitKey(self.waitkey) & 0xFF == 27:
                break

实际运行结果及简单点评

        国际惯例,请上老演员hand_tracking.mp4,效果还不错。

 拓展阅读

如果看了效果想进一步尝试,可以仔细研究下面两个链接

写在最后

 往期TLAIP列表:

【AI十行代码系列】1.手势关键点追踪-MediaPipe Python_朱铭德的博客-CSDN博客

【AI十行代码系列】2.人体关键点定位-MediaPipe Python_朱铭德的博客-CSDN博客

【AI十行代码系列】3.人脸关键点定位-MediaPipe Python_朱铭德的博客-CSDN博客

         AI十行代码仓库地址最重要的→链接 ,后续功能也会在这里更新,欢迎提建议或需求,哈哈,本文相关所有代码和资源也可以在这里下载

        如果文章对你稍微有些帮助的话,麻烦来个三连(疯狂暗示)。

        2021的最后一天了,祝大家明年诸事顺遂~

以上是关于AI十行代码系列4.全身关键点定位-MediaPipe Python的主要内容,如果未能解决你的问题,请参考以下文章

AI十行代码系列4.全身关键点定位-MediaPipe Python

AI十行代码系列4.全身关键点定位-MediaPipe Python

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

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

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

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