#!/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()