Spaces:
Sleeping
Sleeping
File size: 3,350 Bytes
c5eaec9 |
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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
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()
|