frame impact

app/models/swing_analyzer.py

@@ -5,84 +5,125 @@ Swing analysis module for golf swing segmentation and trajectory analysis
 import numpy as np
 from app.models.pose_estimator import calculate_joint_angles
 
-def segment_swing(pose_data, detections, sample_rate=1):
-    swing_phases = {"setup": [], "backswing": [], "downswing": [], "impact": [], "follow_through": []}
-    frame_indices = sorted(pose_data.keys())
-    if not frame_indices:
-        return swing_phases
 
-    angles_by_frame = {}
+def find_top_of_backswing(pose_data):
+    """Helper function to find the peak of backswing"""
+    frame_indices = sorted(pose_data.keys())
+    max_shoulder_angle = -1
+    top_frame = frame_indices[0]
+    
     for idx in frame_indices:
         keypoints = pose_data[idx]
         angles = calculate_joint_angles(keypoints)
-        angles_by_frame[idx] = angles
+        shoulder = angles.get("right_shoulder", 0)
+        if shoulder > max_shoulder_angle:
+            max_shoulder_angle = shoulder
+            top_frame = idx
+    
+    return top_frame
 
-    # --- Dynamic Phase Segmentation ---
-    # 1. Setup: before any significant movement
-    # 2. Backswing: from end of setup to top of backswing
-    # 3. Downswing: from top of backswing to just before impact
-    # 4. Impact: frame(s) where ball first moves
-    # 5. Follow-through: after impact
 
