| 619 | super().__init__() |
| 620 | |
| 621 | def track(self, dets): |
| 622 | self.n_frames += 1 |
| 623 | |
| 624 | trackers = np.zeros((len(self.trackers), 5)) |
| 625 | for i in range(len(trackers)): |
| 626 | trackers[i, :4] = self.trackers[i].predict() |
| 627 | empty = np.isnan(trackers).any(axis=1) |
| 628 | trackers = trackers[~empty] |
| 629 | for ind in np.flatnonzero(empty)[::-1]: |
| 630 | self.trackers.pop(ind) |
| 631 | |
| 632 | matched, unmatched_dets, unmatched_trks = self.match_detections_to_trackers(dets, trackers, self.iou_threshold) |
| 633 | |
| 634 | # update matched trackers with assigned detections |
| 635 | animalindex = [] |
| 636 | for t, trk in enumerate(self.trackers): |
| 637 | if t not in unmatched_trks: |
| 638 | d = matched[np.where(matched[:, 1] == t)[0], 0] |
| 639 | animalindex.append(d[0]) |
| 640 | trk.update(dets[d, :][0]) # update coordinates |
| 641 | else: |
| 642 | animalindex.append("nix") # lost trk! |
| 643 | |
| 644 | # create and initialise new trackers for unmatched detections |
| 645 | for i in unmatched_dets: |
| 646 | trk = BoxTracker(dets[i, :]) |
| 647 | self.trackers.append(trk) |
| 648 | animalindex.append(i) |
| 649 | |
| 650 | i = len(self.trackers) |
| 651 | ret = [] |
| 652 | for trk in reversed(self.trackers): |
| 653 | d = trk.state |
| 654 | if (trk.time_since_update < 1) and (trk.hit_streak >= self.min_hits or self.n_frames <= self.min_hits): |
| 655 | ret.append( |
| 656 | np.concatenate((d, [trk.id, int(animalindex[i - 1])])).reshape(1, -1) |
| 657 | ) # for DLC we also return the original animalid |
| 658 | # +1 as MOT benchmark requires positive >> this is removed for DLC! |
| 659 | i -= 1 |
| 660 | # remove dead tracklet |
| 661 | if trk.time_since_update > self.max_age: |
| 662 | self.trackers.pop(i) |
| 663 | |
| 664 | if len(ret) > 0: |
| 665 | return np.concatenate(ret) |
| 666 | return np.empty((0, 5)) |
| 667 | |
| 668 | @staticmethod |
| 669 | def match_detections_to_trackers(detections, trackers, iou_threshold): |