Spaces:
Sleeping
Sleeping
File size: 4,290 Bytes
40518b9 | 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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | import cv2
import mediapipe as mp
from exercises.hammer_curl import HammerCurl
class PoseEstimator:
def __init__(self):
self.mp_pose = mp.solutions.pose
self.pose = self.mp_pose.Pose()
self.mp_drawing = mp.solutions.drawing_utils
def estimate_pose(self, frame, exercise_type):
# BGR to RGB
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
# Pose estimate
results = self.pose.process(rgb_frame)
# Draw landmarks and specific connections based on exercise type
if results.pose_landmarks:
# Draw specific landmarks and connections based on exercise_type
if exercise_type == "squat":
self.draw_squat_lines(frame, results.pose_landmarks.landmark)
elif exercise_type == "push_up":
self.draw_push_up_lines(frame, results.pose_landmarks.landmark)
elif exercise_type == "hammer_curl":
self.draw_hammerl_curl_lines(frame, results.pose_landmarks.landmark)
return results
def draw_hammerl_curl_lines(self, frame, landmarks):
shoulder_right = [int(landmarks[11].x * frame.shape[1]), int(landmarks[11].y * frame.shape[0])]
elbow_right = [int(landmarks[13].x * frame.shape[1]), int(landmarks[13].y * frame.shape[0])]
hip_right = [int(landmarks[23].x * frame.shape[1]), int(landmarks[23].y * frame.shape[0])]
wrist_right = [int(landmarks[15].x * frame.shape[1]), int(landmarks[15].y * frame.shape[0])]
# Left arm landmarks (shoulder, elbow, hip, wrist)
shoulder_left = [int(landmarks[12].x * frame.shape[1]), int(landmarks[12].y * frame.shape[0])]
elbow_left = [int(landmarks[14].x * frame.shape[1]), int(landmarks[14].y * frame.shape[0])]
hip_left = [int(landmarks[24].x * frame.shape[1]), int(landmarks[24].y * frame.shape[0])]
wrist_left = [int(landmarks[16].x * frame.shape[1]), int(landmarks[16].y * frame.shape[0])]
# Draw lines with improved style
cv2.line(frame, shoulder_left, elbow_left, (0, 0, 255), 4,2)
cv2.line(frame, elbow_left, wrist_left, (0, 0, 255), 4,2)
cv2.line(frame, shoulder_right, elbow_right, (0, 0, 255), 4,2)
cv2.line(frame, elbow_right, wrist_right, (0, 0, 255), 4,2)
def draw_squat_lines(self, frame, landmarks):
# Squat specific lines (hip, knee, shoulder)
hip = [int(landmarks[23].x * frame.shape[1]), int(landmarks[23].y * frame.shape[0])]
knee = [int(landmarks[25].x * frame.shape[1]), int(landmarks[25].y * frame.shape[0])]
shoulder = [int(landmarks[11].x * frame.shape[1]), int(landmarks[11].y * frame.shape[0])]
hip_right = [int(landmarks[24].x * frame.shape[1]), int(landmarks[24].y * frame.shape[0])]
knee_right = [int(landmarks[26].x * frame.shape[1]), int(landmarks[26].y * frame.shape[0])]
shoulder_right = [int(landmarks[12].x * frame.shape[1]), int(landmarks[12].y * frame.shape[0])]
# Draw lines for squat
cv2.line(frame, shoulder, hip, (178, 102, 255), 2)
cv2.line(frame, hip, knee, (178, 102, 255), 2)
cv2.line(frame, shoulder_right, hip_right, (51, 153, 255), 2)
cv2.line(frame, hip_right, knee_right, (51, 153, 255), 2)
def draw_push_up_lines(self, frame, landmarks):
# Push-up specific lines (shoulder, elbow, wrist)
shoulder_left = [int(landmarks[11].x * frame.shape[1]), int(landmarks[11].y * frame.shape[0])]
elbow_left = [int(landmarks[13].x * frame.shape[1]), int(landmarks[13].y * frame.shape[0])]
wrist_left = [int(landmarks[15].x * frame.shape[1]), int(landmarks[15].y * frame.shape[0])]
shoulder_right = [int(landmarks[12].x * frame.shape[1]), int(landmarks[12].y * frame.shape[0])]
elbow_right = [int(landmarks[14].x * frame.shape[1]), int(landmarks[14].y * frame.shape[0])]
wrist_right = [int(landmarks[16].x * frame.shape[1]), int(landmarks[16].y * frame.shape[0])]
# Draw lines for push-up
cv2.line(frame, shoulder_left, elbow_left, (0, 0, 255), 2)
cv2.line(frame, elbow_left, wrist_left, (0, 0, 255), 2)
cv2.line(frame, shoulder_right, elbow_right, (102, 0, 0), 2)
cv2.line(frame, elbow_right, wrist_right, (102, 0, 0), 2)
|