Update app.py

app.py

@@ -1,9 +1,9 @@
 import os
 import cv2
 import time
-import threading
 import numpy as np
 import gradio as gr
+import threading
 from tensorflow.keras import preprocessing
 import pandas as pd
 from calculations import measure_body_sizes
@@ -22,106 +22,118 @@ def save_image(image, filename):
     cv2.imwrite(filepath, image)
     print(f"Image saved as {filepath}")
 
-def draw_bounding_box(frame, box_width, box_height):
-    frame_height, frame_width, _ = frame.shape
-    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_countdown_audio():
-    for i in range(20, 0, -1):
-        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.")
-
-    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 = cap.read()
-        if not ret:
-            print("Failed to capture frame.")
-            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)
-        cv2.imshow("Pose Capture", frame)
-        if cv2.waitKey(1000) & 0xFF == ord('q'):
-            break
-
-    countdown_thread.join()
-    ret, final_frame = cap.read()
-    if ret:
-        filename = f"{label}_image.jpg"
-        save_image(final_frame, filename)
+def process_webcam_image(webcam_image, label, real_height_cm):
+    """Process a webcam image"""
+    if webcam_image is None:
+        raise gr.Error(f"No {label} image captured")
+    
+    # Convert Gradio image to NumPy array
+    frame = webcam_image
 
-    cap.release()
-    cv2.destroyAllWindows()
-    return final_frame
+    filename = f"{label}_image.jpg"
+    save_image(frame, filename)
+    return frame
 
 def process_images(front_img, side_img, real_height_cm):
-    # Ensure images are not None
+    """Process front and side images to extract body measurements"""
+    # Validate inputs
     if front_img is None or side_img is None:
-        raise ValueError("Front or side image is None")
+        raise gr.Error("Both front and side images are required")
 
+    # Convert images to array
     front_image_array = preprocessing.image.img_to_array(front_img)
     side_image_array = preprocessing.image.img_to_array(side_img)
 
+    # Predict body parts
     front_result = bodypix_model.predict_single(front_image_array)
     side_result = bodypix_model.predict_single(side_image_array)
 
+    # Create masks
     front_mask = front_result.get_mask(threshold=0.75)
     side_mask = side_result.get_mask(threshold=0.75)
 
+    # Color the masks
     front_colored_mask = front_result.get_colored_part_mask(front_mask)
     side_colored_mask = side_result.get_colored_part_mask(side_mask)
 
+    # Get poses
     front_poses = front_result.get_poses()
     side_poses = side_result.get_poses()
 
-    # 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)
+    # Calculate body sizes
+    body_sizes = measure_body_sizes(
+        side_colored_mask, 
+        front_colored_mask, 
+        side_poses, 
+        front_poses, 
+        real_height_cm, 
+        rainbow=None
+    )
+    
+    # Convert to DataFrame
     measurements_df = pd.DataFrame(body_sizes)
     return measurements_df
 
-def gradio_interface(person_height_cm):
-    pixels_per_cm = 5
-
-    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)
-
-    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_frame, side_frame, person_height_cm)
-    return measurements_df
-
-def main():
+def create_gradio_interface():
+    """Create Gradio interface for body measurements"""
     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()
+        gr.Markdown("# Body Measurement Tool")
+        
+        with gr.Row():
+            height_input = gr.Number(label="Your Height (cm)")
+            
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("### Front View")
+                front_webcam = gr.Image(source="webcam", type="numpy", label="Front View")
+                front_capture_button = gr.Button("Capture Front View")
+                
+            with gr.Column():
+                gr.Markdown("### Side View")
+                side_webcam = gr.Image(source="webcam", type="numpy", label="Side View")
+                side_capture_button = gr.Button("Capture Side View")
+        
+        # Capture buttons functionality
+        current_front_image = gr.State(None)
+        current_side_image = gr.State(None)
+        
+        def capture_front_image(webcam_input):
+            if webcam_input is None:
+                raise gr.Error("No image in webcam")
+            return webcam_input, webcam_input
+        
+        def capture_side_image(webcam_input):
+            if webcam_input is None:
+                raise gr.Error("No image in webcam")
+            return webcam_input, webcam_input
+        
+        front_capture_button.click(
+            fn=capture_front_image, 
+            inputs=front_webcam, 
+            outputs=[current_front_image, front_webcam]
+        )
+        
+        side_capture_button.click(
+            fn=capture_side_image, 
+            inputs=side_webcam, 
+            outputs=[current_side_image, side_webcam]
+        )
+        
+        # Measurement calculation button
+        calculate_button = gr.Button("Calculate Measurements")
+        output_df = gr.DataFrame(label="Body Measurements")
+        
+        calculate_button.click(
+            fn=lambda height, front, side: process_images(front, side, height),
+            inputs=[height_input, current_front_image, current_side_image],
+            outputs=output_df
+        )
+    
+    return demo
 
-        capture_button.click(fn=gradio_interface, inputs=height_input, outputs=output_df)
-
-    demo.launch()
+def main():
+    demo = create_gradio_interface()
+    demo.launch(share=True, debug=True)
 
 if __name__ == "__main__":
     main()