(self, prob_threshold=0.2, area_threshold=100, vis=False)
| 380 | cv2.imwrite(img_path, img_show) |
| 381 | |
| 382 | def detect(self, prob_threshold=0.2, area_threshold=100, vis=False): |
| 383 | total_dts = 0 |
| 384 | track_instances = None |
| 385 | max_id = 0 |
| 386 | |
| 387 | # we only consider val split (second half images) |
| 388 | for i in tqdm(range((int(self.img_len / 2)), self.img_len)): |
| 389 | # for i in tqdm(range(0, self.img_len)): |
| 390 | img, targets = self.load_img_from_file(self.img_list[i]) |
| 391 | cur_img, ori_img = self.init_img(img) |
| 392 | |
| 393 | # track_instances = None |
| 394 | if track_instances is not None: |
| 395 | track_instances.remove('boxes') |
| 396 | track_instances.remove('labels') |
| 397 | |
| 398 | res = self.detr.inference_single_image(cur_img.cuda().float(), (self.seq_h, self.seq_w), track_instances) |
| 399 | track_instances = res['track_instances'] |
| 400 | max_id = max(max_id, track_instances.obj_idxes.max().item()) |
| 401 | |
| 402 | print("ref points.shape={}".format(res['ref_pts'].shape)) |
| 403 | all_ref_pts = tensor_to_numpy(res['ref_pts'][0, :, :2]) |
| 404 | dt_instances = track_instances.to(torch.device('cpu')) |
| 405 | |
| 406 | # filter det instances by score. |
| 407 | dt_instances = self.filter_dt_by_score(dt_instances, prob_threshold) |
| 408 | dt_instances = self.filter_dt_by_area(dt_instances, area_threshold) |
| 409 | |
| 410 | total_dts += len(dt_instances) |
| 411 | |
| 412 | if vis: |
| 413 | # for visual |
| 414 | cur_vis_img_path = os.path.join(self.save_path, 'frame_{:0>8d}.jpg'.format(i)) |
| 415 | gt_boxes = None |
| 416 | self.visualize_img_with_bbox(cur_vis_img_path, ori_img, dt_instances, ref_pts=all_ref_pts, gt_boxes=gt_boxes) |
| 417 | |
| 418 | tracker_outputs = self.tr_tracker.update(dt_instances) |
| 419 | |
| 420 | self.write_results(txt_path=os.path.join(self.predict_path, 'gt.txt'), |
| 421 | frame_id=(i + 1), |
| 422 | bbox_xyxy=tracker_outputs[:, :4], |
| 423 | identities=tracker_outputs[:, 5]) |
| 424 | print("totally {} dts max_id={}".format(total_dts, max_id)) |
| 425 | |
| 426 | |
| 427 | if __name__ == '__main__': |
no test coverage detected