114 lines
4.5 KiB
Python
114 lines
4.5 KiB
Python
|
import cv2
|
|||
|
|
import mediapipe as mp
|
|||
|
|
import math
|
|||
|
|
|
|||
|
|
|
|||
|
|
class PoseDetector:
|
|||
|
|
"""
|
|||
|
|
人体姿势检测类
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self,
|
|||
|
|
static_image_mode=False,
|
|||
|
|
upper_body_only=False,
|
|||
|
|
smooth_landmarks=True,
|
|||
|
|
# enable_segmentation=False,
|
|||
|
|
min_detection_confidence=0.5,
|
|||
|
|
min_tracking_confidence=0.5):
|
|||
|
|
"""
|
|||
|
|
初始化
|
|||
|
|
:param static_image_mode: 是否是静态图片，默认为否
|
|||
|
|
:param upper_body_only: 是否是上半身，默认为否
|
|||
|
|
:param smooth_landmarks: 设置为True减少抖动
|
|||
|
|
:param min_detection_confidence:人员检测模型的最小置信度值，默认为0.5
|
|||
|
|
:param min_tracking_confidence:姿势可信标记的最小置信度值，默认为0.5
|
|||
|
|
"""
|
|||
|
|
self.static_image_mode = static_image_mode
|
|||
|
|
self.upper_body_only = upper_body_only
|
|||
|
|
self.smooth_landmarks = smooth_landmarks
|
|||
|
|
self.min_detection_confidence = min_detection_confidence
|
|||
|
|
self.min_tracking_confidence = min_tracking_confidence
|
|||
|
|
# 创建一个Pose对象用于检测人体姿势
|
|||
|
|
self.pose = mp.solutions.pose.Pose(self.static_image_mode, self.upper_body_only, self.smooth_landmarks, False,
|
|||
|
|
self.min_detection_confidence, self.min_tracking_confidence)
|
|||
|
|
|
|||
|
|
def find_pose(self, img, draw=True):
|
|||
|
|
"""
|
|||
|
|
检测姿势方法
|
|||
|
|
:param img: 一帧图像
|
|||
|
|
:param draw: 是否画出人体姿势节点和连接图
|
|||
|
|
:return: 处理过的图像
|
|||
|
|
"""
|
|||
|
|
imgRGB = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|||
|
|
# pose.process(imgRGB) 会识别这帧图片中的人体姿势数据，保存到self.results中
|
|||
|
|
self.results = self.pose.process(imgRGB)
|
|||
|
|
if self.results.pose_landmarks:
|
|||
|
|
if draw:
|
|||
|
|
mp.solutions.drawing_utils.draw_landmarks(img, self.results.pose_landmarks,
|
|||
|
|
mp.solutions.pose.POSE_CONNECTIONS)
|
|||
|
|
return img
|
|||
|
|
|
|||
|
|
def find_positions(self, img):
|
|||
|
|
"""
|
|||
|
|
获取人体姿势数据
|
|||
|
|
:param img: 一帧图像
|
|||
|
|
:param draw: 是否画出人体姿势节点和连接图
|
|||
|
|
:return: 人体姿势数据列表
|
|||
|
|
"""
|
|||
|
|
# 人体姿势数据列表，每个成员由3个数字组成：id, x, y
|
|||
|
|
# id代表人体的某个关节点，x和y代表坐标位置数据
|
|||
|
|
self.lmslist = []
|
|||
|
|
if self.results.pose_landmarks:
|
|||
|
|
for id, lm in enumerate(self.results.pose_landmarks.landmark):
|
|||
|
|
h, w, c = img.shape
|
|||
|
|
cx, cy = int(lm.x * w), int(lm.y * h)
|
|||
|
|
self.lmslist.append([id, cx, cy])
|
|||
|
|
|
|||
|
|
return self.lmslist
|
|||
|
|
|
|||
|
|
def find_updown(self, img):
|
|||
|
|
|
|||
|
|
img = self.find_pose(img, draw=True)
|
|||
|
|
# 获取人体姿势列表数据
|
|||
|
|
lmslist = self.find_positions(img)
|
|||
|
|
lfthand = lmslist[16]
|
|||
|
|
lfthandheight = lfthand[2]
|
|||
|
|
rhthand = lmslist[15]
|
|||
|
|
rhthandheight = rhthand[2]
|
|||
|
|
mouthlft = lmslist[10]
|
|||
|
|
mouthlftheight = mouthlft[2]
|
|||
|
|
mouthrht = lmslist[9]
|
|||
|
|
mouthrhtheight = mouthrht[2]
|
|||
|
|
return lfthandheight, rhthandheight, mouthlftheight, mouthrhtheight
|
|||
|
|
|
|||
|
|
def find_angle(self, img, p1, p2, p3, draw=True):
|
|||
|
|
'''
|
|||
|
|
获取人体姿势中3个点p1-p2-p3的角度
|
|||
|
|
:param img: 一帧图像
|
|||
|
|
:param p1: 第1个点
|
|||
|
|
:param p2: 第2个点
|
|||
|
|
:param p3: 第3个点
|
|||
|
|
:param draw: 是否画出3个点的连接图
|
|||
|
|
:return: 角度
|
|||
|
|
'''
|
|||
|
|
x1, y1 = self.lmslist[p1][1], self.lmslist[p1][2]
|
|||
|
|
x2, y2 = self.lmslist[p2][1], self.lmslist[p2][2]
|
|||
|
|
x3, y3 = self.lmslist[p3][1], self.lmslist[p3][2]
|
|||
|
|
|
|||
|
|
# 使用三角函数公式获取3个点p1-p2-p3，以p2为角的角度值，0-180度之间
|
|||
|
|
angle = int(math.degrees(math.atan2(y1 - y2, x1 - x2) - math.atan2(y3 - y2, x3 - x2)))
|
|||
|
|
if angle < 0:
|
|||
|
|
angle = angle + 360
|
|||
|
|
if angle > 180:
|
|||
|
|
angle = 360 - angle
|
|||
|
|
|
|||
|
|
if draw:
|
|||
|
|
cv2.circle(img, (x1, y1), 10, (0, 255, 255), cv2.FILLED)
|
|||
|
|
cv2.circle(img, (x2, y2), 20, (255, 0, 255), cv2.FILLED)
|
|||
|
|
cv2.circle(img, (x3, y3), 10, (0, 255, 255), cv2.FILLED)
|
|||
|
|
cv2.line(img, (x1, y1), (x2, y2), (255, 255, 255, 3))
|
|||
|
|
cv2.line(img, (x2, y2), (x3, y3), (255, 255, 255, 3))
|
|||
|
|
cv2.putText(img, str(angle), (x2 - 50, y2 + 50), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 255, 255), 2)
|
|||
|
|
|
|||
|
|
return angle
|