Spaces:
Runtime error
Runtime error
File size: 8,930 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 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 |
import numpy as np
from scipy.optimize import curve_fit
from scipy.interpolate import make_interp_spline, BSpline
import cv2
class TrajectoryService:
def __init__(self, fps=30):
self.fps = fps
self.points = [] # List of (frame_idx, x, y) tuples
self.trajectory_points = None
self.ball_metrics = None # Store calculated metrics
def add_point(self, frame_idx, x, y):
"""Add a point to the trajectory"""
self.points.append((frame_idx, x, y))
self.points.sort(key=lambda x: x[0]) # Sort by frame index
self.trajectory_points = None # Reset trajectory so it will be recalculated
self.ball_metrics = None # Reset ball metrics so they will be recalculated
def clear_points(self):
"""Clear all trajectory points"""
self.points = []
self.trajectory_points = None
self.ball_metrics = None
def _physics_model(self, t, v0x, v0y, g):
"""
Physics model for projectile motion - used for metrics calculation
t: time array
v0x: initial velocity x
v0y: initial velocity y
g: gravity (pixels/s^2, will be adjusted based on video scale)
"""
x0, y0 = self.points[0][1], self.points[0][2] # Initial position
x = x0 + v0x * t
y = y0 - (v0y * t - 0.5 * g * t * t) # Negative because y is inverted in images
return np.vstack((x, y)).T
def _calculate_ball_metrics(self):
"""Calculate ball speed and trajectory metrics using physics model"""
if len(self.points) < 2:
return None
try:
# Convert frame indices to time
times = np.array([(p[0] / self.fps) for p in self.points])
positions = np.array([(p[1], p[2]) for p in self.points])
# Initial velocity estimates
dt = times[-1] - times[0]
dx = positions[-1, 0] - positions[0, 0]
dy = positions[-1, 1] - positions[0, 1]
v0x_guess = dx / dt
v0y_guess = dy / dt - 0.5 * 9.81 * dt
# Adjust gravity based on video scale
video_height = max(p[2] for p in self.points) - min(
p[2] for p in self.points
)
g_scale = (
video_height / 10
) # Assume trajectory takes about 1/10th of video height
g = 9.81 * g_scale
# Fit trajectory with physics model
def fit_func(t, v0x, v0y):
return self._physics_model(t, v0x, v0y, g).ravel()
params, _ = curve_fit(
fit_func,
times,
positions.ravel(),
p0=[v0x_guess, v0y_guess],
maxfev=10000,
)
# Calculate metrics
v0x, v0y = params
initial_velocity = np.sqrt(v0x**2 + v0y**2)
launch_angle = np.degrees(
np.arctan2(-v0y, v0x)
) # Negative because y is inverted
max_height = max(positions[:, 1]) - positions[0, 1]
total_distance = abs(positions[-1, 0] - positions[0, 0])
# Store metrics
self.ball_metrics = {
"initial_velocity_px": initial_velocity, # pixels per second
"launch_angle_deg": launch_angle,
"max_height_px": max_height,
"total_distance_px": total_distance,
}
except (RuntimeError, ValueError) as e:
print(f"Could not calculate ball metrics: {e}")
self.ball_metrics = None
def _fit_trajectory(self):
"""Fit trajectory through points using interpolation and physics for the end"""
if len(self.points) < 2:
return None
# Sort points by frame index
self.points.sort(key=lambda x: x[0])
# Separate user's final point
final_point = self.points[-1]
model_points = self.points[:-1] if len(self.points) > 2 else self.points
# Extract times and positions for model points
times = np.array([p[0] for p in model_points])
positions = np.array([(p[1], p[2]) for p in model_points])
if len(model_points) > 2:
# Calculate velocities for last few visible points
last_points = positions[-3:] # Use last 3 points
last_times = times[-3:]
# Calculate average velocity
dt = (last_times[-1] - last_times[0]) / self.fps
dx = last_points[-1, 0] - last_points[0, 0]
dy = last_points[-1, 1] - last_points[0, 1]
vx = dx / dt
vy = dy / dt
# Estimate time to reach final point
dx_to_final = final_point[1] - positions[-1, 0]
time_to_final = dx_to_final / vx if abs(vx) > 1e-6 else 1.0
# Generate frames for the gap to final point
gap_frames = np.linspace(
times[-1], final_point[0], int(abs(time_to_final * self.fps))
)
# Calculate trajectory to final point using projectile motion
t = (gap_frames - times[-1]) / self.fps
x_gap = positions[-1, 0] + vx * t
y_gap = positions[-1, 1] + vy * t + 0.5 * 9.81 * t * t # Add gravity effect
# Adjust y_gap to ensure it reaches the final point
y_gap = np.interp(
x_gap, [x_gap[0], final_point[1]], [y_gap[0], final_point[2]]
)
# Combine model points with gap points
frames_smooth = np.concatenate([times, gap_frames])
x_smooth = np.concatenate([positions[:, 0], x_gap])
y_smooth = np.concatenate([positions[:, 1], y_gap])
else:
# If we only have start and end points, use simple interpolation
frames_smooth = np.linspace(
times[0], final_point[0], int((final_point[0] - times[0]) * 2)
)
x_smooth = np.interp(
frames_smooth,
[times[0], final_point[0]],
[positions[0, 0], final_point[1]],
)
y_smooth = np.interp(
frames_smooth,
[times[0], final_point[0]],
[positions[0, 1], final_point[2]],
)
# Apply smoothing to reduce shakiness
def smooth_array(arr, window=5):
kernel = np.ones(window) / window
return np.convolve(arr, kernel, mode="same")
if len(frames_smooth) > 5:
x_smooth = smooth_array(x_smooth)
y_smooth = smooth_array(y_smooth)
# Store interpolated points
self.trajectory_points = list(
zip(frames_smooth.astype(int), x_smooth, y_smooth)
)
# Calculate ball metrics using physics model
self._calculate_ball_metrics()
return self.trajectory_points
def draw_trajectory(self, frame, frame_idx, initial_thickness=8):
"""Draw the trajectory on a frame with clean ball tracking style"""
if not self.trajectory_points:
if len(self.points) >= 2:
self._fit_trajectory()
else:
return frame
# Get points up to current frame
current_points = [
(int(x), int(y)) for f, x, y in self.trajectory_points if f <= frame_idx
]
if len(current_points) > 1:
# Convert points to numpy array for easier manipulation
points_array = np.array(current_points)
# Calculate distance from start for thickness variation
distances = np.sqrt(np.sum((points_array - points_array[0]) ** 2, axis=1))
max_distance = np.max(distances)
# Draw trajectory segments with varying thickness
for i in range(len(points_array) - 1):
p1 = tuple(points_array[i])
p2 = tuple(points_array[i + 1])
# Calculate thickness based on distance (thinner as it goes further)
progress = distances[i] / max_distance if max_distance > 0 else 0
thickness = max(1, int(initial_thickness * (1 - progress * 0.7)))
# Draw single blue line (no white border)
cv2.line(
frame,
p1,
p2,
(255, 144, 30), # BGR: Light blue
thickness,
cv2.LINE_AA,
)
return frame
def get_trajectory(self):
"""Get the computed trajectory points"""
if not self.trajectory_points and len(self.points) >= 2:
return self._fit_trajectory()
return self.trajectory_points
def get_ball_metrics(self):
"""Get the calculated ball metrics"""
if not self.ball_metrics and len(self.points) >= 2:
self._calculate_ball_metrics()
return self.ball_metrics
|