Update app.py

app.py

@@ -1,32 +1,95 @@
 import os
-os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
-os.environ['CUDA_VISIBLE_DEVICES'] = '0'
-
+import cv2
+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 tf_bodypix
 from tf_bodypix.api import download_model, load_model, BodyPixModelPaths
 from tf_bodypix.draw import draw_poses
-from tensorflow.keras import preprocessing
-import cv2
-import json
-from matplotlib import pyplot as plt
-import numpy as np
-from calculations import measure_body_sizes
-import gradio as gr
 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]
+    [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)
+
+def save_image(image, filename):
+    filepath = os.path.join(output_dir, filename)
+    cv2.imwrite(filepath, image)
+    print(f"Image saved as {filepath}")
+
+def draw_bounding_box(frame, person_height_cm, pixels_per_cm):
+    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)
+    bottom_right_x = min(center_x + box_width // 2, frame_width)
+    bottom_right_y = min(center_y + box_height // 2, frame_height)
+    color = (0, 255, 0)
+    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, camera, person_height_cm, pixels_per_cm):
+    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()
+    countdown_thread = threading.Thread(target=play_countdown_audio)
+    countdown_thread.start()
+
+    for i in range(20, 0, -1):
+        ret, frame = camera.read()
+        if not ret:
+            print("Failed to capture frame.")
+            return None
+        draw_bounding_box(frame, person_height_cm, pixels_per_cm)
+        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()
+    if ret:
+        save_image(frame, f"output/{label}_image.jpg")
+    return 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)
@@ -38,18 +101,12 @@ def process_images(front_img, side_img, real_height_cm):
     front_mask = frontresult.get_mask(threshold=0.75)
     side_mask = sideresult.get_mask(threshold=0.75)
 
-    # preprocessing.image.save_img(f'{output_path}/frontbodypix-mask.jpg',front_mask)
-    # preprocessing.image.save_img(f'{output_path}/sidebodypix-mask.jpg',side_mask)
-
     front_colored_mask = frontresult.get_colored_part_mask(front_mask, rainbow)
     side_colored_mask = sideresult.get_colored_part_mask(side_mask, rainbow)
 
-    # preprocessing.image.save_img(f'{output_path}/frontbodypix-colored-mask.jpg',front_colored_mask)
-    # preprocessing.image.save_img(f'{output_path}/sidebodypix-colored-mask.jpg',side_colored_mask)
-
     frontposes = frontresult.get_poses()
     front_image_with_poses = draw_poses(
-        fimage_array.copy(), # create a copy to ensure we are not modifing the source image
+        fimage_array.copy(),
         frontposes,
         keypoints_color=(255, 100, 100),
         skeleton_color=(100, 100, 255)
@@ -57,35 +114,46 @@ def process_images(front_img, side_img, real_height_cm):
 
     sideposes = sideresult.get_poses()
     side_image_with_poses = draw_poses(
-        simage_array.copy(), # create a copy to ensure we are not modifing the source image
+        simage_array.copy(),
         sideposes,
         keypoints_color=(255, 100, 100),
         skeleton_color=(100, 100, 255)
     )
-    # print(np.array(simage).shape)
-    # print(np.array(side_colored_mask).shape)
-
-    # preprocessing.image.save_img(f'{output_path}/frontbodypix-poses.jpg', front_image_with_poses)
-    # preprocessing.image.save_img(f'{output_path}/sidebodypix-poses.jpg', side_image_with_poses)
 
     body_sizes = measure_body_sizes(side_colored_mask, front_colored_mask, sideposes, frontposes, real_height_cm, rainbow)
     measurements_df = pd.DataFrame([body_sizes[0]])
     return measurements_df
 
-# Create the Gradio interface
-interface = gr.Interface(
-    fn=process_images,
-    inputs=[
-        gr.Image(label="Upload Front Pose"),
-        gr.Image(label="Upload Side Pose"),
-        gr.Number(label="Enter Height (cm)")
-    ],
-    outputs=[
-        gr.DataFrame(label="Body Measurements")
-    ],
-    title="Body Sizing System Demo",
-    description="Upload two images: Front View and Side View, and input the height in cm."
-)
-
-# Launch the app
-interface.launch(share=True)
+def gradio_interface(person_height_cm):
+    cap = cv2.VideoCapture(0)
+    if not cap.isOpened():
+        return "Error: Unable to access the camera."
+
+    pixels_per_cm = 5
+    front_frame = capture_image("front", cap, person_height_cm, pixels_per_cm)
+    side_frame = capture_image("side", cap, person_height_cm, pixels_per_cm)
+
+    cap.release()
+    cv2.destroyAllWindows()
+
+    # 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"))
+
+    # 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)
+
+    demo.launch()
+
+if __name__ == "__main__":
+    main()