(self, results_with_low, public_det=None)
| 27 | self.tracks = [] |
| 28 | |
| 29 | def step(self, results_with_low, public_det=None): |
| 30 | |
| 31 | results = [item for item in results_with_low if item['score'] >= self.opt.track_thresh] |
| 32 | |
| 33 | # first association |
| 34 | N = len(results) |
| 35 | M = len(self.tracks) |
| 36 | |
| 37 | dets = np.array( |
| 38 | [det['ct'] + det['tracking'] for det in results], np.float32) # N x 2 |
| 39 | track_size = np.array([((track['bbox'][2] - track['bbox'][0]) * \ |
| 40 | (track['bbox'][3] - track['bbox'][1])) \ |
| 41 | for track in self.tracks], np.float32) # M |
| 42 | track_cat = np.array([track['class'] for track in self.tracks], np.int32) # M |
| 43 | item_size = np.array([((item['bbox'][2] - item['bbox'][0]) * \ |
| 44 | (item['bbox'][3] - item['bbox'][1])) \ |
| 45 | for item in results], np.float32) # N |
| 46 | item_cat = np.array([item['class'] for item in results], np.int32) # N |
| 47 | tracks = np.array( |
| 48 | [pre_det['ct'] for pre_det in self.tracks], np.float32) # M x 2 |
| 49 | dist = (((tracks.reshape(1, -1, 2) - \ |
| 50 | dets.reshape(-1, 1, 2)) ** 2).sum(axis=2)) # N x M |
| 51 | |
| 52 | invalid = ((dist > track_size.reshape(1, M)) + \ |
| 53 | (dist > item_size.reshape(N, 1)) + \ |
| 54 | (item_cat.reshape(N, 1) != track_cat.reshape(1, M))) > 0 |
| 55 | dist = dist + invalid * 1e18 |
| 56 | |
| 57 | if self.opt.hungarian: |
| 58 | assert not self.opt.hungarian, 'we only verify centertrack with greedy_assignment' |
| 59 | item_score = np.array([item['score'] for item in results], np.float32) # N |
| 60 | dist[dist > 1e18] = 1e18 |
| 61 | matched_indices = linear_assignment(dist) |
| 62 | else: |
| 63 | matched_indices = greedy_assignment(copy.deepcopy(dist)) |
| 64 | |
| 65 | unmatched_dets = [d for d in range(dets.shape[0]) \ |
| 66 | if not (d in matched_indices[:, 0])] |
| 67 | unmatched_tracks = [d for d in range(tracks.shape[0]) \ |
| 68 | if not (d in matched_indices[:, 1])] |
| 69 | |
| 70 | if self.opt.hungarian: |
| 71 | assert not self.opt.hungarian, 'we only verify centertrack with greedy_assignment' |
| 72 | matches = [] |
| 73 | for m in matched_indices: |
| 74 | if dist[m[0], m[1]] > 1e16: |
| 75 | unmatched_dets.append(m[0]) |
| 76 | unmatched_tracks.append(m[1]) |
| 77 | else: |
| 78 | matches.append(m) |
| 79 | matches = np.array(matches).reshape(-1, 2) |
| 80 | else: |
| 81 | matches = matched_indices |
| 82 | |
| 83 | ret = [] |
| 84 | for m in matches: |
| 85 | track = results[m[0]] |
| 86 | track['tracking_id'] = self.tracks[m[1]]['tracking_id'] |
nothing calls this directly
no test coverage detected