以下代码示例展示了如何将姿势检测代码模块化,并将其封装为一个名为 'pose_detection.py' 的独立 Python 模块,方便在其他项目中使用。

pose_detection.py

import cv2
import mediapipe as mp
import numpy as np
from Calculate_angle import calculate_angle

mp_pose = mp.solutions.pose

def detect_pose(frame):
    # 设置 mediapipe 实例
    with mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5) as pose:
        # 将图像重新着色为 RGB
        image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        image.flags.writeable = False
    
        # 进行检测
        results = pose.process(image)
    
        # 将图像重新着色回 BGR
        image.flags.writeable = True
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        
        stage = ''
        # 提取地标
        try:
            landmarks = results.pose_landmarks.landmark
            
            # 获取坐标
            left_eye = [landmarks[mp_pose.PoseLandmark.LEFT_EYE.value].x,landmarks[mp_pose.PoseLandmark.LEFT_EYE.value].y]
            left_hip= [landmarks[mp_pose.PoseLandmark.LEFT_HIP.value].x,landmarks[mp_pose.PoseLandmark.LEFT_HIP.value].y]
            left_heel = [landmarks[mp_pose.PoseLandmark.LEFT_HEEL.value].x,landmarks[mp_pose.PoseLandmark.LEFT_HEEL.value].y]
            right_eye = [landmarks[mp_pose.PoseLandmark.RIGHT_EYE.value].x,landmarks[mp_pose.PoseLandmark.RIGHT_EYE.value].y]
            right_hip = [landmarks[mp_pose.PoseLandmark.RIGHT_HIP.value].x,landmarks[mp_pose.PoseLandmark.RIGHT_HIP.value].y]
            right_heel = [landmarks[mp_pose.PoseLandmark.RIGHT_HEEL.value].x,landmarks[mp_pose.PoseLandmark.RIGHT_HEEL.value].y]
            right_index = [landmarks[mp_pose.PoseLandmark.RIGHT_INDEX.value].x,landmarks[mp_pose.PoseLandmark.RIGHT_INDEX.value].y]
            left_index = [landmarks[mp_pose.PoseLandmark.LEFT_INDEX.value].x,landmarks[mp_pose.PoseLandmark.LEFT_INDEX.value].y]
            # 计算角度
            
            angle1 = calculate_angle(left_eye, left_hip, left_heel)
            angle2 = calculate_angle(right_eye, right_hip, right_heel)

            # 可视化角度
            cv2.putText(image, str(angle1), 
                        tuple(np.multiply(left_hip, [640, 480]).astype(int)), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2, cv2.LINE_AA
                                    )
            cv2.putText(image, str(angle2), 
                        tuple(np.multiply(right_hip, [640, 480]).astype(int)), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2, cv2.LINE_AA
                                    )
            # print(f'left_eye={left_eye}')
            # print(f'right_eye={right_eye}')
            
            # print(f'left_hip={left_hip}')
            # print(f'right_hip={right_hip}')

            # print(f'left_heel={left_heel}')
            # print(f'right_heel={right_heel}')

            # print(f'right index={right_index}')
            # print(f'left index={left_index}')

            if ((left_eye[0]>=0.41 and left_eye[0]<=0.43) and (left_hip[0]>=0.44 and left_hip[0]<=0.46) and (left_heel[0]>=0.41 and left_heel[0]<=0.43) or (right_eye[0]>=0.41 and right_eye[0]<=0.43) and (right_hip[0]<=0.43 and right_hip[0]>=0.41) and (right_heel[0]>=0.37 and right_heel[0]<=0.39)):

                if ((left_eye[1]>=0.24 and left_eye[1]<=0.33) and (left_hip[1]<=0.35 and left_hip[1]>=0.45) and (left_heel[1]<=0.74 and left_heel[1]>=0.72) or (right_eye[1]<=0.30 and right_eye[1]>=0.24) and (right_hip[1]<=0.50 and right_hip[1]>=0.32) and (right_heel[1]>=0.71 and right_heel[0]<=0.73)):
                    stage = 'safe :)'
            # 卷曲计数逻辑
            else:
                if angle1 != angle2 and (angle1>170 and angle2>170):
                    if (((right_index[0]<0.70 and right_index[0]>0.20) and (right_index[1]<0.56 and right_index[1]>0.15)) or ((left_index[0]<0.55 and left_index[0]>0.18) and (left_index[1]<0.56 and left_index[1]>0.15))):
                        stage='Hanging on !!'
                    else:
                        stage = 'fallen :('    

                elif angle1 != angle2 and (angle1<140 or angle2<140) :
                    stage = 'Trying to Walk'
                elif angle1!=angle2 and ((angle1<168 and angle1>140) and (angle2<168 and angle2>140)):
                    stage='Barely Walking'
                else:
                    pass
            
        except:
            pass
        
        # 设置状态框
        cv2.rectangle(image, (0,0), (350,125), (245,117,16), -1)
        
        # 重复数据
        cv2.putText(image, 'Condition: ', (15,12), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,0), 1, cv2.LINE_AA)
        cv2.putText(image, str('Calculating Angles'), 
                    (100,12), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 1, cv2.LINE_AA)
        
        # 阶段数据
        cv2.putText(image, 'STAGE: ', (10,50), 
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,0), 1, cv2.LINE_AA)
        cv2.putText(image, stage, 
                    (10,100), 
                    cv2.FONT_HERSHEY_SIMPLEX, 1.5, (255,255,255), 2, cv2.LINE_AA)
        
        # 渲染检测结果
        mp_drawing.draw_landmarks(image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS,
                                    mp_drawing.DrawingSpec(color=(245,117,66), thickness=2, circle_radius=2),
                                    mp_drawing.DrawingSpec(color=(245,66,230), thickness=2, circle_radius=2)
                                    )

        return image

使用模块

现在,你可以像下面这样在你的主程序中导入并使用 'pose_detection.py' 模块:

import cv2
from pose_detection import detect_pose

cap = cv2.VideoCapture(0)

while True:
    ret, frame = cap.read()
    if ret:
        # 使用 detect_pose 函数检测姿势
        processed_frame = detect_pose(frame)
        cv2.imshow('Pose Detection', processed_frame)
    
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

优点

  1. 代码重用性: 你可以将 'pose_detection.py' 模块导入到其他需要姿势检测功能的 Python 程序中,避免重复编写相同的代码。
  2. 模块化: 模块化代码结构,使其更易于理解、维护和扩展。
  3. 提高效率: 将复杂的代码分离成更小的、更易于管理的模块可以提高代码效率。

小结

通过模块化姿势检测代码,我们可以更高效地组织代码结构,并将其轻松集成到其他项目中。这将有助于提高代码的可维护性和可扩展性。

姿势检测模块化:提高代码效率

原文地址: https://www.cveoy.top/t/topic/nZhR 著作权归作者所有。请勿转载和采集!

免费AI点我,无需注册和登录