| import torch |
| import torch.nn as nn |
| import mediapipe as mp |
| import cv2 |
|
|
| class PostureAnalysisModel(nn.Module): |
| def __init__(self): |
| super(PostureAnalysisModel, self).__init__() |
| self.mp_pose = mp.solutions.pose |
| self.pose = self.mp_pose.Pose() |
| self.lmPose = self.mp_pose.PoseLandmark |
|
|
| def findDistance(self, x1, y1, x2, y2): |
| dist = torch.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2) |
| return dist |
|
|
| def findAngle(self, x1, y1, x2, y2): |
| theta = torch.acos((y2 - y1) * (-y1) / (torch.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2) * y1)) |
| degree = (180 / torch.pi) * theta |
| return degree |
|
|
| def forward(self, image): |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) |
| keypoints = self.pose.process(image) |
| image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) |
| lm = keypoints.pose_landmarks |
|
|
| if lm is None: |
| return None |
|
|
| |
| |
| l_shldr_x = int(lm.landmark[self.lmPose.LEFT_SHOULDER].x * image.shape[1]) |
| l_shldr_y = int(lm.landmark[self.lmPose.LEFT_SHOULDER].y * image.shape[0]) |
|
|
| |
| r_shldr_x = int(lm.landmark[self.lmPose.RIGHT_SHOULDER].x * image.shape[1]) |
| r_shldr_y = int(lm.landmark[self.lmPose.RIGHT_SHOULDER].y * image.shape[0]) |
|
|
| |
| l_ear_x = int(lm.landmark[self.lmPose.LEFT_EAR].x * image.shape[1]) |
| l_ear_y = int(lm.landmark[self.lmPose.LEFT_EAR].y * image.shape[0]) |
|
|
| |
| l_hip_x = int(lm.landmark[self.lmPose.LEFT_HIP].x * image.shape[1]) |
| l_hip_y = int(lm.landmark[self.lmPose.LEFT_HIP].y * image.shape[0]) |
|
|
| |
| offset = self.findDistance(l_shldr_x, l_shldr_y, r_shldr_x, r_shldr_y) |
|
|
| |
| neck_inclination = self.findAngle(l_shldr_x, l_shldr_y, l_ear_x, l_ear_y) |
| torso_inclination = self.findAngle(l_hip_x, l_hip_y, l_shldr_x, l_shldr_y) |
|
|
| return offset, neck_inclination, torso_inclination |
