Spaces:
Sleeping
Sleeping
File size: 6,349 Bytes
3b90d9c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 |
"""
Ball Tracking Module using Kalman Filter.
This module implements a Kalman filter-based tracker for smoothing
and predicting tennis ball positions across video frames.
"""
import numpy as np
from typing import Optional, Tuple, List
from filterpy.kalman import KalmanFilter
class BallTracker:
"""
Kalman filter-based tracker for tennis ball position and velocity.
The tracker maintains state estimates for:
- Position (x, y)
- Velocity (vx, vy)
Attributes:
dt (float): Time step between frames (1/fps)
process_noise (float): Process noise covariance
measurement_noise (float): Measurement noise covariance
max_missing_frames (int): Maximum frames without detection before reset
"""
def __init__(
self,
dt: float = 1.0 / 30.0,
process_noise: float = 0.1,
measurement_noise: float = 10.0,
max_missing_frames: int = 10
):
"""
Initialize the ball tracker.
Args:
dt: Time step between frames (seconds)
process_noise: Process noise standard deviation
measurement_noise: Measurement noise standard deviation
max_missing_frames: Max consecutive frames without detection
"""
self.dt = dt
self.process_noise = process_noise
self.measurement_noise = measurement_noise
self.max_missing_frames = max_missing_frames
# Initialize Kalman filter
self.kf = self._create_kalman_filter()
# Tracking state
self.initialized = False
self.missing_frames = 0
self.trajectory = [] # List of (x, y, vx, vy, frame_num)
self.frame_count = 0
def _create_kalman_filter(self) -> KalmanFilter:
"""
Create and configure a Kalman filter for 2D position tracking.
State vector: [x, y, vx, vy]
Measurement vector: [x, y]
Returns:
Configured KalmanFilter instance
"""
kf = KalmanFilter(dim_x=4, dim_z=2)
# State transition matrix (constant velocity model)
kf.F = np.array([
[1, 0, self.dt, 0],
[0, 1, 0, self.dt],
[0, 0, 1, 0],
[0, 0, 0, 1]
])
# Measurement matrix (observe position only)
kf.H = np.array([
[1, 0, 0, 0],
[0, 1, 0, 0]
])
# Measurement noise covariance
kf.R = np.eye(2) * self.measurement_noise
# Process noise covariance
q = self.process_noise
kf.Q = np.array([
[q * self.dt**4 / 4, 0, q * self.dt**3 / 2, 0],
[0, q * self.dt**4 / 4, 0, q * self.dt**3 / 2],
[q * self.dt**3 / 2, 0, q * self.dt**2, 0],
[0, q * self.dt**3 / 2, 0, q * self.dt**2]
])
# Initial state covariance
kf.P = np.eye(4) * 100
return kf
def update(
self,
measurement: Optional[Tuple[float, float]] = None
) -> Optional[Tuple[float, float, float, float]]:
"""
Update tracker with a new measurement or predict if no detection.
Args:
measurement: Ball center position as (x, y), or None if not detected
Returns:
Estimated state as (x, y, vx, vy) or None if tracker not initialized
"""
self.frame_count += 1
if measurement is not None:
# Detection available
if not self.initialized:
# Initialize tracker with first detection
self.kf.x = np.array([
measurement[0],
measurement[1],
0.0,
0.0
])
self.initialized = True
self.missing_frames = 0
else:
# Update with measurement
z = np.array([measurement[0], measurement[1]])
self.kf.predict()
self.kf.update(z)
self.missing_frames = 0
# Record trajectory
x, y, vx, vy = self.kf.x
self.trajectory.append((
float(x), float(y), float(vx), float(vy), self.frame_count
))
return (float(x), float(y), float(vx), float(vy))
else:
# No detection - predict only
if self.initialized:
self.kf.predict()
self.missing_frames += 1
# Reset if too many missing frames
if self.missing_frames > self.max_missing_frames:
self.reset()
return None
# Return prediction
x, y, vx, vy = self.kf.x
self.trajectory.append((
float(x), float(y), float(vx), float(vy), self.frame_count
))
return (float(x), float(y), float(vx), float(vy))
return None
def reset(self):
"""Reset tracker to uninitialized state."""
self.kf = self._create_kalman_filter()
self.initialized = False
self.missing_frames = 0
def get_trajectory(self) -> List[Tuple[float, float, float, float, int]]:
"""
Get the complete trajectory history.
Returns:
List of trajectory points as (x, y, vx, vy, frame_num)
"""
return self.trajectory
def get_speed(self, state: Tuple[float, float, float, float]) -> float:
"""
Calculate speed from velocity components.
Args:
state: Tracker state as (x, y, vx, vy)
Returns:
Speed in pixels per second
"""
_, _, vx, vy = state
speed = np.sqrt(vx**2 + vy**2) / self.dt
return float(speed)
def get_last_n_positions(self, n: int = 20) -> List[Tuple[float, float]]:
"""
Get the last N tracked positions for trail visualization.
Args:
n: Number of recent positions to return
Returns:
List of (x, y) coordinates
"""
if len(self.trajectory) == 0:
return []
recent = self.trajectory[-n:]
return [(x, y) for x, y, _, _, _ in recent]
def is_initialized(self) -> bool:
"""Check if tracker has been initialized with a detection."""
return self.initialized
|