HFup/bodybuilding_pose_analyzer/src/pose_analyzer.py

import cv2
import mediapipe as mp
import numpy as np
from typing import List, Dict, Tuple

class PoseAnalyzer:
    # Add MediaPipe skeleton connections as a class variable
    MP_CONNECTIONS = [
        (11, 13), (13, 15), # Left arm
        (12, 14), (14, 16), # Right arm
        (11, 12),           # Shoulders
        (11, 23), (12, 24), # Torso sides
        (23, 24),           # Hips
        (23, 25), (25, 27), # Left leg
        (24, 26), (26, 28), # Right leg
        (27, 31), (28, 32), # Ankles to feet
        (15, 17), (16, 18), # Wrists to hands
        (15, 19), (16, 20), # Wrists to pinky
        (15, 21), (16, 22), # Wrists to index
        (15, 17), (17, 19), (19, 21), # Left hand
        (16, 18), (18, 20), (20, 22)  # Right hand
    ]
    def __init__(self):
        # Initialize MediaPipe Pose
        self.mp_pose = mp.solutions.pose
        self.pose = self.mp_pose.Pose(
            static_image_mode=False,
            model_complexity=2,  # Using the most accurate model
            min_detection_confidence=0.1,
            min_tracking_confidence=0.1
        )
        self.mp_drawing = mp.solutions.drawing_utils
        
        # Define key angles for bodybuilding poses
        self.key_angles = {
            'front_double_biceps': {
                'shoulder_angle': (90, 120),  # Expected angle range
                'elbow_angle': (80, 100),
                'wrist_angle': (0, 20)
            },
            'side_chest': {
                'shoulder_angle': (45, 75),
                'elbow_angle': (90, 110),
                'wrist_angle': (0, 20)
            },
            'back_double_biceps': {
                'shoulder_angle': (90, 120),
                'elbow_angle': (80, 100),
                'wrist_angle': (0, 20)
            }
        }

    def detect_pose(self, frame: np.ndarray, last_valid_landmarks=None) -> Tuple[np.ndarray, List[Dict]]:
        """
        Detect pose in the given frame and return the frame with pose landmarks drawn
        and the list of detected landmarks. If detection fails, reuse last valid landmarks if provided.
        """
        # Convert the BGR image to RGB
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        
        # Process the frame and detect pose
        results = self.pose.process(rgb_frame)
        
        # Draw the pose landmarks on the frame
        if results.pose_landmarks:
            # Draw all 33 keypoints as bright red, smaller circles, and show index
            for idx, landmark in enumerate(results.pose_landmarks.landmark):
                x = int(landmark.x * frame.shape[1])
                y = int(landmark.y * frame.shape[0])
                if landmark.visibility > 0.1:  # Lowered threshold from 0.3 to 0.1
                    cv2.circle(frame, (x, y), 3, (0, 0, 255), -1)
                    cv2.putText(frame, str(idx), (x+8, y-8), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1)
            # Draw skeleton lines
            # Convert landmarks to pixel coordinates for easier access
            landmark_points = []
            for landmark in results.pose_landmarks.landmark:
                landmark_points.append((int(landmark.x * frame.shape[1]), int(landmark.y * frame.shape[0]), landmark.visibility))
            for pt1, pt2 in self.MP_CONNECTIONS:
                if pt1 < len(landmark_points) and pt2 < len(landmark_points):
                    x1, y1, v1 = landmark_points[pt1]
                    x2, y2, v2 = landmark_points[pt2]
                    if v1 > 0.1 and v2 > 0.1:
                        cv2.line(frame, (x1, y1), (x2, y2), (0, 255, 255), 2)
            # Convert landmarks to a list of dictionaries
            landmarks = []
            for idx, landmark in enumerate(results.pose_landmarks.landmark):
                landmarks.append({
                    'x': landmark.x,
                    'y': landmark.y,
                    'z': landmark.z,
                    'visibility': landmark.visibility
                })
            return frame, landmarks
        # If detection fails, reuse last valid landmarks if provided
        if last_valid_landmarks is not None:
            return frame, last_valid_landmarks
        return frame, []

    def calculate_angle(self, landmarks: List[Dict], joint1: int, joint2: int, joint3: int) -> float:
        """
        Calculate the angle between three joints.
        """
        if len(landmarks) < max(joint1, joint2, joint3):
            return None
            
        # Get the coordinates of the three joints
        p1 = np.array([landmarks[joint1]['x'], landmarks[joint1]['y']])
        p2 = np.array([landmarks[joint2]['x'], landmarks[joint2]['y']])
        p3 = np.array([landmarks[joint3]['x'], landmarks[joint3]['y']])
        
        # Calculate the angle
        v1 = p1 - p2
        v2 = p3 - p2
        
        angle = np.degrees(np.arccos(
            np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))
        ))
        
        return angle

    def analyze_pose(self, landmarks: List[Dict], pose_type: str) -> Dict:
        """
        Analyze the pose and provide feedback based on the pose type.
        Enhanced: Calculates angles for both left and right arms (shoulder, elbow, wrist) for all pose types.
        """
        if not landmarks or pose_type not in self.key_angles:
            return {'error': 'Invalid pose type or no landmarks detected'}
        
        feedback = {
            'pose_type': pose_type,
            'angles': {},
            'corrections': []
        }
        # Indices for MediaPipe 33 keypoints
        LEFT_SHOULDER = 11
        RIGHT_SHOULDER = 12
        LEFT_ELBOW = 13
        RIGHT_ELBOW = 14
        LEFT_WRIST = 15
        RIGHT_WRIST = 16
        LEFT_HIP = 23
        RIGHT_HIP = 24
        LEFT_KNEE = 25
        RIGHT_KNEE = 26
        LEFT_ANKLE = 27
        RIGHT_ANKLE = 28
        # Calculate angles for both arms
        # Shoulder angles (hip-shoulder-elbow)
        l_shoulder_angle = self.calculate_angle(landmarks, LEFT_HIP, LEFT_SHOULDER, LEFT_ELBOW)
        r_shoulder_angle = self.calculate_angle(landmarks, RIGHT_HIP, RIGHT_SHOULDER, RIGHT_ELBOW)
        # Elbow angles (shoulder-elbow-wrist)
        l_elbow_angle = self.calculate_angle(landmarks, LEFT_SHOULDER, LEFT_ELBOW, LEFT_WRIST)
        r_elbow_angle = self.calculate_angle(landmarks, RIGHT_SHOULDER, RIGHT_ELBOW, RIGHT_WRIST)
        # Wrist angles (elbow-wrist-hand index, if available)
        # MediaPipe does not have hand index, so we can use a pseudo point (e.g., extend wrist direction)
        # For now, skip wrist angle or set to None
        # Leg angles (optional)
        l_knee_angle = self.calculate_angle(landmarks, LEFT_HIP, LEFT_KNEE, LEFT_ANKLE)
        r_knee_angle = self.calculate_angle(landmarks, RIGHT_HIP, RIGHT_KNEE, RIGHT_ANKLE)
        # Add angles to feedback
        if l_shoulder_angle:
            feedback['angles']['L Shoulder'] = l_shoulder_angle
            if not self.key_angles[pose_type]['shoulder_angle'][0] <= l_shoulder_angle <= self.key_angles[pose_type]['shoulder_angle'][1]:
                feedback['corrections'].append(
                    f"Adjust L Shoulder to {self.key_angles[pose_type]['shoulder_angle'][0]}-{self.key_angles[pose_type]['shoulder_angle'][1]} deg"
                )
        if r_shoulder_angle:
            feedback['angles']['R Shoulder'] = r_shoulder_angle
            if not self.key_angles[pose_type]['shoulder_angle'][0] <= r_shoulder_angle <= self.key_angles[pose_type]['shoulder_angle'][1]:
                feedback['corrections'].append(
                    f"Adjust R Shoulder to {self.key_angles[pose_type]['shoulder_angle'][0]}-{self.key_angles[pose_type]['shoulder_angle'][1]} deg"
                )
        if l_elbow_angle:
            feedback['angles']['L Elbow'] = l_elbow_angle
            if not self.key_angles[pose_type]['elbow_angle'][0] <= l_elbow_angle <= self.key_angles[pose_type]['elbow_angle'][1]:
                feedback['corrections'].append(
                    f"Adjust L Elbow to {self.key_angles[pose_type]['elbow_angle'][0]}-{self.key_angles[pose_type]['elbow_angle'][1]} deg"
                )
        if r_elbow_angle:
            feedback['angles']['R Elbow'] = r_elbow_angle
            if not self.key_angles[pose_type]['elbow_angle'][0] <= r_elbow_angle <= self.key_angles[pose_type]['elbow_angle'][1]:
                feedback['corrections'].append(
                    f"Adjust R Elbow to {self.key_angles[pose_type]['elbow_angle'][0]}-{self.key_angles[pose_type]['elbow_angle'][1]} deg"
                )
        # Optionally add knee angles
        if l_knee_angle:
            feedback['angles']['L Knee'] = l_knee_angle
        if r_knee_angle:
            feedback['angles']['R Knee'] = r_knee_angle
        return feedback

    def process_frame(self, frame: np.ndarray, pose_type: str = 'front_double_biceps', last_valid_landmarks=None) -> Tuple[np.ndarray, Dict, List[Dict]]:
        """
        Process a single frame, detect pose, and analyze it. Returns frame, analysis, and used landmarks.
        """
        # Detect pose
        frame_with_pose, landmarks = self.detect_pose(frame, last_valid_landmarks=last_valid_landmarks)
        # Analyze pose if landmarks are detected
        analysis = self.analyze_pose(landmarks, pose_type) if landmarks else {'error': 'No pose detected'}
        return frame_with_pose, analysis, landmarks