| 111 | |
| 112 | |
| 113 | class Evaluator(object): |
| 114 | def __init__(self, data_root, seq_name, data_type='mot'): |
| 115 | |
| 116 | self.data_root = data_root |
| 117 | self.seq_name = seq_name |
| 118 | self.data_type = data_type |
| 119 | |
| 120 | self.load_annotations() |
| 121 | self.reset_accumulator() |
| 122 | |
| 123 | def load_annotations(self): |
| 124 | assert self.data_type == 'mot' |
| 125 | |
| 126 | gt_filename = os.path.join(self.data_root, self.seq_name, 'gt', 'gt.txt') |
| 127 | self.gt_frame_dict = read_results(gt_filename, self.data_type, is_gt=True) |
| 128 | self.gt_ignore_frame_dict = read_results(gt_filename, self.data_type, is_ignore=True) |
| 129 | |
| 130 | def reset_accumulator(self): |
| 131 | self.acc = mm.MOTAccumulator(auto_id=True) |
| 132 | |
| 133 | def eval_frame(self, frame_id, trk_tlwhs, trk_ids, rtn_events=False): |
| 134 | # results |
| 135 | trk_tlwhs = np.copy(trk_tlwhs) |
| 136 | trk_ids = np.copy(trk_ids) |
| 137 | |
| 138 | # gts |
| 139 | gt_objs = self.gt_frame_dict.get(frame_id, []) |
| 140 | gt_tlwhs, gt_ids = unzip_objs(gt_objs)[:2] |
| 141 | |
| 142 | # ignore boxes |
| 143 | ignore_objs = self.gt_ignore_frame_dict.get(frame_id, []) |
| 144 | ignore_tlwhs = unzip_objs(ignore_objs)[0] |
| 145 | # remove ignored results |
| 146 | keep = np.ones(len(trk_tlwhs), dtype=bool) |
| 147 | iou_distance = mm.distances.iou_matrix(ignore_tlwhs, trk_tlwhs, max_iou=0.5) |
| 148 | if len(iou_distance) > 0: |
| 149 | match_is, match_js = mm.lap.linear_sum_assignment(iou_distance) |
| 150 | match_is, match_js = map(lambda a: np.asarray(a, dtype=int), [match_is, match_js]) |
| 151 | match_ious = iou_distance[match_is, match_js] |
| 152 | |
| 153 | match_js = np.asarray(match_js, dtype=int) |
| 154 | match_js = match_js[np.logical_not(np.isnan(match_ious))] |
| 155 | keep[match_js] = False |
| 156 | trk_tlwhs = trk_tlwhs[keep] |
| 157 | trk_ids = trk_ids[keep] |
| 158 | |
| 159 | # get distance matrix |
| 160 | iou_distance = mm.distances.iou_matrix(gt_tlwhs, trk_tlwhs, max_iou=0.5) |
| 161 | |
| 162 | # acc |
| 163 | self.acc.update(gt_ids, trk_ids, iou_distance) |
| 164 | |
| 165 | if rtn_events and iou_distance.size > 0 and hasattr(self.acc, 'last_mot_events'): |
| 166 | events = self.acc.last_mot_events # only supported by https://github.com/longcw/py-motmetrics |
| 167 | else: |
| 168 | events = None |
| 169 | return events |
| 170 | |