reab5555's picture
Update pose_analysis.py
e8aec56 verified
raw
history blame
6.44 kB
import math
import cv2
import mediapipe as mp
mp_pose = mp.solutions.pose
mp_drawing = mp.solutions.drawing_utils
pose = mp_pose.Pose(static_image_mode=False, min_detection_confidence=0.8, min_tracking_confidence=0.8)
def calculate_posture_score(frame):
image_height, image_width, _ = frame.shape
results = pose.process(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
if not results.pose_landmarks:
return None, None
landmarks = results.pose_landmarks.landmark
# Use only body landmarks
left_shoulder = landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value]
right_shoulder = landmarks[mp_pose.PoseLandmark.RIGHT_SHOULDER.value]
left_hip = landmarks[mp_pose.PoseLandmark.LEFT_HIP.value]
right_hip = landmarks[mp_pose.PoseLandmark.RIGHT_HIP.value]
left_knee = landmarks[mp_pose.PoseLandmark.LEFT_KNEE.value]
right_knee = landmarks[mp_pose.PoseLandmark.RIGHT_KNEE.value]
# Calculate angles
shoulder_angle = abs(math.degrees(math.atan2(right_shoulder.y - left_shoulder.y, right_shoulder.x - left_shoulder.x)))
hip_angle = abs(math.degrees(math.atan2(right_hip.y - left_hip.y, right_hip.x - left_hip.x)))
knee_angle = abs(math.degrees(math.atan2(right_knee.y - left_knee.y, right_knee.x - left_knee.x)))
# Calculate vertical alignment
shoulder_hip_alignment = abs((left_shoulder.y + right_shoulder.y) / 2 - (left_hip.y + right_hip.y) / 2)
hip_knee_alignment = abs((left_hip.y + right_hip.y) / 2 - (left_knee.y + right_knee.y) / 2)
# Add head landmarks
nose = landmarks[mp_pose.PoseLandmark.NOSE.value]
left_ear = landmarks[mp_pose.PoseLandmark.LEFT_EAR.value]
right_ear = landmarks[mp_pose.PoseLandmark.RIGHT_EAR.value]
# Calculate head tilt
head_tilt = abs(math.degrees(math.atan2(right_ear.y - left_ear.y, right_ear.x - left_ear.x)))
# Calculate head position relative to shoulders
head_position = abs((nose.y - (left_shoulder.y + right_shoulder.y) / 2) /
((left_shoulder.y + right_shoulder.y) / 2 - (left_hip.y + right_hip.y) / 2))
# Combine metrics into a single posture score (you may need to adjust these weights)
posture_score = (
(1 - abs(shoulder_angle - hip_angle) / 90) * 0.3 +
(1 - abs(hip_angle - knee_angle) / 90) * 0.2 +
(1 - shoulder_hip_alignment) * 0.1 +
(1 - hip_knee_alignment) * 0.1 +
(1 - abs(head_tilt - 90) / 90) * 0.15 +
(1 - head_position) * 0.15
)
return posture_score, results.pose_landmarks
def draw_pose_landmarks(frame, landmarks):
annotated_frame = frame.copy()
# Include relevant landmarks for head position and body
body_landmarks = [
mp_pose.PoseLandmark.NOSE,
mp_pose.PoseLandmark.LEFT_SHOULDER,
mp_pose.PoseLandmark.RIGHT_SHOULDER,
mp_pose.PoseLandmark.LEFT_EAR,
mp_pose.PoseLandmark.RIGHT_EAR,
mp_pose.PoseLandmark.LEFT_ELBOW,
mp_pose.PoseLandmark.RIGHT_ELBOW,
mp_pose.PoseLandmark.LEFT_WRIST,
mp_pose.PoseLandmark.RIGHT_WRIST,
mp_pose.PoseLandmark.LEFT_HIP,
mp_pose.PoseLandmark.RIGHT_HIP,
mp_pose.PoseLandmark.LEFT_KNEE,
mp_pose.PoseLandmark.RIGHT_KNEE,
mp_pose.PoseLandmark.LEFT_ANKLE,
mp_pose.PoseLandmark.RIGHT_ANKLE
]
# Connections for head position and body
body_connections = [
(mp_pose.PoseLandmark.LEFT_EAR, mp_pose.PoseLandmark.LEFT_SHOULDER),
(mp_pose.PoseLandmark.RIGHT_EAR, mp_pose.PoseLandmark.RIGHT_SHOULDER),
(mp_pose.PoseLandmark.NOSE, mp_pose.PoseLandmark.LEFT_SHOULDER),
(mp_pose.PoseLandmark.NOSE, mp_pose.PoseLandmark.RIGHT_SHOULDER),
(mp_pose.PoseLandmark.LEFT_SHOULDER, mp_pose.PoseLandmark.RIGHT_SHOULDER),
(mp_pose.PoseLandmark.LEFT_SHOULDER, mp_pose.PoseLandmark.LEFT_ELBOW),
(mp_pose.PoseLandmark.RIGHT_SHOULDER, mp_pose.PoseLandmark.RIGHT_ELBOW),
(mp_pose.PoseLandmark.LEFT_ELBOW, mp_pose.PoseLandmark.LEFT_WRIST),
(mp_pose.PoseLandmark.RIGHT_ELBOW, mp_pose.PoseLandmark.RIGHT_WRIST),
(mp_pose.PoseLandmark.LEFT_SHOULDER, mp_pose.PoseLandmark.LEFT_HIP),
(mp_pose.PoseLandmark.RIGHT_SHOULDER, mp_pose.PoseLandmark.RIGHT_HIP),
(mp_pose.PoseLandmark.LEFT_HIP, mp_pose.PoseLandmark.RIGHT_HIP),
(mp_pose.PoseLandmark.LEFT_HIP, mp_pose.PoseLandmark.LEFT_KNEE),
(mp_pose.PoseLandmark.RIGHT_HIP, mp_pose.PoseLandmark.RIGHT_KNEE),
(mp_pose.PoseLandmark.LEFT_KNEE, mp_pose.PoseLandmark.LEFT_ANKLE),
(mp_pose.PoseLandmark.RIGHT_KNEE, mp_pose.PoseLandmark.RIGHT_ANKLE)
]
# Draw landmarks
for landmark in body_landmarks:
if landmark in landmarks.landmark:
lm = landmarks.landmark[landmark]
h, w, _ = annotated_frame.shape
cx, cy = int(lm.x * w), int(lm.y * h)
cv2.circle(annotated_frame, (cx, cy), 5, (245, 117, 66), -1)
# Draw connections
for connection in body_connections:
start_lm = landmarks.landmark[connection[0]]
end_lm = landmarks.landmark[connection[1]]
h, w, _ = annotated_frame.shape
start_point = (int(start_lm.x * w), int(start_lm.y * h))
end_point = (int(end_lm.x * w), int(end_lm.y * h))
cv2.line(annotated_frame, start_point, end_point, (245, 66, 230), 2)
# Highlight head tilt
left_ear = landmarks.landmark[mp_pose.PoseLandmark.LEFT_EAR]
right_ear = landmarks.landmark[mp_pose.PoseLandmark.RIGHT_EAR]
nose = landmarks.landmark[mp_pose.PoseLandmark.NOSE]
h, w, _ = annotated_frame.shape
left_ear_point = (int(left_ear.x * w), int(left_ear.y * h))
right_ear_point = (int(right_ear.x * w), int(right_ear.y * h))
nose_point = (int(nose.x * w), int(nose.y * h))
# Draw a line between ears to show head tilt
cv2.line(annotated_frame, left_ear_point, right_ear_point, (0, 255, 0), 2)
# Draw a line from nose to the midpoint between shoulders to show head forward/backward tilt
left_shoulder = landmarks.landmark[mp_pose.PoseLandmark.LEFT_SHOULDER]
right_shoulder = landmarks.landmark[mp_pose.PoseLandmark.RIGHT_SHOULDER]
shoulder_mid_x = (left_shoulder.x + right_shoulder.x) / 2
shoulder_mid_y = (left_shoulder.y + right_shoulder.y) / 2
shoulder_mid_point = (int(shoulder_mid_x * w), int(shoulder_mid_y * h))
cv2.line(annotated_frame, nose_point, shoulder_mid_point, (0, 255, 0), 2)
return annotated_frame