compare with pro

app/streamlit_app.py

@@ -17,12 +17,13 @@ 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_downloader import download_youtube_video, download_pro_reference
 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, check_llm_services
 from app.utils.visualizer import create_annotated_video
+from app.utils.comparison import create_frame_by_frame_comparison
 
 # Set page config
 st.set_page_config(page_title="Par-ity Project: Golf Swing Analysis 🏌️‍♀️",
@@ -102,6 +103,8 @@ def main():
             'trajectory_data': None,
             'sample_rate': None
         }
+    if 'pro_reference_path' not in st.session_state:
+        st.session_state.pro_reference_path = None
 
     # Sidebar for configuration
     st.sidebar.title("Configuration")
@@ -164,6 +167,23 @@ def main():
         value=5,
         help=
         "Process every Nth frame. Higher values = faster but less accurate.")
+        
+    # Pro reference toggle
+    enable_pro_comparison = st.sidebar.checkbox(
+        "Enable Pro Comparison",
+        value=True,
+        help="Compare your swing with a professional golfer reference"
+    )
+    
+    # Pro reference URL input
+    if enable_pro_comparison:
+        pro_url = st.sidebar.text_input(
+            "Pro Golfer Reference URL",
+            value="https://www.youtube.com/shorts/geR666LWSHg",
+            help="YouTube URL of professional golfer swing (default provided)"
+        )
+    else:
+        pro_url = None
 
     # Video input options
     st.header("Video Input")
@@ -208,6 +228,18 @@ def main():
                     st.error(f"Error processing video: {str(e)}")
                     st.session_state.video_analyzed = False
                     return
+    
+    # Download pro reference if enabled
+    if enable_pro_comparison and (video_path or st.session_state.video_analyzed):
+        if not st.session_state.pro_reference_path:
+            with st.spinner("Downloading professional golfer reference..."):
+                try:
+                    pro_path = download_pro_reference(pro_url)
+                    st.session_state.pro_reference_path = pro_path
+                    st.success("Professional reference downloaded successfully!")
+                except Exception as e:
+                    st.error(f"Error downloading pro reference: {str(e)}")
+                    st.session_state.pro_reference_path = None
 
     # Process video if available and analyze button was clicked
     if video_path and analyze_clicked:
@@ -255,9 +287,9 @@ def main():
                 'prompt': prompt
             }
 
-            # Present the two options after analysis
+            # Present the options after analysis
             st.subheader("What would you like to do next?")
