trajectory_service.py

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