Spaces:
Runtime error
Runtime error
File size: 4,422 Bytes
013216e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 |
import numpy as np
from segment_anything import sam_model_registry, SamPredictor
import cv2
from scipy.optimize import curve_fit
class GolfTrajectoryPredictor:
def __init__(self, sam_checkpoint):
# Initialize SAM model
self.sam = sam_model_registry["vit_h"](checkpoint=sam_checkpoint)
self.predictor = SamPredictor(self.sam)
def get_club_metrics(self, frame, point):
"""Extract golf club angle and position using SAM"""
self.predictor.set_image(frame)
masks, scores, _ = self.predictor.predict(
point_coords=np.array([point]),
point_labels=np.array([1])
)
# Get the best mask
club_mask = masks[np.argmax(scores)]
# Calculate club angle from mask
coords = np.column_stack(np.where(club_mask))
if len(coords) < 2:
return None, None
# Fit line to get club angle
vx, vy, x0, y0 = cv2.fitLine(coords, cv2.DIST_L2, 0, 0.01, 0.01)
angle = np.arctan2(vy, vx)
return angle, (x0[0], y0[0])
def physics_trajectory(self, t, v0, theta, h0, g=9.81):
"""Model the physics of projectile motion"""
# Convert angle to radians
theta_rad = np.radians(theta)
# Initial velocities
v0x = v0 * np.cos(theta_rad)
v0y = v0 * np.sin(theta_rad)
# Position equations
x = v0x * t
y = h0 + v0y * t - 0.5 * g * t**2
return np.column_stack((x, y))
def fit_trajectory(self, points, initial_height, club_angle=None):
"""Fit trajectory to user-selected points"""
times = np.linspace(0, 1, len(points))
# Initial guess for parameters
# Use club angle if available to better estimate initial velocity direction
initial_theta = club_angle if club_angle is not None else 45
# Calculate approximate initial velocity from first two points
if len(points) >= 2:
dx = points[1][0] - points[0][0]
dy = points[1][1] - points[0][1]
initial_v0 = np.sqrt(dx**2 + dy**2) / (times[1] - times[0])
else:
initial_v0 = 50 # Default initial guess
# Fit physics model to points
try:
params, _ = curve_fit(
lambda t, v0, theta: self.physics_trajectory(t, v0, theta, initial_height),
times,
points,
p0=[initial_v0, initial_theta]
)
return params
except RuntimeError:
return None
def predict_full_trajectory(self, video_path, user_selected_points):
"""Main function to predict full trajectory"""
cap = cv2.VideoCapture(video_path)
fps = cap.get(cv2.CAP_PROP_FPS)
# Get initial frame for club analysis
ret, first_frame = cap.read()
if not ret:
return None
# Get club metrics from first frame
club_angle, club_pos = self.get_club_metrics(first_frame, user_selected_points[0])
# Convert pixel coordinates to physical space
# This requires camera calibration in practice
scale_factor = 0.01 # meters per pixel
physical_points = np.array(user_selected_points) * scale_factor
# Fit trajectory
params = self.fit_trajectory(
physical_points,
initial_height=physical_points[0][1],
club_angle=club_angle
)
if params is None:
return None
# Generate full trajectory
t = np.linspace(0, len(user_selected_points)/fps, 100)
full_trajectory = self.physics_trajectory(t, params[0], params[1], physical_points[0][1])
# Convert back to pixel coordinates
pixel_trajectory = full_trajectory / scale_factor
return pixel_trajectory, params
def visualize_trajectory(self, frame, trajectory, user_points):
"""Visualize the predicted trajectory and user-selected points"""
vis_frame = frame.copy()
# Draw predicted trajectory
trajectory = trajectory.astype(np.int32)
for i in range(len(trajectory)-1):
cv2.line(vis_frame, tuple(trajectory[i]), tuple(trajectory[i+1]),
(0, 255, 0), 2)
# Draw user-selected points
for point in user_points:
cv2.circle(vis_frame, tuple(map(int, point)), 5, (255, 0, 0), -1)
return vis_frame |