Function run_each_dataset

tutorials/ctracker/test_byte.py:47–104

(model_dir, retinanet, dataset_path, subset, cur_dataset)

                f.write(line)
                
def run_each_dataset(model_dir, retinanet, dataset_path, subset, cur_dataset):
    print(cur_dataset)

    img_list = os.listdir(os.path.join(dataset_path, subset, cur_dataset, 'img1'))
    img_list = [os.path.join(dataset_path, subset, cur_dataset, 'img1', _) for _ in img_list if ('jpg' in _) or ('png' in _)]
    img_list = sorted(img_list)

    img_len = len(img_list)
    last_feat = None

    confidence_threshold = 0.6
    IOU_threshold = 0.5
    retention_threshold = 10

    det_list_all = []
    tracklet_all = []
    results = []
    max_id = 0
    max_draw_len = 100
    draw_interval = 5
    img_width = 1920
    img_height = 1080
    fps = 30

    tracker = BYTETracker()

    for idx in range((int(img_len / 2)), img_len + 1):
        i = idx - 1
        print('tracking: ', i)
        with torch.no_grad():
            data_path1 = img_list[min(idx, img_len - 1)]
            img_origin1 = skimage.io.imread(data_path1)
            img_h, img_w, _ = img_origin1.shape
            img_height, img_width = img_h, img_w
            resize_h, resize_w = math.ceil(img_h / 32) * 32, math.ceil(img_w / 32) * 32
            img1 = np.zeros((resize_h, resize_w, 3), dtype=img_origin1.dtype)
            img1[:img_h, :img_w, :] = img_origin1
            img1 = (img1.astype(np.float32) / 255.0 - np.array([[RGB_MEAN]])) / np.array([[RGB_STD]])
            img1 = torch.from_numpy(img1).permute(2, 0, 1).view(1, 3, resize_h, resize_w)
            scores, transformed_anchors, last_feat = retinanet(img1.cuda().float(), last_feat=last_feat)
            
            if idx > (int(img_len / 2)):
                idxs = np.where(scores > 0.1)
                # run tracking
                online_targets = tracker.update(transformed_anchors[idxs[0], :4], scores[idxs[0]])
                online_tlwhs = []
                online_ids = []
                online_scores = []
                for t in online_targets:
                    tlwh = t.tlwh
                    tid = t.track_id
                    online_tlwhs.append(tlwh)
                    online_ids.append(tid)
                    online_scores.append(t.score)
                results.append((idx, online_tlwhs, online_ids, online_scores))
            
    fout_tracking = os.path.join(model_dir, 'results', cur_dataset + '.txt')
    write_results(fout_tracking, results)