| 401 | super().__init__() |
| 402 | |
| 403 | def track(self, poses, identities=None): |
| 404 | self.n_frames += 1 |
| 405 | |
| 406 | trackers = np.zeros((len(self.trackers), 6)) |
| 407 | for i in range(len(trackers)): |
| 408 | trackers[i, :5] = self.trackers[i].predict() |
| 409 | empty = np.isnan(trackers).any(axis=1) |
| 410 | trackers = trackers[~empty] |
| 411 | for ind in np.flatnonzero(empty)[::-1]: |
| 412 | self.trackers.pop(ind) |
| 413 | |
| 414 | ellipses = [] |
| 415 | pred_ids = [] |
| 416 | for i, pose in enumerate(poses): |
| 417 | el = self.fitter.fit(pose) |
| 418 | if el is not None: |
| 419 | ellipses.append(el) |
| 420 | if identities is not None: |
| 421 | pred_ids.append(mode(identities[i])[0][0]) |
| 422 | if not len(trackers): |
| 423 | matches = np.empty((0, 2), dtype=int) |
| 424 | unmatched_detections = np.arange(len(ellipses)) |
| 425 | unmatched_trackers = np.empty((0, 6), dtype=int) |
| 426 | else: |
| 427 | ellipses_trackers = [Ellipse(*t[:5]) for t in trackers] |
| 428 | cost_matrix = np.zeros((len(ellipses), len(ellipses_trackers))) |
| 429 | for i, el in enumerate(ellipses): |
| 430 | for j, el_track in enumerate(ellipses_trackers): |
| 431 | cost = el.calc_similarity_with(el_track) |
| 432 | if identities is not None: |
| 433 | match = 2 if pred_ids[i] == self.trackers[j].id_ else 1 |
| 434 | cost *= match |
| 435 | cost_matrix[i, j] = cost |
| 436 | row_indices, col_indices = linear_sum_assignment(cost_matrix, maximize=True) |
| 437 | unmatched_detections = [i for i, _ in enumerate(ellipses) if i not in row_indices] |
| 438 | unmatched_trackers = [j for j, _ in enumerate(trackers) if j not in col_indices] |
| 439 | matches = [] |
| 440 | for row, col in zip(row_indices, col_indices, strict=False): |
| 441 | val = cost_matrix[row, col] |
| 442 | # diff = val - cost_matrix |
| 443 | # diff[row, col] += val |
| 444 | # if ( |
| 445 | # val < self.iou_threshold |
| 446 | # or np.any(diff[row] <= 0.2) |
| 447 | # or np.any(diff[:, col] <= 0.2) |
| 448 | # ): |
| 449 | if val < self.iou_threshold: |
| 450 | unmatched_detections.append(row) |
| 451 | unmatched_trackers.append(col) |
| 452 | else: |
| 453 | matches.append([row, col]) |
| 454 | if not len(matches): |
| 455 | matches = np.empty((0, 2), dtype=int) |
| 456 | else: |
| 457 | matches = np.stack(matches) |
| 458 | unmatched_trackers = np.asarray(unmatched_trackers) |
| 459 | unmatched_detections = np.asarray(unmatched_detections) |
| 460 | |