Update pages/1_Live_Stream.py

pages/1_Live_Stream.py

@@ -1,12 +1,7 @@
-import av
-import os
-import sys
+import cv2
 import streamlit as st
-from streamlit_webrtc import VideoHTMLAttributes, webrtc_streamer
-from aiortc.contrib.media import MediaRecorder
-
-BASE_DIR = os.path.abspath(os.path.join(__file__, '../../'))
-sys.path.append(BASE_DIR)
+import tempfile
+import os
 
 from utils import get_mediapipe_pose
 from process_frame import ProcessFrame
@@ -14,62 +9,65 @@ from thresholds import get_thresholds_beginner, get_thresholds_pro
 
 st.title('AI Fitness Trainer: Squats Analysis')
 
+# Select mode
 mode = st.radio('Select Mode', ['Beginner', 'Pro'], horizontal=True)
 
-thresholds = None 
+thresholds = None
 if mode == 'Beginner':
     thresholds = get_thresholds_beginner()
 elif mode == 'Pro':
     thresholds = get_thresholds_pro()
 
+# Initialize pose estimation and processing
 live_process_frame = ProcessFrame(thresholds=thresholds, flip_frame=True)
-
-# Initialize face mesh solution
 pose = get_mediapipe_pose()
 
-if 'download' not in st.session_state:
-    st.session_state['download'] = False
+# Temporary file to save processed video
+temp_dir = tempfile.TemporaryDirectory()
+output_video_file = os.path.join(temp_dir.name, 'output_live.avi')
 
-output_video_file = f'output_live.flv'
+# Initialize video writer
+fourcc = cv2.VideoWriter_fourcc(*'XVID')
+video_writer = None
 
-# RTSP stream URL (replace with your RTSP stream URL)
-RTSP_URL = st.text_input("Enter RTSP Stream URL", "rtsp://your-stream-url")
+# Start video capture
+cap = cv2.VideoCapture(0)
+if not cap.isOpened():
+    st.error("Unable to access the camera")
+else:
+    stframe = st.empty()
 
-def video_frame_callback(frame: av.VideoFrame):
-    frame = frame.to_ndarray(format="rgb24")  # Decode and get RGB frame
-    frame, _ = live_process_frame.process(frame, pose)  # Process frame
-    return av.VideoFrame.from_ndarray(frame, format="rgb24")  # Encode and return BGR frame
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            st.error("Failed to capture video")
+            break
 
-def out_recorder_factory() -> MediaRecorder:
-        return MediaRecorder(output_video_file)
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        processed_frame, _ = live_process_frame.process(frame, pose)
+
+        # Write to video file
+        if video_writer is None:
+            height, width, _ = processed_frame.shape
+            video_writer = cv2.VideoWriter(output_video_file, fourcc, 30, (width, height))
+        video_writer.write(cv2.cvtColor(processed_frame, cv2.COLOR_RGB2BGR))
+
+        # Display in Streamlit
+        stframe.image(processed_frame, channels="RGB")
+
+        # Stop button
+        if st.button("Stop"):
+            break
+
+    cap.release()
+    if video_writer:
+        video_writer.release()
+
+# Download processed video
+if os.path.exists(output_video_file):
+    with open(output_video_file, 'rb') as f:
+        st.download_button("Download Processed Video", f, file_name="output_live.avi")
 
-if RTSP_URL:
-    ctx = webrtc_streamer(
-        key="Squats-pose-analysis",
-        video_frame_callback=video_frame_callback,
-        rtc_configuration={"iceServers": [{"urls": "stun:stun.l.google.com:19302"}]},  # Add this config
-        media_stream_constraints={
-            "video": {"frameRate": 30, "deviceId": RTSP_URL},
-            "audio": False
-        },
-        video_html_attrs=VideoHTMLAttributes(autoPlay=True, controls=False, muted=False),
-        out_recorder_factory=out_recorder_factory
-    )
-
-    download_button = st.empty()
-
-    if os.path.exists(output_video_file):
-        with open(output_video_file, 'rb') as op_vid:
-            download = download_button.download_button('Download Video', data=op_vid, file_name='output_live.flv')
-            if download:
-                st.session_state['download'] = True
-
-    if os.path.exists(output_video_file) and st.session_state['download']:
-        os.remove(output_video_file)
-        st.session_state['download'] = False
-        download_button.empty()
-else:
-    st.warning("Please provide a valid RTSP stream URL.")
 
 
 '''