| 154 | |
| 155 | |
| 156 | class BYTETracker(object): |
| 157 | def __init__(self, opt, frame_rate=30): |
| 158 | self.opt = opt |
| 159 | if opt.gpus[0] >= 0: |
| 160 | opt.device = torch.device('cuda') |
| 161 | else: |
| 162 | opt.device = torch.device('cpu') |
| 163 | print('Creating model...') |
| 164 | self.model = create_model(opt.arch, opt.heads, opt.head_conv) |
| 165 | self.model = load_model(self.model, opt.load_model) |
| 166 | self.model = self.model.to(opt.device) |
| 167 | self.model.eval() |
| 168 | |
| 169 | self.tracked_stracks = [] # type: list[STrack] |
| 170 | self.lost_stracks = [] # type: list[STrack] |
| 171 | self.removed_stracks = [] # type: list[STrack] |
| 172 | |
| 173 | self.frame_id = 0 |
| 174 | #self.det_thresh = opt.conf_thres |
| 175 | self.det_thresh = opt.conf_thres + 0.1 |
| 176 | self.buffer_size = int(frame_rate / 30.0 * opt.track_buffer) |
| 177 | self.max_time_lost = self.buffer_size |
| 178 | self.max_per_image = opt.K |
| 179 | self.mean = np.array(opt.mean, dtype=np.float32).reshape(1, 1, 3) |
| 180 | self.std = np.array(opt.std, dtype=np.float32).reshape(1, 1, 3) |
| 181 | |
| 182 | self.kalman_filter = KalmanFilter() |
| 183 | |
| 184 | def post_process(self, dets, meta): |
| 185 | dets = dets.detach().cpu().numpy() |
| 186 | dets = dets.reshape(1, -1, dets.shape[2]) |
| 187 | dets = ctdet_post_process( |
| 188 | dets.copy(), [meta['c']], [meta['s']], |
| 189 | meta['out_height'], meta['out_width'], self.opt.num_classes) |
| 190 | for j in range(1, self.opt.num_classes + 1): |
| 191 | dets[0][j] = np.array(dets[0][j], dtype=np.float32).reshape(-1, 5) |
| 192 | return dets[0] |
| 193 | |
| 194 | def merge_outputs(self, detections): |
| 195 | results = {} |
| 196 | for j in range(1, self.opt.num_classes + 1): |
| 197 | results[j] = np.concatenate( |
| 198 | [detection[j] for detection in detections], axis=0).astype(np.float32) |
| 199 | |
| 200 | scores = np.hstack( |
| 201 | [results[j][:, 4] for j in range(1, self.opt.num_classes + 1)]) |
| 202 | if len(scores) > self.max_per_image: |
| 203 | kth = len(scores) - self.max_per_image |
| 204 | thresh = np.partition(scores, kth)[kth] |
| 205 | for j in range(1, self.opt.num_classes + 1): |
| 206 | keep_inds = (results[j][:, 4] >= thresh) |
| 207 | results[j] = results[j][keep_inds] |
| 208 | return results |
| 209 | |
| 210 | def update(self, im_blob, img0): |
| 211 | self.frame_id += 1 |
| 212 | activated_starcks = [] |
| 213 | refind_stracks = [] |
nothing calls this directly
no outgoing calls
no test coverage detected