Spaces:
Running
Running
| import gradio as gr | |
| import cv2 | |
| import numpy as np | |
| import mediapipe as mp | |
| import time | |
| import traceback # For detailed error logging | |
| # Import your exercise classes | |
| from exercises.hammer_curl import HammerCurl | |
| from exercises.push_up import PushUp | |
| from exercises.squat import Squat | |
| # Initialize MediaPipe Pose | |
| mp_pose = mp.solutions.pose | |
| pose = mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5) | |
| mp_drawing = mp.solutions.drawing_utils | |
| mp_drawing_styles = mp.solutions.drawing_styles | |
| # --- State Variables --- | |
| exercise_trackers = { | |
| "Hammer Curl": HammerCurl(), | |
| "Push Up": PushUp(), | |
| "Squat": Squat() | |
| } | |
| current_exercise_tracker = None | |
| selected_exercise_name = "Hammer Curl" # Default exercise | |
| # Target and progress tracking | |
| target_reps = 10 # Default target reps | |
| target_sets = 3 # Default target sets | |
| current_set_count = 1 | |
| workout_complete_message = "" | |
| def process_frame(video_frame_np, exercise_name_choice, target_reps_input, target_sets_input): | |
| global current_exercise_tracker, selected_exercise_name | |
| global target_reps, target_sets, current_set_count, workout_complete_message | |
| default_h, default_w = 480, 640 | |
| if video_frame_np is not None: | |
| default_h, default_w, _ = video_frame_np.shape | |
| annotated_image = video_frame_np.copy() | |
| else: | |
| blank_frame = np.zeros((default_h, default_w, 3), dtype=np.uint8) | |
| cv2.putText(blank_frame, "No Camera Input", (50, default_h // 2), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255), 2) | |
| return blank_frame, f"0/{target_reps}", f"{current_set_count}/{target_sets}", "No frame", "No camera", "Error: No Frame" | |
| try: | |
| new_target_reps_val = target_reps | |
| try: | |
| new_target_reps_val = int(target_reps_input) | |
| if new_target_reps_val > 0 and target_reps != new_target_reps_val: | |
| target_reps = new_target_reps_val | |
| if current_exercise_tracker: current_exercise_tracker.reset_reps() | |
| current_set_count = 1 | |
| workout_complete_message = "" | |
| except ValueError: pass | |
| new_target_sets_val = target_sets | |
| try: | |
| new_target_sets_val = int(target_sets_input) | |
| if new_target_sets_val > 0 and target_sets != new_target_sets_val: | |
| target_sets = new_target_sets_val | |
| if current_exercise_tracker: current_exercise_tracker.reset_reps() | |
| current_set_count = 1 | |
| workout_complete_message = "" | |
| except ValueError: pass | |
| if selected_exercise_name != exercise_name_choice: | |
| selected_exercise_name = exercise_name_choice | |
| if selected_exercise_name in exercise_trackers: | |
| if selected_exercise_name == "Hammer Curl": exercise_trackers[selected_exercise_name] = HammerCurl() | |
| elif selected_exercise_name == "Push Up": exercise_trackers[selected_exercise_name] = PushUp() | |
| elif selected_exercise_name == "Squat": exercise_trackers[selected_exercise_name] = Squat() | |
| current_set_count = 1 | |
| workout_complete_message = "" | |
| else: current_exercise_tracker = None | |
| current_exercise_tracker = exercise_trackers.get(selected_exercise_name) | |
| image_rgb = cv2.cvtColor(video_frame_np, cv2.COLOR_BGR2RGB) | |
| image_rgb.flags.writeable = False | |
| results = pose.process(image_rgb) | |
| image_rgb.flags.writeable = True | |
| reps_display = f"0/{target_reps}" | |
| sets_display = f"{current_set_count}/{target_sets}" | |
| angle_display = "N/A" | |
| feedback_display = "Initializing..." | |
| temp_workout_message = workout_complete_message | |
| if results.pose_landmarks and current_exercise_tracker and not workout_complete_message: | |
| landmarks_mp = results.pose_landmarks.landmark | |
| frame_height, frame_width, _ = annotated_image.shape | |
| actual_reps_this_set = 0 | |
| try: | |
| if selected_exercise_name == "Hammer Curl": | |
| r_count, r_angle, l_count, l_angle, warn_r, warn_l, _, _, r_stage, l_stage = current_exercise_tracker.track_hammer_curl(landmarks_mp, annotated_image) | |
| actual_reps_this_set = r_count | |
| reps_display = f"R: {r_count}, L: {l_count} (Target: {target_reps} for R)" | |
| angle_display = f"R Ang: {int(r_angle)}, L Ang: {int(l_angle)}" | |
| feedback_list = [] | |
| if warn_r: feedback_list.append(f"R: {warn_r}") | |
| if warn_l: feedback_list.append(f"L: {warn_l}") | |
| feedback_display = " | ".join(feedback_list) if feedback_list else "Good form!" | |
| elif selected_exercise_name == "Push Up": | |
| exercise_data = current_exercise_tracker.track_push_up(landmarks_mp, frame_width, frame_height) | |
| actual_reps_this_set = exercise_data.get("counter", 0) | |
| angle_display = f"L: {int(exercise_data.get('angle_left',0))}, R: {int(exercise_data.get('angle_right',0))}" | |
| feedback_display = str(exercise_data.get("feedback", "No feedback")) | |
| if 'get_drawing_annotations' in dir(current_exercise_tracker): | |
| annotations_to_draw = current_exercise_tracker.get_drawing_annotations(landmarks_mp, frame_width, frame_height, exercise_data) | |
| for ann in annotations_to_draw: | |
| if ann["type"] == "line": cv2.line(annotated_image, tuple(ann["start_point"]), tuple(ann["end_point"]), ann["color_bgr"], ann["thickness"]) | |
| elif ann["type"] == "circle": cv2.circle(annotated_image, tuple(ann["center_point"]), ann["radius"], ann["color_bgr"], -1 if ann.get("filled", False) else ann["thickness"]) | |
| elif ann["type"] == "text": cv2.putText(annotated_image, ann["text_content"], tuple(ann["position"]), cv2.FONT_HERSHEY_SIMPLEX, ann["font_scale"], ann["color_bgr"], ann["thickness"]) | |
| elif selected_exercise_name == "Squat": | |
| exercise_data = current_exercise_tracker.track_squat(landmarks_mp, frame_width, frame_height) | |
| actual_reps_this_set = exercise_data.get("counter", 0) | |
| angle_display = f"L: {int(exercise_data.get('angle_left',0))}, R: {int(exercise_data.get('angle_right',0))}" | |
| feedback_display = str(exercise_data.get("feedback", "No feedback")) | |
| if 'get_drawing_annotations' in dir(current_exercise_tracker): | |
| annotations_to_draw = current_exercise_tracker.get_drawing_annotations(landmarks_mp, frame_width, frame_height, exercise_data) | |
| for ann in annotations_to_draw: | |
| if ann["type"] == "line": cv2.line(annotated_image, tuple(ann["start_point"]), tuple(ann["end_point"]), ann["color_bgr"], ann["thickness"]) | |
| elif ann["type"] == "circle": cv2.circle(annotated_image, tuple(ann["center_point"]), ann["radius"], ann["color_bgr"], -1 if ann.get("filled", False) else ann["thickness"]) | |
| elif ann["type"] == "text": cv2.putText(annotated_image, ann["text_content"], tuple(ann["position"]), cv2.FONT_HERSHEY_SIMPLEX, ann["font_scale"], ann["color_bgr"], ann["thickness"]) | |
| if selected_exercise_name != "Hammer Curl": | |
| reps_display = f"{actual_reps_this_set}/{target_reps}" | |
| if actual_reps_this_set >= target_reps: | |
| if current_set_count < target_sets: | |
| current_set_count += 1 | |
| current_exercise_tracker.reset_reps() | |
| feedback_display = f"Set {current_set_count-1} complete! Starting set {current_set_count}." | |
| if selected_exercise_name == "Hammer Curl": reps_display = f"R: 0, L: 0 (Target: {target_reps} for R)" | |
| else: reps_display = f"0/{target_reps}" | |
| elif current_set_count >= target_sets: | |
| feedback_display = "Workout Complete!" | |
| workout_complete_message = "Workout Complete! Change targets or exercise to restart." | |
| if selected_exercise_name == "Hammer Curl": reps_display = f"R: {target_reps}, L: {target_reps} (Target: {target_reps} for R)" | |
| else: reps_display = f"{target_reps}/{target_reps}" | |
| temp_workout_message = workout_complete_message | |
| except Exception as e_exercise: | |
| print(f"PROCESS_FRAME: Error during exercise '{selected_exercise_name}' logic: {e_exercise}") | |
| print(traceback.format_exc()) | |
| cv2.putText(annotated_image, f"Error in {selected_exercise_name}", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2) | |
| feedback_display = f"Error in {selected_exercise_name} processing." | |
| elif workout_complete_message: | |
| feedback_display = workout_complete_message | |
| reps_display = f"{target_reps}/{target_reps}" if selected_exercise_name != "Hammer Curl" else f"R: {target_reps}, L: {target_reps} (Target: {target_reps} for R)" | |
| sets_display = f"{target_sets}/{target_sets}" | |
| if results.pose_landmarks: | |
| mp_drawing.draw_landmarks(annotated_image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS, landmark_drawing_spec=mp_drawing_styles.get_default_pose_landmarks_style()) | |
| else: | |
| feedback_display = "No person detected or exercise not selected properly." | |
| if results and results.pose_landmarks : | |
| mp_drawing.draw_landmarks(annotated_image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS, landmark_drawing_spec=mp_drawing_styles.get_default_pose_landmarks_style()) | |
| else: | |
| cv2.putText(annotated_image, "No person detected", (50, default_h // 2), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255),2) | |
| sets_display = f"{current_set_count}/{target_sets}" | |
| if not isinstance(annotated_image, np.ndarray) or annotated_image.ndim != 3 or annotated_image.shape[2] != 3: | |
| annotated_image = np.zeros((default_h, default_w, 3), dtype=np.uint8) | |
| cv2.putText(annotated_image, "Display Error", (50, default_h // 2), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255),2) | |
| return annotated_image, reps_display, sets_display, angle_display, feedback_display, temp_workout_message | |
| except Exception as e_main: | |
| print(f"PROCESS_FRAME: CRITICAL error in process_frame: {e_main}") | |
| print(traceback.format_exc()) | |
| error_frame = np.zeros((default_h, default_w, 3), dtype=np.uint8) | |
| cv2.putText(error_frame, f"Error: {e_main}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2) | |
| return error_frame, "Error", "Error", "Error", "Critical Error", "Critical Error" | |
| def trigger_reset_workout(): | |
| global current_set_count, workout_complete_message, selected_exercise_name, target_reps, target_sets, current_exercise_tracker | |
| current_set_count = 1 | |
| workout_complete_message = "" | |
| if current_exercise_tracker: | |
| current_exercise_tracker.reset_reps() | |
| reset_reps_display = f"0/{target_reps}" | |
| if selected_exercise_name == "Hammer Curl": | |
| reset_reps_display = f"R: 0, L: 0 (Target: {target_reps} for R)" | |
| reset_sets_display = f"1/{target_sets}" | |
| reset_angle_display = "N/A" | |
| reset_feedback_display = "Workout Reset. Ready to start." | |
| reset_workout_status = "" | |
| return reset_reps_display, reset_sets_display, reset_angle_display, reset_feedback_display, reset_workout_status | |
| # --- Custom CSS for gradient background and styling --- | |
| # Note: Applying gradient to the root/body might be overridden by Gradio's specific block styling. | |
| # It's often better to target .gradio-container or specific blocks if possible. | |
| # However, for broad effect, body is a start. More specific selectors might be needed for full coverage. | |
| # --- Custom CSS for gradient background and styling --- | |
| custom_css = """ | |
| @font-face { | |
| font-family: 'Sprintura'; | |
| src: url('https://res.cloudinary.com/dgj1gzq0l/raw/upload/v1748875223/sprintura-demo.regular_ttcewg.otf') format('opentype'); | |
| font-weight: normal; | |
| font-style: normal; | |
| font-display: swap; /* Tells the browser to use a fallback until the font loads */ | |
| } | |
| html, body { | |
| /* Ensure html and body take up full viewport height and remove default margins */ | |
| min-height: 100vh; | |
| /* REMOVE THE LINE BELOW: height: 100%; */ /* This can sometimes restrict scrolling when combined with other elements */ | |
| margin: 0; | |
| padding: 0; | |
| overflow-x: hidden; /* Prevent horizontal scrollbar */ | |
| overflow-y: auto; /* ADD THIS LINE: Allow vertical scrolling when content exceeds height */ | |
| width: 100vw; /* Ensure it takes full viewport width */ | |
| } | |
| .gradio-app { | |
| max-width: unset !important; /* Remove any default max-width */ | |
| width: 100vw !important; /* Force it to take full viewport width */ | |
| margin: 0 !important; /* Remove any default auto margins */ | |
| padding: 0 !important; /* Remove any default padding */ | |
| box-shadow: none !important; /* Remove any default shadows that might create a border effect */ | |
| border-radius: 0 !important; /* Remove any default border radius */ | |
| } | |
| body { | |
| /* Apply the background image properties from your outer div */ | |
| background-image: url('https://res.cloudinary.com/dgj1gzq0l/image/upload/v1747821491/new_bg_bz1uqj.svg') !important; | |
| background-size: cover !important; /* Equivalent to bg-cover */ | |
| background-position: center !important; /* Equivalent to bg-center */ | |
| background-repeat: no-repeat !important; /* Equivalent to bg-no-repeat */ | |
| } | |
| .gradio-container { | |
| /* Apply the semi-transparent black overlay from your inner div */ | |
| background-color: rgba(0, 0, 0, 0.60) !important; | |
| /* Make it span the full viewport */ | |
| min-height: 100vh !important; /* Forces full height */ | |
| width: 100vw !important; /* Forces full width */ | |
| margin: 0 !important; /* Remove any auto margins */ | |
| padding: 0 !important; /* IMPORTANT: Remove internal padding so header can go edge-to-edge */ | |
| border-radius: 0 !important; /* Remove rounded corners for full screen */ | |
| box-shadow: none !important; /* Remove shadow for full screen */ | |
| display: flex; /* Helps in making content fill vertical space */ | |
| flex-direction: column; /* Stacks content vertically within the container */ | |
| } | |
| .header-container { | |
| width: 100%; /* Now takes 100% of the 100vw parent */ | |
| height: 160px; | |
| position: relative; | |
| display: flex; /* For centering the h1 */ | |
| flex-direction: column; /* For vertical alignment */ | |
| justify-content: center; /* Centers the h1 vertically */ | |
| align-items: center; /* Centers the h1 horizontally */ | |
| padding: 2rem 1rem; /* Keep padding for text inside */ | |
| overflow: hidden !important; /* Explicitly ensure overflow is hidden */ | |
| border-bottom-left-radius: 1rem; | |
| border-bottom-right-radius: 1rem; | |
| background: linear-gradient(to bottom, rgba(30,18,51,0.9) 0%, rgba(30,18,51,0.7) 60%, rgba(0,0,0,0) 100%); | |
| box-sizing: border-box; /* Include padding in the element's total width/height */ | |
| /* Removed negative margins and calc width, as parent is now full width and has no padding */ | |
| } | |
| label, .gr-checkbox-label span { /* Target labels and checkbox labels */ | |
| color: #E8E8E8 !important; /* Slightly brighter light gray for labels */ | |
| font-weight: bold !important; | |
| } | |
| .header-container h1 { | |
| color: #E6E6FA !important; /* Equivalent to text-lavender-100 */ | |
| font-size: 2.5rem !important; /* Equivalent to text-3xl, but made a bit larger for a main heading */ | |
| font-weight: 400 !important; /* Equivalent to font-bold */ | |
| letter-spacing: 0.05em !important; /* Equivalent to tracking-wide */ | |
| width: 100%; /* Equivalent to w-full */ | |
| text-align: center; /* Equivalent to text-center */ | |
| margin: 0 auto; /* Equivalent to mx-auto */ | |
| /* IMPORTANT: For 'Sprintura' font, you still need to host the font files and | |
| add an @font-face rule to your CSS that points to their public URLs. | |
| Otherwise, browsers will use a generic sans-serif font. */ | |
| font-family: 'Sprintura', sans-serif !important; | |
| } | |
| /* Styling for Controls and Progress Headings */ | |
| #controls-heading h3, | |
| #progress-heading h3 { | |
| font-size: 1.2rem !important; /* Increased font size - adjust this value as needed */ | |
| color: #E6E6FA !important; /* Keeping text color consistent with your main header */ | |
| font-weight: 600 !important; /* Semibold, consistent with your main header */ | |
| margin-top: 0.5rem !important; /* Add some top margin for better spacing */ | |
| margin-bottom: 0.1rem !important; /* Add some bottom margin for better spacing */ | |
| font-family: 'Sprintura', sans-serif !important; | |
| } | |
| .prose { /* Markdown text */ | |
| color: #F0F0F0 !important; | |
| font-size: 1.1rem !important; | |
| } | |
| /* General text within blocks that isn't a label or title */ | |
| .gr-block .gr-box > div > p, .gr-block .gr-box > div > span { | |
| color: #E0E0E0 !important; | |
| } | |
| /* Style for the Reset Button using its specific ID */ | |
| #reset-workout-button { | |
| background-color: #513B8B !important; /* Your desired background color */ | |
| color: #FFFFFF !important; /* White text color */ | |
| border-color: #513B8B !important; /* Match border color to background */ | |
| } | |
| #reset-workout-button:hover { | |
| background-color: #9b87f5 !important; /* Hover background color */ | |
| border-color: #9b87f5 !important; /* Match hover border color */ | |
| } | |
| """ | |
| # --- Gradio Theme --- | |
| # Using a base theme to set font and then overriding some colors. | |
| # For button color, if primary_hue doesn't give desired button color, specific CSS might be needed. | |
| theme = gr.themes.Base( | |
| font=[gr.themes.GoogleFont("Exo 2"), "ui-sans-serif", "system-ui", "sans-serif"], | |
| primary_hue=gr.themes.colors.amber, # For buttons - gives a yellowish/golden hue | |
| secondary_hue=gr.themes.colors.blue, | |
| neutral_hue=gr.themes.colors.gray | |
| ).set( | |
| body_text_color="#E0E0E0", # Light gray for general text (this should also affect input text) | |
| input_background_fill="#4A2A6C", # Darker violet for input backgrounds | |
| input_border_color="#6A3AA2", | |
| # button_primary_text_color="#111111", # Often better to let theme handle this or use CSS | |
| # Ensure other text elements have good contrast automatically or via custom_css | |
| ) | |
| # --- Gradio Interface --- | |
| exercise_choices = ["Hammer Curl", "Push Up", "Squat"] | |
| # Pass the theme and custom_css to gr.Blocks | |
| with gr.Blocks(theme=theme, css=custom_css) as iface: | |
| gr.Markdown("# Live AI Trainer", elem_classes="header-container") # Apply the class directly here | |
| gr.Markdown("Select an exercise, set your targets, and get real-time feedback on your form and reps.") # Styled by .prose | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| webcam_input = gr.Image(sources=["webcam"], streaming=True, type="numpy", label="Your Webcam") | |
| with gr.Column(scale=1): | |
| gr.Markdown("### Controls", elem_id="controls-heading") | |
| exercise_dropdown = gr.Dropdown(choices=exercise_choices, label="Select Exercise", value="Hammer Curl") | |
| target_reps_number = gr.Number(value=target_reps, label="Target Reps per Set", precision=0, minimum=1) | |
| target_sets_number = gr.Number(value=target_sets, label="Target Sets", precision=0, minimum=1) | |
| reset_button = gr.Button("Reset Workout", elem_id="reset-workout-button") | |
| gr.Markdown("### Progress", elem_id="progress-heading") | |
| reps_output = gr.Textbox(label="Reps Progress") | |
| sets_output = gr.Textbox(label="Sets Progress") | |
| angle_output = gr.Textbox(label="Angle Details") | |
| feedback_output = gr.Textbox(label="Feedback", lines=3, max_lines=5) | |
| workout_status_output = gr.Textbox(label="Workout Status", lines=2, interactive=False) | |
| process_frame_inputs = [webcam_input, exercise_dropdown, target_reps_number, target_sets_number] | |
| process_frame_outputs = [webcam_input, reps_output, sets_output, angle_output, feedback_output, workout_status_output] | |
| webcam_input.stream(fn=process_frame, inputs=process_frame_inputs, outputs=process_frame_outputs) | |
| exercise_dropdown.change(fn=process_frame, inputs=process_frame_inputs, outputs=process_frame_outputs) | |
| target_reps_number.change(fn=process_frame, inputs=process_frame_inputs, outputs=process_frame_outputs) | |
| target_sets_number.change(fn=process_frame, inputs=process_frame_inputs, outputs=process_frame_outputs) | |
| reset_button_outputs = [reps_output, sets_output, angle_output, feedback_output, workout_status_output] | |
| reset_button.click(fn=trigger_reset_workout, inputs=None, outputs=reset_button_outputs) | |
| if __name__ == "__main__": | |
| iface.launch(debug=False, share=False) # share=False is default but good to be explicit for Spaces |