app.py

#######################################################################################################################
# Main application of the exercise pose detetion.

# Support four kinds of poses:
# 1. Idle or wrong (inclluding stretch)
# 2. Bicep Curl 
# 3. Shoulder Press
# 4. Squat

#######################################################################################################################

# Import Modules ======================================================================================================
import streamlit as st
import cv2

from tensorflow.keras.models import Model
from tensorflow.keras.layers import (LSTM, Dense, Dropout, Input, Flatten, Bidirectional, Permute, multiply)

import numpy as np
import mediapipe as mp
import math

from streamlit_webrtc import webrtc_streamer, WebRtcMode, RTCConfiguration
import av

import requests

import time

# Build and Load Model ================================================================================================
def attention_block(inputs, time_steps):
    """
    Attention layer for deep neural network
    
    """
    # Attention weights
    a = Permute((2, 1))(inputs)
    a = Dense(time_steps, activation='softmax')(a)
    
    # Attention vector
    a_probs = Permute((2, 1), name='attention_vec')(a)
    
    # Luong's multiplicative score
    output_attention_mul = multiply([inputs, a_probs], name='attention_mul') 
    
    return output_attention_mul

# @st.cache(allow_output_mutation=True)
# ---------------------------------------------------------------------------------------------------------------------
# st.cache was deprecated in version 1.18.0, use st.cache_data or st.cache_resource instead
# st.cache_data is the recommended way to cache computations that return data
# st.cache_resource is the recommended way to cache global resources like ML models or database connections
# refer to https://docs.streamlit.io/library/advanced-features/caching for more details
# ---------------------------------------------------------------------------------------------------------------------
@st.cache_resource(show_spinner="Building Model...")  # 👈 Add the caching decorator using custom text for spinner
def build_model(HIDDEN_UNITS=256, sequence_length=30, num_input_values=33*4, num_classes=4):
    """
    Function used to build the deep neural network model on startup

    Args:
        HIDDEN_UNITS (int, optional): Number of hidden units for each neural network hidden layer. Defaults to 256.
        sequence_length (int, optional): Input sequence length (i.e., number of frames). Defaults to 30.
        num_input_values (_type_, optional): Input size of the neural network model. Defaults to 33*4 (i.e., number of keypoints x number of metrics).
        num_classes (int, optional): Number of classification categories (i.e., model output size). Defaults to 4.

    Returns:
        keras model: neural network with pre-trained weights
    """
    # Input
    inputs = Input(shape=(sequence_length, num_input_values))
    # Bi-LSTM
    lstm_out = Bidirectional(LSTM(HIDDEN_UNITS, return_sequences=True))(inputs)
    # Attention
    attention_mul = attention_block(lstm_out, sequence_length)
    attention_mul = Flatten()(attention_mul)
    # Fully Connected Layer
    x = Dense(2*HIDDEN_UNITS, activation='relu')(attention_mul)
    x = Dropout(0.5)(x)
    # Output
    x = Dense(num_classes, activation='softmax')(x)
    # Bring it all together
    model = Model(inputs=[inputs], outputs=x)

    ## Load Model Weights
    load_dir = "./models/LSTM_Attention.h5"  
    model.load_weights(load_dir)
    
    return model

HIDDEN_UNITS = 256
model = build_model(HIDDEN_UNITS)

# App =================================================================================================================
## UI -----------------------------------------------------------------------------------------------------------------
st.markdown("<h1 style='text-align: center;'>🏋️‍♂️ Exercise Detector 🕵️</h1>", unsafe_allow_html=True)

st.write("## ⚙️Settings")
threshold1 = st.slider("**1. Minimum Keypoint Detection Confidence**", 0.00, 1.00, 0.80)
threshold2 = st.slider("**2. Minimum Tracking Confidence**", 0.00, 1.00, 0.80)
threshold3 = st.slider("**3. Minimum Activity Classification Confidence**", 0.00, 1.00, 0.80)
cool_down_time = st.slider("**4. Cool Down Time**", 3.00, 10.00, 5.00)
st.write("> P.S. 🧊*Cool Down Time*🧊 means if you fail to finish a correct movement within \
    this short period of time, your current state will switch to idle or wrong.")

st.write("## 🚀Launch")

## Mediapipe ----------------------------------------------------------------------------------------------------------
mp_pose = mp.solutions.pose # Pre-trained pose estimation model from Google Mediapipe
mp_drawing = mp.solutions.drawing_utils # Supported Mediapipe visualization tools
pose = mp_pose.Pose(min_detection_confidence=threshold1, min_tracking_confidence=threshold2) # mediapipe pose model

