Add Kalman tracker and distance visualization

app.py

@@ -229,6 +229,14 @@ class AppState:
         self.kick_debug_area: list[float] = []
         self.kick_debug_kick_frame: int | None = None
         self.kick_debug_distance: list[float] = []
+        self.kalman_centers: dict[int, dict[int, tuple[float, float]]] = {}
+        self.kalman_speeds: dict[int, dict[int, float]] = {}
+        self.kalman_residuals: dict[int, dict[int, float]] = {}
+        self.kick_debug_kalman_speeds: list[float] = []
+        self.kalman_centers: dict[int, dict[int, tuple[float, float]]] = {}
+        self.kalman_speeds: dict[int, dict[int, float]] = {}
+        self.kalman_residuals: dict[int, dict[int, float]] = {}
+        self.kick_debug_kalman_speeds: list[float] = []
         # Model selection
         self.model_repo_key: str = "tiny"
         self.model_repo_id: str | None = None
@@ -308,6 +316,18 @@ def init_video_session(GLOBAL_STATE: gr.State, video: str | dict) -> tuple[AppSt
     GLOBAL_STATE.smoothed_centers = {}
     GLOBAL_STATE.ball_speeds = {}
     GLOBAL_STATE.kick_frame = None
+    GLOBAL_STATE.kalman_centers = {}
+    GLOBAL_STATE.kalman_speeds = {}
+    GLOBAL_STATE.kalman_residuals = {}
+    GLOBAL_STATE.kick_debug_kalman_speeds = []
+    GLOBAL_STATE.kick_debug_frames = []
+    GLOBAL_STATE.kick_debug_speeds = []
+    GLOBAL_STATE.kick_debug_threshold = None
+    GLOBAL_STATE.kick_debug_baseline = None
+    GLOBAL_STATE.kick_debug_speed_std = None
+    GLOBAL_STATE.kick_debug_area = []
+    GLOBAL_STATE.kick_debug_kick_frame = None
+    GLOBAL_STATE.kick_debug_distance = []
 
     load_model_if_needed(GLOBAL_STATE)
 
@@ -543,6 +563,69 @@ def _update_centroids_for_frame(state: AppState, frame_idx: int):
     _recompute_motion_metrics(state)
 
 
+def _run_kalman_filter(
+    ordered_items: list[tuple[int, tuple[float, float]]],
+    base_dt: float,
+) -> tuple[dict[int, tuple[float, float]], dict[int, float], dict[int, float]]:
+    if not ordered_items:
+        return {}, {}, {}
+
+    H = np.array([[1, 0, 0, 0], [0, 1, 0, 0]], dtype=float)
+    R = np.eye(2, dtype=float) * 25.0
+
+    state_vec = np.array(
+        [ordered_items[0][1][0], ordered_items[0][1][1], 0.0, 0.0], dtype=float
+    )
+    P = np.eye(4, dtype=float) * 50.0
+
+    positions: dict[int, tuple[float, float]] = {}
+    speeds: dict[int, float] = {}
+    residuals: dict[int, float] = {}
+
+    prev_frame = ordered_items[0][0]
+
+    for frame_idx, (cx, cy) in ordered_items:
+        frame_delta = max(1, frame_idx - prev_frame) if frame_idx != prev_frame else 1
+        dt = frame_delta * base_dt
+        F = np.array(
+            [
+                [1, 0, dt, 0],
+                [0, 1, 0, dt],
+                [0, 0, 1, 0],
+                [0, 0, 0, 1],
+            ],
+            dtype=float,
+        )
+        q = 0.5 * dt**2
+        Q = np.array(
+            [
+                [q, 0, dt, 0],
+                [0, q, 0, dt],
+                [dt, 0, 1, 0],
+                [0, dt, 0, 1],
+            ],
+            dtype=float,
+        ) * 0.05
+
+        state_vec = F @ state_vec
+        P = F @ P @ F.T + Q
+
+        z = np.array([cx, cy], dtype=float)
+        innovation = z - H @ state_vec
+        S = H @ P @ H.T + R
+        K = P @ H.T @ np.linalg.inv(S)
+        state_vec = state_vec + K @ innovation
+        P = (np.eye(4) - K @ H) @ P
+
+        positions[frame_idx] = (state_vec[0], state_vec[1])
+        speeds[frame_idx] = float(math.hypot(state_vec[2], state_vec[3]))
+        residuals[frame_idx] = float(math.hypot(innovation[0], innovation[1]))
+
+        prev_frame = frame_idx
+
+    return positions, speeds, residuals
+
+
 def _build_kick_plot(state: AppState):
     fig = go.Figure()
     if state is None or not state.kick_debug_frames or not state.kick_debug_speeds:
@@ -605,9 +688,18 @@ def _build_kick_plot(state: AppState):
             mode="lines",
             name="Distance from start",
             line=dict(color="#9467bd"),
-            yaxis="y2",
         )
     )
