Spaces:
Running
Running
| """ | |
| 🚴 AI Bike Fitting Analyzer | |
| Analyze cycling posture and get recommendations for bike adjustments. | |
| """ | |
| import gradio as gr | |
| import cv2 | |
| import mediapipe as mp | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| import tempfile | |
| import os | |
| # Initialize MediaPipe | |
| mp_pose = mp.solutions.pose | |
| # Landmark indices | |
| LANDMARKS = { | |
| 'left_shoulder': mp_pose.PoseLandmark.LEFT_SHOULDER, | |
| 'right_shoulder': mp_pose.PoseLandmark.RIGHT_SHOULDER, | |
| 'left_elbow': mp_pose.PoseLandmark.LEFT_ELBOW, | |
| 'right_elbow': mp_pose.PoseLandmark.RIGHT_ELBOW, | |
| 'left_wrist': mp_pose.PoseLandmark.LEFT_WRIST, | |
| 'right_wrist': mp_pose.PoseLandmark.RIGHT_WRIST, | |
| 'left_hip': mp_pose.PoseLandmark.LEFT_HIP, | |
| 'right_hip': mp_pose.PoseLandmark.RIGHT_HIP, | |
| 'left_knee': mp_pose.PoseLandmark.LEFT_KNEE, | |
| 'right_knee': mp_pose.PoseLandmark.RIGHT_KNEE, | |
| 'left_ankle': mp_pose.PoseLandmark.LEFT_ANKLE, | |
| 'right_ankle': mp_pose.PoseLandmark.RIGHT_ANKLE, | |
| 'left_heel': mp_pose.PoseLandmark.LEFT_HEEL, | |
| 'right_heel': mp_pose.PoseLandmark.RIGHT_HEEL, | |
| 'left_foot_index': mp_pose.PoseLandmark.LEFT_FOOT_INDEX, | |
| 'right_foot_index': mp_pose.PoseLandmark.RIGHT_FOOT_INDEX, | |
| } | |
| # Optimal angle ranges for bike fitting | |
| OPTIMAL_RANGES = { | |
| 'torso_angle': (80, 90), | |
| 'hip_angle': (60, 100), | |
| 'knee_angle': (75, 160), | |
| 'ankle_angle': (90, 130), | |
| 'elbow_angle': (150, 175), | |
| } | |
| ANGLE_DESCRIPTIONS = { | |
| 'torso_angle': 'Torso (elbow-shoulder-hip)', | |
| 'hip_angle': 'Hip (shoulder-hip-knee)', | |
| 'knee_angle': 'Knee (hip-knee-ankle)', | |
| 'ankle_angle': 'Ankle (knee-ankle-foot)', | |
| 'elbow_angle': 'Elbow (shoulder-elbow-wrist)', | |
| } | |
| def calculate_angle(point1, point2, point3): | |
| """Calculate angle at point2 formed by point1-point2-point3.""" | |
| a = np.array(point1) | |
| b = np.array(point2) | |
| c = np.array(point3) | |
| ba = a - b | |
| bc = c - b | |
| cosine_angle = np.dot(ba, bc) / (np.linalg.norm(ba) * np.linalg.norm(bc) + 1e-6) | |
| cosine_angle = np.clip(cosine_angle, -1.0, 1.0) | |
| return np.degrees(np.arccos(cosine_angle)) | |
| def get_landmark_coords(landmarks, landmark_name, image_shape): | |
| """Extract pixel coordinates for a landmark.""" | |
| landmark = landmarks.landmark[LANDMARKS[landmark_name]] | |
| h, w = image_shape[:2] | |
| return (int(landmark.x * w), int(landmark.y * h)) | |
| def compute_angles(landmarks, image_shape, side='right'): | |
| """Compute all bike fitting angles.""" | |
| prefix = side + '_' | |
| shoulder = get_landmark_coords(landmarks, prefix + 'shoulder', image_shape) | |
| elbow = get_landmark_coords(landmarks, prefix + 'elbow', image_shape) | |
| wrist = get_landmark_coords(landmarks, prefix + 'wrist', image_shape) | |
| hip = get_landmark_coords(landmarks, prefix + 'hip', image_shape) | |
| knee = get_landmark_coords(landmarks, prefix + 'knee', image_shape) | |
| ankle = get_landmark_coords(landmarks, prefix + 'ankle', image_shape) | |
| foot = get_landmark_coords(landmarks, prefix + 'foot_index', image_shape) | |
| angles = { | |
| 'torso_angle': calculate_angle(elbow, shoulder, hip), | |
| 'hip_angle': calculate_angle(shoulder, hip, knee), | |
| 'knee_angle': calculate_angle(hip, knee, ankle), | |
| 'ankle_angle': calculate_angle(knee, ankle, foot), | |
| 'elbow_angle': calculate_angle(shoulder, elbow, wrist), | |
| '_coords': { | |
| 'shoulder': shoulder, 'elbow': elbow, 'wrist': wrist, | |
| 'hip': hip, 'knee': knee, 'ankle': ankle, 'foot': foot, | |
| } | |
| } | |
| return angles | |
| def get_status_color(angle_name, value): | |
| """Get color based on whether angle is in optimal range.""" | |
| if angle_name not in OPTIMAL_RANGES: | |
| return (255, 255, 255) | |
| min_val, max_val = OPTIMAL_RANGES[angle_name] | |
| if min_val <= value <= max_val: | |
| return (0, 255, 0) # Green | |
| elif value < min_val - 10 or value > max_val + 10: | |
| return (0, 0, 255) # Red | |
| else: | |
| return (0, 165, 255) # Orange | |
| def draw_overlay(image, angles): | |
| """Draw skeleton and angle annotations on image.""" | |
| annotated = image.copy() | |
| coords = angles['_coords'] | |
| skeleton_color = (0, 255, 0) | |
| # Draw skeleton | |
| cv2.line(annotated, coords['shoulder'], coords['elbow'], skeleton_color, 3) | |
| cv2.line(annotated, coords['elbow'], coords['wrist'], skeleton_color, 3) | |
| cv2.line(annotated, coords['shoulder'], coords['hip'], skeleton_color, 3) | |
| cv2.line(annotated, coords['hip'], coords['knee'], skeleton_color, 3) | |
| cv2.line(annotated, coords['knee'], coords['ankle'], skeleton_color, 3) | |
| cv2.line(annotated, coords['ankle'], coords['foot'], skeleton_color, 3) | |
| # Draw angle labels | |
| angle_positions = [ | |
| ('torso_angle', coords['shoulder'], (-60, -30)), | |
| ('hip_angle', coords['hip'], (40, -10)), | |
| ('knee_angle', coords['knee'], (-80, 0)), | |
| ('ankle_angle', coords['ankle'], (10, 30)), | |
| ] | |
| for angle_name, position, offset in angle_positions: | |
| value = angles[angle_name] | |
| color = get_status_color(angle_name, value) | |
| text_pos = (position[0] + offset[0], position[1] + offset[1]) | |
| # Background | |
| text = f"{value:.0f}" | |
| (tw, th), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2) | |
| cv2.rectangle(annotated, (text_pos[0]-5, text_pos[1]-th-5), | |
| (text_pos[0]+tw+15, text_pos[1]+5), (0,0,0), -1) | |
| cv2.rectangle(annotated, (text_pos[0]-5, text_pos[1]-th-5), | |
| (text_pos[0]+tw+15, text_pos[1]+5), color, 2) | |
| # Text | |
| cv2.putText(annotated, text, text_pos, cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255,255,255), 2) | |
| cv2.putText(annotated, "o", (text_pos[0]+tw, text_pos[1]-th+5), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.3, (255,255,255), 1) | |
| # Optimal range | |
| if angle_name in OPTIMAL_RANGES: | |
| opt = OPTIMAL_RANGES[angle_name] | |
| cv2.putText(annotated, f"{opt[0]}-{opt[1]}", | |
| (text_pos[0], text_pos[1]+20), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.4, (200,200,200), 1) | |
| # Draw joints | |
| for coord in coords.values(): | |
| cv2.circle(annotated, coord, 6, skeleton_color, -1) | |
| cv2.circle(annotated, coord, 8, (255,255,255), 2) | |
| return annotated | |
| def create_plots(angle_history): | |
| """Create time-series plots for angles.""" | |
| if not angle_history: | |
| return None | |
| fig, axes = plt.subplots(2, 2, figsize=(12, 8)) | |
| axes = axes.flatten() | |
| angle_names = ['torso_angle', 'hip_angle', 'knee_angle', 'ankle_angle'] | |
| times = [a['time'] for a in angle_history] | |
| for idx, name in enumerate(angle_names): | |
| ax = axes[idx] | |
| values = [a[name] for a in angle_history] | |
| ax.plot(times, values, 'b-', linewidth=1.5) | |
| if name in OPTIMAL_RANGES: | |
| opt_min, opt_max = OPTIMAL_RANGES[name] | |
| ax.axhspan(opt_min, opt_max, alpha=0.2, color='green', label=f'Optimal ({opt_min}-{opt_max}°)') | |
| ax.axhline(y=np.mean(values), color='red', linestyle='--', label=f'Mean: {np.mean(values):.1f}°') | |
| ax.set_xlabel('Time (s)') | |
| ax.set_ylabel('Angle (°)') | |
| ax.set_title(ANGLE_DESCRIPTIONS[name]) | |
| ax.legend(loc='upper right', fontsize=8) | |
| ax.grid(True, alpha=0.3) | |
| plt.suptitle('🚴 Bike Fitting Angle Analysis', fontsize=14, fontweight='bold') | |
| plt.tight_layout() | |
| # Save to temp file | |
| plot_path = tempfile.mktemp(suffix='.png') | |
| plt.savefig(plot_path, dpi=150, bbox_inches='tight') | |
| plt.close() | |
| return plot_path | |
| def process_video(video_path, side, progress=gr.Progress()): | |
| """Main processing function.""" | |
| if video_path is None: | |
| return None, None | |
| cap = cv2.VideoCapture(video_path) | |
| if not cap.isOpened(): | |
| return None, None | |
| fps = int(cap.get(cv2.CAP_PROP_FPS)) or 30 | |
| width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| # Output video | |
| output_path = tempfile.mktemp(suffix='.mp4') | |
| fourcc = cv2.VideoWriter_fourcc(*'mp4v') | |
| out = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) | |
| angle_history = [] | |
| with mp_pose.Pose( | |
| static_image_mode=False, | |
| model_complexity=2, | |
| min_detection_confidence=0.5, | |
| min_tracking_confidence=0.5 | |
| ) as pose: | |
| frame_idx = 0 | |
| while True: | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| # Update progress | |
| progress(frame_idx / total_frames, desc=f"Processing frame {frame_idx}/{total_frames}") | |
| # Process frame | |
| rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) | |
| results = pose.process(rgb) | |
| if results.pose_landmarks: | |
| angles = compute_angles(results.pose_landmarks, frame.shape, side) | |
| annotated = draw_overlay(frame, angles) | |
| # Store for plotting | |
| angle_data = {k: v for k, v in angles.items() if not k.startswith('_')} | |
| angle_data['time'] = frame_idx / fps | |
| angle_history.append(angle_data) | |
| else: | |
| annotated = frame | |
| out.write(annotated) | |
| frame_idx += 1 | |
| cap.release() | |
| out.release() | |
| # Convert to H.264 for browser compatibility | |
| web_output = tempfile.mktemp(suffix='.mp4') | |
| os.system(f'ffmpeg -y -i "{output_path}" -vcodec libx264 -acodec aac "{web_output}" -hide_banner -loglevel error') | |
| # Generate plot | |
| plot_path = create_plots(angle_history) | |
| return web_output, plot_path | |
| # Build Gradio interface | |
| with gr.Blocks(title="🚴 AI Bike Fitting Analyzer", theme=gr.themes.Soft()) as demo: | |
| gr.Markdown(""" | |
| # 🚴 AI Bike Fitting Analyzer | |
| Upload a video of a cyclist on a stationary trainer (side view) to analyze their position | |
| and get recommendations for bike adjustments. | |
| **Tips for best results:** | |
| - Film from the side (perpendicular to the bike) | |
| - Ensure good lighting | |
| - Keep the full body in frame | |
| - 10-30 seconds of pedaling is ideal | |
| """) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| video_input = gr.Video(label="📹 Upload Video") | |
| side_select = gr.Radio( | |
| choices=["left", "right"], | |
| value="right", | |
| label="Which side of the cyclist faces the camera?" | |
| ) | |
| analyze_btn = gr.Button("🔍 Analyze", variant="primary", size="lg") | |
| with gr.Column(scale=1): | |
| video_output = gr.Video(label="🎬 Analyzed Video") | |
| with gr.Row(): | |
| plot_output = gr.Image(label="📊 Angle Analysis Over Time") | |
| analyze_btn.click( | |
| fn=process_video, | |
| inputs=[video_input, side_select], | |
| outputs=[video_output, plot_output], | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |