Update app.py

app.py

@@ -1,22 +1,35 @@
-#!/usr/bin/env python
-
-from __future__ import annotations
-
-import pathlib
-import math
-
-import gradio as gr
 import cv2
 import mediapipe as mp
 import numpy as np
+import gradio as gr
 
-mp_drawing = mp.solutions.drawing_utils
-mp_drawing_styles = mp.solutions.drawing_styles
+# Initialize MediaPipe Pose class.
 mp_pose = mp.solutions.pose
+pose = mp_pose.Pose(static_image_mode=True, min_detection_confidence=0.5)
 
-TITLE = "MediaPipe Human Pose Estimation"
-DESCRIPTION = "https://google.github.io/mediapipe/"
+# Initialize drawing class
+mp_drawing = mp.solutions.drawing_utils
 
+def calculateAngle(landmark1, landmark2, landmark3):
+    '''
+    This function calculates the angle between three landmarks.
+    Args:
+        landmark1: The first landmark containing the x, y, and z coordinates.
+        landmark2: The second landmark containing the x, y, and z coordinates.
+        landmark3: The third landmark containing the x, y, and z coordinates.
+    Returns:
+        angle: The calculated angle between the three landmarks.
+    '''
+    x1, y1, _ = landmark1
+    x2, y2, _ = landmark2
+    x3, y3, _ = landmark3
+    
+    angle = np.degrees(np.arctan2(y3 - y2, x3 - x2) - np.arctan2(y1 - y2, x1 - x2))
+    
+    if angle < 0:
+        angle += 360
+    
+    return angle
 
 def classifyPose(landmarks, output_image, display=False):
     '''
@@ -30,44 +43,33 @@ def classifyPose(landmarks, output_image, display=False):
         output_image: The image with the detected pose landmarks drawn and pose label written.
         label: The classified pose label of the person in the output_image.
     '''
-    # Initialize the label of the pose. It is not known at this stage.
     label = 'Unknown Pose'
-
-    # Specify the color (Red) with which the label will be written on the image.
     color = (0, 0, 255)
 
-    # Calculate the required angles.
-    # Get the angle between the left shoulder, elbow and wrist points.
     left_elbow_angle = calculateAngle(landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value],
                                       landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value],
                                       landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value])
 
-    # Get the angle between the right shoulder, elbow and wrist points.
     right_elbow_angle = calculateAngle(landmarks[mp_pose.PoseLandmark.RIGHT_SHOULDER.value],
                                        landmarks[mp_pose.PoseLandmark.RIGHT_ELBOW.value],
                                        landmarks[mp_pose.PoseLandmark.RIGHT_WRIST.value])
 
-    # Get the angle between the left elbow, shoulder and hip points.
     left_shoulder_angle = calculateAngle(landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value],
                                          landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value],
                                          landmarks[mp_pose.PoseLandmark.LEFT_HIP.value])
 
-    # Get the angle between the right hip, shoulder and elbow points.
     right_shoulder_angle = calculateAngle(landmarks[mp_pose.PoseLandmark.RIGHT_HIP.value],
                                           landmarks[mp_pose.PoseLandmark.RIGHT_SHOULDER.value],
                                           landmarks[mp_pose.PoseLandmark.RIGHT_ELBOW.value])
 
-    # Get the angle between the left hip, knee and ankle points.
     left_knee_angle = calculateAngle(landmarks[mp_pose.PoseLandmark.LEFT_HIP.value],
                                      landmarks[mp_pose.PoseLandmark.LEFT_KNEE.value],
                                      landmarks[mp_pose.PoseLandmark.LEFT_ANKLE.value])
 
-    # Get the angle between the right hip, knee and ankle points
     right_knee_angle = calculateAngle(landmarks[mp_pose.PoseLandmark.RIGHT_HIP.value],
                                       landmarks[mp_pose.PoseLandmark.RIGHT_KNEE.value],
                                       landmarks[mp_pose.PoseLandmark.RIGHT_ANKLE.value])
 
-    # Print calculated angles for debugging
     print(f"Left Elbow Angle: {left_elbow_angle}")
     print(f"Right Elbow Angle: {right_elbow_angle}")
     print(f"Left Shoulder Angle: {left_shoulder_angle}")
@@ -75,14 +77,12 @@ def classifyPose(landmarks, output_image, display=False):
     print(f"Left Knee Angle: {left_knee_angle}")
     print(f"Right Knee Angle: {right_knee_angle}")
 
