Update app.py

app.py

@@ -59,17 +59,20 @@ class VideoProcessor:
         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())
-        # Now we can open the video file using cv2.VideoCapture()
+
+        # Process the video and save the processed video to a new file
+        output_filename = "processed_video.mp4"
         cap = cv2.VideoCapture(filename)
-        output_frames = []
+        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:
@@ -77,11 +80,15 @@ class VideoProcessor:
             frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
             results = self.pose.process(frame_rgb)
             processed_frame = self.process_frame(frame, results)
-            output_frames.append(processed_frame)
+            out.write(processed_frame)
         cap.release()
+        out.release()
+
         # Remove the temporary file
         os.remove(filename)
-        return output_frames
+
+        # Return the path to the processed video file
+        return output_filename
     
     def process_frame(self, frame, results):
         # Process the frame using the `process` function
@@ -89,15 +96,7 @@ class VideoProcessor:
         return processed_frame
     
     def process(self, image):
-        """
-        Function to process the video frame and run the fitness trainer AI
-
-        Args:
-            image (numpy array): input image from the video
-
-        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)
@@ -118,20 +117,25 @@ class VideoProcessor:
             
             self.current_action = self.actions[np.argmax(res)]
             confidence = np.max(res)
+            print("confidence", confidence)  # Debug print statement
+            print("current action" , self.current_action)
             
             # 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
-            try:
-                landmarks = results.pose_landmarks.landmark
-                self.count_reps(image, landmarks, mp_pose)
-            except:
-                pass
+            
+            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)
@@ -157,90 +161,105 @@ class VideoProcessor:
     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
@@ -258,15 +277,48 @@ class VideoProcessor:
             
         return output_frame
 
+    def get_coordinates(self, landmarks, mp_pose, side, part):
+        
+        
+        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] 
+    
+    
+   
+       
+
+
+
+    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
+
+    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()
-        output_frames = video_processor.process_video(video_file)
-        for frame in output_frames:
-            st.image(frame, channels="BGR")
+        
+        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()