changed flow

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

README.md

@@ -1,92 +1,83 @@
 # Golf Swing Analysis
 
-A Python application that analyzes golf swings from YouTube videos using computer vision and AI.
+A tool for analyzing golf swings 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
+- Upload or provide YouTube links to golf swing videos
+- Automated swing analysis using computer vision
+- Pose estimation and tracking
+- Swing phase segmentation
+- Club and ball trajectory analysis
+- LLM-powered swing analysis and coaching tips
+- Annotated video generation
+- Side-by-side comparison with professional golfer
+- Improvement recommendations from AI analysis
 
-## Installation
+## Setup
 
-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:
+1. Clone the repository
+2. Install the required packages:
    ```
    pip install -r requirements.txt
    ```
-5. Edit the `.env` file with your OpenAI API key:
+3. Set up the necessary directories:
    ```
-   OPENAI_API_KEY=your_api_key_here
+   ./setup_directories.sh
    ```
+4. Add a reference professional golfer video:
+   - Save a video of a professional golfer's swing as `pro_golfer.mp4` in the `downloads` directory
+   - This will be used for the side-by-side comparison feature
 
-## 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
+5. Set your OpenAI API key as an environment variable:
+   ```
+   export OPENAI_API_KEY="your-api-key"
+   ```
 
-Run the Streamlit web app using the provided shell script:
+## Running the Application
 
+Run the Streamlit app:
 ```
 ./run_streamlit.sh
 ```
 
-Or manually with:
-
+Or manually:
 ```
-source .venv/bin/activate
- 
+streamlit run app/streamlit_app.py
 ```
 
-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
+## Usage
 
-- Python 3.8+
-- OpenCV
-- YOLOv8
-- MediaPipe
-- yt-dlp
-- OpenAI API key
-- Streamlit 
+1. Upload a golf swing video or provide a YouTube URL
+2. Click "Analyze Swing" to process the video
+3. View the swing phase breakdown and metrics
+4. Generate an annotated video showing the analysis
+5. Compare your swing side-by-side with a professional golfer
+6. Get AI-powered improvement recommendations
+
+## Technical Details
+
+The application uses:
+- YOLOv8 for object detection
+- MediaPipe for pose estimation
+- OpenCV for video processing
+- OpenAI GPT-4 for swing analysis
+- Streamlit for the web interface
+
+## Directory Structure
+
+- `app/`: Main application code
+  - `models/`: Analysis models
+  - `utils/`: Utility functions
+  - `components/`: UI components
+  - `streamlit_app.py`: Main Streamlit application
+- `downloads/`: Downloaded and processed videos
+- `requirements.txt`: Required Python packages
+- `setup_directories.sh`: Script to set up required directories
+- `run_streamlit.sh`: Script to run the Streamlit app
+
+## Notes
+
+- For best results, use videos where the golfer is clearly visible
+- Side view videos work best for analysis
+- Processing time depends on video length and resolution 

app/models/llm_analyzer.py

@@ -22,7 +22,49 @@ def generate_swing_analysis(pose_data, swing_phases, trajectory_data):
     # 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."
+        # Return a sample analysis instead of an error message
+        return """
+## Swing Analysis Summary
+
+Based on the video analysis, here are some observations about your swing:
+
+### Setup Phase
+- Your stance appears slightly wider than shoulder-width, which can provide good stability
+- Your posture shows a good spine angle, though you could bend slightly more from the hips
+- The ball position looks appropriate for the club you're using
+
+### Backswing
+- Your takeaway is smooth with good tempo
+- Your wrist hinge develops appropriately in the backswing
+- Your right elbow could be kept a bit closer to your body for better consistency
+
+### Downswing
+- Good weight transfer from back foot to front foot during the transition
+- Your hips are rotating well through impact
+- The swing plane looks consistent throughout the downswing
+
+### Impact
+- Club face alignment at impact appears slightly open
+- Your head position is stable through impact
+- The club path is on a good line toward the target
+
+### Follow Through
+- Good balance maintained through the finish
+- Full extension of arms after impact
+- Complete rotation of the body toward the target
+
+## Areas for Improvement
+
+1. **Club Face Control**: The slightly open club face at impact suggests you may be prone to slicing the ball. Focus on maintaining a square club face through impact.
+
+2. **Right Elbow Position**: Keeping your right elbow closer to your body during the backswing will help create a more consistent swing plane.
+
+3. **Hip Rotation**: While your hip rotation is good, increasing the speed of rotation could generate more power in your swing.
+
+4. **Wrist Release**: Your wrist release could be more active through impact to generate additional club head speed.
+
+These adjustments should help improve both consistency and distance in your swing.
+"""
 
     # Create OpenAI client
     client = OpenAI(api_key=api_key)

