Update app.py

app.py

@@ -9,7 +9,6 @@ import pandas as pd
 from calculations import measure_body_sizes
 import tf_bodypix
 from tf_bodypix.api import download_model, load_model, BodyPixModelPaths
-from tf_bodypix.draw import draw_poses
 
 # Load BodyPix model
 bodypix_model = load_model(download_model(BodyPixModelPaths.MOBILENET_FLOAT_50_STRIDE_16))
@@ -59,7 +58,7 @@ def capture_image_with_timer(label, person_height_cm, pixels_per_cm):
         ret, frame = cap.read()
         if not ret:
             print("Failed to capture frame.")
-            break
+            return None
         draw_bounding_box(frame, box_width, box_height)
         cv2.putText(frame, f"{label.capitalize()} pose: {i} seconds remaining", (10, 30),
                     cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2)
@@ -78,6 +77,10 @@ def capture_image_with_timer(label, person_height_cm, pixels_per_cm):
     return final_frame
 
 def process_images(front_img, side_img, real_height_cm):
+    # Ensure images are not None
+    if front_img is None or side_img is None:
+        raise ValueError("Front or side image is None")
+
     front_image_array = preprocessing.image.img_to_array(front_img)
     side_image_array = preprocessing.image.img_to_array(side_img)
 
@@ -93,8 +96,9 @@ def process_images(front_img, side_img, real_height_cm):
     front_poses = front_result.get_poses()
     side_poses = side_result.get_poses()
 
-    body_sizes = measure_body_sizes(side_colored_mask, front_colored_mask, side_poses, front_poses, real_height_cm)
-    measurements_df = pd.DataFrame([body_sizes[0]])
+    # Added a placeholder for the rainbow parameter
+    body_sizes = measure_body_sizes(side_colored_mask, front_colored_mask, side_poses, front_poses, real_height_cm, rainbow=None)
+    measurements_df = pd.DataFrame(body_sizes)
     return measurements_df
 
 def gradio_interface(person_height_cm):
@@ -103,13 +107,10 @@ def gradio_interface(person_height_cm):
     front_frame = capture_image_with_timer("front", person_height_cm, pixels_per_cm)
     side_frame = capture_image_with_timer("side", person_height_cm, pixels_per_cm)
 
-    front_img_path = os.path.join(output_dir, "front_image.jpg")
-    side_img_path = os.path.join(output_dir, "side_image.jpg")
-
-    front_img = cv2.imread(front_img_path)
-    side_img = cv2.imread(side_img_path)
+    if front_frame is None or side_frame is None:
+        return pd.DataFrame()  # Return empty DataFrame if image capture fails
 
-    measurements_df = process_images(front_img, side_img, person_height_cm)
+    measurements_df = process_images(front_frame, side_frame, person_height_cm)
     return measurements_df
 
 def main():
@@ -123,4 +124,4 @@ def main():
     demo.launch()
 
 if __name__ == "__main__":
-    main()
+    main()