Update app.py

app.py

@@ -1,4 +1,3 @@
-import gc
 import cv2
 import tempfile
 import numpy as np
@@ -18,61 +17,80 @@ POSE_CONNECTIONS = [
     (30, 32)
 ]
 
+
 @st.cache_resource
 def load_model():
     return tf.saved_model.load("Models/ssd_mobilenet/saved_model")
 
+
 model = load_model()
 mp_pose = mp.solutions.pose
 
 labels = {1: 'person'}
 
+
 def detect_persons(image):
     tensor_img = tf.convert_to_tensor(image)
     tensor_img = tensor_img[tf.newaxis, ...]
+
     detections = model(tensor_img)
+
     boxes = detections['detection_boxes'][0].numpy()
     scores = detections['detection_scores'][0].numpy()
     classes = detections['detection_classes'][0].numpy().astype(np.int32)
+
     return boxes, scores, classes
 
+
 def draw_landmarks(img, landmarks):
     height, width, _ = img.shape
     for lm in landmarks.landmark:
         cx, cy = int(lm.x * width), int(lm.y * height)
         cv2.circle(img, (cx, cy), 8, (0, 0, 255), -1)
+
     for connection in POSE_CONNECTIONS:
         start_idx, end_idx = connection
         if landmarks.landmark[start_idx] and landmarks.landmark[end_idx]:
             start_point = landmarks.landmark[start_idx]
             end_point = landmarks.landmark[end_idx]
+
             start_coordinates = (int(start_point.x * width), int(start_point.y * height))
             end_coordinates = (int(end_point.x * width), int(end_point.y * height))
+
             cv2.line(img, start_coordinates, end_coordinates, (0, 255, 0), 3)
+
     return img
 
+
 def draw_bounding_box(img, box, width, height):
     y_min, x_min, y_max, x_max = box
     left, right, top, bottom = x_min * width, x_max * width, y_min * height, y_max * height
     cv2.rectangle(img, (int(left), int(top)), (int(right), int(bottom)), (255, 0, 0), 2)
 
+
 def process_frame(frame, pose, draw_box):
     image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
     boxes, scores, classes = detect_persons(image_rgb)
+
     height, width, _ = frame.shape
     for i in range(len(scores)):
         if scores[i] > 0.5 and classes[i] == 1:
             y_min, x_min, y_max, x_max = boxes[i]
             left, right, top, bottom = x_min * width, x_max * width, y_min * height, y_max * height
             person_roi = frame[int(top):int(bottom), int(left):int(right)]
+
             results = pose.process(cv2.cvtColor(person_roi, cv2.COLOR_BGR2RGB))
+
             if results.pose_landmarks:
                 person_roi = draw_landmarks(person_roi, results.pose_landmarks)
+
             frame[int(top):int(bottom), int(left):int(right)] = person_roi
             if draw_box:
                 draw_bounding_box(frame, boxes[i], width, height)
+
     return frame
 
+
 def main():
     st.markdown(
         """
@@ -153,9 +171,7 @@ def main():
                     frame = process_frame(frame, pose, draw_box)
 
                     st_frame.image(frame, channels='BGR', use_column_width=True)
-
-                    # Ensure proper synchronization and frame display
-                    st_frame.empty()
+                    st.empty()
 
                 st.text("Completed")
                 cam.release()
@@ -166,8 +182,6 @@ def main():
 
                 st.image(processed_image, channels='BGR', use_column_width=True)
 
-            gc.collect()
-
     elif operation_type == "Demo":
         st.empty()
         st.markdown("<p class='intro'>Demo video will be shown below:</p>", unsafe_allow_html=True)
@@ -185,11 +199,11 @@ def main():
             frame = process_frame(frame, pose, draw_box=False)
 
             st_frame.image(frame, channels='BGR', use_column_width=True)
-            st_frame.empty()
+            st.empty()
 
         st.text("Completed")
         cam.release()
-        gc.collect()
+
 
 if __name__ == "__main__":
     main()