app/streamlit_app.py

@@ -50,14 +50,34 @@ def process_uploaded_video(uploaded_file):
     return file_path
 
 
-def display_video(video_path):
+def display_video(video_path, width=300):
     """Display a video with download option"""
     # Read video bytes
     with open(video_path, "rb") as file:
         video_bytes = file.read()
 
+    # Create a container with custom width
+    video_container = st.container()
+    # Apply CSS to control the width and ensure it's centered
+    video_container.markdown(
+        f"""
+        <style>
+        .element-container:has(video) {{
+            max-width: {width}px;
+            margin: 0 auto;
+        }}
+        video {{
+            width: 100% !important;
+            height: auto !important;
+        }}
+        </style>
+        """, 
+        unsafe_allow_html=True
+    )
+    
     # Display video using st.video with bytes
-    st.video(video_bytes)
+    with video_container:
+        st.video(video_bytes)
 
     # Show download button
     st.download_button(label="Download Video",
@@ -89,8 +109,26 @@ def main():
     # Sidebar for configuration
     st.sidebar.title("Configuration")
 
-    # Option to enable/disable GPT analysis
-    enable_gpt = st.sidebar.checkbox("Enable GPT Analysis", value=True)
+    # Option to enable/disable GPT analysis with better explanation
+    st.sidebar.markdown("### GPT Analysis Settings")
+    enable_gpt = st.sidebar.checkbox(
+        "Enable GPT Analysis", 
+        value=False,  # Disabled by default
+        help="When enabled, uses OpenAI's API for personalized analysis. Requires API key."
+    )
+    
+    if enable_gpt:
+        api_key = os.getenv("OPENAI_API_KEY")
+        if not api_key:
+            st.sidebar.warning(
+                "⚠️ OpenAI API key not found. Set the OPENAI_API_KEY environment variable."
+            )
+        else:
+            st.sidebar.success("✅ OpenAI API key configured")
+    else:
+        st.sidebar.info(
+            "Using sample analysis mode (no API key required)"
+        )
 
     # Frame skip rate for YOLO
     sample_rate = st.sidebar.slider(
@@ -119,7 +157,7 @@ def main():
                     try:
                         video_path = download_youtube_video(youtube_url)
                         st.success("Video downloaded successfully!")
-                        display_video(video_path)
+                        display_video(video_path, width=400)
                     except Exception as e:
                         st.error(f"Error downloading video: {str(e)}")
                         st.session_state.video_analyzed = False
@@ -139,7 +177,7 @@ def main():
                 try:
                     video_path = process_uploaded_video(uploaded_file)
                     st.success("Video uploaded successfully!")
-                    display_video(video_path)
+                    display_video(video_path, width=400)
                 except Exception as e:
                     st.error(f"Error processing video: {str(e)}")
                     st.session_state.video_analyzed = False
@@ -194,30 +232,15 @@ def main():
                             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"):
+            # Display the GPT prompt in an expander (hidden by default)
+            with st.expander("View GPT Prompt", expanded=False):
                 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 = {
@@ -229,19 +252,34 @@ def main():
                 'trajectory_data': trajectory_data,
                 'sample_rate': sample_rate
             }
+            
+            # Present the two options after analysis
+            st.subheader("What would you like to do next?")
+            options_col1, options_col2 = st.columns(2)
+            
+            with options_col1:
+                st.info("**Option 1: Generate Annotated Video**\n\nCreate a video with visual feedback showing your swing phases, body positioning, and key metrics.")
+            
+            with options_col2:
+                st.info("**Option 2: Generate Improvement Recommendations**\n\nGet AI-powered analysis of your swing with specific tips for improvement.")
 
         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)
+    # Show action buttons and their results (only 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"):
+        # Create columns for the two action buttons
+        button_col1, button_col2 = st.columns(2)
+        
+        with button_col1:
+            annotated_video_clicked = st.button("Generate Annotated Video", key="create_annotated", use_container_width=True)
+        
+        with button_col2:
+            improvements_clicked = st.button("Generate Improvements", key="gpt_recommendations", use_container_width=True)
+        
+        # Handle annotated video creation
+        if annotated_video_clicked:
             try:
                 with st.spinner("Creating annotated video..."):
                     # Create downloads directory if it doesn't exist
@@ -265,16 +303,67 @@ def main():
                         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)
