initial commit

Signed-off-by: Emily Chen <emilychen@Emilys-iMac.lan>

README.md

@@ -0,0 +1,92 @@
+# Golf Swing Analysis
+
+A Python application that analyzes golf swings from YouTube videos using computer vision and AI.
+
+## Features
+
+- YouTube video retrieval and processing using yt-dlp
+- Golfer, club, and ball detection using YOLOv8
+- Pose estimation for swing analysis
+- Swing phase segmentation (setup, backswing, downswing, impact, follow-through)
+- Trajectory and speed analysis
+- AI-powered swing evaluation and coaching tips
+- Visual feedback with annotations
+- Streamlit web interface
+
+## Installation
+
+1. Clone this repository
+2. Run the setup script to create necessary directories:
+   ```
+   chmod +x setup_directories.sh
+   ./setup_directories.sh
+   ```
+3. Create a virtual environment:
+   ```
+   python -m venv .venv
+   source .venv/bin/activate  # On Windows: .venv\Scripts\activate
+   ```
+4. Install dependencies:
+   ```
+   pip install -r requirements.txt
+   ```
+5. Edit the `.env` file with your OpenAI API key:
+   ```
+   OPENAI_API_KEY=your_api_key_here
+   ```
+
+## Usage
+
+### Command Line Interface
+
+Run the main application:
+
+```
+python app/main.py
+```
+
+Follow the prompts to input a YouTube URL containing a golf swing recording.
+
+### Streamlit Web Interface
+
+Run the Streamlit web app using the provided shell script:
+
+```
+./run_streamlit.sh
+```
+
+Or manually with:
+
+```
+source .venv/bin/activate
+ 
+```
+
+The web interface provides:
+- Options to upload a video or use a YouTube URL
+- Control over frame skip rate for YOLO detection
+- Toggle for enabling/disabling GPT analysis
+- Interactive display of analysis results
+- Option to create and view annotated videos
+
+## File Organization
+
+- **downloads/**: Contains both downloaded YouTube videos and annotated videos
+- All videos (both original and annotated) are stored in the same directory for easy access
+
+## Troubleshooting
+
+If you encounter issues with the "Create Annotated Video" button:
+1. Make sure you've run the setup script to create the downloads directory
+2. Check that the `downloads` directory has write permissions
+3. Try restarting the Streamlit app
+
+## Requirements
+
+- Python 3.8+
+- OpenCV
+- YOLOv8
+- MediaPipe
+- yt-dlp
+- OpenAI API key
+- Streamlit 

app/main.py

@@ -0,0 +1,102 @@
+#!/usr/bin/env python3
+"""
+Golf Swing Analysis - Main Application
+"""
+
+import os
+import sys
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+
+# Add the app directory to the path
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from app.utils.video_downloader import download_youtube_video
+from app.utils.video_processor import process_video
+from app.models.pose_estimator import analyze_pose
+from app.models.swing_analyzer import segment_swing, analyze_trajectory
+from app.models.llm_analyzer import generate_swing_analysis
+from app.utils.visualizer import create_annotated_video
+
+
+def main():
+    """Main application function"""
+    print("\n===== Golf Swing Analysis =====\n")
+
+    # Step 1: Get YouTube URL from user
+    youtube_url = input("Enter YouTube URL of golf swing: ")
+
+    # Step 2: Configure analysis options
+    enable_gpt = input(
+        "\nEnable GPT analysis? (y/n, default: y): ").lower() != 'n'
+
+    sample_rate_input = input(
+        "\nFrame skip rate for YOLO (1-10, default: 5): ")
+    sample_rate = 5  # Default value
+    if sample_rate_input.isdigit():
+        sample_rate = max(1, min(10, int(sample_rate_input)))
+
+    try:
+        # Step 3: Download the video
+        print("\nDownloading video...")
+        video_path = download_youtube_video(youtube_url)
+        print(f"Video downloaded to: {video_path}")
+
+        # Step 4: Process video and detect golfer, club, and ball
+        print("\nProcessing video and detecting objects...")
+        frames, detections = process_video(video_path, sample_rate=sample_rate)
+
+        # Step 5: Analyze pose throughout the swing
+        print("\nAnalyzing golfer's pose...")
+        pose_data = analyze_pose(frames)
+
+        # Step 6: Segment swing into phases
+        print("\nSegmenting swing phases...")
+        swing_phases = segment_swing(pose_data,
+                                     detections,
+                                     sample_rate=sample_rate)
+
+        # Step 7: Analyze trajectory and speed
+        print("\nAnalyzing trajectory and speed...")
+        trajectory_data = analyze_trajectory(frames,
+                                             detections,
+                                             swing_phases,
+                                             sample_rate=sample_rate)
+
+        # Step 8: Generate swing analysis using LLM (if enabled)
+        if enable_gpt:
+            print("\nGenerating swing analysis and coaching tips...")
+            analysis = generate_swing_analysis(pose_data, swing_phases,
+                                               trajectory_data)
+
+            # Display results
+            print("\n===== Swing Analysis Results =====\n")
+            print(analysis)
+        else:
+            print("\nGPT analysis disabled. Skipping swing evaluation.")
+
+        # Step 9: Create annotated video (optional)
+        create_video = input(
+            "\nCreate annotated video? (y/n): ").lower() == 'y'
+        if create_video:
+            print("\nCreating annotated video...")
+            output_path = create_annotated_video(video_path,
+                                                 frames,
+                                                 detections,
+                                                 pose_data,
+                                                 swing_phases,
+                                                 trajectory_data,
+                                                 sample_rate=sample_rate)
+            print(f"Annotated video saved to: {output_path}")
+
+        print("\nAnalysis complete!")
+
+    except Exception as e:
+        print(f"\nError: {str(e)}")
+        return
+
+
+if __name__ == "__main__":
+    main()

app/models/llm_analyzer.py

@@ -0,0 +1,172 @@
+"""
+LLM-based golf swing analysis module
+"""
+
+import os
+import json
+from openai import OpenAI
+
+
+def generate_swing_analysis(pose_data, swing_phases, trajectory_data):
+    """
+    Generate swing analysis and coaching tips using LLM
+    
+    Args:
+        pose_data (dict): Dictionary mapping frame indices to pose keypoints
+        swing_phases (dict): Dictionary mapping phase names to lists of frame indices
+        trajectory_data (dict): Dictionary mapping frame indices to trajectory data
+        
+    Returns:
+        str: Detailed swing analysis and coaching tips
+    """
+    # Check if OpenAI API key is available
+    api_key = os.getenv("OPENAI_API_KEY")
+    if not api_key:
+        return "Error: OpenAI API key not found. Please set the OPENAI_API_KEY environment variable."
+
+    # Create OpenAI client
+    client = OpenAI(api_key=api_key)
+
+    # Prepare data for LLM
+    analysis_data = prepare_data_for_llm(pose_data, swing_phases,
+                                         trajectory_data)
+
+    # Generate prompt for LLM
+    prompt = create_llm_prompt(analysis_data)
+
+    try:
+        # Call OpenAI API
+        response = client.chat.completions.create(
+            model="gpt-4",
+            messages=[{
+                "role":
+                "system",
+                "content":
+                "You are a professional golf coach with expertise in analyzing golf swings. Provide detailed, actionable feedback based on the swing data provided."
+            }, {
+                "role": "user",
+                "content": prompt
+            }],
+            temperature=0.7,
+            max_tokens=1000)
+
+        # Extract and return analysis
+        analysis = response.choices[0].message.content
+        return analysis
+
+    except Exception as e:
+        return f"Error generating swing analysis: {str(e)}"
+
+
+def prepare_data_for_llm(pose_data, swing_phases, trajectory_data):
+    """
+    Prepare swing data for LLM analysis
+    
+    Args:
+        pose_data (dict): Dictionary mapping frame indices to pose keypoints
+        swing_phases (dict): Dictionary mapping phase names to lists of frame indices
+        trajectory_data (dict): Dictionary mapping frame indices to trajectory data
+        
+    Returns:
+        dict: Processed data for LLM analysis
+    """
+    analysis_data = {"swing_phases": {}, "joint_angles": {}, "trajectory": {}}
+
+    # Process swing phases
+    for phase, frames in swing_phases.items():
+        if frames:
+            # Get a representative frame for each phase
+            mid_frame = frames[len(frames) // 2]
+
+            # Get joint angles for the representative frame
+            if mid_frame in pose_data:
+                keypoints = pose_data[mid_frame]
+
+                # Calculate key metrics for each phase
+                analysis_data["swing_phases"][phase] = {
+                    "frame_index": mid_frame,
+                    "duration_frames": len(frames)
+                }
+
+    # Process trajectory data
+    impact_frames = swing_phases.get("impact", [])
+    if impact_frames:
+        impact_frame = impact_frames[len(impact_frames) // 2]
+        if impact_frame in trajectory_data:
+            impact_data = trajectory_data[impact_frame]
+            if "club_speed" in impact_data and impact_data["club_speed"]:
+                analysis_data["trajectory"]["club_speed_mph"] = impact_data[
+                    "club_speed"]
+
+    # Add additional metrics that would be calculated in a real implementation
+    # These are placeholder values for demonstration
+    analysis_data["metrics"] = {
+        "tempo_ratio": 3.0,  # Backswing to downswing time ratio
+        "swing_plane_consistency": 0.85,  # 0-1 scale
+        "weight_shift": 0.7,  # 0-1 scale
+        "hip_rotation": 45,  # degrees
+        "shoulder_rotation": 90,  # degrees
+        "wrist_hinge": 80,  # degrees
+        "posture_score": 0.8  # 0-1 scale
+    }
+
+    return analysis_data
+
+
+def create_llm_prompt(analysis_data):
+    """
+    Create a prompt for the LLM based on swing analysis data
+    
+    Args:
+        analysis_data (dict): Processed swing analysis data
+        
+    Returns:
+        str: Prompt for LLM
+    """
+    prompt = """
+I've analyzed a golf swing and extracted the following data:
+
+## Swing Phases
+"""
+
+    # Add swing phases information
+    for phase, data in analysis_data["swing_phases"].items():
+        prompt += f"- {phase.capitalize()}: Frame {data['frame_index']}, Duration: {data['duration_frames']} frames\n"
+
+    # Add trajectory information
+    prompt += "\n## Trajectory Data\n"
+    if "trajectory" in analysis_data and "club_speed_mph" in analysis_data[
+            "trajectory"]:
+        prompt += f"- Club Speed: {analysis_data['trajectory']['club_speed_mph']:.1f} mph\n"
+
+    # Add metrics
+    prompt += "\n## Swing Metrics\n"
+    for metric, value in analysis_data["metrics"].items():
+        # Format metric name for readability
+        metric_name = metric.replace("_", " ").title()
+
+        # Format value based on type
+        if isinstance(value, float):
+            if 0 <= value <= 1:
+                # Format as percentage for 0-1 scale metrics
+                formatted_value = f"{value * 100:.0f}%"
+            else:
+                # Format as decimal for other floats
+                formatted_value = f"{value:.1f}"
+        else:
+            # Use as is for integers and other types
+            formatted_value = str(value)
+
+        prompt += f"- {metric_name}: {formatted_value}\n"
+
+    prompt += """
+Based on this data, please provide:
+1. A detailed analysis of the golf swing
+2. Key strengths and weaknesses
+3. Specific recommendations for improvement
+4. Drills or exercises that could help address the identified issues
+
+Please be specific and actionable in your feedback.
+"""
+
+    return prompt

app/models/pose_estimator.py

@@ -0,0 +1,156 @@
+"""
+Pose estimation module for golf swing analysis
+"""
+
+import cv2
+import numpy as np
+import mediapipe as mp
+from tqdm import tqdm
+
+
+class PoseEstimator:
+    """MediaPipe-based pose estimator for golf swing analysis"""
+
+    def __init__(self):
+        """Initialize the pose estimator"""
+        self.mp_pose = mp.solutions.pose
+        self.pose = self.mp_pose.Pose(static_image_mode=False,
+                                      model_complexity=2,
+                                      enable_segmentation=False,
+                                      min_detection_confidence=0.5,
+                                      min_tracking_confidence=0.5)
+
+    def process_frame(self, frame):
+        """
+        Process a single frame and extract pose landmarks
+        
+        Args:
+            frame (numpy.ndarray): Input frame
+            
+        Returns:
+            list: List of keypoints [x, y, visibility] or None if not detected
+        """
+        # Convert BGR to RGB
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+
+        # Process the frame
+        results = self.pose.process(frame_rgb)
+
+        if not results.pose_landmarks:
+            return None
+
+        # Extract keypoints
+        keypoints = []
+        for landmark in results.pose_landmarks.landmark:
+            # Convert normalized coordinates to pixel coordinates
+            h, w, _ = frame.shape
+            x, y = int(landmark.x * w), int(landmark.y * h)
+            visibility = landmark.visibility
+            keypoints.append([x, y, visibility])
+
+        return keypoints
+
+    def close(self):
+        """Release resources"""
+        self.pose.close()
+
+
+def analyze_pose(frames):
+    """
+    Analyze pose in video frames
+    
+    Args:
+        frames (list): List of video frames
+        
+    Returns:
+        dict: Dictionary mapping frame indices to pose keypoints
+    """
+    pose_estimator = PoseEstimator()
+    pose_data = {}
+
+    for i, frame in enumerate(tqdm(frames, desc="Analyzing pose")):
+        keypoints = pose_estimator.process_frame(frame)
+        if keypoints:
+            pose_data[i] = keypoints
+
+    pose_estimator.close()
+
+    return pose_data
+
+
+def calculate_joint_angles(keypoints):
+    """
+    Calculate joint angles from pose keypoints
+    
+    Args:
+        keypoints (list): List of keypoints [x, y, visibility]
+        
+    Returns:
+        dict: Dictionary of joint angles in degrees
+    """
+    # Define joint connections for angle calculation
+    joint_connections = {
+        "right_shoulder": [
+            mp.solutions.pose.PoseLandmark.RIGHT_ELBOW.value,
+            mp.solutions.pose.PoseLandmark.RIGHT_SHOULDER.value,
+            mp.solutions.pose.PoseLandmark.RIGHT_HIP.value
+        ],
+        "left_shoulder": [
+            mp.solutions.pose.PoseLandmark.LEFT_ELBOW.value,
+            mp.solutions.pose.PoseLandmark.LEFT_SHOULDER.value,
+            mp.solutions.pose.PoseLandmark.LEFT_HIP.value
+        ],
+        "right_elbow": [
+            mp.solutions.pose.PoseLandmark.RIGHT_WRIST.value,
+            mp.solutions.pose.PoseLandmark.RIGHT_ELBOW.value,
+            mp.solutions.pose.PoseLandmark.RIGHT_SHOULDER.value
+        ],
+        "left_elbow": [
+            mp.solutions.pose.PoseLandmark.LEFT_WRIST.value,
+            mp.solutions.pose.PoseLandmark.LEFT_ELBOW.value,
+            mp.solutions.pose.PoseLandmark.LEFT_SHOULDER.value
+        ],
+        "right_hip": [
+            mp.solutions.pose.PoseLandmark.RIGHT_KNEE.value,
+            mp.solutions.pose.PoseLandmark.RIGHT_HIP.value,
+            mp.solutions.pose.PoseLandmark.RIGHT_SHOULDER.value
+        ],
+        "left_hip": [
+            mp.solutions.pose.PoseLandmark.LEFT_KNEE.value,
+            mp.solutions.pose.PoseLandmark.LEFT_HIP.value,
+            mp.solutions.pose.PoseLandmark.LEFT_SHOULDER.value
+        ],
+        "right_knee": [
+            mp.solutions.pose.PoseLandmark.RIGHT_ANKLE.value,
+            mp.solutions.pose.PoseLandmark.RIGHT_KNEE.value,
+            mp.solutions.pose.PoseLandmark.RIGHT_HIP.value
+        ],
+        "left_knee": [
+            mp.solutions.pose.PoseLandmark.LEFT_ANKLE.value,
+            mp.solutions.pose.PoseLandmark.LEFT_KNEE.value,
+            mp.solutions.pose.PoseLandmark.LEFT_HIP.value
+        ]
+    }
+
+    angles = {}
+
+    for joint_name, landmarks in joint_connections.items():
+        # Get the three points that form the angle
+        if all(landmarks[i] < len(keypoints) for i in range(3)):
+            p1 = np.array(keypoints[landmarks[0]][:2])
+            p2 = np.array(keypoints[landmarks[1]][:2])
+            p3 = np.array(keypoints[landmarks[2]][:2])
+
+            # Calculate vectors
+            v1 = p1 - p2
+            v2 = p3 - p2
+
+            # Calculate angle
+            cosine_angle = np.dot(
+                v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))
+            angle = np.arccos(np.clip(cosine_angle, -1.0, 1.0))
+            angle_degrees = np.degrees(angle)
+
+            angles[joint_name] = angle_degrees
+
+    return angles

app/models/swing_analyzer.py

@@ -0,0 +1,174 @@
+"""
+Swing analysis module for golf swing segmentation and trajectory analysis
+"""
+
+import numpy as np
+import cv2
+from app.models.pose_estimator import calculate_joint_angles
+
+
+def segment_swing(pose_data, detections, sample_rate=5):
+    """
+    Segment the golf swing into key phases
+    
+    Args:
+        pose_data (dict): Dictionary mapping frame indices to pose keypoints
+        detections (list): List of Detection objects
+        sample_rate (int): The frame sampling rate used during processing
+        
+    Returns:
+        dict: Dictionary mapping phase names to lists of frame indices
+    """
+    # Initialize swing phases
+    swing_phases = {
+        "setup": [],
+        "backswing": [],
+        "downswing": [],
+        "impact": [],
+        "follow_through": []
+    }
+
+    # Get frame indices with pose data
+    frame_indices = sorted(pose_data.keys())
+
+    if not frame_indices:
+        return swing_phases
+
+    # Calculate joint angles for each frame
+    angles_by_frame = {}
+    for idx in frame_indices:
+        keypoints = pose_data[idx]
+        angles = calculate_joint_angles(keypoints)
+        angles_by_frame[idx] = angles
+
+    # Analyze shoulder rotation to identify swing phases
+    # This is a simplified approach - a more sophisticated algorithm would be needed for production
+
+    # Find the frame with the maximum right shoulder angle (top of backswing)
+    max_shoulder_angle = -1
+    top_backswing_frame = frame_indices[0]
+
+    for idx in frame_indices:
+        angles = angles_by_frame[idx]
+        if "right_shoulder" in angles and angles[
+                "right_shoulder"] > max_shoulder_angle:
+            max_shoulder_angle = angles["right_shoulder"]
+            top_backswing_frame = idx
+
+    # Find impact frame (when club meets ball)
+    # In a real implementation, this would use club and ball detection
+    impact_frame = None
+    person_positions = {}
+
+    # Extract person positions from detections
+    for detection in detections:
+        if detection.class_name == "person":
+            frame_idx = detection.frame_idx // sample_rate  # Convert to processed frame index
+            if frame_idx in frame_indices:
+                person_positions[frame_idx] = detection.bbox
+
+    # Find the frame with the most forward position (impact)
+    if person_positions:
+        min_x = float('inf')
+        for idx, bbox in person_positions.items():
+            if idx > top_backswing_frame and bbox[0] < min_x:
+                min_x = bbox[0]
+                impact_frame = idx
+
+    # If impact frame not found, estimate it as 2/3 between top of backswing and end
+    if impact_frame is None:
+        impact_frame = frame_indices[0] + int(
+            (frame_indices[-1] - top_backswing_frame) * 2 / 3)
+
+    # Assign frames to phases
+    for idx in frame_indices:
+        if idx < frame_indices[len(frame_indices) // 5]:
+            # First 20% of frames are setup
+            swing_phases["setup"].append(idx)
+        elif idx < top_backswing_frame:
+            # Frames before top of backswing are backswing
+            swing_phases["backswing"].append(idx)
+        elif idx < impact_frame:
+            # Frames between top of backswing and impact are downswing
+            swing_phases["downswing"].append(idx)
+        elif idx < impact_frame + 5:
+            # Frames around impact
+            swing_phases["impact"].append(idx)
+        else:
+            # Remaining frames are follow-through
+            swing_phases["follow_through"].append(idx)
+
+    return swing_phases
+
+
+def analyze_trajectory(frames, detections, swing_phases, sample_rate=5):
+    """
+    Analyze club and ball trajectory and speed
+    
+    Args:
+        frames (list): List of video frames
+        detections (list): List of Detection objects
+        swing_phases (dict): Dictionary mapping phase names to lists of frame indices
+        sample_rate (int): The frame sampling rate used during processing
+        
+    Returns:
+        dict: Dictionary mapping frame indices to trajectory data
+    """
+    trajectory_data = {}
+
+    # Extract ball detections
+    ball_detections = [d for d in detections if d.class_name == "sports ball"]
+
+    # Get impact frame index
+    impact_frames = swing_phases.get("impact", [])
+    if not impact_frames:
+        return trajectory_data
+
+    impact_frame_idx = impact_frames[len(impact_frames) // 2]
+
+    # Track ball trajectory after impact
+    ball_trajectory = []
+    ball_positions = {}
+
+    for detection in ball_detections:
+        frame_idx = detection.frame_idx // sample_rate  # Convert to processed frame index
+        if frame_idx >= impact_frame_idx:
+            # Calculate ball center
+            x1, y1, x2, y2 = detection.bbox
+            center_x = (x1 + x2) / 2
+            center_y = (y1 + y2) / 2
+            ball_positions[frame_idx] = (center_x, center_y)
+
+    # Sort ball positions by frame index
+    sorted_frames = sorted(ball_positions.keys())
+    for idx in sorted_frames:
+        ball_trajectory.append(ball_positions[idx])
+
+    # Estimate club speed at impact
+    # In a real implementation, this would use more sophisticated tracking
+    club_speed = None
+    if len(swing_phases.get("downswing", [])) >= 2:
+        # Simplified club speed calculation
+        # In reality, this would require tracking the club head specifically
+        downswing_frames = swing_phases["downswing"]
+        time_diff = (downswing_frames[-1] -
+                     downswing_frames[0]) / 30  # Assuming 30 fps
+        if time_diff > 0:
+            # Simplified speed calculation (just an example)
+            club_speed = 100 * (1 / time_diff)  # Arbitrary scaling
+
+    # Populate trajectory data
+    for idx in sorted(swing_phases.keys()):
+        frames_in_phase = swing_phases[idx]
+        for frame_idx in frames_in_phase:
+            trajectory_data[frame_idx] = {
+                "phase":
+                idx,
+                "club_speed":
+                club_speed if idx == "impact" else None,
+                "ball_trajectory":
+                ball_trajectory
+                if idx == "impact" or idx == "follow_through" else None
+            }
+
+    return trajectory_data

app/streamlit_app.py

@@ -0,0 +1,281 @@
+"""
+Streamlit web UI for Golf Swing Analysis
+"""
+
+import os
+import sys
+import tempfile
+import streamlit as st
+from dotenv import load_dotenv
+import base64
+from pathlib import Path
+import shutil
+
+# Load environment variables
+load_dotenv()
+
+# Add the app directory to the path
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from app.utils.video_downloader import download_youtube_video
+from app.utils.video_processor import process_video
+from app.models.pose_estimator import analyze_pose
+from app.models.swing_analyzer import segment_swing, analyze_trajectory
+from app.models.llm_analyzer import generate_swing_analysis, create_llm_prompt, prepare_data_for_llm
+from app.utils.visualizer import create_annotated_video
+
+# Set page config
+st.set_page_config(page_title="Golf Swing Analysis",
+                   page_icon="🏌️",
+                   layout="wide",
+                   initial_sidebar_state="expanded")
+
+
+# Define functions
+def validate_youtube_url(url):
+    """Validate if the URL is a YouTube URL"""
+    return "youtube.com" in url or "youtu.be" in url
+
+
+def process_uploaded_video(uploaded_file):
+    """Process an uploaded video file"""
+    # Create downloads directory if it doesn't exist
+    os.makedirs("downloads", exist_ok=True)
+
+    # Save uploaded file to the downloads directory
+    file_path = os.path.join("downloads", uploaded_file.name)
+    with open(file_path, "wb") as f:
+        f.write(uploaded_file.getvalue())
+
+    return file_path
+
+
+def display_video(video_path):
+    """Display a video with download option"""
+    # Read video bytes
+    with open(video_path, "rb") as file:
+        video_bytes = file.read()
+
+    # Display video using st.video with bytes
+    st.video(video_bytes)
+
+    # Show download button
+    st.download_button(label="Download Video",
+                       data=video_bytes,
+                       file_name=os.path.basename(video_path),
+                       mime="video/mp4")
+
+
+# Main app
+def main():
+    """Main Streamlit application"""
+    st.title("🏌️ Golf Swing Analysis")
+    st.write("Analyze your golf swing using computer vision and AI")
+
+    # Initialize session state for storing analysis results
+    if 'video_analyzed' not in st.session_state:
+        st.session_state.video_analyzed = False
+    if 'analysis_data' not in st.session_state:
+        st.session_state.analysis_data = {
+            'video_path': None,
+            'frames': None,
+            'detections': None,
+            'pose_data': None,
+            'swing_phases': None,
+            'trajectory_data': None,
+            'sample_rate': None
+        }
+
+    # Sidebar for configuration
+    st.sidebar.title("Configuration")
+
+    # Option to enable/disable GPT analysis
+    enable_gpt = st.sidebar.checkbox("Enable GPT Analysis", value=True)
+
+    # Frame skip rate for YOLO
+    sample_rate = st.sidebar.slider(
+        "Frame Skip Rate (YOLO)",
+        min_value=1,
+        max_value=10,
+        value=5,
+        help=
+        "Process every Nth frame. Higher values = faster but less accurate.")
+
+    # Video input options
+    st.header("Video Input")
+    input_option = st.radio("Choose input method:",
+                            ["YouTube URL", "Upload Video"])
+
+    video_path = None
+    analyze_clicked = False
+
+    if input_option == "YouTube URL":
+        youtube_url = st.text_input("Enter YouTube URL of golf swing:")
+
+        analyze_clicked = st.button("Analyze Swing", key="analyze_youtube")
+        if youtube_url and analyze_clicked:
+            if validate_youtube_url(youtube_url):
+                with st.spinner("Downloading video..."):
+                    try:
+                        video_path = download_youtube_video(youtube_url)
+                        st.success("Video downloaded successfully!")
+                        display_video(video_path)
+                    except Exception as e:
+                        st.error(f"Error downloading video: {str(e)}")
+                        st.session_state.video_analyzed = False
+                        return
+            else:
+                st.error("Please enter a valid YouTube URL")
+                st.session_state.video_analyzed = False
+                return
+
+    else:  # Upload Video
+        uploaded_file = st.file_uploader("Upload a golf swing video",
+                                         type=["mp4", "mov", "avi"])
+
+        analyze_clicked = st.button("Analyze Swing", key="analyze_upload")
+        if uploaded_file and analyze_clicked:
+            with st.spinner("Processing uploaded video..."):
+                try:
+                    video_path = process_uploaded_video(uploaded_file)
+                    st.success("Video uploaded successfully!")
+                    display_video(video_path)
+                except Exception as e:
+                    st.error(f"Error processing video: {str(e)}")
+                    st.session_state.video_analyzed = False
+                    return
+
+    # Process video if available and analyze button was clicked
+    if video_path and analyze_clicked:
+        try:
+            # Step 1: Process video and detect objects
+            with st.spinner("Processing video and detecting objects..."):
+                frames, detections = process_video(video_path,
+                                                   sample_rate=sample_rate)
+                st.success(f"Processed {len(frames)} frames")
+
+            # Step 2: Analyze pose
+            with st.spinner("Analyzing golfer's pose..."):
+                pose_data = analyze_pose(frames)
+                st.success("Pose analysis complete")
+
+            # Step 3: Segment swing into phases
+            with st.spinner("Segmenting swing phases..."):
+                swing_phases = segment_swing(pose_data,
+                                             detections,
+                                             sample_rate=sample_rate)
+
+                # Display swing phases
+                st.subheader("Swing Phases")
+                phase_cols = st.columns(5)
+                for i, (phase,
+                        frames_in_phase) in enumerate(swing_phases.items()):
+                    with phase_cols[i]:
+                        st.metric(label=phase.capitalize(),
+                                  value=f"{len(frames_in_phase)} frames")
+
+            # Step 4: Analyze trajectory and speed
+            with st.spinner("Analyzing trajectory and speed..."):
+                trajectory_data = analyze_trajectory(frames,
+                                                     detections,
+                                                     swing_phases,
+                                                     sample_rate=sample_rate)
+
+                # Display club speed if available
+                impact_frames = swing_phases.get("impact", [])
+                if impact_frames:
+                    impact_frame = impact_frames[len(impact_frames) // 2]
+                    if impact_frame in trajectory_data and trajectory_data[
+                            impact_frame].get("club_speed"):
+                        st.subheader("Club Speed")
+                        st.metric(
+                            label="Estimated Club Speed",
+                            value=
+                            f"{trajectory_data[impact_frame]['club_speed']:.1f} mph"
+                        )
+
+            # Step 5: Generate swing analysis using LLM (if enabled)
+            # Prepare data for LLM regardless of whether GPT is enabled
+            analysis_data = prepare_data_for_llm(pose_data, swing_phases,
+                                                 trajectory_data)
+            prompt = create_llm_prompt(analysis_data)
+
+            # Display the GPT prompt
+            with st.expander("View GPT Prompt"):
+                st.code(prompt, language="text")
+
+            if enable_gpt:
+                with st.spinner(
+                        "Generating swing analysis and coaching tips..."):
+                    analysis = generate_swing_analysis(pose_data, swing_phases,
+                                                       trajectory_data)
+
+                    # Display analysis
+                    st.subheader("Swing Analysis")
+                    st.write(analysis)
+            else:
+                st.info(
+                    "GPT Analysis is disabled. Enable it in the sidebar to generate coaching tips."
+                )
+
+            # Store analysis data in session state
+            st.session_state.video_analyzed = True
+            st.session_state.analysis_data = {
+                'video_path': video_path,
+                'frames': frames,
+                'detections': detections,
+                'pose_data': pose_data,
+                'swing_phases': swing_phases,
+                'trajectory_data': trajectory_data,
+                'sample_rate': sample_rate
+            }
+
+        except Exception as e:
+            st.error(f"Error during analysis: {str(e)}")
+            st.session_state.video_analyzed = False
+
+    # Create annotated video button (only show if analysis is complete)
+    if st.session_state.video_analyzed:
+        st.header("Create Annotated Video")
+        st.write(
+            "Create a video with annotations showing the analysis results")
+
+        # Create a separate section for the annotated video
+        if st.button("Generate Annotated Video", key="create_annotated"):
+            try:
+                with st.spinner("Creating annotated video..."):
+                    # Create downloads directory if it doesn't exist
+                    os.makedirs("downloads", exist_ok=True)
+
+                    # Get data from session state
+                    data = st.session_state.analysis_data
+
+                    # Create the annotated video
+                    output_path = create_annotated_video(
+                        data['video_path'],
+                        data['frames'],
+                        data['detections'],
+                        data['pose_data'],
+                        data['swing_phases'],
+                        data['trajectory_data'],
+                        sample_rate=data['sample_rate'])
+
+                    # Verify the file exists
+                    if not os.path.exists(output_path):
+                        raise FileNotFoundError(
+                            f"Annotated video file not found at {output_path}")
+
+                    st.success("Annotated video created successfully!")
+
+                    # Display the video with download option
+                    display_video(output_path)
+
+            except Exception as e:
+                st.error(f"Error creating annotated video: {str(e)}")
+                st.error(
+                    "Please check if the downloads directory exists and is writable"
+                )
+
+
+if __name__ == "__main__":
+    main()

app/utils/video_downloader.py

@@ -0,0 +1,79 @@
+"""
+YouTube video downloader module using yt-dlp
+"""
+
+import os
+import yt_dlp
+
+
+def download_youtube_video(url, output_dir="downloads"):
+    """
+    Download a YouTube video from the provided URL using yt-dlp
+    
+    Args:
+        url (str): YouTube video URL
+        output_dir (str): Directory to save the downloaded video
+        
+    Returns:
+        str: Path to the downloaded video file
+        
+    Raises:
+        ValueError: If the URL is invalid or video is unavailable
+    """
+    # Create output directory if it doesn't exist
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Set output template for the downloaded file
+    output_template = os.path.join(output_dir, "%(title)s.%(ext)s")
+
+    # Configure yt-dlp options
+    ydl_opts = {
+        'format': 'best[ext=mp4]/best',  # Prefer mp4 format
+        'outtmpl': output_template,
+        'noplaylist': True,
+        'quiet': False,
+        'no_warnings': False,
+        'ignoreerrors': False,
+    }
+
+    try:
+        # Create yt-dlp object and download the video
+        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+            info = ydl.extract_info(url, download=True)
+
+            # Get the downloaded file path
+            if 'entries' in info:
+                # Playlist (should not happen with noplaylist=True)
+                raise ValueError("Playlists are not supported")
+
+            # Get video title and extension
+            title = info.get('title', 'video')
+            ext = info.get('ext', 'mp4')
+
+            # Construct the file path
+            video_path = os.path.join(output_dir, f"{title}.{ext}")
+
+            # Check if file exists
+            if not os.path.exists(video_path):
+                # Try with sanitized filename
+                sanitized_title = ''.join(c for c in title
+                                          if c.isalnum() or c in ' ._-')
+                video_path = os.path.join(output_dir,
+                                          f"{sanitized_title}.{ext}")
+
+                if not os.path.exists(video_path):
+                    # If still not found, look for any mp4 file in the directory
+                    mp4_files = [
+                        f for f in os.listdir(output_dir) if f.endswith('.mp4')
+                    ]
+                    if mp4_files:
+                        video_path = os.path.join(output_dir, mp4_files[0])
+                    else:
+                        raise ValueError("Downloaded file not found")
+
+            return video_path
+
+    except yt_dlp.utils.DownloadError as e:
+        raise ValueError(f"Error downloading video: {str(e)}")
+    except Exception as e:
+        raise ValueError(f"Error: {str(e)}")

app/utils/video_processor.py

@@ -0,0 +1,91 @@
+"""
+Video processing and object detection module
+"""
+
+import cv2
+import numpy as np
+from tqdm import tqdm
+from ultralytics import YOLO
+
+
+class Detection:
+    """Class to store detection results"""
+
+    def __init__(self, frame_idx, class_id, class_name, bbox, confidence):
+        self.frame_idx = frame_idx
+        self.class_id = class_id
+        self.class_name = class_name
+        self.bbox = bbox  # [x1, y1, x2, y2]
+        self.confidence = confidence
+
+
+def process_video(video_path, sample_rate=5):
+    """
+    Process video and detect golfer, club, and ball
+    
+    Args:
+        video_path (str): Path to the video file
+        sample_rate (int): Process every nth frame
+        
+    Returns:
+        tuple: (frames, detections)
+            - frames: List of processed frames
+            - detections: List of Detection objects
+    """
+    # Load YOLOv8 model
+    model = YOLO("yolov8n.pt")
+
+    # Custom class names for golf-specific objects
+    class_names = model.names
+
+    # Open video file
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        raise ValueError("Error opening video file")
+
+    # Get video properties
+    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+
+    frames = []
+    detections = []
+
+    # Process frames
+    for frame_idx in tqdm(range(0, frame_count, sample_rate),
+                          desc="Processing frames"):
+        # Set frame position
+        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
+
+        # Read frame
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        # Store original frame
+        frames.append(frame)
+
+        # Run YOLOv8 detection
+        results = model(frame)
+
+        # Process detection results
+        for result in results:
+            boxes = result.boxes
+            for box in boxes:
+                # Get detection information
+                class_id = int(box.cls.item())
+                class_name = class_names[class_id]
+
+                # Filter for relevant objects (person, sports ball)
+                if class_name in ["person", "sports ball"]:
+                    bbox = box.xyxy[0].tolist()  # [x1, y1, x2, y2]
+                    confidence = box.conf.item()
+
+                    # Create Detection object
+                    detection = Detection(frame_idx, class_id, class_name,
+                                          bbox, confidence)
+                    detections.append(detection)
+
+    # Release video capture
+    cap.release()
+
+    return frames, detections

app/utils/visualizer.py

@@ -0,0 +1,225 @@
+"""
+Visualization module for creating annotated videos
+"""
+
+import os
+import cv2
+import numpy as np
+from tqdm import tqdm
+import logging
+import mediapipe as mp
+
+# Define body part groups and their colors
+BODY_PART_COLORS = {
+    "head": (255, 0, 0),  # Blue
+    "torso": (0, 255, 0),  # Green
+    "arms": (255, 165, 0),  # Orange
+    "hands": (255, 0, 255),  # Magenta
+    "legs": (0, 255, 255),  # Cyan
+    "feet": (255, 255, 0)  # Yellow
+}
+
+# Define which landmarks belong to which body part groups
+BODY_PARTS_MAPPING = {
+    "head": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10],  # Nose, eyes, ears, mouth
+    "torso": [11, 12, 23, 24],  # Shoulders and hips
+    "arms": [11, 12, 13, 14],  # Shoulders and elbows
+    "hands": [15, 16, 17, 18, 19, 20, 21,
+              22],  # Wrists, pinkies, indices, thumbs
+    "legs": [23, 24, 25, 26],  # Hips and knees
+    "feet": [27, 28, 29, 30, 31, 32]  # Ankles, heels, foot indices
+}
+
+
+def create_annotated_video(video_path,
+                           frames,
+                           detections,
+                           pose_data,
+                           swing_phases,
+                           trajectory_data,
+                           output_dir="downloads",
+                           sample_rate=5):
+    """
+    Create an annotated video with swing analysis visualizations
+    
+    Args:
+        video_path (str): Path to the original video
+        frames (list): List of video frames
+        detections (list): List of Detection objects
+        pose_data (dict): Pose estimation data
+        swing_phases (dict): Swing phase segmentation data
+        trajectory_data (dict): Trajectory and speed analysis data
+        output_dir (str): Directory to save the output video
+        sample_rate (int): The frame sampling rate used during processing
+        
+    Returns:
+        str: Path to the annotated video
+    """
+    try:
+        # Create output directory if it doesn't exist
+        os.makedirs(output_dir, exist_ok=True)
+
+        # Get original video filename without extension
+        video_name = os.path.splitext(os.path.basename(video_path))[0]
+        output_path = os.path.join(output_dir, f"{video_name}_annotated.mp4")
+
+        # Get video properties
+        if not frames or len(frames) == 0:
+            raise ValueError("No frames provided for annotation")
+
+        height, width = frames[0].shape[:2]
+        fps = 30  # Default fps
+
+        # Create video writer
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+
+        if not out.isOpened():
+            raise IOError(
+                f"Failed to create video writer for {output_path}. Check directory permissions."
+            )
+
+        # Process each frame
+        for i, frame in enumerate(tqdm(frames,
+                                       desc="Creating annotated video")):
+            # Create a copy of the frame for annotations
+            annotated_frame = frame.copy()
+
+            # Draw detections
+            frame_detections = [
+                d for d in detections if d.frame_idx == i * sample_rate
+            ]
+            for detection in frame_detections:
+                x1, y1, x2, y2 = map(int, detection.bbox)
+
+                # Draw bounding box
+                color = (0, 255,
+                         0) if detection.class_name == "person" else (0, 0,
+                                                                      255)
+                cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), color, 2)
+
+                # Draw label
+                label = f"{detection.class_name}: {detection.confidence:.2f}"
+                cv2.putText(annotated_frame, label, (x1, y1 - 10),
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+
+            # Draw pose keypoints with different colors for different body parts
+            if i in pose_data:
+                keypoints = pose_data[i]
+
+                # Draw each keypoint with its corresponding body part color
+                for part_name, part_indices in BODY_PARTS_MAPPING.items():
+                    color = BODY_PART_COLORS[part_name]
+                    for idx in part_indices:
+                        if idx < len(keypoints
+                                     ) and keypoints[idx] is not None and len(
+                                         keypoints[idx]) >= 2:
+                            x, y = int(keypoints[idx][0]), int(
+                                keypoints[idx][1])
+                            cv2.circle(annotated_frame, (x, y), 5, color, -1)
+
+                # Draw connections between keypoints
+                mp_pose = mp.solutions.pose
+                connections = mp_pose.POSE_CONNECTIONS
+
+                for connection in connections:
+                    start_idx, end_idx = connection
+
+                    if (start_idx < len(keypoints) and end_idx < len(keypoints)
+                            and keypoints[start_idx] is not None
+                            and keypoints[end_idx] is not None
+                            and len(keypoints[start_idx]) >= 2
+                            and len(keypoints[end_idx]) >= 2):
+
+                        # Determine the color based on the body part of the start point
+                        color = None
+                        for part_name, part_indices in BODY_PARTS_MAPPING.items(
+                        ):
+                            if start_idx in part_indices:
+                                color = BODY_PART_COLORS[part_name]
+                                break
+
+                        # If no color found, use white
+                        if color is None:
+                            color = (255, 255, 255)
+
+                        start_point = (int(keypoints[start_idx][0]),
+                                       int(keypoints[start_idx][1]))
+                        end_point = (int(keypoints[end_idx][0]),
+                                     int(keypoints[end_idx][1]))
+
+                        cv2.line(annotated_frame, start_point, end_point,
+                                 color, 2)
+
+            # Draw swing phase information
+            phase = None
+            for phase_name, phase_frames in swing_phases.items():
+                if i in phase_frames:
+                    phase = phase_name
+                    break
+
+            if phase:
+                cv2.putText(annotated_frame, f"Phase: {phase}", (10, 30),
+                            cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
+
+            # Draw trajectory information if available
+            if i in trajectory_data:
+                traj_info = trajectory_data[i]
+                if "club_speed" in traj_info and traj_info["club_speed"]:
+                    cv2.putText(
+                        annotated_frame,
+                        f"Club Speed: {traj_info['club_speed']:.1f} mph",
+                        (10, 70), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0),
+                        2)
+
+                if "ball_trajectory" in traj_info and traj_info[
+                        "ball_trajectory"]:
+                    points = traj_info["ball_trajectory"]
+                    for j in range(1, len(points)):
+                        pt1 = (int(points[j - 1][0]), int(points[j - 1][1]))
+                        pt2 = (int(points[j][0]), int(points[j][1]))
+                        cv2.line(annotated_frame, pt1, pt2, (0, 255, 255), 2)
+
+            # Add legend for body part colors
+            legend_y_start = 110
+            legend_y_spacing = 30
+            legend_x = 10
+            legend_box_size = 20
+
+            # Draw legend title
+            cv2.putText(annotated_frame, "Body Parts Legend:",
+                        (legend_x, legend_y_start - 10),
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
+
+            # Draw color boxes and labels for each body part
+            for idx, (part_name, color) in enumerate(BODY_PART_COLORS.items()):
+                y_pos = legend_y_start + idx * legend_y_spacing
+
+                # Draw color box
+                cv2.rectangle(annotated_frame,
+                              (legend_x, y_pos - legend_box_size + 5),
+                              (legend_x + legend_box_size, y_pos + 5), color,
+                              -1)
+
+                # Draw part name
+                cv2.putText(annotated_frame, part_name.capitalize(),
+                            (legend_x + legend_box_size + 10, y_pos + 5),
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
+
+            # Write the annotated frame to the output video
+            out.write(annotated_frame)
+
+        # Release video writer
+        out.release()
+
+        # Verify the file was created
+        if not os.path.exists(output_path) or os.path.getsize(
+                output_path) == 0:
+            raise IOError(f"Failed to create video file at {output_path}")
+
+        print(f"Annotated video saved to: {output_path}")
+        return output_path
+
+    except Exception as e:
+        print(f"Error creating annotated video: {str(e)}")
+        raise

requirements.txt

@@ -0,0 +1,12 @@
+opencv-python==4.8.1.78
+yt-dlp==2025.2.19
+ultralytics==8.1.0
+mediapipe==0.10.13
+numpy==1.26.4
+matplotlib==3.8.0
+torch==2.2.0
+torchvision==0.17.0
+openai==1.6.0
+python-dotenv==1.0.0
+tqdm==4.66.1
+streamlit==1.30.0 

run_streamlit.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+
+# Activate the virtual environment
+source .venv/bin/activate
+
+# Run the Streamlit app
+streamlit run app/streamlit_app.py
+
+# Deactivate the virtual environment when done
+deactivate 

setup_directories.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+# Create downloads directory for the application
+mkdir -p downloads
+
+# Set permissions
+chmod 755 downloads
+
+echo "Directory created and permissions set:"
+echo "- downloads: for storing downloaded YouTube videos and annotated videos"
+
+# Create .env file if it doesn't exist
+if [ ! -f .env ]; then
+    echo "Creating .env file template..."
+    echo "OPENAI_API_KEY=your_api_key_here" > .env
+    echo ".env file created. Please edit it to add your OpenAI API key."
+fi
+
+echo "Setup complete!"