import streamlit as st import cv2 import mediapipe as mp import numpy as np import tempfile import os # Initialize MediaPipe Pose mp_pose = mp.solutions.pose pose = mp_pose.Pose(static_image_mode=False, model_complexity=1, enable_segmentation=True, min_detection_confidence=0.5, min_tracking_confidence=0.5) mp_drawing = mp.solutions.drawing_utils def calculate_angle_between_vectors(v1, v2): unit_vector_1 = v1 / np.linalg.norm(v1) unit_vector_2 = v2 / np.linalg.norm(v2) dot_product = np.dot(unit_vector_1, unit_vector_2) angle = np.arccos(dot_product) return np.degrees(angle) def process_video(video_path): cap = cv2.VideoCapture(video_path) output_dir = tempfile.mkdtemp() current_phase = "Not Setup phase" prev_wrist_left_y = None prev_wrist_right_y = None top_backswing_detected = False mid_downswing_detected = False ball_impact_detected = False top_backswing_frame = -2 mid_downswing_frame = -2 ball_impact_frame = -2 BALL_IMPACT_DURATION = 2 # Duration in frames to display Ball Impact phase MIN_MOVEMENT_THRESHOLD = 0.01 HIP_NEAR_THRESHOLD = 0.05 MID_SWING_THRESHOLD = 0.05 saved_phases = set() while cap.isOpened(): ret, frame = cap.read() if not ret: break frame_no = int(cap.get(cv2.CAP_PROP_POS_FRAMES)) image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) result = pose.process(image_rgb) h, w, c = frame.shape if result.pose_landmarks: mp_drawing.draw_landmarks( frame, result.pose_landmarks, mp_pose.POSE_CONNECTIONS, mp_drawing.DrawingSpec(color=(255, 0, 0), thickness=2, circle_radius=2), mp_drawing.DrawingSpec(color=(255, 0, 255), thickness=2, circle_radius=2) ) landmarks = result.pose_landmarks.landmark wrist_left_y = landmarks[mp_pose.PoseLandmark.LEFT_WRIST].y wrist_right_y = landmarks[mp_pose.PoseLandmark.RIGHT_WRIST].y hip_left_y = landmarks[mp_pose.PoseLandmark.LEFT_HIP].y hip_right_y = landmarks[mp_pose.PoseLandmark.RIGHT_HIP].y shoulder_left_y = landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER].y shoulder_right_y = landmarks[mp_pose.PoseLandmark.RIGHT_SHOULDER].y hip_y_avg = (hip_left_y + hip_right_y) / 2 shoulder_y_avg = (shoulder_left_y + shoulder_right_y) / 2 mid_swing_y = (shoulder_y_avg + hip_y_avg) / 2 # Ensure the current phase persists for a few more milliseconds if it's Ball Impact if ball_impact_detected and frame_no <= ball_impact_frame + BALL_IMPACT_DURATION: current_phase = "Ball impact phase" elif (abs(wrist_left_y - hip_y_avg) < HIP_NEAR_THRESHOLD and abs(wrist_right_y - hip_y_avg) < HIP_NEAR_THRESHOLD): if prev_wrist_left_y is not None and prev_wrist_right_y is not None: if (abs(wrist_left_y - prev_wrist_left_y) < MIN_MOVEMENT_THRESHOLD and abs(wrist_right_y - prev_wrist_right_y) < MIN_MOVEMENT_THRESHOLD): if mid_downswing_detected and frame_no > mid_downswing_frame: current_phase = "Ball impact phase" ball_impact_detected = True ball_impact_frame = frame_no else: current_phase = "Setup phase" top_backswing_detected = False mid_downswing_detected = False else: current_phase = "" else: if mid_downswing_detected and frame_no > mid_downswing_frame: current_phase = "Ball impact phase" ball_impact_detected = True ball_impact_frame = frame_no else: current_phase = "Setup phase" top_backswing_detected = False mid_downswing_detected = False elif (abs(wrist_left_y - mid_swing_y) < MID_SWING_THRESHOLD and abs(wrist_right_y - mid_swing_y) < MID_SWING_THRESHOLD and not top_backswing_detected and not ball_impact_detected): current_phase = "Mid backswing phase" elif (wrist_left_y < shoulder_left_y and wrist_right_y < shoulder_right_y and not mid_downswing_detected and not ball_impact_detected): current_phase = "Top backswing phase" top_backswing_detected = True top_backswing_frame = frame_no elif (abs(wrist_left_y - mid_swing_y) < MID_SWING_THRESHOLD and abs(wrist_right_y - mid_swing_y) < MID_SWING_THRESHOLD and top_backswing_detected and frame_no > top_backswing_frame): current_phase = "Mid downswing phase" mid_downswing_detected = True mid_downswing_frame = frame_no elif (wrist_left_y < shoulder_left_y and wrist_right_y < shoulder_right_y and ball_impact_detected and frame_no > ball_impact_frame): current_phase = "Follow through phase" else: current_phase = "" prev_wrist_left_y = wrist_left_y prev_wrist_right_y = wrist_right_y cv2.putText(frame, f"Phase: {current_phase}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv2.LINE_AA) # Save the frame for each detected phase if current_phase and current_phase not in saved_phases: phase_filename = os.path.join(output_dir, f"{current_phase.replace(' ', '_')}.png") cv2.imwrite(phase_filename, frame) saved_phases.add(current_phase) cap.release() cv2.destroyAllWindows() pose.close() return output_dir st.title("Golf Swing Phase Detection") st.write("Upload a video to detect different phases of a golf swing.") video_file = st.file_uploader("Upload Video", type=["mp4", "avi", "mov", "mkv"]) if video_file is not None: tfile = tempfile.NamedTemporaryFile(delete=False) tfile.write(video_file.read()) tfile_path = tfile.name st.write("Processing video...") output_dir = process_video(tfile_path) st.write("Detected phases saved to:", output_dir) st.write("Example frames from detected phases:") for phase_image in os.listdir(output_dir): st.image(os.path.join(output_dir, phase_image), caption=phase_image)