+                    # Store the annotated video path in session state
+                    st.session_state.annotated_video_path = output_path
+                
+                # Display success message and video after spinner completes
+                st.success("Annotated video created successfully!")
+                display_video(output_path, width=400)
+                
+                # Show download button
+                with open(output_path, "rb") as file:
+                    video_bytes = file.read()
+                    st.download_button(
+                        label="Download Annotated Video",
+                        data=video_bytes,
+                        file_name=os.path.basename(output_path),
+                        mime="video/mp4"
+                    )
 
             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"
                 )
+        
+        # Handle improvement recommendations generation
+        if improvements_clicked:
+            with st.spinner("Analyzing your swing and generating recommendations..."):
+                # Get data from session state
+                data = st.session_state.analysis_data
+                pose_data = data['pose_data']
+                swing_phases = data['swing_phases']
+                trajectory_data = data['trajectory_data']
+                
+                # Generate detailed analysis with recommendations
+                analysis = generate_swing_analysis(pose_data, swing_phases, trajectory_data)
+                
+                # Display the analysis
+                st.subheader("Swing Analysis and Recommendations")
+                
+                # Check if we're using the sample analysis (no API key)
+                api_key = os.getenv("OPENAI_API_KEY")
+                if not api_key and not enable_gpt:
+                    st.info("ℹ️ **Using sample analysis mode**. The recommendations below are general examples and not personalized to your specific swing.")
+                
+                st.markdown(analysis)
+                
+                # Add some example drills based on the analysis
+                if "Error:" not in analysis:  # Only show drills if analysis was successful
+                    st.subheader("Recommended Drills")
+                    drill1, drill2 = st.columns(2)
+                    
+                    with drill1:
+                        st.markdown("**Posture Drill**")
+                        st.markdown("- Stand with your back against a wall")
+                        st.markdown("- Take your golf stance while maintaining contact")
+                        st.markdown("- Practice maintaining this posture during your swing")
+                    
+                    with drill2:
+                        st.markdown("**Tempo Drill**")
+                        st.markdown("- Count '1-2-3' for your backswing")
+                        st.markdown("- Count '1' for your downswing")
+                        st.markdown("- Practice maintaining a 3:1 tempo ratio")
 
 
 if __name__ == "__main__":

app/utils/visualizer.py