## Send Request -------------------------------------------------------------------------------------------------------
class RequestSender:
    def __init__(self):
        self.domain = 'https://wombat-meet-antelope.ngrok-free.app/'
        self.notify_domain = 'https://ab4b-128-54-47-14.ngrok-free.app/'
        self.action_sent = ''
        self.count_sent = 0
        self.can_send_count = False
            
    
    def send_activity_request(self):
        try:
            postObject = {'user': 'user_1',
                    'device': 'video',
                    'value': self.action_sent,
                    'type': 'activity',}
            print(postObject)
            t = time.localtime()
            current_time = time.strftime("%H:%M:%S", t)
            print(current_time)
            #x = requests.post(self.domain + 'send-data', json = postObject)
            # print(x.text)
        except:
            print("Activity Request Failed!")        
            
    def send_count_request(self):
        try:
            postObject = {'activity': self.action_sent,
                        'reps': self.count_sent,}
            print(postObject)
            t = time.localtime()
            current_time = time.strftime("%H:%M:%S", t)
            print(current_time)
            #x = requests.post(self.notify_domain + 'notify', json = postObject)
            # print(x.text)
        except:
            print("Count Request Failed!") 
        
# Instantiate a request sender
sender = RequestSender()

## Real Time Machine Learning and Computer Vision Processes -----------------------------------------------------------
class VideoProcessor:
    def __init__(self):
        # Parameters
        self.actions = np.array(['curl', 'press', 'squat', 'idle or wrong'])
        self.sequence_length = 30
        self.colors = [(240,160,80), (200,80,90), (80,180,240), (0,0,0)]
        self.threshold = threshold3
        
        # Detection variables
        self.sequence = []
        self.previous_action = 'idle or wrong'
        self.current_action = 'idle or wrong'
        self.detected_action = ''

        # Rep counter logic variables
        self.curl_counter = 0
        self.press_counter = 0
        self.squat_counter = 0
        self.curl_stage = None
        self.press_stage = None
        self.squat_stage = None
        self.cool_down_start = 0
    
    
    def draw_landmarks(self, image, results):
        """
        This function draws keypoints and landmarks detected by the human pose estimation model
        
        """
        mp_drawing.draw_landmarks(image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS,
                                    mp_drawing.DrawingSpec(color=(245,117,66), thickness=2, circle_radius=2), 
                                    mp_drawing.DrawingSpec(color=(245,66,230), thickness=2, circle_radius=2) 
                                    )
        return
    
    
    def extract_keypoints(self, results):
        """
        Processes and organizes the keypoints detected from the pose estimation model 
        to be used as inputs for the exercise decoder models
        
        """
        pose = np.array([[res.x, res.y, res.z, res.visibility] for res in results.pose_landmarks.landmark]).flatten() if results.pose_landmarks else np.zeros(33*4)
        return pose
    
    
    def calculate_angle(self, a,b,c):
        """
        Computes 3D joint angle inferred by 3 keypoints and their relative positions to one another
        
        """
        a = np.array(a) # First
        b = np.array(b) # Mid
        c = np.array(c) # End
        
        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 
    
    
    def get_coordinates(self, landmarks, mp_pose, side, joint):
        """
        Retrieves x and y coordinates of a particular keypoint from the pose estimation model
            
        Args:
            landmarks: processed keypoints from the pose estimation model
            mp_pose: Mediapipe pose estimation model
            side: 'left' or 'right'. Denotes the side of the body of the landmark of interest.
            joint: 'shoulder', 'elbow', 'wrist', 'hip', 'knee', or 'ankle'. Denotes which body joint is associated with the landmark of interest.
        
        """
        coord = getattr(mp_pose.PoseLandmark,side.upper()+"_"+joint.upper())
        x_coord_val = landmarks[coord.value].x
        y_coord_val = landmarks[coord.value].y
        return [x_coord_val, y_coord_val] 
    
    
    def viz_joint_angle(self, image, angle, joint):
        """
        Displays the joint angle value near the joint within the image frame
        
        """
        cv2.putText(image, str(int(angle)), 
                    tuple(np.multiply(joint, [640, 480]).astype(int)), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2, cv2.LINE_AA
                    )
        return
    
    
    def count_reps(self, image, landmarks, mp_pose):
        """
        Counts repetitions of each exercise. Global count and stage (i.e., state) variables are updated within this function.
        
        """
        
        if self.detected_action == 'curl':
            # Get coords
            shoulder = self.get_coordinates(landmarks, mp_pose, 'left', 'shoulder')
            elbow = self.get_coordinates(landmarks, mp_pose, 'left', 'elbow')
            wrist = self.get_coordinates(landmarks, mp_pose, 'left', 'wrist')
            
            # calculate elbow angle
            angle = self.calculate_angle(shoulder, elbow, wrist)
            
            # curl counter logic
            if angle < 30:
                self.curl_stage = "up" 
            if angle > 140 and self.curl_stage =='up':
                self.curl_stage="down"  
                self.curl_counter +=1
                self.current_action = 'curl'
                self.cool_down_start = time.time()
                if (self.curl_counter != 0 and self.curl_counter % 10 == 0):
                    sender.count_sent = self.curl_counter
                    sender.can_send_count = True
            self.press_stage = None
            self.squat_stage = None
                
            # Viz joint angle
            self.viz_joint_angle(image, angle, elbow)
            
        elif self.detected_action == 'press':           
            # Get coords
            shoulder = self.get_coordinates(landmarks, mp_pose, 'left', 'shoulder')
            elbow = self.get_coordinates(landmarks, mp_pose, 'left', 'elbow')
            wrist = self.get_coordinates(landmarks, mp_pose, 'left', 'wrist')

            # Calculate elbow angle
            elbow_angle = self.calculate_angle(shoulder, elbow, wrist)
            
            # Compute distances between joints
            shoulder2elbow_dist = abs(math.dist(shoulder,elbow))
            shoulder2wrist_dist = abs(math.dist(shoulder,wrist))
            
            # Press counter logic
            if (elbow_angle > 130) and (shoulder2elbow_dist < shoulder2wrist_dist):
                self.press_stage = "up"
            if (elbow_angle < 50) and (shoulder2elbow_dist > shoulder2wrist_dist) and (self.press_stage =='up'):
                self.press_stage='down'
                self.press_counter += 1
                self.current_action = 'press'
                self.cool_down_start = time.time()
                if (self.press_counter != 0 and self.press_counter % 10 == 0):
                    sender.count_sent = self.press_counter
                    sender.can_send_count = True
            self.curl_stage = None
            self.squat_stage = None
                
            # Viz joint angle
            self.viz_joint_angle(image, elbow_angle, elbow)
            
        elif self.detected_action == 'squat':
            # Get coords
            # left side
            left_shoulder = self.get_coordinates(landmarks, mp_pose, 'left', 'shoulder')
            left_hip = self.get_coordinates(landmarks, mp_pose, 'left', 'hip')
            left_knee = self.get_coordinates(landmarks, mp_pose, 'left', 'knee')
            left_ankle = self.get_coordinates(landmarks, mp_pose, 'left', 'ankle')
            # right side
            right_shoulder = self.get_coordinates(landmarks, mp_pose, 'right', 'shoulder')
            right_hip = self.get_coordinates(landmarks, mp_pose, 'right', 'hip')
            right_knee = self.get_coordinates(landmarks, mp_pose, 'right', 'knee')
            right_ankle = self.get_coordinates(landmarks, mp_pose, 'right', 'ankle')
            
            # Calculate knee angles
            left_knee_angle = self.calculate_angle(left_hip, left_knee, left_ankle)
            right_knee_angle = self.calculate_angle(right_hip, right_knee, right_ankle)
            
            # Calculate hip angles
            left_hip_angle = self.calculate_angle(left_shoulder, left_hip, left_knee)
            right_hip_angle = self.calculate_angle(right_shoulder, right_hip, right_knee)
            
            # Squat counter logic
            thr = 165
            if (left_knee_angle < thr) and (right_knee_angle < thr) and (left_hip_angle < thr) and (right_hip_angle < thr):
                self.squat_stage = "down"
            if (left_knee_angle > thr) and (right_knee_angle > thr) and (left_hip_angle > thr) and (right_hip_angle > thr) and (self.squat_stage =='down'):
                self.squat_stage='up'
                self.squat_counter += 1
                self.current_action = 'squat'
                self.cool_down_start = time.time()
                if (self.squat_counter != 0 and self.squat_counter % 10 == 0):
                    sender.count_sent = self.squat_counter
                    sender.can_send_count = True
            self.curl_stage = None
            self.press_stage = None
                
            # Viz joint angles
            self.viz_joint_angle(image, left_knee_angle, left_knee)
            self.viz_joint_angle(image, left_hip_angle, left_hip)
            
        else:
            if round(time.time() - self.cool_down_start > cool_down_time):
                self.current_action = 'idle or wrong'
                sender.count_sent = 0
                sender.can_send_count = False
            else:    
                pass
        return
    
    
    def prob_viz(self, res, input_frame):
        """
        This function displays the model prediction probability distribution over the set of exercise classes
        as a horizontal bar graph
        
        """
        output_frame = input_frame.copy()
        for num, prob in enumerate(res):        
            # rectangle
            if num == 3:
                # idle or wrong
                cv2.rectangle(output_frame, (0,60+num*40), (int(prob*220), 95+num*40), self.colors[num], -1)
            else:
                # one of the three exercises
                cv2.rectangle(output_frame, (0,60+num*40), (int(prob*100), 95+num*40), self.colors[num], -1)
            # text
            cv2.putText(output_frame, self.actions[num], (0, 85+num*40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255,255,255), 2, cv2.LINE_AA)
            
        return output_frame
    
    
    def process(self, image):
        """
        Function to process the video frame from the user's webcam and run the fitness trainer AI

        Args:
            image (numpy array): input image from the webcam

        Returns:
            numpy array: processed image with keypoint detection and fitness activity classification visualized
        """
        # Pose detection model
        image.flags.writeable = False
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        results = pose.process(image)

        # Draw the hand annotations on the image.
        image.flags.writeable = True
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        self.draw_landmarks(image, results) 
        
        # Prediction logic
        keypoints = self.extract_keypoints(results)        
        self.sequence.append(keypoints.astype('float32',casting='same_kind'))      
        self.sequence = self.sequence[-self.sequence_length:]
        
        if len(self.sequence) == self.sequence_length:
            res = model.predict(np.expand_dims(self.sequence, axis=0), verbose=0)[0]
            
            self.detected_action = self.actions[np.argmax(res)]
            confidence = np.max(res)
            
            # Erase current action variable if no probability is above threshold
            if confidence < self.threshold:
                self.detected_action = 'idle or wrong'

            # Viz probabilities
            image = self.prob_viz(res, image)
            
            # Count reps
            try:
                landmarks = results.pose_landmarks.landmark
                self.count_reps(image, landmarks, mp_pose)
            except:
                pass

            # Display graphical information
            cv2.rectangle(image, (0,0), (640, 40), self.colors[np.argmax(res)], -1)
            if self.current_action == 'idle or wrong':
                cv2.putText(image, 'idle, stretch, or wrong move', (80,30), cv2.FONT_HERSHEY_TRIPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
            else:
                cv2.putText(image, 'curl ' + str(self.curl_counter), (3,30), cv2.FONT_HERSHEY_TRIPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
                cv2.putText(image, 'press ' + str(self.press_counter), (230,30), cv2.FONT_HERSHEY_TRIPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
                cv2.putText(image, 'squat ' + str(self.squat_counter), (490,30), cv2.FONT_HERSHEY_TRIPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
        
            # Send requests
            ##  activity request
            ##  1. if it's the first frame and nothing has been sent: send
            ##  2. if it's different from the previous action
            ##     2.1 if precious action is "idle or wrong": send
            ##     2.2 if precious action is a valid pose: delayed send
            if sender.action_sent == '' or self.current_action != self.previous_action:
                if sender.action_sent == '':
                    sender.action_sent = 'idle or wrong'
                else:
                    sender.action_sent = self.current_action
                sender.send_activity_request()
                
            ## count request
            if sender.count_sent != 0 and sender.can_send_count == True:
                sender.send_count_request()
                sender.can_send_count = False


            # Update previous action
            self.previous_action = self.current_action
            
        # return cv2.flip(image, 1)
        return image
    
    def recv(self, frame):
        """
        Receive and process video stream from webcam

        Args:
            frame: current video frame

        Returns:
            av.VideoFrame: processed video frame
        """
        img = frame.to_ndarray(format="bgr24")
        img = self.process(img)
        return av.VideoFrame.from_ndarray(img, format="bgr24")
        
## Stream Webcam Video and Run Model ----------------------------------------------------------------------------------
# Options
RTC_CONFIGURATION = RTCConfiguration(
    {"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]}
)
# Streamer
webrtc_ctx = webrtc_streamer(
    key="Exercise Detector",
    mode=WebRtcMode.SENDRECV,
    rtc_configuration=RTC_CONFIGURATION,
    media_stream_constraints={"video": True, "audio": False},
    video_processor_factory=VideoProcessor,
    async_processing=True,
)