Update interface_pages/yoga_position_from_stream.py

interface_pages/yoga_position_from_stream.py

@@ -6,6 +6,7 @@ from mediapipe.python.solutions import drawing_utils as mp_drawing
 from PoseClassification.pose_embedding import FullBodyPoseEmbedding
 from PoseClassification.pose_classifier import PoseClassifier
 from PoseClassification.utils import EMADictSmoothing
+import time
 
 # Initialize components
 mp_pose = mp.solutions.pose
@@ -61,9 +62,11 @@ def yoga_position_from_stream():
     last_update_time = 0
     recording = False
     recorded_frames = []
+    start_time = 0
+    frame_count = 0
 
     def classify_pose(frame):
-        nonlocal current_position, position_timer, last_update_time, recording, recorded_frames
+        nonlocal current_position, position_timer, last_update_time, recording, recorded_frames, start_time, frame_count
         if frame is None:
             return (
                 None,
@@ -112,6 +115,9 @@ def yoga_position_from_stream():
 
         if recording:
             recorded_frames.append(processed_frame)
+            frame_count += 1
+            if frame_count == 1:
+                start_time = time.time()
 
         return (
             frame,
@@ -128,9 +134,11 @@ def yoga_position_from_stream():
         ]
 
     def start_recording():
-        nonlocal recording, recorded_frames
+        nonlocal recording, recorded_frames, start_time, frame_count
         recording = True
         recorded_frames = []
+        start_time = 0
+        frame_count = 0
         return "Recording started"
 
     def stop_recording():
@@ -139,46 +147,76 @@ def yoga_position_from_stream():
         return "Recording stopped"
 
     def save_video():
-        nonlocal recorded_frames
+        nonlocal recorded_frames, start_time, frame_count
         if not recorded_frames:
             return None, "No recorded frames available"
 
         output_path = "recorded_yoga_session.mp4"
         height, width, _ = recorded_frames[0].shape
+
+        # Calculate the actual frame rate
+        elapsed_time = time.time() - start_time
+        fps = frame_count / elapsed_time if elapsed_time > 0 else 30.0
+
         fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-        out = cv2.VideoWriter(output_path, fourcc, 30.0, (width, height))
+        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
 
         for frame in recorded_frames:
-            out.write(frame)
+            # Convert frame to BGR color space before writing
+            frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
+            out.write(frame_bgr)
         out.release()
 
-        return output_path, "Video saved successfully"
+        return output_path, f"Video saved successfully at {fps:.2f} FPS"
 
     with gr.Column() as yoga_stream:
-        gr.Markdown("# Yoga Position Classifier")
-        gr.Markdown("Stream live yoga sessions and get real-time pose classification.")
-
-        debug_toggle = gr.Checkbox(label="Debug Mode", value=False)
+        gr.Markdown("# Yoga Position Classifier", elem_classes=["custom-title"])
+        gr.Markdown(
+            "Stream live yoga sessions and get real-time pose classification.",
+            elem_classes=["custom-subtitle"],
+        )
 
-        with gr.Column(visible=True) as normal_view:
-            video_feed = gr.Webcam(streaming=True, elem_classes="fullscreen")
-            pose_output = gr.Textbox(label="Current Pose")
-            timer_output = gr.Textbox(label="Pose Duration")
+        with gr.Row():
+            with gr.Column(scale=3):
+                video_feed = gr.Webcam(streaming=True, elem_classes=["custom-webcam"])
+
+            with gr.Column(scale=2):
+                pose_output = gr.Textbox(
+                    label="Current Pose", elem_classes=["custom-textbox"]
+                )
+                timer_output = gr.Textbox(
+                    label="Pose Duration", elem_classes=["custom-textbox"]
+                )
+                debug_toggle = gr.Checkbox(
+                    label="Debug Mode", value=False, elem_classes=["custom-checkbox"]
+                )
 
         with gr.Column(visible=False) as debug_view:
-            classified_video = gr.Image(label="Classified Video Feed")
+            classified_video = gr.Image(
+                label="Classified Video Feed", elem_classes=["custom-image"]
+            )
             with gr.Row():
-                start_button = gr.Button("Start Recording")
-                stop_button = gr.Button("Stop Recording")
-            save_button = gr.Button("Save Recording")
-            recording_status = gr.Textbox(label="Recording Status")
-            recorded_video = gr.Video(label="Recorded Video")
-            download_button = gr.Button("Download Recorded Video")
+                start_button = gr.Button(
+                    "Start Recording", elem_classes=["custom-button"]
+                )
+                stop_button = gr.Button(
+                    "Stop Recording", elem_classes=["custom-button"]
+                )
+            save_button = gr.Button("Save Recording", elem_classes=["custom-button"])
+            recording_status = gr.Textbox(
+                label="Recording Status", elem_classes=["custom-textbox"]
+            )
+            recorded_video = gr.Video(
+                label="Recorded Video", elem_classes=["custom-video"]
+            )
+            download_button = gr.Button(
+                "Download Recorded Video", elem_classes=["custom-button"]
+            )
 
         debug_toggle.change(
             toggle_debug,
             inputs=[debug_toggle],
-            outputs=[debug_view, normal_view, classified_video],
+            outputs=[debug_view, video_feed, classified_video],
         )
 
         video_feed.stream(
@@ -197,6 +235,17 @@ def yoga_position_from_stream():
 
 
 if __name__ == "__main__":
-    with gr.Blocks() as demo:
+    with gr.Blocks(
+        css="""
+        .custom-title { font-size: 36px; font-weight: bold; margin-bottom: 10px; }
+        .custom-subtitle { font-size: 18px; margin-bottom: 20px; }
+        .custom-webcam { height: 480px; }
+        .custom-textbox input { font-size: 24px; }
+        .custom-checkbox label { font-size: 18px; }
+        .custom-button { font-size: 18px; }
+        .custom-image img { max-height: 400px; }
+        .custom-video video { max-height: 400px; }
+        """
+    ) as demo:
         yoga_position_from_stream()
     demo.launch()