| 309 | |
| 310 | |
| 311 | class BoxTracker(BaseTracker): |
| 312 | def __init__(self, bbox): |
| 313 | super().__init__(dim=4, dim_z=4) |
| 314 | self.kf = KalmanFilter(dim_x=7, dim_z=4) |
| 315 | self.kf.F = np.array( |
| 316 | [ |
| 317 | [1, 0, 0, 0, 1, 0, 0], |
| 318 | [0, 1, 0, 0, 0, 1, 0], |
| 319 | [0, 0, 1, 0, 0, 0, 1], |
| 320 | [0, 0, 0, 1, 0, 0, 0], |
| 321 | [0, 0, 0, 0, 1, 0, 0], |
| 322 | [0, 0, 0, 0, 0, 1, 0], |
| 323 | [0, 0, 0, 0, 0, 0, 1], |
| 324 | ] |
| 325 | ) |
| 326 | self.kf.H = np.array( |
| 327 | [ |
| 328 | [1, 0, 0, 0, 0, 0, 0], |
| 329 | [0, 1, 0, 0, 0, 0, 0], |
| 330 | [0, 0, 1, 0, 0, 0, 0], |
| 331 | [0, 0, 0, 1, 0, 0, 0], |
| 332 | ] |
| 333 | ) |
| 334 | self.kf.R[2:, 2:] *= 10.0 |
| 335 | # Give high uncertainty to the unobservable initial velocities |
| 336 | self.kf.P[4:, 4:] *= 1000.0 |
| 337 | self.kf.P *= 10.0 |
| 338 | self.kf.Q[-1, -1] *= 0.01 |
| 339 | self.kf.Q[4:, 4:] *= 0.01 |
| 340 | self.state = bbox |
| 341 | |
| 342 | def update(self, bbox): |
| 343 | super().update(self.convert_bbox_to_z(bbox)) |
| 344 | |
| 345 | def predict(self): |
| 346 | if (self.kf.x[6] + self.kf.x[2]) <= 0: |
| 347 | self.kf.x[6] *= 0.0 |
| 348 | return super().predict() |
| 349 | |
| 350 | @property |
| 351 | def state(self): |
| 352 | return self.convert_x_to_bbox(self.kf.x)[0] |
| 353 | |
| 354 | @state.setter |
| 355 | def state(self, bbox): |
| 356 | state = self.convert_bbox_to_z(bbox) |
| 357 | super(BoxTracker, type(self)).state.fset(self, state) |
| 358 | |
| 359 | @staticmethod |
| 360 | def convert_x_to_bbox(x, score=None): |
| 361 | """Takes a bounding box in the centre form [x,y,s,r] and returns it in the form |
| 362 | [x1,y1,x2,y2] where x1,y1 is the top left and x2,y2 is the bottom right.""" |
| 363 | w = np.sqrt(x[2] * x[3]) |
| 364 | h = x[2] / w |
| 365 | if score is None: |
| 366 | return np.array([x[0] - w / 2.0, x[1] - h / 2.0, x[0] + w / 2.0, x[1] + h / 2.0]).reshape((1, 4)) |
| 367 | else: |
| 368 | return np.array([x[0] - w / 2.0, x[1] - h / 2.0, x[0] + w / 2.0, x[1] + h / 2.0, score]).reshape((1, 5)) |