LISHUZUOXUN_yangjiang/MCamera/mp_algorithm.py

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)