Update app.py

app.py

@@ -4,28 +4,16 @@ import time
 import threading
 import numpy as np
 import gradio as gr
-from gtts import gTTS
-import pygame
 from tensorflow.keras import preprocessing
-import tensorflow as tf
+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
-import pandas as pd
-from calculations import measure_body_sizes
 
 # Load BodyPix model
 bodypix_model = load_model(download_model(BodyPixModelPaths.MOBILENET_FLOAT_50_STRIDE_16))
 
-rainbow = [
-    [110, 64, 170], [143, 61, 178], [178, 60, 178], [210, 62, 167],
-    [238, 67, 149], [255, 78, 125], [255, 94, 99],  [255, 115, 75],
-    [255, 140, 56], [239, 167, 47], [217, 194, 49], [194, 219, 64],
-    [175, 240, 91], [135, 245, 87], [96, 247, 96],  [64, 243, 115],
-    [40, 234, 141], [28, 219, 169], [26, 199, 194], [33, 176, 213],
-    [47, 150, 224], [65, 125, 224], [84, 101, 214], [99, 81, 195]
-]
-
 # Create output directory
 output_dir = "output_images"
 os.makedirs(output_dir, exist_ok=True)
@@ -35,10 +23,8 @@ def save_image(image, filename):
     cv2.imwrite(filepath, image)
     print(f"Image saved as {filepath}")
 
-def draw_bounding_box(frame, person_height_cm, pixels_per_cm):
+def draw_bounding_box(frame, box_width, box_height):
     frame_height, frame_width, _ = frame.shape
-    box_height = int(2 * person_height_cm * pixels_per_cm)
-    box_width = int(2 * person_height_cm * pixels_per_cm)
     center_x, center_y = frame_width // 2, frame_height // 2
     top_left_x = max(center_x - box_width // 2, 0)
     top_left_y = max(center_y - box_height // 2, 0)
@@ -48,105 +34,88 @@ def draw_bounding_box(frame, person_height_cm, pixels_per_cm):
     thickness = 2
     cv2.rectangle(frame, (top_left_x, top_left_y), (bottom_right_x, bottom_right_y), color, thickness)
 
-def play_audio_instruction(text, filename):
-    tts = gTTS(text=text, lang='en')
-    tts.save(filename)
-    pygame.mixer.init()
-    pygame.mixer.music.load(filename)
-    pygame.mixer.music.play()
-
 def play_countdown_audio():
     for i in range(20, 0, -1):
-        tts = gTTS(text=str(i), lang='en')
-        filename = f"input1/countdown_{i}.mp3"
-        tts.save(filename)
-        pygame.mixer.music.load(filename)
-        pygame.mixer.music.play()
-        while pygame.mixer.music.get_busy():
-            time.sleep(0.1)
-
-def capture_image(label, person_height_cm, pixels_per_cm):
-    camera = gr.Button("Capture Images")
+        audio_file = f"input1/countdown_{i}.mp3"
+        if os.path.exists(audio_file):
+            os.system(f"mpg123 {audio_file}")
+        time.sleep(1)
+
+def capture_image_with_timer(label, person_height_cm, pixels_per_cm):
+    cap = cv2.VideoCapture(0)
+    if not cap.isOpened():
+        print("Error: Unable to access the camera.")
+        return None
+
     print(f"Please position yourself for the {label} pose.")
-    instruction_text = f"Please stand straight in the box for the {label} pose."
-    audio_file = f"input1/instruction_{label}.mp3"
-    threading.Thread(target=play_audio_instruction, args=(instruction_text, audio_file)).start()
+
+    box_height = int(2 * person_height_cm * pixels_per_cm)
+    box_width = int(2 * person_height_cm * pixels_per_cm)
+
     countdown_thread = threading.Thread(target=play_countdown_audio)
     countdown_thread.start()
 
     for i in range(20, 0, -1):
-        ret, frame = camera.read()
+        ret, frame = cap.read()
         if not ret:
             print("Failed to capture frame.")
-            return None
-        draw_bounding_box(frame, person_height_cm, pixels_per_cm)
+            break
+        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)
         cv2.imshow("Pose Capture", frame)
         if cv2.waitKey(1000) & 0xFF == ord('q'):
             break
+
     countdown_thread.join()
-    ret, frame = camera.read()
+    ret, final_frame = cap.read()
     if ret:
-        save_image(frame, f"output/{label}_image.jpg")
-    return frame
+        filename = f"{label}_image.jpg"
+        save_image(final_frame, filename)
+
+    cap.release()
+    cv2.destroyAllWindows()
+    return final_frame
 
 def process_images(front_img, side_img, real_height_cm):
-    fimage_array = preprocessing.image.img_to_array(front_img)
-    simage_array = preprocessing.image.img_to_array(side_img)
-
-    # bodypix prediction
-    frontresult = bodypix_model.predict_single(fimage_array)
-    sideresult = bodypix_model.predict_single(simage_array)
-
-    front_mask = frontresult.get_mask(threshold=0.75)
-    side_mask = sideresult.get_mask(threshold=0.75)
-
-    front_colored_mask = frontresult.get_colored_part_mask(front_mask, rainbow)
-    side_colored_mask = sideresult.get_colored_part_mask(side_mask, rainbow)
-
-    frontposes = frontresult.get_poses()
-    front_image_with_poses = draw_poses(
-        fimage_array.copy(),
-        frontposes,
-        keypoints_color=(255, 100, 100),
-        skeleton_color=(100, 100, 255)
-    )
-
-    sideposes = sideresult.get_poses()
-    side_image_with_poses = draw_poses(
-        simage_array.copy(),
-        sideposes,
-        keypoints_color=(255, 100, 100),
-        skeleton_color=(100, 100, 255)
-    )
-
-    body_sizes = measure_body_sizes(side_colored_mask, front_colored_mask, sideposes, frontposes, real_height_cm, rainbow)
+    front_image_array = preprocessing.image.img_to_array(front_img)
+    side_image_array = preprocessing.image.img_to_array(side_img)
+
+    front_result = bodypix_model.predict_single(front_image_array)
+    side_result = bodypix_model.predict_single(side_image_array)
+
+    front_mask = front_result.get_mask(threshold=0.75)
+    side_mask = side_result.get_mask(threshold=0.75)
+
+    front_colored_mask = front_result.get_colored_part_mask(front_mask)
+    side_colored_mask = side_result.get_colored_part_mask(side_mask)
+
+    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]])
     return measurements_df
 
 def gradio_interface(person_height_cm):
-    
-#    if not cap.isOpened():
-#        return "Error: Unable to access the camera."
-
     pixels_per_cm = 5
-    front_frame = capture_image("front", person_height_cm, pixels_per_cm)
-    side_frame = capture_image("side", person_height_cm, pixels_per_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")
 
-    # Load the captured images to process for body measurement
-    front_img = cv2.imread(os.path.join(output_dir, "output/front_image.jpg"))
-    side_img = cv2.imread(os.path.join(output_dir, "output/side_image.jpg"))
+    front_img = cv2.imread(front_img_path)
+    side_img = cv2.imread(side_img_path)
 
-    # Process images and return measurement DataFrame
     measurements_df = process_images(front_img, side_img, person_height_cm)
     return measurements_df
 
 def main():
-    # Updated Gradio interface with separate button trigger
     with gr.Blocks() as demo:
         height_input = gr.Number(label="Enter your height (in cm)")
+        capture_button = gr.Button("Capture Images")
         output_df = gr.DataFrame()
 
         capture_button.click(fn=gradio_interface, inputs=height_input, outputs=output_df)