import cv2 as cv
# MediaPipe是一个用于构建机器学习管道的框架,用于处理视频、音频等时间序列数据。这个跨平台框架适用于桌面/服务器、Android、iOS和嵌入式设备,如Raspberry Pi和Jetson Nano。
import mediapipe as mp
import time
import math

class poseDetector():
 def __init__(self,static_image_mode=False, # 静态图模式,False代表置信度高时继续跟踪,True代表实时跟踪检测新的结果
 upper_body_only=False, # 是否只检测上半身
 model_complexity=1,
 smooth_landmarks=True, # 平滑,一般为True
 min_detection_confidence=0.5, # 检测置信度# 检测置信度大于0.5代表检测到了,若此时跟踪置信度大于0.5就继续跟踪,小于就沿用上一次,避免一次又一次重复使用模型
 min_tracking_confidence=0.5): # 跟踪置信度
 self.static_image_mode = static_image_mode
 self.upper_body_only = upper_body_only
 self.model_complexity=model_complexity
 self.smooth_landmarks = smooth_landmarks
 self.min_detection_confidence = min_detection_confidence
 self.min_tracking_confidence = min_tracking_confidence

 # mediapipe.solutions.drawing_utils.draw_landmarks()绘制关键点的连线
 self.mpDraw = mp.solutions.drawing_utils
 # mediapipe.solutions.pose.Pose()姿态关键点检测函数
 self.mpPose = mp.solutions.pose
 self.pose = self.mpPose.Pose(self.static_image_mode,self.upper_body_only,self.model_complexity,self.smooth_landmarks,self.min_detection_confidence,self.min_tracking_confidence)

 # 检测关键点方法 
 def findPose(self, frame, draw=True):
 self.frame_RGB = cv.cvtColor(frame, cv.COLOR_BGR2RGB)
 # 3.将图像传给姿态识别模型
 self.res = self.pose.process(self.frame_RGB)
 if self.res.pose_landmarks:
 if draw:
 # 绘制姿态坐标点,img为画板,传入姿态点坐标,坐标连线
 # mediapipe.solutions.drawing_utils.draw_landmarks()绘制手部关键点的连线
 # mpDraw.draw_landmarks(frame,res.pose_landmarks)
 self.mpDraw.draw_landmarks(frame, self.res.pose_landmarks, self.mpPose.POSE_CONNECTIONS)
 return frame

 
 #关键点信息,是否绘制关键点连线 
 def findPosition(self, frame, draw=True):
 self.lmList = []
 if self.res.pose_landmarks:
 for id,lm in enumerate(self.res.pose_landmarks.landmark):
 h, w, c = frame.shape
 cx, cy = int(lm.x * w), int(lm.y * h)
 self.lmList.append([id,cx,cy])
 if draw:
 cv.circle(frame, (cx, cy), 10, (255, 0, 0), cv.FILLED)
 return self.lmList
 
 #计算运动时关键点角度 
 def findAngle(self,frame,p1,p2,p3,draw=True):
 # 关键点坐标
 x1,y1 = self.lmList[p1][1:]
 x2,y2 = self.lmList[p2][1:]
 x3,y3 = self.lmList[p3][1:]

 # 关键点角度 atan2 方法返回一个 -pi 到 pi 之间的数值,表示点 (x, y) 对应的偏移角度。这是一个逆时针角度,以弧度为单位,正X轴和点 (x, y) 与原点连线 之间
 angle = math.degrees(math.atan2(y3-y2,x3-x2)-math.atan2(y1-y2,x1-x2))

 # 绘制关键点
 if draw:
 cv.line(frame,(x1,y1),(x2,y2),(255,0,255),3)
 cv.line(frame,(x3,y3),(x2,y2),(255,0,255),3)
 cv.circle(frame,(x1,y1),10,(0,0,255),cv.FILLED)
 cv.circle(frame,(x1,y1),15,(0,0,255),2)
 cv.circle(frame,(x2,y2),10,(0,0,255),cv.FILLED)
 cv.circle(frame,(x2,y2),15,(0,0,255),2)
 cv.circle(frame,(x3,y3),10,(0,0,255),cv.FILLED)
 cv.circle(frame,(x3,y3),15,(0,0,255),2)
 return angle