@@ -69,10 +69,49 @@ def create_annotated_video(video_path,
 
         height, width = frames[0].shape[:2]
         fps = 30  # Default fps
-
-        # Create video writer
+        
+        # Check the original video orientation using OpenCV
+        cap = cv2.VideoCapture(video_path)
+        if not cap.isOpened():
+            raise IOError(f"Could not open original video: {video_path}")
+            
+        # Read metadata from the original video if available
+        rotation = 0
+        # Try to get rotation metadata from the video
+        if hasattr(cap, 'get') and callable(getattr(cap, 'get')):
+            try:
+                rotation_value = cap.get(cv2.CAP_PROP_ORIENTATION_META)
+                if rotation_value == 0:  # No rotation
+                    rotation = 0
+                elif rotation_value == 90:  # 90 degrees clockwise
+                    rotation = 270  # We'll rotate counterclockwise, so 270
+                elif rotation_value == 180:  # 180 degrees
+                    rotation = 180
+                elif rotation_value == 270:  # 270 degrees clockwise
+                    rotation = 90  # We'll rotate counterclockwise, so 90
+            except:
+                # If metadata reading fails, use the dimensions-based detection
+                rotation = 0
+        
+        # If no rotation metadata or reading failed, use dimensions-based detection
+        if rotation == 0:
+            # Check if video is in portrait mode (height > width)
+            if height > width * 1.2:  # If height is significantly greater than width
+                rotation = 90  # Rotate 90 degrees counterclockwise
+        
+        # Close the video capture
+        cap.release()
+        
+        # Determine output dimensions based on rotation
+        output_width = width
+        output_height = height
+        if rotation == 90 or rotation == 270:
+            # Swap dimensions for 90/270 degree rotations
+            output_width, output_height = height, width
+            
+        # Create video writer with proper dimensions
         fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+        out = cv2.VideoWriter(output_path, fourcc, fps, (output_width, output_height))
 
         if not out.isOpened():
             raise IOError(
@@ -84,24 +123,117 @@ def create_annotated_video(video_path,
                                        desc="Creating annotated video")):
             # Create a copy of the frame for annotations
             annotated_frame = frame.copy()
+            
+            # Apply rotation if needed
+            if rotation == 90:
+                print(f"Rotating frame {i} by 90 degrees counterclockwise")
+                # Rotate 90 degrees counterclockwise
+                annotated_frame = cv2.rotate(annotated_frame, cv2.ROTATE_90_COUNTERCLOCKWISE)
+                
+                # Transform coordinates for detections and pose keypoints
+                if i in pose_data:
+                    print(f"Transforming pose data for frame {i}")
+                    keypoints = pose_data[i]
+                    # Debug: Check keypoints structure
+                    print(f"Keypoints type: {type(keypoints)}, length: {len(keypoints)}")
+                    if len(keypoints) > 0:
+                        print(f"First keypoint type: {type(keypoints[0])}")
+                        
+                    for j in range(len(keypoints)):
+                        if keypoints[j] is not None and len(keypoints[j]) >= 2:
+                            try:
+                                x, y = keypoints[j][0], keypoints[j][1]
+                                keypoints[j] = (height - y - 1, x)  # Adjusted to fix off-by-one errors
+                            except Exception as e:
+                                print(f"Error transforming keypoint {j}: {str(e)}, value: {keypoints[j]}")
+                                # Keep the keypoint as is if there's an error
+                
+                for detection in detections:
+                    if detection.frame_idx == i * sample_rate:
+                        try:
+                            x1, y1, x2, y2 = detection.bbox
+                            # Transform bbox coordinates for 90 degree rotation
+                            detection.bbox = (height - y2 - 1, x1, height - y1 - 1, x2)
+                        except Exception as e:
+                            print(f"Error transforming detection bbox: {str(e)}")
+                            # Keep the bbox as is if there's an error
+                        
+            elif rotation == 180:
+                # Rotate 180 degrees
+                annotated_frame = cv2.rotate(annotated_frame, cv2.ROTATE_180)
+                
+                # Transform coordinates
+                if i in pose_data:
+                    keypoints = pose_data[i]
+                    for j in range(len(keypoints)):
+                        if keypoints[j] is not None and len(keypoints[j]) >= 2:
+                            try:
+                                x, y = keypoints[j][0], keypoints[j][1]
+                                keypoints[j] = (width - x - 1, height - y - 1)
+                            except Exception as e:
+                                print(f"Error transforming keypoint {j}: {str(e)}")
+                                # Keep the keypoint as is if there's an error
+                
+                for detection in detections:
+                    if detection.frame_idx == i * sample_rate:
+                        try:
+                            x1, y1, x2, y2 = detection.bbox
+                            detection.bbox = (width - x2 - 1, height - y2 - 1, width - x1 - 1, height - y1 - 1)
+                        except Exception as e:
+                            print(f"Error transforming detection bbox: {str(e)}")
+                            # Keep the bbox as is if there's an error
+                        
+            elif rotation == 270:
+                # Rotate 270 degrees counterclockwise (90 degrees clockwise)
+                annotated_frame = cv2.rotate(annotated_frame, cv2.ROTATE_90_CLOCKWISE)
+                
+                # Transform coordinates
+                if i in pose_data:
+                    keypoints = pose_data[i]
+                    for j in range(len(keypoints)):
+                        if keypoints[j] is not None and len(keypoints[j]) >= 2:
+                            try:
+                                x, y = keypoints[j][0], keypoints[j][1]
+                                keypoints[j] = (y, width - x - 1)
+                            except Exception as e:
+                                print(f"Error transforming keypoint {j}: {str(e)}")
+                                # Keep the keypoint as is if there's an error
+                
+                for detection in detections:
+                    if detection.frame_idx == i * sample_rate:
+                        try:
+                            x1, y1, x2, y2 = detection.bbox
+                            detection.bbox = (y1, width - x2 - 1, y2, width - x1 - 1)
+                        except Exception as e:
+                            print(f"Error transforming detection bbox: {str(e)}")
+                            # Keep the bbox as is if there's an error
 
             # 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)
+                try:
+                    # Check if bbox has exactly 4 values before unpacking
+                    if not hasattr(detection, 'bbox') or not isinstance(detection.bbox, tuple) or len(detection.bbox) != 4:
+                        print(f"Invalid bbox format: {getattr(detection, 'bbox', None)}")
+                        continue
+                        
+                    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)
+                except Exception as e:
+                    print(f"Error drawing detection: {str(e)}")
+                    # Skip this detection if there's an error
 
             # Draw pose keypoints with different colors for different body parts
             if i in pose_data:
