Create trajectory_predictor.py

trajectory_predictor.py

@@ -0,0 +1,63 @@
+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"
+    }