new_pose.py

import math
import cv2
import numpy as np
from time import time
import mediapipe as mp
import matplotlib.pyplot as plt
import socket
import time
from collections import deque

# This line was added by Trevor Evetts.
# I believe there's been an update to the syntax since this code was established, and so this is now required to model the figure in 3D space.
from mpl_toolkits.mplot3d import Axes3D

# Create a socket object
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

# Define the address and port to send the data
address = 'localhost'
port = 12345

# Connect to the receiver script
sock.connect((address, port))

# Initializing mediapipe pose class.
mp_pose = mp.solutions.pose
label_history = deque(maxlen=4)

# Setting up the Pose function.
pose = mp_pose.Pose(static_image_mode=True, min_detection_confidence=0.3, model_complexity=2)

# Initializing mediapipe drawing class, useful for annotation.
mp_drawing = mp.solutions.drawing_utils

def detectPose(image, pose, display=True):
    '''
    This function performs pose detection on an image.
    Args:
        image: The input image with a prominent person whose pose landmarks needs to be detected.
        pose: The pose setup function required to perform the pose detection.
        display: A boolean value that is if set to true the function displays the original input image, the resultant image, 
                 and the pose landmarks in 3D plot and returns nothing.
    Returns:
        output_image: The input image with the detected pose landmarks drawn.
        landmarks: A list of detected landmarks converted into their original scale.
    '''
    
    # Create a copy of the input image.
    output_image = image.copy()
    
    # Convert the image from BGR into RGB format.
    imageRGB = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
    # Perform the Pose Detection.
    results = pose.process(imageRGB)
    
    # Retrieve the height and width of the input image.
    height, width, _ = image.shape
    
    # Initialize a list to store the detected landmarks.
    landmarks = []
    
    # Check if any landmarks are detected.
    if results.pose_landmarks:
    
        # Draw Pose landmarks on the output image.
        mp_drawing.draw_landmarks(image=output_image, landmark_list=results.pose_landmarks,
                                  connections=mp_pose.POSE_CONNECTIONS)
        
        # Iterate over the detected landmarks.
        for landmark in results.pose_landmarks.landmark:
            
            # Append the landmark into the list.
            landmarks.append((int(landmark.x * width), int(landmark.y * height),
                                  (landmark.z * width)))
    
    # Check if the original input image and the resultant image are specified to be displayed.
    if display:
    
        # Display the original input image and the resultant image.
        plt.figure(figsize=[22,22])
        plt.subplot(121);plt.imshow(image[:,:,::-1]);plt.title("Original Image");plt.axis('off');
        plt.subplot(122);plt.imshow(output_image[:,:,::-1]);plt.title("Output Image");plt.axis('off');
        
        # Also Plot the Pose landmarks in 3D.
        mp_drawing.plot_landmarks(results.pose_world_landmarks, mp_pose.POSE_CONNECTIONS)
        
    # Otherwise
    else:
        
        # Return the output image and the found landmarks.
        return output_image, landmarks
    
def calculateAngle(landmark1, landmark2, landmark3):
    '''
    This function calculates angle between three different 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.

    '''

    # Get the required landmarks coordinates.
    x1, y1, _ = landmark1
    x2, y2, _ = landmark2
    x3, y3, _ = landmark3

    # Calculate the angle between the three points
    angle = math.degrees(math.atan2(y3 - y2, x3 - x2) - math.atan2(y1 - y2, x1 - x2))
    
    # Check if the angle is less than zero.
    if angle < 0:

        # Add 360 to the found angle.
        angle += 360
    
    # Return the calculated angle.
    return angle