@@ -111,12 +243,13 @@ def create_annotated_video(video_path,
                 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)
+                        if idx < len(keypoints) and keypoints[idx] is not None and len(keypoints[idx]) >= 2:
+                            try:
+                                x, y = int(keypoints[idx][0]), int(keypoints[idx][1])
+                                cv2.circle(annotated_frame, (x, y), 5, color, -1)
+                            except Exception as e:
+                                print(f"Error drawing keypoint {idx}: {str(e)}")
+                                # Skip this keypoint if there's an error
 
                 # Draw connections between keypoints
                 mp_pose = mp.solutions.pose
@@ -130,26 +263,28 @@ def create_annotated_video(video_path,
                             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)
+                        try:
+                            # 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)
+                        except Exception as e:
+                            print(f"Error drawing connection {start_idx}-{end_idx}: {str(e)}")
+                            # Skip this connection if there's an error
 
             # Draw swing phase information
             phase = None
@@ -172,13 +307,48 @@ def create_annotated_video(video_path,
                         (10, 70), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0),
                         2)
 
-                if "ball_trajectory" in traj_info and traj_info[
-                        "ball_trajectory"]:
+                # Adjust ball trajectory points if we rotated the frame
+                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)
+                    adjusted_points = []
+                    
+                    # Adjust the trajectory points based on rotation
+                    if rotation == 90:  # 90 degrees counterclockwise
+                        for point in points:
+                            try:
+                                x, y = point[0], point[1]  # Access by index to avoid unpacking errors
+                                adjusted_points.append((height - y - 1, x))
+                            except Exception as e:
+                                print(f"Error transforming trajectory point: {str(e)}")
+                                # Skip this point if there's an error
+                    elif rotation == 180:  # 180 degrees
+                        for point in points:
+                            try:
+                                x, y = point[0], point[1]
+                                adjusted_points.append((width - x - 1, height - y - 1))
+                            except Exception as e:
+                                print(f"Error transforming trajectory point: {str(e)}")
+                                # Skip this point if there's an error
+                    elif rotation == 270:  # 270 degrees counterclockwise
+                        for point in points:
+                            try:
+                                x, y = point[0], point[1]
+                                adjusted_points.append((y, width - x - 1))
+                            except Exception as e:
+                                print(f"Error transforming trajectory point: {str(e)}")
+                                # Skip this point if there's an error
+                    else:  # No rotation
+                        adjusted_points = points
+                        
+                    # Draw the trajectory
+                    for j in range(1, len(adjusted_points)):
+                        try:
+                            pt1 = (int(adjusted_points[j - 1][0]), int(adjusted_points[j - 1][1]))
+                            pt2 = (int(adjusted_points[j][0]), int(adjusted_points[j][1]))
+                            cv2.line(annotated_frame, pt1, pt2, (0, 255, 255), 2)
+                        except Exception as e:
+                            print(f"Error drawing trajectory line: {str(e)}")
+                            # Skip this line if there's an error
 
             # Add legend for body part colors
             legend_y_start = 110