app.py

"""
🚴 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()