app.py

#!/usr/bin/env python

from __future__ import annotations

import pathlib

import gradio as gr
import mediapipe as mp
import numpy as np
import cv2
mp_drawing = mp.solutions.drawing_utils
mp_drawing_styles = mp.solutions.drawing_styles
mp_pose = mp.solutions.pose

TITLE = "MediaPipe Human Pose Estimation"

# Function to calculate the angle between three points
def calculate_angle(a, b, c):
    a = np.array([a.x, a.y])  # First point
    b = np.array([b.x, b.y])  # Mid point
    c = np.array([c.x, c.y])  # End point
    
    radians = np.arctan2(c[1] - b[1], c[0] - b[0]) - np.arctan2(a[1] - b[1], a[0] - b[0])
    angle = np.abs(radians * 180.0 / np.pi)
    
    if angle > 180.0:
        angle = 360 - angle
        
    return angle

# Define a function to classify yoga poses
def classify_pose(landmarks, output_image):
    label = 'Unknown Pose'

    # Calculate the required angles
    left_elbow_angle = calculate_angle(
        landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value],
        landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value],
        landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value])

    right_elbow_angle = calculate_angle(
        landmarks[mp_pose.PoseLandmark.RIGHT_SHOULDER.value],
        landmarks[mp_pose.PoseLandmark.RIGHT_ELBOW.value],
        landmarks[mp_pose.PoseLandmark.RIGHT_WRIST.value])

    left_shoulder_angle = calculate_angle(
        landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value],
        landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value],
        landmarks[mp_pose.PoseLandmark.LEFT_HIP.value])

    right_shoulder_angle = calculate_angle(
        landmarks[mp_pose.PoseLandmark.RIGHT_HIP.value],
        landmarks[mp_pose.PoseLandmark.RIGHT_SHOULDER.value],
        landmarks[mp_pose.PoseLandmark.RIGHT_ELBOW.value])

    left_knee_angle = calculate_angle(
        landmarks[mp_pose.PoseLandmark.LEFT_HIP.value],
        landmarks[mp_pose.PoseLandmark.LEFT_KNEE.value],
        landmarks[mp_pose.PoseLandmark.LEFT_ANKLE.value])

    right_knee_angle = calculate_angle(
        landmarks[mp_pose.PoseLandmark.RIGHT_HIP.value],
        landmarks[mp_pose.PoseLandmark.RIGHT_KNEE.value],
        landmarks[mp_pose.PoseLandmark.RIGHT_ANKLE.value])

    # Check for Five-Pointed Star Pose
    if abs(landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value][1] - landmarks[mp_pose.PoseLandmark.LEFT_HIP.value][1]) < 100 and \
       abs(landmarks[mp_pose.PoseLandmark.RIGHT_WRIST.value][1] - landmarks[mp_pose.PoseLandmark.RIGHT_HIP.value][1]) < 100 and \
       abs(landmarks[mp_pose.PoseLandmark.LEFT_ANKLE.value][0] - landmarks[mp_pose.PoseLandmark.RIGHT_ANKLE.value][0]) > 200 and \
       abs(landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value][0] - landmarks[mp_pose.PoseLandmark.RIGHT_WRIST.value][0]) > 200:
        label = "Five-Pointed Star Pose"  
    
    # Check for Warrior II pose
    if 165 < left_elbow_angle < 195 and 165 < right_elbow_angle < 195 and \
       80 < left_shoulder_angle < 110 and 80 < right_shoulder_angle < 110:
        if (165 < left_knee_angle < 195 or 165 < right_knee_angle < 195) and \
           (90 < left_knee_angle < 120 or 90 < right_knee_angle < 120):
            label = 'Warrior II Pose'

    # Check for T pose
    if 165 < left_elbow_angle < 195 and 165 < right_elbow_angle < 195 and \
       80 < left_shoulder_angle < 110 and 80 < right_shoulder_angle < 110 and \
       160 < left_knee_angle < 195 and 160 < right_knee_angle < 195:
        label = 'T Pose'

    # Check for Tree Pose
    if (165 < left_knee_angle < 195 or 165 < right_knee_angle < 195) and \
       (315 < left_knee_angle < 335 or 25 < right_knee_angle < 45):
        label = 'Tree Pose'
    
    # Check for Upward Salute Pose
    if abs(landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value][0] - landmarks[mp_pose.PoseLandmark.LEFT_HIP.value][0]) < 100 and \
       abs(landmarks[mp_pose.PoseLandmark.RIGHT_WRIST.value][0] - landmarks[mp_pose.PoseLandmark.RIGHT_HIP.value][0]) < 100 and \
       landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value][1] < landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value][1] and \
       landmarks[mp_pose.PoseLandmark.RIGHT_WRIST.value][1] < landmarks[mp_pose.PoseLandmark.RIGHT_SHOULDER.value][1] and \
       abs(landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value][1] - landmarks[mp_pose.PoseLandmark.RIGHT_SHOULDER.value][1]) < 50:
        label = "Upward Salute Pose"

   # Check for Hands Under Feet Pose
    if landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value][1] > landmarks[mp_pose.PoseLandmark.LEFT_KNEE.value][1] and \
       landmarks[mp_pose.PoseLandmark.RIGHT_WRIST.value][1] > landmarks[mp_pose.PoseLandmark.RIGHT_KNEE.value][1] and \
       abs(landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value][0] - landmarks[mp_pose.PoseLandmark.LEFT_ANKLE.value][0]) < 50 and \
       abs(landmarks[mp_pose.PoseLandmark.RIGHT_WRIST.value][0] - landmarks[mp_pose.PoseLandmark.RIGHT_ANKLE.value][0]) < 50:
        label = "Hands Under Feet Pose"      
        

    # Check for Plank Pose
    if 160 < left_shoulder_angle < 200 and 160 < right_shoulder_angle < 200 and \
       160 < left_knee_angle < 200 and 160 < right_knee_angle < 200:
        label = "Plank Pose"

    # Write the label on the output image
    cv2.putText(output_image, label, (10, 30), cv2.FONT_HERSHEY_PLAIN, 2, (0, 255, 0), 2)
    
    return output_image, label

    
def run(
    image: np.ndarray,
    model_complexity: int,
    enable_segmentation: bool,
    min_detection_confidence: float,
    background_color: str,
) -> np.ndarray:
    with mp_pose.Pose(
        static_image_mode=True,
        model_complexity=model_complexity,
        enable_segmentation=enable_segmentation,
        min_detection_confidence=min_detection_confidence,
    ) as pose:
        results = pose.process(image)

    res = image[:, :, ::-1].copy()
    if enable_segmentation:
        if background_color == "white":
            bg_color = 255
        elif background_color == "black":
            bg_color = 0
        elif background_color == "green":
            bg_color = (0, 255, 0)  # type: ignore
        else:
            raise ValueError

        if results.segmentation_mask is not None:
            res[results.segmentation_mask <= 0.1] = bg_color
        else:
            res[:] = bg_color

    mp_drawing.draw_landmarks(
        res,
        results.pose_landmarks,
        mp_pose.POSE_CONNECTIONS,
        landmark_drawing_spec=mp_drawing_styles.get_default_pose_landmarks_style(),
    )

    return res[:, :, ::-1]


model_complexities = list(range(3))
background_colors = ["white", "black", "green"]

image_paths = sorted(pathlib.Path("images").rglob("*.jpg"))
examples = [[path, model_complexities[1], True, 0.5, background_colors[0]] for path in image_paths]

demo = gr.Interface(
    fn=run,
    inputs=[
        gr.Image(label="Input", type="numpy"),
        gr.Radio(label="Model Complexity", choices=model_complexities, type="index", value=model_complexities[1]),
        gr.Checkbox(label="Enable Segmentation", value=True),
        gr.Slider(label="Minimum Detection Confidence", minimum=0, maximum=1, step=0.05, value=0.5),
        gr.Radio(label="Background Color", choices=background_colors, type="value", value=background_colors[0]),
    ],
    outputs=gr.Image(label="Output"),
    examples=examples,
    title=TITLE,
)

if __name__ == "__main__":
    demo.queue().launch()