-            options_col1, options_col2 = st.columns(2)
+            options_col1, options_col2, options_col3 = st.columns(3)
 
             with options_col1:
                 st.info(
@@ -268,6 +300,12 @@ def main():
                 st.info(
                     "**Option 2: Generate Improvement Recommendations**\n\nGet AI-powered analysis of your swing with specific tips for improvement."
                 )
+                
+            with options_col3:
+                if enable_pro_comparison and st.session_state.pro_reference_path:
+                    st.info(
+                        "**Option 3: Compare With Pro**\n\nSee a side-by-side frame-by-frame comparison with a professional golfer's swing."
+                    )
 
         except Exception as e:
             st.error(f"Error during analysis: {str(e)}")
@@ -307,8 +345,12 @@ def main():
                 st.code(st.session_state.analysis_data['prompt'],
                         language="text")
 
-        # Create columns for the two action buttons
-        button_col1, button_col2 = st.columns(2)
+        # Create columns for the action buttons
+        if enable_pro_comparison and st.session_state.pro_reference_path:
+            button_col1, button_col2, button_col3 = st.columns(3)
+        else:
+            button_col1, button_col2 = st.columns(2)
+            button_col3 = None
 
         with button_col1:
             annotated_video_clicked = st.button("Generate Annotated Video",
@@ -319,6 +361,15 @@ def main():
             improvements_clicked = st.button("Generate Improvements",
                                              key="gpt_recommendations",
                                              use_container_width=True)
+        
+        # Add pro comparison button if enabled
+        if enable_pro_comparison and st.session_state.pro_reference_path and button_col3:
+            with button_col3:
+                comparison_clicked = st.button("Compare With Pro",
+                                              key="pro_comparison",
+                                              use_container_width=True)
+        else:
+            comparison_clicked = False
 
         # Handle annotated video creation
         if annotated_video_clicked:
@@ -425,6 +476,65 @@ def main():
                         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")
+        
+        # Handle pro comparison video creation
+        if comparison_clicked and st.session_state.pro_reference_path:
+            try:
+                with st.spinner("Creating frame-by-frame comparison..."):
+                    # Get data from session state
+                    user_video_path = st.session_state.analysis_data['video_path']
+                    pro_video_path = st.session_state.pro_reference_path
+                    
+                    # Create the comparison video
+                    comparison_path = create_frame_by_frame_comparison(
+                        user_video_path,
+                        pro_video_path
+                    )
+                    
+                    # Verify the file exists
+                    if not os.path.exists(comparison_path):
+                        raise FileNotFoundError(
+                            f"Comparison video file not found at {comparison_path}")
+                    
+                    # Store the comparison video path in session state
+                    st.session_state.comparison_video_path = comparison_path
+                
+                # Display success message and video
+                st.success("Frame-by-frame comparison created successfully!")
+                st.subheader("Side-by-Side Comparison with Pro Golfer")
+                
+                # Display video with larger width
+                display_video(comparison_path, width=800)
+                
+                # Show download button
+                with open(comparison_path, "rb") as file:
+                    video_bytes = file.read()
+                    st.download_button(
+                        label="Download Comparison Video",
+                        data=video_bytes,
+                        file_name=os.path.basename(comparison_path),
+                        mime="video/mp4"
+                    )
+                
+                # Add some guidance for interpreting the comparison
+                with st.expander("How to use this comparison", expanded=True):
+                    st.markdown("""
+                    ### How to Interpret This Comparison
+                    
+                    This side-by-side comparison allows you to see how your swing compares to a professional golfer's swing frame by frame. Look for these key differences:
+                    
+                    1. **Posture and Setup**: Compare your stance, grip, and alignment at address
+                    2. **Backswing Rotation**: Note how much shoulder and hip rotation occurs
+                    3. **Top of Swing Position**: Observe club position and body alignment
+                    4. **Downswing Sequence**: Watch how the pro initiates the downswing
+                    5. **Impact Position**: Compare body positioning at impact
+                    6. **Follow-through**: Note how weight transfers and body rotates after impact
+                    
+                    Try pausing the video at key positions to analyze differences in detail.
+                    """)
+            
+            except Exception as e:
+                st.error(f"Error creating comparison video: {str(e)}")
 
 
 if __name__ == "__main__":

app/utils/comparison.py

@@ -0,0 +1,218 @@
+"""
+Comparison module for frame-by-frame analysis between user and pro swings
+"""
+
+import os
+import cv2
+import numpy as np
+from tqdm import tqdm
+
+
+def extract_frames(video_path, max_frames=100):
+    """
+    Extract frames from a video
+    
+    Args:
+        video_path (str): Path to the video file
+        max_frames (int): Maximum number of frames to extract
+        
+    Returns:
+        list: List of extracted frames as numpy arrays
+    """
+    frames = []
+    
+    if not os.path.exists(video_path):
+        raise ValueError(f"Video file not found: {video_path}")
+        
+    cap = cv2.VideoCapture(video_path)
+    
+    if not cap.isOpened():
+        raise ValueError(f"Could not open video: {video_path}")
+    
+    # Get total frame count
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    
+    # Calculate step to get approximately max_frames
+    step = max(1, total_frames // max_frames)
+    
+    current_frame = 0
+    
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+            
+        if current_frame % step == 0:
+            frames.append(frame)
+            
+        current_frame += 1
+        
+    cap.release()
+    
+    return frames
+
+
+def normalize_frames(frames, target_height=480):
+    """
+    Normalize frames to a consistent size while maintaining aspect ratio
+    
+    Args:
+        frames (list): List of frames
+        target_height (int): Target height for normalized frames
+        
+    Returns:
+        list: List of normalized frames
+    """
+    normalized_frames = []
+    
+    for frame in frames:
+        # Get current dimensions
+        h, w = frame.shape[:2]
+        
+        # Calculate new width to maintain aspect ratio
+        target_width = int(w * (target_height / h))
+        
+        # Resize the frame
+        resized = cv2.resize(frame, (target_width, target_height))
+        normalized_frames.append(resized)
+        
+    return normalized_frames
+
+
+def create_side_by_side_comparison(user_frames, pro_frames, output_path, fps=30):
+    """
+    Create a side-by-side comparison video
+    
+    Args:
+        user_frames (list): List of user swing frames
+        pro_frames (list): List of pro swing frames
+        output_path (str): Path to save the comparison video
+        fps (int): Frames per second for output video
+        
+    Returns:
+        str: Path to the comparison video
+    """
+    if not user_frames or not pro_frames:
+        raise ValueError("Both user and pro frames must be provided")
+        
+    # Normalize frames to same height
+    user_normalized = normalize_frames(user_frames)
+    pro_normalized = normalize_frames(pro_frames)
+    
+    # Ensure we have the same number of frames by duplicating the last frame if needed
+    max_frames = max(len(user_normalized), len(pro_normalized))
+    
+    while len(user_normalized) < max_frames:
+        user_normalized.append(user_normalized[-1])
+        
+    while len(pro_normalized) < max_frames:
+        pro_normalized.append(pro_normalized[-1])
+    
+    # Create output directory if it doesn't exist
+    os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
+    
+    # Get dimensions for the combined frame
+    user_h, user_w = user_normalized[0].shape[:2]
+    pro_h, pro_w = pro_normalized[0].shape[:2]
+    
+    # Create a combined frame with padding
+    padding = 20  # Pixels between the two videos
+    combined_width = user_w + pro_w + padding
+    combined_height = max(user_h, pro_h)
+    
+    # Create video writer
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(output_path, fourcc, fps, (combined_width, combined_height))
+    
+    if not out.isOpened():
+        raise IOError(f"Failed to create video writer for {output_path}")
+    
+    # Create the combined video
+    for i in tqdm(range(min(len(user_normalized), len(pro_normalized))), desc="Creating comparison video"):
+        # Create a blank canvas
+        combined = np.ones((combined_height, combined_width, 3), dtype=np.uint8) * 255
+        
+        # Add title text
+        font = cv2.FONT_HERSHEY_SIMPLEX
+        cv2.putText(combined, "Your Swing", (user_w//2 - 60, 30), font, 1, (0, 0, 0), 2)
+        cv2.putText(combined, "Pro Swing", (user_w + padding + pro_w//2 - 60, 30), font, 1, (0, 0, 0), 2)
+        
+        # Add frame number
+        cv2.putText(combined, f"Frame: {i+1}/{min(len(user_normalized), len(pro_normalized))}", 
+                   (10, combined_height - 10), font, 0.5, (0, 0, 0), 1)
+        
+        # Paste user frame
+        combined[0:user_h, 0:user_w] = user_normalized[i]
+        
+        # Paste pro frame
+        combined[0:pro_h, user_w+padding:user_w+padding+pro_w] = pro_normalized[i]
+        
+        # Draw vertical line between frames
+        cv2.line(combined, (user_w + padding//2, 0), (user_w + padding//2, combined_height), (0, 0, 0), 2)
+        
+        # Write to video
+        out.write(combined)
+    
+    out.release()
+    
+    return output_path
+
+
+def align_swings(user_frames, pro_frames, method="manual"):
+    """
+    Align user and pro swings based on swing phases
+    
+    Args:
+        user_frames (list): List of user swing frames
+        pro_frames (list): List of pro swing frames
+        method (str): Alignment method ('manual' or 'auto')
+        
+    Returns:
+        tuple: Aligned user frames and pro frames
+    """
+    # For now, we'll use a simple frame stretching approach
+    # In the future, this could be enhanced with ML-based swing phase detection
+    
+    # Get frame counts
+    user_count = len(user_frames)
+    pro_count = len(pro_frames)
+    
+    # If almost equal, return as-is
+    if abs(user_count - pro_count) <= 5:
+        return user_frames, pro_frames
+    
+    # If user has more frames, subsample
+    if user_count > pro_count:
+        indices = np.linspace(0, user_count - 1, pro_count, dtype=int)
+        return [user_frames[i] for i in indices], pro_frames
+    
+    # If pro has more frames, subsample
+    indices = np.linspace(0, pro_count - 1, user_count, dtype=int)
+    return user_frames, [pro_frames[i] for i in indices]
+
+
+def create_frame_by_frame_comparison(user_video_path, pro_video_path, output_dir="downloads"):
+    """
+    Create a frame-by-frame comparison between user and pro golfer swings
+    
+    Args:
+        user_video_path (str): Path to the user's golf swing video
+        pro_video_path (str): Path to the professional golfer's swing video
+        output_dir (str): Directory to save the comparison video
+        
+    Returns:
+        str: Path to the comparison video
+    """
+    # Extract frames
+    user_frames = extract_frames(user_video_path)
+    pro_frames = extract_frames(pro_video_path)
+    
+    # Align swings
+    aligned_user, aligned_pro = align_swings(user_frames, pro_frames)
+    
+    # Create output path
+    video_name = os.path.splitext(os.path.basename(user_video_path))[0]
+    output_path = os.path.join(output_dir, f"{video_name}_comparison.mp4")
+    
+    # Create side-by-side comparison
+    return create_side_by_side_comparison(aligned_user, aligned_pro, output_path) 

app/utils/video_downloader.py

@@ -77,3 +77,61 @@ def download_youtube_video(url, output_dir="downloads"):
         raise ValueError(f"Error downloading video: {str(e)}")
     except Exception as e:
         raise ValueError(f"Error: {str(e)}")
+
+
+def download_pro_reference(url="https://www.youtube.com/shorts/geR666LWSHg", output_dir="downloads"):
+    """
+    Download a professional golfer reference video
+    
+    Args:
+        url (str): YouTube video URL of professional golfer (default: provided reference)
+        output_dir (str): Directory to save the downloaded video
+        
+    Returns:
+        str: Path to the downloaded pro reference video file
+    """
+    try:
+        # Create a specific filename for the pro reference
+        os.makedirs(output_dir, exist_ok=True)
+        
+        # Check if pro reference already exists to avoid re-downloading
+        pro_file_path = os.path.join(output_dir, "pro_reference.mp4")
+        if os.path.exists(pro_file_path):
+            return pro_file_path
+            
+        # Set output template for the downloaded file with fixed name
+        output_template = os.path.join(output_dir, "pro_reference.%(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,
+        }
+        
+        # Create yt-dlp object and download the video
+        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+            ydl.extract_info(url, download=True)
+            
+        # Check if file exists with mp4 extension
+        if os.path.exists(pro_file_path):
+            return pro_file_path
+        else:
+            # Try other extensions
+            for ext in ['webm', 'mkv']:
+                alt_path = os.path.join(output_dir, f"pro_reference.{ext}")
+                if os.path.exists(alt_path):
+                    return alt_path
+                    
+            # If still not found, download as normal video and rename
+            video_path = download_youtube_video(url, output_dir)
+            ext = os.path.splitext(video_path)[1]
+            new_path = os.path.join(output_dir, f"pro_reference{ext}")
+            os.rename(video_path, new_path)
+            return new_path
+            
+    except Exception as e:
+        raise ValueError(f"Error downloading pro reference: {str(e)}")

requirements.txt

@@ -1,5 +1,5 @@
 opencv-python-headless
-yt-dlp==2025.2.19
+yt-dlp==2025.05.22
 ultralytics
 mediapipe
 numpy