-    # Check for Five-Pointed Star Pose
     if abs(landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].y - landmarks[mp_pose.PoseLandmark.LEFT_HIP.value].y) < 0.1 and \
        abs(landmarks[mp_pose.PoseLandmark.RIGHT_WRIST.value].y - landmarks[mp_pose.PoseLandmark.RIGHT_HIP.value].y) < 0.1 and \
        abs(landmarks[mp_pose.PoseLandmark.LEFT_ANKLE.value].x - landmarks[mp_pose.PoseLandmark.RIGHT_ANKLE.value].x) > 0.2 and \
        abs(landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].x - landmarks[mp_pose.PoseLandmark.RIGHT_WRIST.value].x) > 0.2:
-        label = "Five-Pointed Star Pose"  
-    
-    # Check if it is the warrior II pose or the T pose.
+        label = "Five-Pointed Star Pose"
+
     if left_elbow_angle > 165 and left_elbow_angle < 195 and right_elbow_angle > 165 and right_elbow_angle < 195:
         if left_shoulder_angle > 80 and left_shoulder_angle < 110 and right_shoulder_angle > 80 and right_shoulder_angle < 110:
             if left_knee_angle > 165 and left_knee_angle < 195 or right_knee_angle > 165 and right_knee_angle < 195:
@@ -91,12 +91,10 @@ def classifyPose(landmarks, output_image, display=False):
             if left_knee_angle > 160 and left_knee_angle < 195 and right_knee_angle > 160 and right_knee_angle < 195:
                 label = 'T Pose'
     
-    # Check if it is the tree pose.
     if left_knee_angle > 165 and left_knee_angle < 195 or right_knee_angle > 165 and right_knee_angle < 195:
         if left_knee_angle > 315 and left_knee_angle < 335 or right_knee_angle > 25 and right_knee_angle < 45:
             label = 'Tree Pose'
 
-    # Check for Upward Salute Pose
     if abs(landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].x - landmarks[mp_pose.PoseLandmark.LEFT_HIP.value].x) < 0.1 and \
        abs(landmarks[mp_pose.PoseLandmark.RIGHT_WRIST.value].x - landmarks[mp_pose.PoseLandmark.RIGHT_HIP.value].x) < 0.1 and \
        landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].y < landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].y and \
@@ -104,14 +102,12 @@ def classifyPose(landmarks, output_image, display=False):
        abs(landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].y - landmarks[mp_pose.PoseLandmark.RIGHT_SHOULDER.value].y) < 0.05:
         label = "Upward Salute Pose"
 
-    # Check for Hands Under Feet Pose
     if landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].y > landmarks[mp_pose.PoseLandmark.LEFT_KNEE.value].y and \
        landmarks[mp_pose.PoseLandmark.RIGHT_WRIST.value].y > landmarks[mp_pose.PoseLandmark.RIGHT_KNEE.value].y and \
        abs(landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].x - landmarks[mp_pose.PoseLandmark.LEFT_ANKLE.value].x) < 0.05 and \
        abs(landmarks[mp_pose.PoseLandmark.RIGHT_WRIST.value].x - landmarks[mp_pose.PoseLandmark.RIGHT_ANKLE.value].x) < 0.05:
-        label = "Hands Under Feet Pose"      
+        label = "Hands Under Feet Pose"
 
-    # If the pose is classified successfully
     if label != 'Unknown Pose':
         color = (0, 255, 0)
 
@@ -125,72 +121,18 @@ def classifyPose(landmarks, output_image, display=False):
     else:
         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(),
-    )
+def run(image):
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    results = pose.process(image)
 
     if results.pose_landmarks:
-        res, pose_classification = classifyPose(results.pose_landmarks.landmark, res)
-
-    return res[:, :, ::-1]
-
-
-model_complexities = list(range(3))
-background_colors = ["white", "black", "green"]
-
-image_paths = sorted(pathlib.Path("images").rglob("*.jpg"))
-examples = [[str(path), model_complexities[1], True, 0.5, background_colors[0]] for path in image_paths]
+        mp_drawing.draw_landmarks(image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS)
+        image, classification = classifyPose(results.pose_landmarks.landmark, image)
+        return image, classification
+    else:
+        return image, "No Pose Detected"
 
-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,
-    description=DESCRIPTION,
-)
+demo = gr.Interface(fn=run, inputs="image", outputs=["image", "text"], live=True)
 
 if __name__ == "__main__":
-    demo.queue().launch()
+    demo.launch()