| 153 | |
| 154 | |
| 155 | class DeepSort(object): |
| 156 | def __init__(self, model_path, max_dist=0.1, min_confidence=0.3, nms_max_overlap=1.0, max_iou_distance=0.7, max_age=30, n_init=3, nn_budget=100, use_cuda=True): |
| 157 | self.min_confidence = min_confidence |
| 158 | self.nms_max_overlap = nms_max_overlap |
| 159 | |
| 160 | self.extractor = Extractor(model_path, use_cuda=use_cuda) |
| 161 | |
| 162 | max_cosine_distance = max_dist |
| 163 | metric = NearestNeighborDistanceMetric( |
| 164 | "cosine", max_cosine_distance, nn_budget) |
| 165 | self.tracker = Tracker( |
| 166 | metric, max_iou_distance=max_iou_distance, max_age=max_age, n_init=n_init) |
| 167 | |
| 168 | def update(self, output_results, img_info, img_size, img_file_name): |
| 169 | img_file_name = os.path.join(get_yolox_datadir(), 'mot', 'train', img_file_name) |
| 170 | ori_img = cv2.imread(img_file_name) |
| 171 | self.height, self.width = ori_img.shape[:2] |
| 172 | # post process detections |
| 173 | output_results = output_results.cpu().numpy() |
| 174 | confidences = output_results[:, 4] * output_results[:, 5] |
| 175 | |
| 176 | bboxes = output_results[:, :4] # x1y1x2y2 |
| 177 | img_h, img_w = img_info[0], img_info[1] |
| 178 | scale = min(img_size[0] / float(img_h), img_size[1] / float(img_w)) |
| 179 | bboxes /= scale |
| 180 | bbox_xyxy = bboxes |
| 181 | bbox_tlwh = self._xyxy_to_tlwh_array(bbox_xyxy) |
| 182 | remain_inds = confidences > self.min_confidence |
| 183 | bbox_tlwh = bbox_tlwh[remain_inds] |
| 184 | confidences = confidences[remain_inds] |
| 185 | |
| 186 | # generate detections |
| 187 | features = self._get_features(bbox_tlwh, ori_img) |
| 188 | detections = [Detection(bbox_tlwh[i], conf, features[i]) for i, conf in enumerate( |
| 189 | confidences) if conf > self.min_confidence] |
| 190 | classes = np.zeros((len(detections), )) |
| 191 | |
| 192 | # run on non-maximum supression |
| 193 | boxes = np.array([d.tlwh for d in detections]) |
| 194 | scores = np.array([d.confidence for d in detections]) |
| 195 | |
| 196 | # update tracker |
| 197 | self.tracker.predict() |
| 198 | self.tracker.update(detections, classes) |
| 199 | |
| 200 | # output bbox identities |
| 201 | outputs = [] |
| 202 | for track in self.tracker.tracks: |
| 203 | if not track.is_confirmed() or track.time_since_update > 1: |
| 204 | continue |
| 205 | box = track.to_tlwh() |
| 206 | x1, y1, x2, y2 = self._tlwh_to_xyxy_noclip(box) |
| 207 | track_id = track.track_id |
| 208 | class_id = track.class_id |
| 209 | outputs.append(np.array([x1, y1, x2, y2, track_id, class_id], dtype=np.int)) |
| 210 | if len(outputs) > 0: |
| 211 | outputs = np.stack(outputs, axis=0) |
| 212 | return outputs |