Spaces:
Sleeping
Sleeping
File size: 2,372 Bytes
44c2095 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
import av
import os
import sys
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)
from utils import get_mediapipe_pose
from process_frame import ProcessFrame
from thresholds import get_thresholds_beginner, get_thresholds_pro
st.title('AI Fitness Trainer: Squats Analysis')
mode = st.radio('Select Mode', ['Beginner', 'Pro'], horizontal=True)
thresholds = None
if mode == 'Beginner':
thresholds = get_thresholds_beginner()
elif mode == 'Pro':
thresholds = get_thresholds_pro()
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
output_video_file = f'output_live.flv'
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
def out_recorder_factory() -> MediaRecorder:
return MediaRecorder(output_video_file)
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": {"width": {'min':480, 'ideal':480}}, "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()
|