+    if state.kick_debug_kalman_speeds:
+        fig.add_trace(
+            go.Scatter(
+                x=frames,
+                y=state.kick_debug_kalman_speeds,
+                mode="lines",
+                name="Kalman speed",
+                line=dict(color="#8c564b"),
+            )
+        )
     if kick_frame is not None:
         fig.add_trace(
             go.Scatter(
@@ -649,6 +741,10 @@ def _recompute_motion_metrics(state: AppState, target_obj_id: int = 1):
         state.kick_debug_area = []
         state.kick_debug_kick_frame = None
         state.kick_debug_distance = []
+        state.kalman_centers[target_obj_id] = {}
+        state.kalman_speeds[target_obj_id] = {}
+        state.kalman_residuals[target_obj_id] = {}
+        state.kick_debug_kalman_speeds = []
         return
 
     items = sorted(centers.items())
@@ -680,6 +776,10 @@ def _recompute_motion_metrics(state: AppState, target_obj_id: int = 1):
 
     state.smoothed_centers[target_obj_id] = smoothed
     state.ball_speeds[target_obj_id] = speeds
+    kalman_pos, kalman_speed, kalman_res = _run_kalman_filter(items, dt)
+    state.kalman_centers[target_obj_id] = kalman_pos
+    state.kalman_speeds[target_obj_id] = kalman_speed
+    state.kalman_residuals[target_obj_id] = kalman_res
     state.kick_frame = _detect_kick_frame(state, target_obj_id)
 
 
@@ -720,6 +820,8 @@ def _detect_kick_frame(state: AppState, target_obj_id: int) -> int | None:
         math.hypot(smoothed[f][0] - initial_center[0], smoothed[f][1] - initial_center[1])
         for f in frames
     ]
+    kalman_speed_dict = state.kalman_speeds.get(target_obj_id, {})
+    state.kick_debug_kalman_speeds = [kalman_speed_dict.get(f, 0.0) for f in frames]
     state.kick_debug_kick_frame = None
 
     for idx in range(baseline_window, len(frames)):
@@ -988,6 +1090,18 @@ def reset_session(GLOBAL_STATE: gr.State) -> tuple[AppState, Image.Image, int, i
     GLOBAL_STATE.ball_speeds.clear()
     GLOBAL_STATE.kick_frame = None
     GLOBAL_STATE.ball_centers.clear()
+    GLOBAL_STATE.kalman_centers.clear()
+    GLOBAL_STATE.kalman_speeds.clear()
+    GLOBAL_STATE.kalman_residuals.clear()
+    GLOBAL_STATE.kick_debug_frames = []
+    GLOBAL_STATE.kick_debug_speeds = []
+    GLOBAL_STATE.kick_debug_threshold = None
+    GLOBAL_STATE.kick_debug_baseline = None
+    GLOBAL_STATE.kick_debug_speed_std = None
+    GLOBAL_STATE.kick_debug_area = []
+    GLOBAL_STATE.kick_debug_kick_frame = None
+    GLOBAL_STATE.kick_debug_distance = []
+    GLOBAL_STATE.kick_debug_kalman_speeds = []
 
     # Dispose and re-init inference session for current model with existing frames
     try: