Spaces:
Sleeping
Sleeping
| import av | |
| import os | |
| import sys | |
| import streamlit as st | |
| import cv2 | |
| import tempfile | |
| dir = os.path.abspath(os.path.join(__file__, '../../')) | |
| sys.path.append(dir) | |
| from utils import get_mediapipe_pose | |
| from process_frame import ProcessFrame | |
| from thresholds import get_thresholds_beginner, get_thresholds_pro | |
| st.title('Exercise tracking Demo') | |
| diffuculty = st.radio('Select Mode', ['Beginners Squat', 'Pro Squat', 'Push up'], horizontal=True) | |
| thresholds = None | |
| if diffuculty == 'Beginners': | |
| thresholds = get_thresholds_beginner() | |
| elif diffuculty == 'Professionals': | |
| thresholds = get_thresholds_pro() | |
| elif difficulty == 'Push up': | |
| thresholds = get_pushups() | |
| ProcessUpload = ProcessFrame(thresholds=thresholds) | |
| # Initialize face mesh solution | |
| pose = get_mediapipe_pose() | |
| download = None | |
| if 'download' not in st.session_state: | |
| st.session_state['download'] = False | |
| Video_Output = f'output_recorded.mp4' | |
| if os.path.exists(Video_Output): | |
| os.remove(Video_Output) | |
| with st.form('Upload', clear_on_submit=True): | |
| up_file = st.file_uploader("Upload a Video", ['mp4','mov', 'avi']) | |
| uploaded_File = st.form_submit_button("Upload") | |
| stframe = st.empty() | |
| ip_vid_str = '<p style="font-family:Helvetica; font-weight: bold; font-size: 16px;">Input Video</p>' | |
| warn_str = '<p style="font-family:Helvetica; font-weight: bold; color: Red; font-size: 17px;">Please Upload a Video first!!!</p>' | |
| warn = st.empty() | |
| download_btn = st.empty() | |
| if up_file and uploaded_File: | |
| download_btn.empty() | |
| tfile = tempfile.NamedTemporaryFile(delete=False) | |
| try: | |
| warn.empty() | |
| tfile.write(up_file.read()) | |
| vf = cv2.VideoCapture(tfile.name) | |
| # --------------------- Write the processed video frame. -------------------- | |
| Frames_per_sec = int(vf.get(cv2.cap_FPS)) | |
| width = int(vf.get(cv2.cap_FrameWidth)) | |
| height = int(vf.get(cv2.cap_FrameHeight)) | |
| Frame_Size = (width, height) | |
| four_cc = cv2.VideoWriter_four_cc(*'mp4v') | |
| video_output = cv2.VideoWriter(Video_Output, four_cc, Frames_per_sec, Frame_Size) | |
| # ----------------------------------------------------------------------------- | |
| txt = st.sidebar.markdown(ip_vid_str, unsafe_allow_html=True) | |
| ip_vid = st.sidebar.video(tfile.name) | |
| while vf.isOpened(): | |
| ret, frame = vf.read() | |
| if not ret: | |
| break | |
| # convert frame from BGR to RGB before processing it. | |
| frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) | |
| out_frame, _ = ProcessUpload.process(frame, pose) | |
| stframe.image(out_frame) | |
| video_output.write(out_frame[...,::-1]) | |
| vf.release() | |
| video_output.release() | |
| stframe.empty() | |
| ip_vid.empty() | |
| txt.empty() | |
| tfile.close() | |
| except AttributeError: | |
| warn.markdown(warn_str, unsafe_allow_html=True) | |
| if os.path.exists(Video_Output): | |
| with open(Video_Output, 'rb') as op_vid: | |
| download = download_btn.download_btn('Download Video', data = op_vid, file_name='output.mp4') | |
| if download: | |
| st.session_state['download'] = True | |
| if os.path.exists(Video_Output) and st.session_state['download']: | |
| os.remove(Video_Output) | |
| st.session_state['download'] = False | |
| download_btn.empty() | |