Update app.py

app.py

@@ -36,74 +36,62 @@ def build_model(HIDDEN_UNITS=256, sequence_length=30, num_input_values=33*4, num
     model.load_weights(load_dir)
     return model
 
+# Define the VideoProcessor class for real-time video processing
 class VideoProcessor:
     def __init__(self):
         self.actions = np.array(['curl', 'press', 'squat'])
         self.sequence_length = 30
         self.colors = [(245,117,16), (117,245,16), (16,117,245)]
-        self.threshold = 0.50  # Default threshold for activity classification confidence
-        
-        # Detection variables
-        self.sequence = []
-        self.current_action = ''
-        
-        # Initialize pose model
-        self.mp_pose = mp.solutions.pose
-        self.mp_drawing = mp.solutions.drawing_utils
-        self.pose = self.mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5)
-        self.model = build_model()  # Load the LSTM model
+        self.pose = mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5)
+        self.model = build_model()
+    
+    def process_video(self, video_file):
+        # Get the filename from the file object
+        filename = video_file.name
+        # Create a temporary file to write the contents of the uploaded video file
+        temp_file = open(filename, 'wb')
+        temp_file.write(video_file.read())
+        temp_file.close()
+        # Now we can open the video file using cv2.VideoCapture()
+        cap = cv2.VideoCapture(filename)
+        out_frames = []
+        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)
+            frame = self.draw_landmarks(frame, results)
+            out_frames.append(frame)
+        cap.release()
+        # Remove the temporary file
+        os.remove(filename)
+        return out_frames
     
-    @st.cache()
     def draw_landmarks(self, image, results):
-        """
-        This function draws keypoints and landmarks detected by the human pose estimation model
-        
-        """
-        self.mp_drawing.draw_landmarks(image, results.pose_landmarks, self.mp_pose.POSE_CONNECTIONS,
-                                        self.mp_drawing.DrawingSpec(color=(245,117,66), thickness=2, circle_radius=2), 
-                                        self.mp_drawing.DrawingSpec(color=(245,66,230), thickness=2, circle_radius=2) 
-                                        )
+        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
-    
+
     @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
     
     @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
-        
         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 
     
     @st.cache()
     def get_coordinates(self, landmarks, 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
-            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(self.mp_pose.PoseLandmark, side.upper() + "_" + joint.upper())
         x_coord_val = landmarks[coord.value].x
         y_coord_val = landmarks[coord.value].y
@@ -111,100 +99,12 @@ class VideoProcessor:
     
     @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()
-    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]
-            # 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 = ''
-
-            # 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 process_video(self, video_file):
-        # Get the filename from the file object
-        filename = video_file.name
-        # Create a temporary file to write the contents of the uploaded video file
-        temp_file = open(filename, 'wb')
-        temp_file.write(video_file.read())
-        temp_file.close()
-        # Now we can open the video file using cv2.VideoCapture()
-        cap = cv2.VideoCapture(filename)
-        out_frames = []
-        while cap.isOpened():
-            ret, frame = cap.read()
-            if not ret:
-                break
-            frame_processed = self.process(frame)
-            out_frames.append(frame_processed)
-        cap.release()
-        # Remove the temporary file
-        os.remove(filename)
-        return out_frames
-
-
 # Define Streamlit app
 def main():
     st.title("Real-time Exercise Detection")