File size: 2,103 Bytes
096b664 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
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
# Acquire the landmark coordinates.
# Left shoulder.
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])
# Right shoulder
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])
# Left ear.
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])
# Left hip.
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])
# Calculate distance between left shoulder and right shoulder points.
offset = self.findDistance(l_shldr_x, l_shldr_y, r_shldr_x, r_shldr_y)
# Calculate angles.
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 |