|
2 | 2 |
|
3 | 3 | import numpy as np |
4 | 4 |
|
5 | | -from trackers.core.sort.kalman_box_tracker import SORTKalmanBoxTracker |
6 | 5 |
|
7 | | - |
8 | | -class DeepSORTKalmanBoxTracker(SORTKalmanBoxTracker): |
| 6 | +class DeepSORTKalmanBoxTracker: |
9 | 7 | """ |
10 | 8 | The `DeepSORTKalmanBoxTracker` class represents the internals of a single |
11 | 9 | tracked object (bounding box), with a Kalman filter to predict and update |
12 | 10 | its position. It also maintains a feature vector for the object, which is |
13 | 11 | used to identify the object across frames. |
| 12 | +
|
| 13 | + Attributes: |
| 14 | + tracker_id (int): Unique identifier for the tracker. |
| 15 | + number_of_successful_updates (int): Number of times the object has been |
| 16 | + updated successfully. |
| 17 | + time_since_update (int): Number of frames since the last update. |
| 18 | + state (np.ndarray): State vector of the bounding box. |
| 19 | + F (np.ndarray): State transition matrix. |
| 20 | + H (np.ndarray): Measurement matrix. |
| 21 | + Q (np.ndarray): Process noise covariance matrix. |
| 22 | + R (np.ndarray): Measurement noise covariance matrix. |
| 23 | + P (np.ndarray): Error covariance matrix. |
| 24 | + features (list[np.ndarray]): List of feature vectors. |
| 25 | + count_id (int): Class variable to assign unique IDs to each tracker. |
| 26 | +
|
| 27 | + Args: |
| 28 | + bbox (np.ndarray): Initial bounding box in the form [x1, y1, x2, y2]. |
| 29 | + feature (Optional[np.ndarray]): Optional initial feature vector. |
14 | 30 | """ |
15 | 31 |
|
| 32 | + count_id = 0 |
| 33 | + |
| 34 | + @classmethod |
| 35 | + def get_next_tracker_id(cls) -> int: |
| 36 | + """ |
| 37 | + Class method that returns the next available tracker ID. |
| 38 | +
|
| 39 | + Returns: |
| 40 | + int: The next available tracker ID. |
| 41 | + """ |
| 42 | + next_id = cls.count_id |
| 43 | + cls.count_id += 1 |
| 44 | + return next_id |
| 45 | + |
16 | 46 | def __init__(self, bbox: np.ndarray, feature: Optional[np.ndarray] = None): |
17 | | - super().__init__(bbox) |
| 47 | + # Initialize with a temporary ID of -1 |
| 48 | + # Will be assigned a real ID when the track is considered mature |
| 49 | + self.tracker_id = -1 |
| 50 | + |
| 51 | + # Number of hits indicates how many times the object has been |
| 52 | + # updated successfully |
| 53 | + self.number_of_successful_updates = 1 |
| 54 | + # Number of frames since the last update |
| 55 | + self.time_since_update = 0 |
| 56 | + |
| 57 | + # For simplicity, we keep a small state vector: |
| 58 | + # (x, y, x2, y2, vx, vy, vx2, vy2). |
| 59 | + # We'll store the bounding box in "self.state" |
| 60 | + self.state = np.zeros((8, 1), dtype=np.float32) |
| 61 | + |
| 62 | + # Initialize state directly from the first detection |
| 63 | + self.state[0] = bbox[0] |
| 64 | + self.state[1] = bbox[1] |
| 65 | + self.state[2] = bbox[2] |
| 66 | + self.state[3] = bbox[3] |
| 67 | + |
| 68 | + # Basic constant velocity model |
| 69 | + self._initialize_kalman_filter() |
| 70 | + |
| 71 | + # Initialize features list |
18 | 72 | self.features: list[np.ndarray] = [] |
19 | 73 | if feature is not None: |
20 | 74 | self.features.append(feature) |
21 | 75 |
|
| 76 | + def _initialize_kalman_filter(self) -> None: |
| 77 | + """ |
| 78 | + Sets up the matrices for the Kalman filter. |
| 79 | + """ |
| 80 | + # State transition matrix (F): 8x8 |
| 81 | + # We assume a constant velocity model. Positions are incremented by |
| 82 | + # velocity each step. |
| 83 | + self.F = np.eye(8, dtype=np.float32) |
| 84 | + for i in range(4): |
| 85 | + self.F[i, i + 4] = 1.0 |
| 86 | + |
| 87 | + # Measurement matrix (H): we directly measure x1, y1, x2, y2 |
| 88 | + self.H = np.eye(4, 8, dtype=np.float32) # 4x8 |
| 89 | + |
| 90 | + # Process covariance matrix (Q) |
| 91 | + self.Q = np.eye(8, dtype=np.float32) * 0.01 |
| 92 | + |
| 93 | + # Measurement covariance (R): noise in detection |
| 94 | + self.R = np.eye(4, dtype=np.float32) * 0.1 |
| 95 | + |
| 96 | + # Error covariance matrix (P) |
| 97 | + self.P = np.eye(8, dtype=np.float32) |
| 98 | + |
| 99 | + def predict(self) -> None: |
| 100 | + """ |
| 101 | + Predict the next state of the bounding box (applies the state transition). |
| 102 | + """ |
| 103 | + # Predict state |
| 104 | + self.state = self.F @ self.state |
| 105 | + # Predict error covariance |
| 106 | + self.P = self.F @ self.P @ self.F.T + self.Q |
| 107 | + |
| 108 | + # Increase time since update |
| 109 | + self.time_since_update += 1 |
| 110 | + |
| 111 | + def update(self, bbox: np.ndarray) -> None: |
| 112 | + """ |
| 113 | + Updates the state with a new detected bounding box. |
| 114 | +
|
| 115 | + Args: |
| 116 | + bbox (np.ndarray): Detected bounding box in the form [x1, y1, x2, y2]. |
| 117 | + """ |
| 118 | + self.time_since_update = 0 |
| 119 | + self.number_of_successful_updates += 1 |
| 120 | + |
| 121 | + # Kalman Gain |
| 122 | + S = self.H @ self.P @ self.H.T + self.R |
| 123 | + K = self.P @ self.H.T @ np.linalg.inv(S) |
| 124 | + |
| 125 | + # Residual |
| 126 | + measurement = bbox.reshape((4, 1)) |
| 127 | + y = measurement - self.H @ self.state |
| 128 | + |
| 129 | + # Update state |
| 130 | + self.state = self.state + K @ y |
| 131 | + |
| 132 | + # Update covariance |
| 133 | + identity_matrix = np.eye(8, dtype=np.float32) |
| 134 | + self.P = (identity_matrix - K @ self.H) @ self.P |
| 135 | + |
| 136 | + def get_state_bbox(self) -> np.ndarray: |
| 137 | + """ |
| 138 | + Returns the current bounding box estimate from the state vector. |
| 139 | +
|
| 140 | + Returns: |
| 141 | + np.ndarray: The bounding box [x1, y1, x2, y2]. |
| 142 | + """ |
| 143 | + return np.array( |
| 144 | + [ |
| 145 | + self.state[0], # x1 |
| 146 | + self.state[1], # y1 |
| 147 | + self.state[2], # x2 |
| 148 | + self.state[3], # y2 |
| 149 | + ], |
| 150 | + dtype=float, |
| 151 | + ).reshape(-1) |
| 152 | + |
22 | 153 | def update_feature(self, feature: np.ndarray): |
23 | 154 | self.features.append(feature) |
24 | 155 |
|
|
0 commit comments