Update app.py

app.py

@@ -1,103 +1,51 @@
+import os
 import streamlit as st
 import cv2
-
+import mediapipe as mp
+import numpy as np
+import math
 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
+# Load the pose estimation model from Mediapipe
+mp_pose = mp.solutions.pose 
+mp_drawing = mp.solutions.drawing_utils 
+pose = mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5) 
 
-from streamlit_webrtc import webrtc_streamer, WebRtcMode, RTCConfiguration
-import av
-
-## Build and Load Model
+# Define the attention block for the LSTM 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
 
+# Build and load the LSTM model
 @st.cache(allow_output_mutation=True)
 def build_model(HIDDEN_UNITS=256, sequence_length=30, num_input_values=33*4, num_classes=3):
-    """
-    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 3.
-
-    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
-st.write("# AI Personal Fitness Trainer Web App")
-
-st.markdown("❗❗ **Development Note** ❗❗")
-st.markdown("Currently, the exercise recognition model uses the the x, y, and z coordinates of each anatomical landmark from the MediaPipe Pose model. These coordinates are normalized with respect to the image frame (e.g., the top left corner represents (x=0,y=0) and the bottom right corner represents(x=1,y=1)).")
-st.markdown("I'm currently developing and testing two new feature engineering strategies:")
-st.markdown("- Normalizing coordinates by the detected bounding box of the user")
-st.markdown("- Using joint angles rather than keypoint coordaintes as features")            
-st.write("Stay Tuned!")
-
-st.write("## Settings")
-threshold1 = st.slider("Minimum Keypoint Detection Confidence", 0.00, 1.00, 0.50)
-threshold2 = st.slider("Minimum Tracking Confidence", 0.00, 1.00, 0.50)
-threshold3 = st.slider("Minimum Activity Classification Confidence", 0.00, 1.00, 0.50)
-
-st.write("## Activate the AI 🤖🏋️‍♂️")
-
-## 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
-
-## Real Time Machine Learning and Computer Vision Processes
+# Define the VideoProcessor class for real-time video processing
 class VideoProcessor:
     def __init__(self):
         # Parameters
         self.actions = np.array(['curl', 'press', 'squat'])
         self.sequence_length = 30
         self.colors = [(245,117,16), (117,245,16), (16,117,245)]
-        self.threshold = threshold3
+        self.threshold = 0.5
+
+        self.model = build_model(256)
         
         # Detection variables
         self.sequence = []
@@ -110,169 +58,212 @@ class VideoProcessor:
         self.curl_stage = None
         self.press_stage = None
         self.squat_stage = None
+        self.pose = mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5)
+
+    def process_video(self, video_file):
+        # Get the filename from the file object
+        filename = "temp_video.mp4"
+        # Create a temporary file to write the contents of the uploaded video file
+        with open(filename, 'wb') as temp_file:
+            temp_file.write(video_file.read())
+
+        # Process the video and save the processed video to a new file
+        output_filename = "processed_video.mp4"
+        cap = cv2.VideoCapture(filename)
+        frame_width = int(cap.get(3))
+        frame_height = int(cap.get(4))
+        out = cv2.VideoWriter(output_filename, cv2.VideoWriter_fourcc(*'h264'), 30, (frame_width, frame_height))
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
+                break
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            results = self.pose.process(frame_rgb)
+            processed_frame = self.process_frame(frame, results)
+            out.write(processed_frame)
+        cap.release()
+        out.release()
+
+        # Remove the temporary file
+        os.remove(filename)
+
+        # Return the path to the processed video file
+        return output_filename
     
-    @st.cache()    
-    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
-    
-    @st.cache()
-    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 process_frame(self, frame, results):
+        # Process the frame using the `process` function
+        processed_frame = self.process(frame)
+        return processed_frame
     
-    @st.cache()
-    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
+    def process(self, image):
+      
+        # 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) 
         
-        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)
+        # Prediction logic
+        keypoints = self.extract_keypoints(results)        
+        self.sequence.append(keypoints.astype('float32',casting='same_kind'))      
+        self.sequence = self.sequence[-self.sequence_length:]
         
-        if angle > 180.0:
-            angle = 360-angle
+        if len(self.sequence) == self.sequence_length:
+            res = self.model.predict(np.expand_dims(self.sequence, axis=0), verbose=0)[0]
             
-        return angle 
-    
-    @st.cache()
-    def get_coordinates(self, landmarks, mp_pose, side, joint):
-        """
-        Retrieves x and y coordinates of a particular keypoint from the pose estimation model
+            self.current_action = self.actions[np.argmax(res)]
+            confidence = np.max(res)
+            print("confidence", confidence)  # Debug print statement
+            print("current action" , self.current_action)
             
