Update app.py

app.py

@@ -5,27 +5,32 @@ import streamlit as st
 import cv2
 import tempfile
 
-dir = os.path.abspath(os.path.join(__file__, '../../'))
-sys.path.append(dir)
+
+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('Exercise tracking Demo')
 
-diffuculty = st.radio('Select Mode', ['Beginners Squat', 'Pro Squat', 'Push up'], horizontal=True)
+
+st.title('AI Fitness Trainer: Squats Analysis')
+
+mode = st.radio('Select Mode', ['Beginner', 'Pro'], horizontal=True)
+
+
 
 thresholds = None 
 
-if diffuculty == 'Beginners':
+if mode == 'Beginner':
     thresholds = get_thresholds_beginner()
 
-elif diffuculty == 'Professionals':
+elif mode == 'Pro':
     thresholds = get_thresholds_pro()
 
-#elif difficulty == 'Push up':
-    #thresholds = get_pushups()
+
 
 upload_process_frame = ProcessFrame(thresholds=thresholds)
 
@@ -39,29 +44,29 @@ if 'download' not in st.session_state:
     st.session_state['download'] = False
 
 
-Video_Output = f'output_recorded.mp4'
+output_video_file = f'output_recorded.mp4'
 
-if os.path.exists(Video_Output):
-    os.remove(Video_Output)
+if os.path.exists(output_video_file):
+    os.remove(output_video_file)
 
 
 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")
+    uploaded = 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>'
+warning_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()
+download_button = st.empty()
 
-if up_file and uploaded_File:
+if up_file and uploaded:
     
-    download_btn.empty()
+    download_button.empty()
     tfile = tempfile.NamedTemporaryFile(delete=False)
 
     try:
@@ -71,16 +76,17 @@ if up_file and uploaded_File:
         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)
+        fps = int(vf.get(cv2.CAP_PROP_FPS))
+        width = int(vf.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(vf.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        frame_size = (width, height)
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        video_output = cv2.VideoWriter(output_video_file, fourcc, fps, frame_size)
         # -----------------------------------------------------------------------------
+
         
         txt = st.sidebar.markdown(ip_vid_str, unsafe_allow_html=True)   
-        ip_vid = st.sidebar.video(tfile.name) 
+        ip_video = st.sidebar.video(tfile.name) 
 
         while vf.isOpened():
             ret, frame = vf.read()
@@ -92,25 +98,30 @@ if up_file and uploaded_File:
             out_frame, _ = upload_process_frame.process(frame, pose)
             stframe.image(out_frame)
             video_output.write(out_frame[...,::-1])
+
         
         vf.release()
         video_output.release()
         stframe.empty()
-        ip_vid.empty()
+        ip_video.empty()
         txt.empty()
         tfile.close()
     
     except AttributeError:
-        warn.markdown(warn_str, unsafe_allow_html=True)   
+        warn.markdown(warning_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 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_recorded.mp4')
     
     if download:
         st.session_state['download'] = True
 
-if os.path.exists(Video_Output) and st.session_state['download']:
-    os.remove(Video_Output)
+
+
+if os.path.exists(output_video_file) and st.session_state['download']:
+    os.remove(output_video_file)
     st.session_state['download'] = False
-    download_btn.empty()
+    download_button.empty()