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