-        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] 
-    
-    @st.cache()
-    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
-    
-    @st.cache()
+            # Erase current action variable if no probability is above threshold
+            if confidence < self.threshold:
+                self.current_action = ''
+                
+                
+            print("current action" , self.current_action)
+
+
+            # Viz probabilities
+            image = self.prob_viz(res, image)
+            
+            # Count reps
+            
+            landmarks = results.pose_landmarks.landmark
+            self.count_reps(image, landmarks, mp_pose)
+            
+
+            # Display graphical information
+            cv2.rectangle(image, (0,0), (640, 40), self.colors[np.argmax(res)], -1)
+            cv2.putText(image, 'curl ' + str(self.curl_counter), (3,30), 
+                           cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
+            cv2.putText(image, 'press ' + str(self.press_counter), (240,30), 
+                           cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
+            cv2.putText(image, 'squat ' + str(self.squat_counter), (490,30), 
+                           cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
+          
+        return image
+
+    def draw_landmarks(self, image, results):
+        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 image
+
+    def extract_keypoints(self, results):
+        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 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.current_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')
-            
+            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
+            print("Curl Angle:", angle)  # Debug print statement
             if angle < 30:
-                self.curl_stage = "up" 
-            if angle > 140 and self.curl_stage =='up':
-                self.curl_stage="down"  
-                self.curl_counter +=1
+                self.curl_stage = "up"
+            if angle > 140 and self.curl_stage == 'up':
+                self.curl_stage = "down"
+                self.curl_counter += 1
+                print("count:",self.curl_counter)
             self.press_stage = None
             self.squat_stage = None
-                
+
             # Viz joint angle
             self.viz_joint_angle(image, angle, elbow)
-            
-        elif self.current_action == 'press':           
+
+        elif self.current_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')
+            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)
-            
+            print(shoulder, elbow, wrist)
             # Compute distances between joints
-            shoulder2elbow_dist = abs(math.dist(shoulder,elbow))
-            shoulder2wrist_dist = abs(math.dist(shoulder,wrist))
-            
+            shoulder2elbow_dist = abs(math.dist(shoulder, elbow))
+            shoulder2wrist_dist = abs(math.dist(shoulder, wrist))
+
             # Press counter logic
+            print("Press Angle:", elbow_angle)  # Debug print statement
+            print("Shoulder to Elbow Distance:", shoulder2elbow_dist)  # Debug print statement
+            print("Shoulder to Wrist Distance:", shoulder2wrist_dist)  # Debug print statement
             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'
+            if (elbow_angle < 50) and (shoulder2elbow_dist > shoulder2wrist_dist) and (self.press_stage == 'up'):
+                self.press_stage = 'down'
                 self.press_counter += 1
+                
+                
+                print("count:",self.press_counter)
             self.curl_stage = None
             self.squat_stage = None
-                
+
             # Viz joint angle
             self.viz_joint_angle(image, elbow_angle, elbow)
-            
+
         elif self.current_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')
+            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')
-            
+            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):
+            print("Left Knee Angle:", left_knee_angle)  # Debug print statement
+            print("Right Knee Angle:", right_knee_angle)  # Debug print statement
+            print("Left Hip Angle:", left_hip_angle)  # Debug print statement
+            print("Right Hip Angle:", right_hip_angle)  # Debug print statement
+            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'
+            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
+                print("count:",self.squat_counter)
             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:
             pass
         return
-    
-    @st.cache()
+
     def prob_viz(self, res, input_frame):
         """
         This function displays the model prediction probability distribution over the set of exercise classes
@@ -285,94 +276,49 @@ class VideoProcessor:
             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
-    
-    @st.cache()
-    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) 
+    def get_coordinates(self, landmarks, mp_pose, side, part):
         
-        # 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]
-            # interpreter.set_tensor(self.input_details[0]['index'], np.expand_dims(self.sequence, axis=0))
-            # interpreter.invoke()
-            # res = interpreter.get_tensor(self.output_details[0]['index'])
-            
-            self.current_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.current_action = ''
+        coord = getattr(mp_pose.PoseLandmark,side.upper()+"_"+part.upper())
+        x_coord_val = landmarks[coord.value].x
+        y_coord_val = landmarks[coord.value].y
+        return [x_coord_val, y_coord_val] 
+    
+    
+   
+       
 
-            # 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)
-            cv2.putText(image, 'curl ' + str(self.curl_counter), (3,30), 
-                           cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
-            cv2.putText(image, 'press ' + str(self.press_counter), (240,30), 
-                           cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
-            cv2.putText(image, 'squat ' + str(self.squat_counter), (490,30), 
-                           cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
-          
-        # return cv2.flip(image, 1)
-        return image
-    
-    def recv(self, frame):
-        """
-        Receive and process video stream from webcam
 
-        Args:
-            frame: current video frame
+    def calculate_angle(self, a, b, c):
+        a = np.array(a) 
+        b = np.array(b) 
+        c = np.array(c)
+        radians = math.atan2(c[1]-b[1], c[0]-b[0]) - math.atan2(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
 
-        Returns:
-            av.VideoFrame: processed video frame
-        """
-        img = frame.to_ndarray(format="bgr24")
-        img = self.process(img)
-        return av.VideoFrame.from_ndarray(img, format="bgr24")
+    def viz_joint_angle(self, image, angle, joint):
+        cv2.putText(image, str(round(angle, 2)), 
+                    tuple(np.multiply(joint, [640, 480]).astype(int)), 
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 2, cv2.LINE_AA)
+
+# Define Streamlit app
+def main():
+    st.title("Real-time Exercise Detection")
+    video_file = st.file_uploader("Upload a video file", type=["mp4", "avi"])
+    if video_file is not None:
+        video_processor = VideoProcessor()
         
-## Stream Webcam Video and Run Model
-# Options
-RTC_CONFIGURATION = RTCConfiguration(
-    {"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]}
-)
-# Streamer
-webrtc_ctx = webrtc_streamer(
-    key="AI trainer",
-    mode=WebRtcMode.SENDRECV,
-    rtc_configuration=RTC_CONFIGURATION,
-    media_stream_constraints={"video": True, "audio": False},
-    video_processor_factory=VideoProcessor,
-    async_processing=True,
-)
+        output_video = video_processor.process_video(video_file)
+        
+        
+        video_file = open(output_video, 'rb')
+        video_bytes = video_file.read()
+        st.video(video_bytes)
+
+if __name__ == "__main__":
+    main()