File size: 6,777 Bytes
690e476 |
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 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 |
import cv2
import mediapipe as mp
import numpy as np
import gradio as gr
# Initialize MediaPipe Pose
mp_drawing = mp.solutions.drawing_utils
mp_pose = mp.solutions.pose
# Global variables for rep counting
counter = 0
stage = None
# Function to calculate angle between three points
def calculate_angle(a, b, c):
a = np.array(a) # First point
b = np.array(b) # Mid point (pivot)
c = np.array(c) # End point
radians = np.arctan2(c[1] - b[1], c[0] - b[0]) - np.arctan2(a[1] - b[1], a[0] - b[0])
angle = np.abs(radians * 180.0 / np.pi)
if angle > 180.0:
angle = 360 - angle
return angle
# Function to process frames (used for both webcam and uploaded video)
def process_frames(cap):
global counter, stage
with mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5) as pose:
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
image.flags.writeable = False
results = pose.process(image)
image.flags.writeable = True
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
try:
landmarks = results.pose_landmarks.landmark
shoulder = [landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].x,
landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].y]
elbow = [landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value].x,
landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value].y]
wrist = [landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].x,
landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].y]
angle = calculate_angle(shoulder, elbow, wrist)
cv2.putText(image, str(round(angle, 2)),
tuple(np.multiply(elbow, [frame.shape[1], frame.shape[0]]).astype(int)),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2, cv2.LINE_AA)
if angle > 160:
stage = "down"
if angle < 30 and stage == "down":
stage = "up"
counter += 1
except:
pass
cv2.rectangle(image, (0, 0), (225, 73), (245, 117, 16), -1)
cv2.putText(image, 'REPS', (15, 12),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
cv2.putText(image, str(counter), (10, 60),
cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 255, 255), 2, cv2.LINE_AA)
cv2.putText(image, 'STAGE', (65, 12),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
cv2.putText(image, stage if stage else "", (60, 60),
cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 255, 255), 2, cv2.LINE_AA)
mp_drawing.draw_landmarks(image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS,
mp_drawing.DrawingSpec(color=(245, 117, 66), thickness=2, circle_radius=2),
mp_drawing.DrawingSpec(color=(245, 66, 230), thickness=2, circle_radius=2))
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
yield image_rgb, f"Reps: {counter}", f"Stage: {stage if stage else 'None'}"
cap.release()
# Function to process webcam feed
def process_webcam():
global counter, stage
counter = 0
stage = None
cap = cv2.VideoCapture(0) # Use default webcam
for frame, reps, stage in process_frames(cap):
yield frame, reps, stage
# Function to process uploaded video
def process_uploaded_video(video_path):
global counter, stage
counter = 0
stage = None
cap = cv2.VideoCapture(video_path)
for frame, reps, stage in process_frames(cap):
yield frame, reps, stage
# Gradio interface using gr.Blocks
with gr.Blocks() as demo:
gr.Markdown("# Exercise Pose Detection")
gr.Markdown("Choose an input source to track bicep curls and pose.")
# Main Layout: Split into Left and Right Columns
with gr.Row():
# Left Side: Input Source and Controls (Made Smaller)
with gr.Column(scale=1):
input_source = gr.Dropdown(choices=["Use Webcam", "Upload Video"], label="Select Input Source", value="Use Webcam", elem_classes="small-dropdown")
start_webcam_btn = gr.Button("Start Webcam", visible=True, elem_classes="small-button")
video_input = gr.Video(label="Upload Video", visible=False, elem_classes="small-video-input")
# Right Side: Pose Detection Feed and Outputs
with gr.Column(scale=3):
with gr.Row():
video_output = gr.Image(label="Pose Detection Feed", streaming=True, elem_classes="large-video")
with gr.Column():
rep_count = gr.Textbox(label="Rep Count", elem_classes="small-textbox")
stage_output = gr.Textbox(label="Stage", elem_classes="small-textbox")
# Custom CSS for styling
demo.css = """
.small-dropdown {
max-width: 200px !important;
}
.small-button {
max-width: 200px !important;
padding: 5px !important;
font-size: 14px !important;
}
.small-video-input {
max-width: 200px !important;
}
.large-video {
width: 100% !important;
max-width: 800px !important;
margin: 0 auto !important;
}
.small-textbox {
max-width: 150px !important;
height: 50px !important;
margin: 10px 0 !important;
}
"""
# Function to toggle visibility of input components based on dropdown
def toggle_inputs(source):
if source == "Use Webcam":
return gr.update(visible=True), gr.update(visible=False), gr.update(value=None), gr.update(value=None), gr.update(value=None)
else:
return gr.update(visible=False), gr.update(visible=True), gr.update(value=None), gr.update(value=None), gr.update(value=None)
# Update visibility and clear outputs when dropdown changes
input_source.change(
toggle_inputs,
inputs=[input_source],
outputs=[start_webcam_btn, video_input, video_output, rep_count, stage_output]
)
# Start webcam feed when the button is clicked
start_webcam_btn.click(
process_webcam,
inputs=None,
outputs=[video_output, rep_count, stage_output],
queue=True
)
# Process uploaded video when a file is uploaded
video_input.change(
process_uploaded_video,
inputs=[video_input],
outputs=[video_output, rep_count, stage_output],
queue=True
)
demo.launch() |