| 135 | |
| 136 | |
| 137 | class BYTETracker(object): |
| 138 | def __init__(self, opt, frame_rate=30): |
| 139 | self.opt = opt |
| 140 | self.model = Darknet(opt.cfg, nID=14455) |
| 141 | # load_darknet_weights(self.model, opt.weights) |
| 142 | self.model.load_state_dict(torch.load(opt.weights, map_location='cpu')['model'], strict=False) |
| 143 | self.model.cuda().eval() |
| 144 | |
| 145 | self.tracked_stracks = [] # type: list[STrack] |
| 146 | self.lost_stracks = [] # type: list[STrack] |
| 147 | self.removed_stracks = [] # type: list[STrack] |
| 148 | |
| 149 | self.frame_id = 0 |
| 150 | self.det_thresh = opt.conf_thres |
| 151 | self.init_thresh = self.det_thresh + 0.2 |
| 152 | self.low_thresh = 0.3 |
| 153 | self.buffer_size = int(frame_rate / 30.0 * opt.track_buffer) |
| 154 | self.max_time_lost = self.buffer_size |
| 155 | |
| 156 | self.kalman_filter = KalmanFilter() |
| 157 | |
| 158 | def update(self, im_blob, img0): |
| 159 | """ |
| 160 | Processes the image frame and finds bounding box(detections). |
| 161 | |
| 162 | Associates the detection with corresponding tracklets and also handles lost, removed, refound and active tracklets |
| 163 | |
| 164 | Parameters |
| 165 | ---------- |
| 166 | im_blob : torch.float32 |
| 167 | Tensor of shape depending upon the size of image. By default, shape of this tensor is [1, 3, 608, 1088] |
| 168 | |
| 169 | img0 : ndarray |
| 170 | ndarray of shape depending on the input image sequence. By default, shape is [608, 1080, 3] |
| 171 | |
| 172 | Returns |
| 173 | ------- |
| 174 | output_stracks : list of Strack(instances) |
| 175 | The list contains information regarding the online_tracklets for the recieved image tensor. |
| 176 | |
| 177 | """ |
| 178 | |
| 179 | self.frame_id += 1 |
| 180 | activated_starcks = [] # for storing active tracks, for the current frame |
| 181 | refind_stracks = [] # Lost Tracks whose detections are obtained in the current frame |
| 182 | lost_stracks = [] # The tracks which are not obtained in the current frame but are not removed.(Lost for some time lesser than the threshold for removing) |
| 183 | removed_stracks = [] |
| 184 | |
| 185 | t1 = time.time() |
| 186 | ''' Step 1: Network forward, get detections & embeddings''' |
| 187 | with torch.no_grad(): |
| 188 | pred = self.model(im_blob) |
| 189 | # pred is tensor of all the proposals (default number of proposals: 54264). Proposals have information associated with the bounding box and embeddings |
| 190 | pred = pred[pred[:, :, 4] > self.low_thresh] |
| 191 | # pred now has lesser number of proposals. Proposals rejected on basis of object confidence score |
| 192 | if len(pred) > 0: |
| 193 | dets = non_max_suppression(pred.unsqueeze(0), self.low_thresh, self.opt.nms_thres)[0].cpu() |
| 194 | # Final proposals are obtained in dets. Information of bounding box and embeddings also included |
no outgoing calls
no test coverage detected