#!/usr/bin/env python3 """ Workout Buddy - AI Personal Trainer Main application """ import cv2 import json import time import argparse from pathlib import Path from camera import ESP32Camera, FPSCounter from pose_detector import PoseDetector from form_analyzer import create_analyzer import sys class WorkoutBuddy: """ Main workout tracking application """ def __init__(self, config_path: str = "config.json"): """ Initialize workout buddy Args: config_path: Path to configuration JSON file """ # Load configuration with open(config_path, 'r') as f: self.config = json.load(f) # Initialize camera camera_url = self.config['camera']['url'] print(f"Connecting to camera at {camera_url}...") self.camera = ESP32Camera( url=camera_url, timeout=self.config['camera']['timeout'] ) print("Camera connected!") # Initialize pose detector print("Loading pose detection model...") pose_config = self.config['pose'] self.detector = PoseDetector( model_complexity=pose_config['model_complexity'], min_detection_confidence=pose_config['min_detection_confidence'], min_tracking_confidence=pose_config['min_tracking_confidence'], enable_segmentation=pose_config['enable_segmentation'] ) print("Pose detector ready!") # FPS counter self.fps_counter = FPSCounter() # Current exercise self.current_exercise = None self.analyzer = None # Display settings self.display_config = self.config['display'] self.window_name = "Workout Buddy" def set_exercise(self, exercise_name: str): """ Set the current exercise Args: exercise_name: Name of exercise (must be in config) """ if exercise_name not in self.config['exercises']: raise ValueError(f"Unknown exercise: {exercise_name}") self.current_exercise = exercise_name exercise_config = self.config['exercises'][exercise_name] print(f"\nStarting exercise: {exercise_config['name']}") self.analyzer = create_analyzer(exercise_name, exercise_config) def draw_ui(self, frame, analysis_results=None): """ Draw UI overlay on frame Args: frame: Video frame analysis_results: Results from form analyzer Returns: Frame with UI overlay """ height, width = frame.shape[:2] # Draw FPS if self.display_config['show_fps']: fps = self.fps_counter.get_fps() cv2.putText(frame, f"FPS: {fps:.1f}", (width - 150, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) if analysis_results: # Exercise name exercise_name = self.config['exercises'][self.current_exercise]['name'] cv2.putText(frame, exercise_name, (10, 40), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (255, 255, 255), 3) # Rep counter reps = analysis_results['rep_count'] cv2.putText(frame, f"Reps: {reps}", (10, 90), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 255, 255), 3) # Form score form_score = analysis_results['form_score'] score_percent = int(form_score * 100) # Color based on score (red = bad, yellow = ok, green = good) if form_score >= 0.85: score_color = (0, 255, 0) # Green elif form_score >= 0.7: score_color = (0, 255, 255) # Yellow else: score_color = (0, 0, 255) # Red cv2.putText(frame, f"Form: {score_percent}%", (10, 140), cv2.FONT_HERSHEY_SIMPLEX, 1.0, score_color, 2) # Form issues y_offset = 190 if analysis_results['form_issues']: cv2.putText(frame, "Issues:", (10, y_offset), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2) y_offset += 35 for issue in analysis_results['form_issues'][:3]: # Show max 3 issues cv2.putText(frame, f"- {issue}", (10, y_offset), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2) y_offset += 30 # Debug: Show angles if enabled if self.config['feedback'].get('show_angle_debug', False): y_offset = height - 100 for angle_name, angle_value in analysis_results.get('angles', {}).items(): text = f"{angle_name}: {angle_value:.1f}°" cv2.putText(frame, text, (10, y_offset), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2) y_offset += 25 # Instructions else: cv2.putText(frame, "Press number key to select exercise:", (10, height - 80), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2) cv2.putText(frame, "1=Squat 2=Push-up 3=Deadlift", (10, height - 50), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2) return frame def run(self): """ Main application loop """ # Create window cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL) if self.display_config['fullscreen']: cv2.setWindowProperty(self.window_name, cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN) else: cv2.resizeWindow(self.window_name, self.display_config['width'], self.display_config['height']) print("\n=== Workout Buddy Started ===") print("Controls:") print(" 1 - Squat") print(" 2 - Push-up") print(" 3 - Deadlift") print(" r - Reset rep counter") print(" q - Quit") print("===========================\n") try: while True: # Capture frame frame = self.camera.capture() if frame is None: print("Failed to capture frame") time.sleep(0.1) continue # Update FPS self.fps_counter.update() # Detect pose pose_detected = self.detector.detect(frame) analysis_results = None if pose_detected: # Draw landmarks if enabled if self.display_config['show_landmarks']: frame = self.detector.draw_landmarks(frame) # Analyze form if exercise selected if self.analyzer: analysis_results = self.analyzer.analyze(self.detector) # Draw UI overlay frame = self.draw_ui(frame, analysis_results) # Show frame cv2.imshow(self.window_name, frame) # Handle keyboard input key = cv2.waitKey(1) & 0xFF if key == ord('q'): print("\nQuitting...") break elif key == ord('1'): self.set_exercise('squat') elif key == ord('2'): self.set_exercise('pushup') elif key == ord('3'): self.set_exercise('deadlift') elif key == ord('r'): if self.analyzer: self.analyzer.reset() print("Rep counter reset") finally: # Cleanup cv2.destroyAllWindows() print("\nWorkout session ended.") if self.analyzer and self.analyzer.rep_count > 0: print(f"\nSession summary:") print(f" Exercise: {self.current_exercise}") print(f" Total reps: {self.analyzer.rep_count}") print(f" Final form score: {int(self.analyzer.form_score * 100)}%") def main(): """Entry point""" parser = argparse.ArgumentParser(description="Workout Buddy - AI Personal Trainer") parser.add_argument('--config', type=str, default='config.json', help='Path to configuration file') parser.add_argument('--exercise', type=str, choices=['squat', 'pushup', 'deadlift'], help='Start with this exercise') args = parser.parse_args() try: app = WorkoutBuddy(config_path=args.config) if args.exercise: app.set_exercise(args.exercise) app.run() except KeyboardInterrupt: print("\nInterrupted by user") except Exception as e: print(f"\nError: {e}") import traceback traceback.print_exc() sys.exit(1) if __name__ == "__main__": main()