Create tracking/ball_tracker.py

tracking/ball_tracker.py

@@ -0,0 +1,20 @@
+import numpy as np
+from filterpy.kalman import KalmanFilter
+
+class BallTracker:
+    def __init__(self):
+        self.kf = KalmanFilter(dim_x=4, dim_z=2)
+        self.kf.F = np.array([[1,0,1,0], [0,1,0,1], [0,0,1,0], [0,0,0,1]])
+        self.kf.H = np.array([[1,0,0,0], [0,1,0,0]])
+        self.kf.P *= 10.
+        self.kf.R *= 0.5
+        self.kf.Q *= 0.01
+        self.initialized = False
+
+    def update(self, measurement):
+        if not self.initialized:
+            self.kf.x[:2] = np.array([[measurement[0]], [measurement[1]]])
+            self.initialized = True
+        self.kf.predict()
+        self.kf.update(measurement)
+        return self.kf.x[:2].flatten()