-    # --- 1. Find end of setup (first significant movement) ---
+def detect_impact_frame(pose_data, detections, sample_rate=2):
+    """
+    Simple impact detection: ball movement first, wrist speed fallback
+    """
+    frame_indices = sorted(pose_data.keys())
+    if len(frame_indices) < 10:
+        return None
+    
+    top_backswing = find_top_of_backswing(pose_data)
+    downswing_frames = [f for f in frame_indices if f > top_backswing]
+    
+    if not downswing_frames:
+        return None
+    
+    # Method 1: Ball movement (if we have ball detections)
+    if detections:
+        ball_detections = [d for d in detections if d.class_name == "sports ball"]
+        ball_positions = {}
+        
+        for detection in ball_detections:
+            frame_idx = detection.frame_idx // sample_rate
+            if frame_idx > top_backswing:
+                x1, y1, x2, y2 = detection.bbox
+                center_x, center_y = (x1 + x2) / 2, (y1 + y2) / 2
+                ball_positions[frame_idx] = (center_x, center_y)
+        
+        # Find first significant ball movement
+        if len(ball_positions) >= 2:
+            sorted_frames = sorted(ball_positions.keys())
+            for i in range(1, len(sorted_frames)):
+                curr_pos = ball_positions[sorted_frames[i]]
+                prev_pos = ball_positions[sorted_frames[i-1]]
+                movement = np.sqrt((curr_pos[0] - prev_pos[0])**2 + (curr_pos[1] - prev_pos[1])**2)
+                
+                if movement > 15:  # Significant movement threshold
+                    print(f"Impact detected via ball movement at frame {sorted_frames[i]}")
+                    return sorted_frames[i]
+    
+    # Method 2: Wrist speed fallback (simple and reliable)
+    max_wrist_speed = 0
+    impact_frame = None
+    
+    for i in range(1, len(downswing_frames)):
+        curr_frame = downswing_frames[i]
+        prev_frame = downswing_frames[i-1]
+        
+        curr_angles = calculate_joint_angles(pose_data[curr_frame])
+        prev_angles = calculate_joint_angles(pose_data[prev_frame])
+        
+        curr_wrist = curr_angles.get("right_wrist", 0)
+        prev_wrist = prev_angles.get("right_wrist", 0)
+        wrist_speed = abs(curr_wrist - prev_wrist)
+        
+        if wrist_speed > max_wrist_speed:
+            max_wrist_speed = wrist_speed
+            impact_frame = curr_frame
+    
+    print(f"Impact detected via wrist speed at frame {impact_frame}")
+    return impact_frame or downswing_frames[len(downswing_frames) // 3]
+
+
+def segment_swing_pose_based(pose_data, detections=None, sample_rate=2):
+    """
+    Simple swing segmentation with clean impact detection
+    """
+    swing_phases = {"setup": [], "backswing": [], "downswing": [], "impact": [], "follow_through": []}
+    frame_indices = sorted(pose_data.keys())
+    
+    if not frame_indices:
+        return swing_phases
+
+    # 1. Find setup end (first significant movement)
     setup_end = frame_indices[0]
-    initial_angles = angles_by_frame[frame_indices[0]]
-    initial_shoulder = initial_angles.get("right_shoulder")
-    initial_wrist = initial_angles.get("right_elbow")
-    movement_threshold = 8  # degrees, can be tuned
+    initial_angles = calculate_joint_angles(pose_data[frame_indices[0]])
+    initial_shoulder = initial_angles.get("right_shoulder", 0)
+    
     for idx in frame_indices[1:]:
-        angles = angles_by_frame[idx]
-        shoulder = angles.get("right_shoulder")
-        wrist = angles.get("right_elbow")
-        if (shoulder and initial_shoulder and abs(shoulder - initial_shoulder) > movement_threshold) or \
-           (wrist and initial_wrist and abs(wrist - initial_wrist) > movement_threshold):
-            setup_end = idx - 1
+        angles = calculate_joint_angles(pose_data[idx])
+        shoulder = angles.get("right_shoulder", 0)
+        if abs(shoulder - initial_shoulder) > 10:
+            setup_end = max(frame_indices[0], idx - 2)
             break
 
-    # --- 2. Top of backswing (max shoulder angle after setup) ---
-    max_shoulder_angle = -1
-    top_backswing_frame = setup_end + 1
-    for idx in frame_indices:
-        if idx <= setup_end:
-            continue
-        shoulder = angles_by_frame[idx].get("right_shoulder")
-        if shoulder and shoulder > max_shoulder_angle:
-            max_shoulder_angle = shoulder
-            top_backswing_frame = idx
+    # 2. Find top of backswing
+    top_backswing = find_top_of_backswing(pose_data)
 
-    # --- 3. Impact: first significant ball movement ---
-    impact_frame = None
-    ball_detections = [d for d in detections if d.class_name == "sports ball"]
-    ball_detections.sort(key=lambda x: x.frame_idx)
-    movement_threshold_px = 2
-    prev_x = prev_y = None
-    prev_frame = None
-    for detection in ball_detections:
-        frame_idx = detection.frame_idx
-        if frame_idx < top_backswing_frame:
-            continue
-        x1, y1, x2, y2 = detection.bbox
-        ball_x = (x1 + x2) / 2
-        ball_y = (y1 + y2) / 2
-        if prev_x is not None and prev_y is not None:
-            dx = abs(ball_x - prev_x)
-            dy = abs(ball_y - prev_y)
-            if dx > movement_threshold_px or dy > movement_threshold_px:
-                impact_frame = frame_idx
-                break
-        prev_x = ball_x
-        prev_y = ball_y
-        prev_frame = frame_idx
-    if impact_frame is None and prev_frame is not None:
-        impact_frame = prev_frame
-    if impact_frame is None:
-        impact_frame = frame_indices[-1]
-
-    # --- 4. Assign phases dynamically ---
+    # 3. Find impact frame
+    impact_frame = detect_impact_frame(pose_data, detections, sample_rate)
+    
+    # Simple validation and fallback
+    if not impact_frame or impact_frame <= top_backswing:
+        downswing_frames = [f for f in frame_indices if f > top_backswing]
+        impact_frame = downswing_frames[len(downswing_frames) // 3] if downswing_frames else top_backswing + 1
+
+    print(f"Swing phases: Setup end={setup_end}, Top backswing={top_backswing}, Impact={impact_frame}")
+
+    # 4. Assign phases
     for idx in frame_indices:
         if idx <= setup_end:
             swing_phases["setup"].append(idx)
-        elif idx <= top_backswing_frame:
+        elif idx <= top_backswing:
             swing_phases["backswing"].append(idx)
         elif idx < impact_frame:
             swing_phases["downswing"].append(idx)
@@ -93,7 +134,19 @@ def segment_swing(pose_data, detections, sample_rate=1):
 
     return swing_phases
 
-def analyze_trajectory(frames, detections, swing_phases, sample_rate=1):
+
+# Wrapper function to maintain compatibility with existing Streamlit app
+def segment_swing(pose_data, detections, sample_rate=2):
+    """
+    Main swing segmentation function (wrapper for pose-based approach)
+    """
+    return segment_swing_pose_based(pose_data, detections, sample_rate)
+
+
+def analyze_trajectory(frames, detections, swing_phases, sample_rate=2):
+    """
+    Analyze ball trajectory and calculate club speed
+    """
     trajectory_data = {}
     if len(frames) < 150:
         sample_rate = 1
@@ -108,7 +161,7 @@ def analyze_trajectory(frames, detections, swing_phases, sample_rate=1):
     ball_positions = {}
 
     for detection in ball_detections:
-        frame_idx = detection.frame_idx
+        frame_idx = detection.frame_idx // sample_rate
         if frame_idx >= impact_frame_idx:
             x1, y1, x2, y2 = detection.bbox
             center_x = (x1 + x2) / 2
@@ -122,7 +175,7 @@ def analyze_trajectory(frames, detections, swing_phases, sample_rate=1):
     club_speed = None
     downswing_frames = swing_phases.get("downswing", [])
     if len(downswing_frames) >= 2:
-        actual_frames_elapsed = (downswing_frames[-1] - downswing_frames[0])
+        actual_frames_elapsed = (downswing_frames[-1] - downswing_frames[0]) * sample_rate
         time_diff = actual_frames_elapsed / 30
         if time_diff > 0:
             club_speed = 100 * (1 / time_diff)

app/streamlit_app.py

@@ -166,7 +166,7 @@ def main():
         "Frame Skip Rate (YOLO)",
         min_value=1,
         max_value=10,
-        value=5,
+        value=2,
         help=
         "Process every Nth frame. Higher values = faster but less accurate.")
         
@@ -476,7 +476,8 @@ def main():
                 with st.spinner("Extracting key frames from your swing..."):
                     user_video_path = st.session_state.analysis_data['video_path']
                     user_swing_phases = st.session_state.analysis_data['swing_phases']
-                    key_frames = extract_key_swing_frames(user_video_path, user_swing_phases)
+                    frames = st.session_state.analysis_data['frames']
+                    key_frames = extract_key_swing_frames(user_video_path, frames, user_swing_phases)
 
                 st.success("Key frame analysis complete!")
                 st.subheader("Key Frame Analysis: Your Swing's Critical Positions")

app/utils/comparison.py

@@ -120,133 +120,83 @@ def extract_frames(video_path, max_frames=100):
     return frames
 
 
-def extract_key_swing_frames(video_path, swing_phases=None):
+def extract_key_swing_frames(video_path, frames, swing_phases=None):
     """
-    Extract 3 key frames from a golf swing video:
+    Extract 3 key frames from a list of processed frames.
     1. First setup frame
     2. Last backswing frame (top of backswing)
     3. First impact frame
     
     Args:
-        video_path (str): Path to the video file
+        video_path (str): Path to the original video file (used for rotation metadata).
+        frames (list): List of processed video frames.
         swing_phases (dict): Dictionary mapping phase names to lists of frame indices
+                             relative to the 'frames' list.
         
     Returns:
         dict: Dictionary mapping phase names to frames
     """
-    if not os.path.exists(video_path):
-        raise ValueError(f"Video file not found: {video_path}")
-    
-    print(f"Extracting key frames from: {video_path}")
-    
-    # Use basic OpenCV VideoCapture
-    cap = cv2.VideoCapture(video_path)
-    
-    if not cap.isOpened():
-        raise ValueError(f"Could not open video: {video_path}")
+    key_frames = {'setup': None, 'backswing': None, 'impact': None}
     
-    try:
-        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-        if total_frames <= 0:
-            raise ValueError(f"Invalid video: no frames found in {video_path}")
-        
-        print(f"Total frames in video: {total_frames}")
+    if not frames:
+        print("Warning: No frames provided to extract_key_swing_frames.")
+        return key_frames
         
-        # Check for rotation metadata
-        rotation_angle = 0
-        try:
-            # Try to get orientation metadata if available
-            orientation = cap.get(cv2.CAP_PROP_ORIENTATION_META)
-            if orientation == 90:
-                rotation_angle = 270  # Rotate counterclockwise
-            elif orientation == 180:
-                rotation_angle = 180
-            elif orientation == 270:
-                rotation_angle = 90   # Rotate counterclockwise
-            print(f"Video orientation metadata: {orientation}, applying rotation: {rotation_angle}")
-        except:
-            print("No orientation metadata available")
-        
-        key_frames = {}
-        
-        # Determine frame indices based on swing phases
-        if swing_phases:
-            # Get first setup frame
-            setup_frames = swing_phases.get('setup', [])
-            setup_idx = setup_frames[0] if setup_frames else 0
-            
-            # Get last backswing frame (top of backswing)
-            backswing_frames = swing_phases.get('backswing', [])
-            backswing_idx = backswing_frames[-1] if backswing_frames else total_frames//3
-            
-            # Get first impact frame
-            impact_frames = swing_phases.get('impact', [])
-            impact_idx = impact_frames[0] if impact_frames else total_frames//2
-        else:
-            # Fallback to default indices if no swing phases provided
-            setup_idx = 0
-            backswing_idx = total_frames // 3
-            impact_idx = int(total_frames * 0.6)
-        
-        print(f"Frame indices - Setup: {setup_idx}, Backswing: {backswing_idx}, Impact: {impact_idx}")
-        
-        # Extract frames for each phase
-        phases = [
-            ('setup', setup_idx),
-            ('backswing', backswing_idx),
-            ('impact', impact_idx)
-        ]
-        
-        for phase_name, frame_idx in phases:
-            frame = _extract_single_frame(cap, frame_idx, total_frames, rotation_angle, phase_name)
-            if frame is not None:
-                key_frames[phase_name] = frame
-                print(f"Successfully extracted {phase_name} frame")
-            else:
-                print(f"Failed to extract {phase_name} frame")
+    # Determine frame indices based on swing phases
+    if swing_phases:
+        # Get first setup frame
+        setup_frames = swing_phases.get('setup', [])
+        setup_idx = setup_frames[0] if setup_frames else 0
         
-        return key_frames
+        # Get last backswing frame (top of backswing)
+        backswing_frames = swing_phases.get('backswing', [])
+        backswing_idx = backswing_frames[-1] if backswing_frames else len(frames) // 3
         
-    except Exception as e:
-        raise ValueError(f"Error extracting frames from {video_path}: {str(e)}")
-    finally:
-        cap.release()
+        # Get first impact frame
+        impact_frames = swing_phases.get('impact', [])
+        impact_idx = impact_frames[0] if impact_frames else len(frames) // 2
+    else:
+        # Fallback to default indices if no swing phases provided
+        setup_idx = 0
+        backswing_idx = len(frames) // 3
+        impact_idx = len(frames) // 2
+    
+    print(f"Key frame indices (relative to processed frames) - Setup: {setup_idx}, Backswing: {backswing_idx}, Impact: {impact_idx}")
+    
+    # Get rotation angle from the original video file
+    rotation_angle = 0
+    if os.path.exists(video_path):
+        cap = cv2.VideoCapture(video_path)
+        if cap.isOpened():
+            try:
+                orientation = int(cap.get(cv2.CAP_PROP_ORIENTATION_META))
+                if orientation == 90:
+                    rotation_angle = 270  # Rotate counterclockwise
+                elif orientation == 180:
+                    rotation_angle = 180
+                elif orientation == 270:
+                    rotation_angle = 90   # Rotate counterclockwise
+                print(f"Video orientation metadata: {orientation}, applying rotation: {rotation_angle}")
+            except Exception as e:
+                print(f"Could not read orientation metadata: {e}")
+            finally:
+                cap.release()
+    else:
+        print(f"Warning: Video path {video_path} not found for rotation check.")
 
+    phase_indices = {'setup': setup_idx, 'backswing': backswing_idx, 'impact': impact_idx}
 
-def _extract_single_frame(cap, target_idx, total_frames, rotation_angle, phase_name):
-    """
-    Extract a single frame from video with validation and rotation correction
-    """
-    # Try the target frame first
-    for attempt_idx in [target_idx, target_idx + 1, target_idx - 1, target_idx + 2, target_idx - 2]:
-        if attempt_idx < 0 or attempt_idx >= total_frames:
-            continue
-            
-        cap.set(cv2.CAP_PROP_POS_FRAMES, attempt_idx)
-        ret, frame = cap.read()
-        
-        if not ret or frame is None:
-            print(f"Failed to read frame at index {attempt_idx} for {phase_name}")
-            continue
-        
-        # Validate frame has 3 channels (color)
-        if len(frame.shape) != 3 or frame.shape[2] != 3:
-            print(f"Frame at index {attempt_idx} for {phase_name} is not in color format: {frame.shape}")
-            continue
-        
-        print(f"Successfully read frame at index {attempt_idx} for {phase_name}, shape: {frame.shape}")
-        
-        # Apply rotation correction if needed
-        if rotation_angle != 0:
-            print(f"Before rotation: {frame.shape}")
-            frame = _apply_rotation(frame, rotation_angle)
-            print(f"After {rotation_angle}° rotation: {frame.shape}")
-            print(f"Applied {rotation_angle}° rotation to {phase_name} frame")
-        
-        return frame.copy()
-    
-    print(f"Could not extract valid frame for {phase_name} after trying multiple indices")
-    return None
+    for phase_name, frame_idx in phase_indices.items():
+        if 0 <= frame_idx < len(frames):
+            frame = frames[frame_idx].copy()
+            if rotation_angle != 0:
+                frame = _apply_rotation(frame, rotation_angle)
+            key_frames[phase_name] = frame
+            print(f"Successfully extracted {phase_name} frame from memory.")
+        else:
+            print(f"Failed to extract {phase_name} frame: index {frame_idx} is out of bounds for {len(frames)} frames.")
+
+    return key_frames
 
 
 def _apply_rotation(frame, rotation_angle):
@@ -432,13 +382,13 @@ def create_key_frame_comparison(user_video_path, pro_video_path=None, user_swing
               'user_image_path', 'pro_image_path', 'title', and 'comments' as values
     """
     # Extract key frames from user video
-    user_frames = extract_key_swing_frames(user_video_path, user_swing_phases)
+    user_frames = extract_key_swing_frames(user_video_path, user_frames, user_swing_phases)
     
     # Get pro frames either from provided images or video
     if use_pro_images:
         pro_frames = load_pro_reference_images()
     else:
-        pro_frames = extract_key_swing_frames(pro_video_path, pro_swing_phases)
+        pro_frames = extract_key_swing_frames(pro_video_path, pro_frames, pro_swing_phases)
     
     # Create output directory with absolute path
     output_dir = os.path.abspath(output_dir)

app/utils/video_processor.py

@@ -3,6 +3,7 @@ Video processing and object detection module
 """
 
 import cv2
+import platform
 import numpy as np
 from tqdm import tqdm
 from ultralytics import YOLO
@@ -35,9 +36,17 @@ def process_video(video_path, sample_rate=5):
     model = YOLO("yolov8n.pt")
     class_names = model.names
 
-    cap = cv2.VideoCapture(video_path)
+    # On macOS ("Darwin"), the AVFoundation backend is often more reliable.
+    # For other systems, FFMPEG is a good choice.
+    backend = cv2.CAP_AVFOUNDATION if platform.system() == "Darwin" else cv2.CAP_FFMPEG
+    cap = cv2.VideoCapture(video_path, backend)
+    
     if not cap.isOpened():
-        raise ValueError("Error opening video file")
+        backend_name = "AVFoundation" if platform.system() == "Darwin" else "FFMPEG"
+        print(f"Warning: Could not open video with {backend_name} backend. Trying default.")
+        cap = cv2.VideoCapture(video_path) # Fallback to default
+        if not cap.isOpened():
+            raise ValueError("Error opening video file with any available backend.")
 
     frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))