updated requirements, added streamlit to the stack

Aditya Raj · Aditya Raj · commit d16f12f4ac69 · 2024-10-25T15:36:16.000+05:30
diff --git a/.gitignore b/.gitignore
@@ -1,3 +1,5 @@
 .venv/
 __pycache__/
-.DS_Store
+.DS_Store
+videos/
+processed_videos/
diff --git a/requirements.txt b/requirements.txt
@@ -1,4 +1,7 @@
+altair==5.4.1
+attrs==24.2.0
 blinker==1.8.2
+cachetools==5.5.0
 certifi==2024.8.30
 charset-normalizer==3.4.0
 click==8.1.7
@@ -10,36 +13,55 @@ Flask==3.0.3
 Flask-Cors==5.0.0
 fonttools==4.54.1
 fsspec==2024.10.0
+gitdb==4.0.11
+GitPython==3.1.43
 idna==3.10
 imageio==2.36.0
 imageio-ffmpeg==0.5.1
 itsdangerous==2.2.0
 Jinja2==3.1.4
+jsonschema==4.23.0
+jsonschema-specifications==2024.10.1
 kiwisolver==1.4.7
+markdown-it-py==3.0.0
 MarkupSafe==3.0.2
 matplotlib==3.9.2
+mdurl==0.1.2
 moviepy==1.0.3
 mpmath==1.3.0
+narwhals==1.10.0
 networkx==3.4.2
 numpy==1.26.4
 opencv-python==4.10.0.84
 packaging==24.1
 pandas==2.2.3
 pillow==10.3.0
 proglog==0.1.10
+protobuf==5.28.3
 psutil==6.1.0
 py-cpuinfo==9.0.0
+pyarrow==17.0.0
+pydeck==0.9.1
+Pygments==2.18.0
 pyparsing==3.2.0
 python-dateutil==2.9.0.post0
 pytz==2024.2
 PyYAML==6.0.2
+referencing==0.35.1
 requests==2.32.3
+rich==13.9.3
+rpds-py==0.20.0
 scipy==1.14.1
 seaborn==0.13.2
 six==1.16.0
+smmap==5.0.1
+streamlit==1.39.0
 sympy==1.13.1
+tenacity==9.0.0
+toml==0.10.2
 torch==2.5.0
 torchvision==0.20.0
+tornado==6.4.1
 tqdm==4.66.5
 typing_extensions==4.12.2
 tzdata==2024.2
