import cv2 import math import mediapipe as mp from mediapipe.python.solutions.pose_connections import POSE_CONNECTIONS class MediapipeAlgorithmPlugin: detect_corner = None config = (True, 1, True, False, True, 0.7, 0.7) mp_pose = mp.solutions.pose pose = None @classmethod def class_init(cls): MediapipeAlgorithmPlugin.pose = mp.solutions.pose.Pose( *MediapipeAlgorithmPlugin.config ) def __init__(self): super().__init__() self.mpDraw = mp.solutions.drawing_utils self.lm_list = [] if MediapipeAlgorithmPlugin.pose is None: MediapipeAlgorithmPlugin.class_init() @classmethod def get_config(cls): return MediapipeAlgorithmPlugin.config @classmethod def set_config(cls, config): if config != MediapipeAlgorithmPlugin.config: # 强制修改为静态识别 new_config = list(config) new_config[0] = True MediapipeAlgorithmPlugin.config = new_config MediapipeAlgorithmPlugin.pose = MediapipeAlgorithmPlugin.mp_pose.Pose(*new_config) @classmethod def get_corner(cls): return MediapipeAlgorithmPlugin.detect_corner @classmethod def set_corner(cls, corner): if corner != MediapipeAlgorithmPlugin.detect_corner: MediapipeAlgorithmPlugin.detect_corner = corner def set_result(self, lm_list): self.lm_list = lm_list @classmethod def find_pose(cls, img): dc = MediapipeAlgorithmPlugin.detect_corner if dc: img_cropped = img[dc[0]:dc[1], dc[2]:dc[3]] else: img_cropped = img img_rgb = cv2.cvtColor(img_cropped, cv2.COLOR_BGR2RGB) results = MediapipeAlgorithmPlugin.pose.process(img_rgb) return results, img_cropped def find_pose_with_drawing(self, img, draw=True): img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) results = MediapipeAlgorithmPlugin.pose.process(img_rgb) if results.pose_landmarks: if draw: self.mpDraw.draw_landmarks( img, results.pose_landmarks, POSE_CONNECTIONS ) return img def drawPoint(self, img, p1, p2, p3, bias_x=0, bias_y=0): if len(self.lm_list) == 0: return x1, y1 = self.lm_list[p1][1], self.lm_list[p1][2] x2, y2 = self.lm_list[p2][1], self.lm_list[p2][2] x3, y3 = self.lm_list[p3][1], self.lm_list[p3][2] # print((x1, y1), (x2, y2), (x3, y3)) cv2.circle(img, (x1 + bias_x, y1 + bias_y), 5, (102, 106, 233), cv2.FILLED) cv2.circle(img, (x2 + bias_x, y2 + bias_y), 5, (50, 80, 4), cv2.FILLED) cv2.circle(img, (x3 + bias_x, y3 + bias_y), 5, (20, 95, 104), cv2.FILLED) def drawPoint_more(self, img, queue, bias_x=0, bias_y=0): if len(self.lm_list) == 0: return color_list = [(0, 255, 255), (255, 255, 0), (255, 0, 255), (0.0, 255), (0, 255, 0), (255, 255, 255)] xy_list = [] for i in queue: x, y = self.lm_list[i][1], self.lm_list[i][2] xy_list.append([x, y]) for i in range(len(xy_list)): x = xy_list[i][0] y = xy_list[i][1] c = i % 6 cv2.circle(img, (x + bias_x, y + bias_y), 5, color_list[c], cv2.FILLED) def Screen_occupation_left(self): if self.lm_list: ocp = self.lm_list[29][1] - self.lm_list[23][1] else: ocp = 0 return ocp def Screen_occupation_right(self): if self.lm_list: ocp = self.lm_list[30][1] - self.lm_list[24][1] else: ocp = 0 return ocp def findPosition(self, img, draw=True): if self.lm_list: if draw: for _id, cx, cy, cz, vis in self.lm_list: cv2.circle(img, (cx, cy), 5, (255, 0, 0), cv2.FILLED) return self.lm_list def findConcave(self, p1, p2, p3): x1, y1 = self.lm_list[p1][1], self.lm_list[p1][2] x2, y2 = self.lm_list[p2][1], self.lm_list[p2][2] x3, y3 = self.lm_list[p3][1], self.lm_list[p3][2] # Calculate Angle Concave_angle = math.degrees(math.atan2(y3 - y2, x3 - x2) - math.atan2(y1 - y2, x1 - x2)) if Concave_angle < 0: Concave_angle += 360 return Concave_angle def findAngle(self, img, p1, p2, p3, draw=True): # Get the landmarks x1, y1 = self.lm_list[p1][1], self.lm_list[p1][2] x2, y2 = self.lm_list[p2][1], self.lm_list[p2][2] x3, y3 = self.lm_list[p3][1], self.lm_list[p3][2] # Calculate Angle angle = math.degrees(math.atan2(y3 - y2, x3 - x2) - math.atan2(y1 - y2, x1 - x2)) if angle < 0: angle += 360 if angle > 180: angle = 360 - angle elif angle > 180: angle = 360 - angle # Draw if draw: cv2.line(img, (x1, y1), (x2, y2), (255, 255, 255), 2) cv2.line(img, (x3, y3), (x2, y2), (255, 255, 255), 2) cv2.circle(img, (x1, y1), 3, (0, 0, 255), cv2.FILLED) cv2.circle(img, (x1, y1), 10, (0, 0, 255), 2) cv2.circle(img, (x2, y2), 3, (0, 0, 255), cv2.FILLED) cv2.circle(img, (x2, y2), 10, (0, 0, 255), 2) cv2.circle(img, (x3, y3), 3, (0, 0, 255), cv2.FILLED) cv2.circle(img, (x3, y3), 10, (0, 0, 255), 2) cv2.putText(img, str(int(angle)), (x2 - 20, y2 + 20), cv2.FONT_HERSHEY_PLAIN, 1, (0, 0, 255), 2) return angle def findIncludedAngle(self, x1, y1, x2, y2): IncludedAngle = math.asin((y1 - y2) / math.hypot((y1 - y2), (x1 - x2))) * 180 / math.pi return IncludedAngle def findRange(self, x1, y1, x2, y2): Range = math.sqrt(math.pow(x1 - x2, 2) + math.pow(y1 - y2, 2)) return Range def draw_dashed_line(self, img, x1, y1, x2, y2, color, thickness=1, dash_length=15): dist_x = abs(x2 - x1) dist_y = abs(y2 - y1) # 斜率和截距 slope = float(y2 - y1) / (x2 - x1) if x2 != x1 else float('inf') if slope == 0: # 水平虚线 step_length = dash_length else: # 竖直或倾斜虚线 step_length = int(max(dist_x, dist_y) * (dash_length / max(dist_x, dist_y))) num = 0 for i in range(0, max(dist_x, dist_y), step_length): num += 1 if num % 2 == 0: # 计算当前线段的起点和终点 interm_x1 = x1 + i if x2 >= x1 else x1 - i interm_y1 = y1 + (i * slope) if x2 >= x1 else y1 - (i * slope) interm_x2 = x1 + (i + dash_length) if x2 >= x1 else x1 - (i + dash_length) interm_y2 = y1 + ((dash_length + i) * slope) if x2 >= x1 else y1 - ((dash_length + i) * slope) # 画出当前线段 cv2.line(img, (int(interm_x1), int(interm_y1)), (int(interm_x2), int(interm_y2)), color, thickness)