def classifyPose(landmarks, output_image, display=False):
    '''
    This function classifies yoga poses depending upon the angles of various body joints.
    Args:
        landmarks: A list of detected landmarks of the person whose pose needs to be classified.
        output_image: A image of the person with the detected pose landmarks drawn.
        display: A boolean value that is if set to true the function displays the resultant image with the pose label 
        written on it and returns nothing.
    Returns:
        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.
    text_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])
    
    #----------------------------------------------------------------------------------------------------------------
    # Below are a collection of poses.
    #----------------------------------------------------------------------------------------------------------------

    # Check if it is the "Go Right" Pose.
    #----------------------------------------------------------------------------------------------------------------

    # Check if the both arms are straight.
    if left_elbow_angle > 155 and left_elbow_angle < 195 and right_elbow_angle > 155 and right_elbow_angle < 195:

        # Check if both legs are straight
        if left_knee_angle > 160 and left_knee_angle < 195 and right_knee_angle > 160 and right_knee_angle < 195:
                
            # Check if shoulders are at the required angle.
            if right_shoulder_angle > 0 and right_shoulder_angle < 50 and left_shoulder_angle > 70 and left_shoulder_angle < 120:

                # Specify the label of the pose that is tree pose.
                label = 'D1va, Strafe Right'
        
            # Check if shoulders are at the required angle.
            if left_shoulder_angle < 50 and left_shoulder_angle > 0 and right_shoulder_angle < 120 and right_shoulder_angle > 70:

                # Specify the label of the pose that is tree pose.
                label = 'D1va, Strafe Left'
    
    # Checks to see if legs are straight.
    if left_knee_angle > 160 and left_knee_angle < 195 and right_knee_angle > 160 and right_knee_angle < 195:

        # Checks to see if left shoulder is up and other is down.
        if right_shoulder_angle > 0 and right_shoulder_angle < 50 and left_shoulder_angle > 70 and left_shoulder_angle < 120:

            #Checks to see if left elbow is bent.
            if right_elbow_angle > 155 and right_elbow_angle < 195 and left_elbow_angle > 70 and left_elbow_angle < 120:

                label = 'D1va, Spin Right'

        # Checks to see if right shoulder is up and other is down.
        if left_shoulder_angle > 0 and left_shoulder_angle < 50 and right_shoulder_angle > 70 and right_shoulder_angle < 120:

            #Checks t osee if right elbow is bent.
            if left_elbow_angle > 155 and left_elbow_angle < 195 and right_elbow_angle > 250 and right_elbow_angle < 300:

                label = 'D1va, Spin Left'
        
        # Checks to see if both shoulders are up.
        if left_shoulder_angle > 70 and left_shoulder_angle < 120 and right_shoulder_angle > 70 and right_shoulder_angle < 120:

            # Checks to see if both elbows are bent.
            if left_elbow_angle > 70 and left_elbow_angle < 120 and right_elbow_angle > 250 and right_elbow_angle < 300:

                label = 'D1va, Turn Around'

        # Checks to see if shoulders are slightly up.
        if left_shoulder_angle > 30 and left_shoulder_angle < 70 and right_shoulder_angle > 30 and right_shoulder_angle < 70:

            # Checks to see if elbows are slightly bent.
            if left_elbow_angle > 250 and left_elbow_angle < 300 and right_elbow_angle > 70 and right_elbow_angle < 120:

                # This is a "hands on hips" position
                label = 'D1va, Go Forward'

        # Checks to see if shoulders are up.
        if left_shoulder_angle > 70 and left_shoulder_angle < 140 and right_shoulder_angle > 70 and right_shoulder_angle < 140:

            # Checks to see if elbows are bent further.
            if left_elbow_angle > 0 and left_elbow_angle < 60 and right_elbow_angle > 300 and right_elbow_angle < 360:

                # This is a "hands on head" position.
                label = 'D1va, Go Backward'

        # Checks to see if arms are up.
        if right_shoulder_angle > 70 and right_shoulder_angle < 120 and left_shoulder_angle > 70 and left_shoulder_angle < 120:

            # Checks to see if elbows are straight.
            if right_elbow_angle > 155 and right_elbow_angle < 195 and left_elbow_angle > 155 and left_elbow_angle < 195:

                label = 'D1va, Lay Down'
        
        # Checks to see if arms are higher than usual.
        if right_shoulder_angle > 110 and right_shoulder_angle < 150 and left_shoulder_angle > 110 and left_shoulder_angle < 150:

            # Checks to see if elbows are straight.
            if right_elbow_angle > 155 and right_elbow_angle < 195 and left_elbow_angle > 155 and left_elbow_angle < 195:

                label = 'D1va, Stand Up'

    #----------------------------------------------------------------------------------------------------------------
    # Check if the pose is classified successfully
    # if label != 'Unknown Pose':
        
    #     # Update the color (to green) with which the label will be written on the image.
    #     color = (0, 255, 0)
    #     message = label
    #     # Send the message
    #     sock.sendall(message.encode())

    label_history.append(label)  # Add current label to the history

    
    if label != 'Unknown Pose':
            if len(label_history) == 4 and len(set(label_history)) == 1:
                # Check if the last 5 labels are all equal
                label = label_history[-1]
                text_color = (0, 255, 0)  # Update the color to green for recognized poses
                cv2.putText(output_image, label, (10, 60), cv2.FONT_HERSHEY_PLAIN, 2, text_color, 2)
                # Send the message
                sock.sendall(label.encode())
    
    # Write the label on the output image. 
    #cv2.putText(output_image, label, (10, 30),cv2.FONT_HERSHEY_PLAIN, 2, color, 2)
    else:
        cv2.putText(output_image, label, (10, 60), cv2.FONT_HERSHEY_PLAIN, 2, text_color, 2)
    
    # Check if the resultant image is specified to be displayed.
    if display:
    
        # Display the resultant image.
        plt.figure(figsize=[10,10])
        plt.imshow(output_image[:,:,::-1]);plt.title("Output Image");plt.axis('off');
        
    else:
        
        # Return the output image and the classified label.
        return output_image, label

# Setup Pose function for video.
pose_video = mp_pose.Pose(static_image_mode=False, min_detection_confidence=0.5, model_complexity=1)

# Initialize the VideoCapture object to read from the webcam.
camera_video = cv2.VideoCapture(0)
camera_video.set(3,1280)
camera_video.set(4,960)

# Initialize a resizable window.
cv2.namedWindow('Pose Classification', cv2.WINDOW_NORMAL)

# Iterate until the webcam is accessed successfully.
while camera_video.isOpened():
    
    # Read a frame.
    ok, frame = camera_video.read()
    
    # Check if frame is not read properly.
    if not ok:
        
        # Continue to the next iteration to read the next frame and ignore the empty camera frame.
        continue
    
    # Flip the frame horizontally for natural (selfie-view) visualization.
    frame = cv2.flip(frame, 1)
    
    # Get the width and height of the frame
    frame_height, frame_width, _ =  frame.shape
    
    # Resize the frame while keeping the aspect ratio.
    frame = cv2.resize(frame, (int(frame_width * (640 / frame_height)), 640))
    
    # Perform Pose landmark detection.
    frame, landmarks = detectPose(frame, pose_video, display=False)
    
    # Check if the landmarks are detected.
    if landmarks:
        
        # Perform the Pose Classification.
        frame, _ = classifyPose(landmarks, frame, display=False)
    
    # Display the frame. 
    cv2.imshow('Pose Classification', frame)
    
    # Wait until a key is pressed.
    # Retreive the ASCII code of the key pressed
    k = cv2.waitKey(1) & 0xFF
    
    # Check if 'ESC' is pressed.
    if(k == 27):
        
        # Break the loop.
        break

# Release the VideoCapture object and close the windows.
camera_video.release()
cv2.destroyAllWindows()

# Close the socket
sock.close()