trajectory_predictor.py

# trajectory_predictor.py
import numpy as np
from sklearn.linear_model import LinearRegression

def predict_trajectory(ball_positions, future_frames=10):
    """
    Predicts future trajectory based on current ball positions using polynomial regression.
    Returns extrapolated list of (x, y) points.
    """
    if len(ball_positions) < 5:
        return []  # not enough data

    frames = np.array([p[0] for p in ball_positions])
    xs = np.array([p[1] for p in ball_positions])
    ys = np.array([p[2] for p in ball_positions])

    # Fit 2nd-degree polynomial (quadratic) to x and y separately
    x_poly = np.poly1d(np.polyfit(frames, xs, 2))
    y_poly = np.poly1d(np.polyfit(frames, ys, 2))

    last_frame = frames[-1]
    future_frame_ids = np.arange(last_frame, last_frame + future_frames)

    trajectory = [(int(x_poly(f)), int(y_poly(f))) for f in future_frame_ids]
    return trajectory