trajectory_predictor.py

import numpy as np

def predict_trajectory(detection_data, pitch_height=720, stump_zone=(280, 360)):
    """
    Uses polynomial regression to predict post-impact ball trajectory.
    
    Args:
        detection_data: output from `detect_lbw_event`
        pitch_height: total frame height (in pixels) to simulate stumps
        stump_zone: x-coordinate range for stumps (min_x, max_x)

    Returns:
        dict with:
            - trajectory_points: [(x, y), ...] actual + predicted
            - decision: "OUT" or "NOT OUT"
    """
    ball_positions = detection_data["ball_positions"]
    impact_frame = detection_data["impact_frame"]

    if not ball_positions or impact_frame == -1:
        return {
            "trajectory_points": [],
            "decision": "NOT ENOUGH DATA"
        }

    # Extract coordinates pre-impact
    xs = []
    ys = []
    for idx, x, y in ball_positions:
        if idx <= impact_frame:
            xs.append(x)
            ys.append(y)

    if len(xs) < 5:
        return {
            "trajectory_points": [],
            "decision": "NOT ENOUGH POINTS"
        }

    # Fit polynomial regression (degree 2 for parabolic path)
    coeffs = np.polyfit(xs, ys, deg=2)
    poly = np.poly1d(coeffs)

    # Predict future trajectory
    last_x = xs[-1]
    future_xs = list(range(last_x, last_x + 60, 5))  # simulate 60px ahead
    future_ys = [int(poly(x)) for x in future_xs]

    trajectory_points = list(zip(xs, ys)) + list(zip(future_xs, future_ys))

    # OUT logic: predicted y crosses stump plane and x within stump zone
    for x, y in zip(future_xs, future_ys):
        if y >= pitch_height - 150:  # near stump base
            if stump_zone[0] <= x <= stump_zone[1]:
                return {
                    "trajectory_points": trajectory_points,
                    "decision": "OUT"
                }

    return {
        "trajectory_points": trajectory_points,
        "decision": "NOT OUT"
    }