diff --git a/streamlit.py b/streamlit.py
@@ -0,0 +1,179 @@
+import os
+import logging
+import cv2
+from ultralytics import YOLO
+import streamlit as st
+
+# Set up logging
+logging.basicConfig(level=logging.DEBUG)
+
+# Load YOLO models
+yolo_models = {
+    'regular_deadlift': YOLO("muscleAi_weights/best.pt"),
+    'sumo_deadlift': YOLO("muscleAi_weights/sumo_best.pt"),
+    'squat': YOLO("muscleAi_weights/squats_best.pt"),
+    'romanian_deadlift': YOLO("muscleAi_weights/best_romanian.pt"),
+    "zercher_squat": YOLO("muscleAi_weights/zercher_best.pt"),
+    "front_squat": YOLO("muscleAi_weights/front_squats_best.pt")
+}
+
+# Function to check for injury risk
+def check_injury_risk(labels, exercise_type):
+    if exercise_type in ['regular_deadlift', 'squat']:
+        ibw_value = labels.get('ibw', 1.0)
+        down_value = labels.get('down', 1.0)
+    else:
+        ibw_value = labels.get('up', 1.0)
+        down_value = labels.get('down', 1.0)
+
+    return "stop right now to prevent injury" if ibw_value < 0.80 or down_value < 0.70 else "No significant risk"
+
+# Function to draw keypoints on the frame
+def draw_keypoints(frame, keypoints):
+    for point in keypoints:
+        x, y = int(point[0]), int(point[1])
+        cv2.circle(frame, (x, y), 5, (0, 255, 0), -1)
+    return frame
+
+# Function to process video with YOLO
+def process_video_with_yolo(video_path, exercise_type):
+    processed_frames = []
+    
+    yolo_model = yolo_models[exercise_type]
+    
+    cap = cv2.VideoCapture(video_path)
+    
+    if not cap.isOpened():
+        st.error("Error opening video file")
+        return None
+
+    last_ibw_label = None
+    rep_count = 0
+    rep_started = False
+
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        
+        results = yolo_model(source=frame, stream=True, conf=0.3)
+        
+        for result in results:
+            frame = result.orig_img
+            
+            labels = {result.names[int(box.cls)]: float(box.conf) for box in result.boxes} if result.boxes is not None else {}
+            injury_risk = check_injury_risk(labels, exercise_type)
+
+            current_ibw_label = labels.get('ibw') if exercise_type in ['regular_deadlift', 'squat'] else labels.get('up')
+            
+            if last_ibw_label is not None and current_ibw_label is not None:
+                if not rep_started:
+                    if last_ibw_label > 0.89 and current_ibw_label <= 0.89:
+                        rep_started = True
+                else:
+                    if last_ibw_label <= 0.89 and current_ibw_label > 0.89:
+                        rep_count += 1
+                        rep_started = False
+
+            last_ibw_label = current_ibw_label
+
+            # Draw keypoints on the frame if available
+            if hasattr(result, 'keypoints') and result.keypoints is not None:
+                keypoints = result.keypoints.xy[0]
+                frame = draw_keypoints(frame, keypoints)
+
+            cv2.putText(frame, f"Injury Risk: {injury_risk}", (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
+            cv2.putText(frame, f"Repetitions: {rep_count}", (50, 100), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
+
+            processed_frames.append(frame)
+
+    cap.release()
+    
+    return processed_frames
+
+# Function to process live video stream
+def process_live_video(exercise_type):
+    cap = cv2.VideoCapture(0)
+    
+    if not cap.isOpened():
+        st.error("Error opening webcam")
+        return
+
+    stframe = st.empty()
+    
+    last_ibw_label = None
+    rep_count = 0
+    rep_started = False
+
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        
+        results = yolo_models[exercise_type](source=frame, stream=True, conf=0.3)
+
+        for result in results:
+            frame = result.orig_img
+            
+            labels = {result.names[int(box.cls)]: float(box.conf) for box in result.boxes} if result.boxes is not None else {}
+            injury_risk = check_injury_risk(labels, exercise_type)
+
+            current_ibw_label = labels.get('ibw') if exercise_type in ['regular_deadlift', 'squat'] else labels.get('up')
+            
+            if last_ibw_label is not None and current_ibw_label is not None:
+                if not rep_started:
+                    if last_ibw_label > 0.89 and current_ibw_label <= 0.89:
+                        rep_started = True
+                else:
+                    if last_ibw_label <= 0.89 and current_ibw_label > 0.89:
+                        rep_count += 1
+                        rep_started = False
+
+            last_ibw_label = current_ibw_label
+
+            # Draw keypoints on the frame if available
+            if hasattr(result, 'keypoints') and result.keypoints is not None:
+                keypoints = result.keypoints.xy[0]
+                frame = draw_keypoints(frame, keypoints)
+
+            cv2.putText(frame, f"Injury Risk: {injury_risk}", (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
+            cv2.putText(frame, f"Repetitions: {rep_count}", (50, 100), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
+
+            # Update the Streamlit frame with the processed frame
+            stframe.image(frame, channels="BGR")
+
+    cap.release()
+
+# Streamlit UI
+st.title("Exercise Video Analysis")
+
+exercise_type = st.selectbox("Select Exercise Type", list(yolo_models.keys()))
+
+uploaded_file = st.file_uploader("Upload a Video", type=["mp4", "mov"])
+
+if uploaded_file is not None:
+    video_path = os.path.join('./videos', uploaded_file.name)
+    
+    # Save uploaded video
+    with open(video_path, "wb") as f:
+        f.write(uploaded_file.getbuffer())
+
+    st.video(uploaded_file)
+
+    if st.button("Process Video"):
+        processed_frames = process_video_with_yolo(video_path, exercise_type)
+
+        if processed_frames is not None:
+            output_video_path = os.path.join('./processed_videos', f'processed_{uploaded_file.name}')
+            out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'XVID'), 30,
+                                  (processed_frames[0].shape[1], processed_frames[0].shape[0]))
+
+            for frame in processed_frames:
+                out.write(frame)
+            out.release()
+
+            st.success(f"Video processed successfully! You can download it [here](./processed_videos/processed_{uploaded_file.name})")
+
+if st.button("Start Live Stream"):
+    process_live_video(exercise_type)
+
