Update app.py

app.py

@@ -11,6 +11,9 @@ import math
 
 from streamlit_webrtc import webrtc_streamer, WebRtcMode, RTCConfiguration
 import av
+from io import BytesIO
+import av
+from PIL import Image
 
 ## Build and Load Model
 def attention_block(inputs, time_steps):
@@ -286,87 +289,148 @@ class VideoProcessor:
             
         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
+    # @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
+    #     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)
+    #     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) 
+    #     # 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:]
+    #     # 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'])
+    #     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)
+    #         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 = ''
+    #         # 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)
+    #         # 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
+    #         # 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)
+    #         # 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
+    #     # return cv2.flip(image, 1)
+    #     return image
     
-    def recv(self, frame):
-        """
-        Receive and process video stream from webcam
+    # 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")
+    def process_uploaded_file(self, file):
+        """
+        Function to process an uploaded image or video file and run the fitness trainer AI
         Args:
-            frame: current video frame
+            file (BytesIO): uploaded image or video file
+        Returns:
+            numpy array: processed image with keypoint detection and fitness activity classification visualized
+        """
+        # Initialize an empty list to store processed frames
+        processed_frames = []
+
+        # Check if the uploaded file is a video
+        is_video = hasattr(file, 'name') and file.name.endswith(('.mp4', '.avi', '.mov'))
 
+        if is_video:
+            container = av.open(file)
+            for frame in container.decode(video=0):
+                # Convert the frame to OpenCV format
+                image = frame.to_image().convert("RGB")
+                image = np.array(image)
+                
+                # Process the frame
+                processed_frame = self.process(image)
+                
+                # Append the processed frame to the list
+                processed_frames.append(processed_frame)
+            
+            # Close the video file container
+            container.close()
+        else:
+            # If the uploaded file is an image
+            # Load the image from the BytesIO object
+            image = Image.open(file)
+            image = np.array(image)
+            
+            # Process the image
+            processed_frame = self.process(image)
+            
+            # Append the processed frame to the list
+            processed_frames.append(processed_frame)
+        
+        return processed_frames
+
+    def recv_uploaded_file(self, file):
+        """
+        Receive and process an uploaded video file
+        Args:
+            file (BytesIO): uploaded video file
         Returns:
-            av.VideoFrame: processed video frame
+            List[av.VideoFrame]: list of processed video frames
         """
-        img = frame.to_ndarray(format="bgr24")
-        img = self.process(img)
-        return av.VideoFrame.from_ndarray(img, format="bgr24")
+        # Process the uploaded file
+        processed_frames = self.process_uploaded_file(file)
+        
+        # Convert processed frames to av.VideoFrame objects
+        av_frames = []
+        for frame in processed_frames:
+            av_frame = av.VideoFrame.from_ndarray(frame, format="bgr24")
+            av_frames.append(av_frame)
+        
+        return av_frames
         
-## 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",