| 149 | |
| 150 | |
| 151 | class BYTETracker(object): |
| 152 | def __init__(self, opt, frame_rate=30): |
| 153 | self.opt = opt |
| 154 | if int(opt.gpus[0]) >= 0: |
| 155 | opt.device = torch.device('cuda') |
| 156 | else: |
| 157 | opt.device = torch.device('cpu') |
| 158 | print('Creating model...') |
| 159 | |
| 160 | ckpt = torch.load(opt.weights, map_location=opt.device) # load checkpoint |
| 161 | self.model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=1).to(opt.device) # create |
| 162 | exclude = ['anchor'] if opt.cfg else [] # exclude keys |
| 163 | if type(ckpt['model']).__name__ == "OrderedDict": |
| 164 | state_dict = ckpt['model'] |
| 165 | else: |
| 166 | state_dict = ckpt['model'].float().state_dict() # to FP32 |
| 167 | state_dict = intersect_dicts(state_dict, self.model.state_dict(), exclude=exclude) # intersect |
| 168 | self.model.load_state_dict(state_dict, strict=False) # load |
| 169 | self.model.cuda().eval() |
| 170 | total_params = sum(p.numel() for p in self.model.parameters()) |
| 171 | print(f'{total_params:,} total parameters.') |
| 172 | |
| 173 | |
| 174 | self.tracked_stracks = [] # type: list[STrack] |
| 175 | self.lost_stracks = [] # type: list[STrack] |
| 176 | self.removed_stracks = [] # type: list[STrack] |
| 177 | |
| 178 | self.frame_id = 0 |
| 179 | self.det_thresh = opt.conf_thres |
| 180 | self.buffer_size = int(frame_rate / 30.0 * opt.track_buffer) |
| 181 | self.max_time_lost = self.buffer_size |
| 182 | self.mean = np.array(opt.mean, dtype=np.float32).reshape(1, 1, 3) |
| 183 | self.std = np.array(opt.std, dtype=np.float32).reshape(1, 1, 3) |
| 184 | |
| 185 | self.kalman_filter = KalmanFilter() |
| 186 | self.low_thres = 0.1 |
| 187 | self.high_thres = self.opt.conf_thres + 0.1 |
| 188 | |
| 189 | def update(self, im_blob, img0,seq_num, save_dir): |
| 190 | self.frame_id += 1 |
| 191 | activated_starcks = [] |
| 192 | refind_stracks = [] |
| 193 | lost_stracks = [] |
| 194 | removed_stracks = [] |
| 195 | dets = [] |
| 196 | |
| 197 | ''' Step 1: Network forward, get detections & embeddings''' |
| 198 | with torch.no_grad(): |
| 199 | output = self.model(im_blob, augment=False) |
| 200 | pred, train_out = output[1] |
| 201 | |
| 202 | pred = pred[pred[:, :, 4] > self.low_thres] |
| 203 | detections = [] |
| 204 | if len(pred) > 0: |
| 205 | dets,x_inds,y_inds = non_max_suppression_and_inds(pred[:,:6].unsqueeze(0), 0.1, self.opt.nms_thres,method='cluster_diou') |
| 206 | dets = dets.numpy() |
| 207 | if len(dets) != 0: |
| 208 | scale_coords(self.opt.img_size, dets[:, :4], img0.shape).round() |
nothing calls this directly
no outgoing calls
no test coverage detected