Update trajectory_predictor.py

trajectory_predictor.py

@@ -1,63 +1,25 @@
+# trajectory_predictor.py
 import numpy as np
+from sklearn.linear_model import LinearRegression
 
-def predict_trajectory(detection_data, pitch_height=720, stump_zone=(280, 360)):
+def predict_trajectory(ball_positions, future_frames=10):
     """
-    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"
+    Predicts future trajectory based on current ball positions using polynomial regression.
+    Returns extrapolated list of (x, y) points.
     """
-    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)
+    if len(ball_positions) < 5:
+        return []  # not enough data
 
-    # 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]
+    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])
 
-    trajectory_points = list(zip(xs, ys)) + list(zip(future_xs, future_ys))
+    # 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))
 
-    # 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"
-                }
+    last_frame = frames[-1]
+    future_frame_ids = np.arange(last_frame, last_frame + future_frames)
 
-    return {
-        "trajectory_points": trajectory_points,
-        "decision": "NOT OUT"
-    }
+    trajectory = [(int(x_poly(f)), int(y_poly(f))) for f in future_frame_ids]
+    return trajectory