(self, results_with_low, public_det=None)
| 37 | self.alive = [] |
| 38 | |
| 39 | def step(self, results_with_low, public_det=None): |
| 40 | results = [item for item in results_with_low if item['score'] >= self.opt.track_thresh] |
| 41 | |
| 42 | # first association |
| 43 | N = len(results) |
| 44 | M = len(self.tracks) |
| 45 | self.alive = [] |
| 46 | |
| 47 | track_boxes = np.array([[track['bbox'][0], track['bbox'][1], |
| 48 | track['bbox'][2], track['bbox'][3]] for track in self.tracks], np.float32) # M x 4 |
| 49 | det_boxes = np.array([[item['bbox'][0], item['bbox'][1], |
| 50 | item['bbox'][2], item['bbox'][3]] for item in results], np.float32) # N x 4 |
| 51 | box_ious = self.bbox_overlaps_py(det_boxes, track_boxes) |
| 52 | |
| 53 | dets = np.array( |
| 54 | [det['ct'] + det['tracking'] for det in results], np.float32) # N x 2 |
| 55 | track_size = np.array([((track['bbox'][2] - track['bbox'][0]) * \ |
| 56 | (track['bbox'][3] - track['bbox'][1])) \ |
| 57 | for track in self.tracks], np.float32) # M |
| 58 | track_cat = np.array([track['class'] for track in self.tracks], np.int32) # M |
| 59 | item_size = np.array([((item['bbox'][2] - item['bbox'][0]) * \ |
| 60 | (item['bbox'][3] - item['bbox'][1])) \ |
| 61 | for item in results], np.float32) # N |
| 62 | item_cat = np.array([item['class'] for item in results], np.int32) # N |
| 63 | tracks = np.array( |
| 64 | [pre_det['ct'] for pre_det in self.tracks], np.float32) # M x 2 |
| 65 | dist = (((tracks.reshape(1, -1, 2) - \ |
| 66 | dets.reshape(-1, 1, 2)) ** 2).sum(axis=2)) # N x M |
| 67 | |
| 68 | if self.opt.dataset == 'youtube_vis': |
| 69 | invalid = ((dist > track_size.reshape(1, M)) + \ |
| 70 | (dist > item_size.reshape(N, 1)) + (box_ious < self.opt.overlap_thresh)) > 0 |
| 71 | else: |
| 72 | invalid = ((dist > track_size.reshape(1, M)) + \ |
| 73 | (dist > item_size.reshape(N, 1)) + \ |
| 74 | (item_cat.reshape(N, 1) != track_cat.reshape(1, M)) + (box_ious < self.opt.overlap_thresh)) > 0 |
| 75 | dist = dist + invalid * 1e18 |
| 76 | |
| 77 | if self.opt.hungarian: |
| 78 | item_score = np.array([item['score'] for item in results], np.float32) # N |
| 79 | dist[dist > 1e18] = 1e18 |
| 80 | matched_indices = linear_assignment(dist) |
| 81 | else: |
| 82 | matched_indices = greedy_assignment(copy.deepcopy(dist)) |
| 83 | unmatched_dets = [d for d in range(dets.shape[0]) \ |
| 84 | if not (d in matched_indices[:, 0])] |
| 85 | unmatched_tracks = [d for d in range(tracks.shape[0]) \ |
| 86 | if not (d in matched_indices[:, 1])] |
| 87 | |
| 88 | if self.opt.hungarian: |
| 89 | matches = [] |
| 90 | for m in matched_indices: |
| 91 | if dist[m[0], m[1]] > 1e16: |
| 92 | unmatched_dets.append(m[0]) |
| 93 | unmatched_tracks.append(m[1]) |
| 94 | else: |
| 95 | matches.append(m) |
| 96 | matches = np.array(matches).reshape(-1, 2) |
no test coverage detected