Пример #1
0
    samples = sample_generator(target_bbox, opts['n_samples'])
    sample_scores = forward_samples(model, image, samples, out_layer='fc6')

    top_scores, top_idx = sample_scores[:, 1].topk(5)
    top_idx = top_idx.cpu()
    target_score = top_scores.mean()
    target_bbox = samples[top_idx]
    if top_idx.shape[0] > 1:
        target_bbox = target_bbox.mean(axis=0)
    success = target_score > 0

    # Expand search area at failure
    if success:
        sample_generator.set_trans(opts['trans'])
    else:
        sample_generator.expand_trans(opts['trans_limit'])

    # Bbox regression
    if success:
        bbreg_samples = samples[top_idx]
        if top_idx.shape[0] == 1:
            bbreg_samples = bbreg_samples[None, :]
        bbreg_feats = forward_samples(model, image, bbreg_samples)
        bbreg_samples = bbreg.predict(bbreg_feats, bbreg_samples)
        bbreg_bbox = bbreg_samples.mean(axis=0)
    else:
        bbreg_bbox = target_bbox

    pred_poly = Rectangle(bbreg_bbox[0], bbreg_bbox[1], bbreg_bbox[2],
                          bbreg_bbox[3])
    confidence = target_score.item()
Пример #2
0
def run_mdnet(img_list, init_bbox, gt=None, savefig_dir='', display=False):

    # Init bbox
    target_bbox = np.array(init_bbox)
    result = np.zeros((len(img_list), 4))
    result_bb = np.zeros((len(img_list), 4))
    result[0] = target_bbox
    result_bb[0] = target_bbox

    if gt is not None:
        overlap = np.zeros(len(img_list))
        overlap[0] = 1

    # Init model
    model = MDNet(opts['model_path'])
    if opts['use_gpu']:
        model = model.cuda()

    # Init criterion and optimizer
    criterion = BCELoss()
    model.set_learnable_params(opts['ft_layers'])
    init_optimizer = set_optimizer(model, opts['lr_init'], opts['lr_mult'])
    update_optimizer = set_optimizer(model, opts['lr_update'], opts['lr_mult'])

    tic = time.time()
    # Load first image
    image = Image.open(img_list[0]).convert('RGB')

    # Draw pos/neg samples
    pos_examples = SampleGenerator('gaussian', image.size, opts['trans_pos'],
                                   opts['scale_pos'])(target_bbox,
                                                      opts['n_pos_init'],
                                                      opts['overlap_pos_init'])

    neg_examples = np.concatenate([
        SampleGenerator('uniform', image.size, opts['trans_neg_init'],
                        opts['scale_neg_init'])(target_bbox,
                                                int(opts['n_neg_init'] * 0.5),
                                                opts['overlap_neg_init']),
        SampleGenerator('whole', image.size)(target_bbox,
                                             int(opts['n_neg_init'] * 0.5),
                                             opts['overlap_neg_init'])
    ])
    neg_examples = np.random.permutation(neg_examples)

    # Extract pos/neg features
    pos_feats = forward_samples(model, image, pos_examples)
    neg_feats = forward_samples(model, image, neg_examples)

    # Initial training
    train(model, criterion, init_optimizer, pos_feats, neg_feats,
          opts['maxiter_init'])
    del init_optimizer, neg_feats
    torch.cuda.empty_cache()

    # Train bbox regressor
    bbreg_examples = SampleGenerator(
        'uniform', image.size, opts['trans_bbreg'], opts['scale_bbreg'],
        opts['aspect_bbreg'])(target_bbox, opts['n_bbreg'],
                              opts['overlap_bbreg'])
    bbreg_feats = forward_samples(model, image, bbreg_examples)
    bbreg = BBRegressor(image.size)
    bbreg.train(bbreg_feats, bbreg_examples, target_bbox)
    del bbreg_feats
    torch.cuda.empty_cache()

    # Init sample generators for update:这三个是【生成器】,不是产生的数据
    sample_generator = SampleGenerator('gaussian', image.size, opts['trans'],
                                       opts['scale'])
    pos_generator = SampleGenerator('gaussian', image.size, opts['trans_pos'],
                                    opts['scale_pos'])
    neg_generator = SampleGenerator('uniform', image.size, opts['trans_neg'],
                                    opts['scale_neg'])

    # Init pos/neg features for update
    neg_examples = neg_generator(target_bbox, opts['n_neg_update'],
                                 opts['overlap_neg_init'])
    neg_feats = forward_samples(model, image, neg_examples)
    pos_feats_all = [pos_feats]
    neg_feats_all = [neg_feats]

    spf_total = time.time() - tic

    # Display
    savefig = savefig_dir != ''
    if display or savefig:
        dpi = 80.0
        figsize = (image.size[0] / dpi, image.size[1] / dpi)

        fig = plt.figure(frameon=False, figsize=figsize, dpi=dpi)
        ax = plt.Axes(fig, [0., 0., 1., 1.])
        ax.set_axis_off()
        fig.add_axes(ax)
        im = ax.imshow(image, aspect='auto')

        if gt is not None:
            gt_rect = plt.Rectangle(tuple(gt[0, :2]),
                                    gt[0, 2],
                                    gt[0, 3],
                                    linewidth=3,
                                    edgecolor="#00ff00",
                                    zorder=1,
                                    fill=False)
            ax.add_patch(gt_rect)

        rect = plt.Rectangle(tuple(result_bb[0, :2]),
                             result_bb[0, 2],
                             result_bb[0, 3],
                             linewidth=3,
                             edgecolor="#ff0000",
                             zorder=1,
                             fill=False)
        ax.add_patch(rect)

        if display:
            plt.pause(.01)
            plt.draw()
        if savefig:
            fig.savefig(os.path.join(savefig_dir, '0000.jpg'), dpi=dpi)

    # Main loop
    for i in range(1, len(img_list)):

        tic = time.time()
        # Load image
        image = Image.open(img_list[i]).convert('RGB')

        # Estimate target bbox:指的是下一次的target_bbox,用于迭代
        samples = sample_generator(target_bbox, opts['n_samples'])
        sample_scores = forward_samples(
            model, image, samples, out_layer='fc6'
        )  #forward_samples是根据当前的给当前的多个sample打分数,【相当于执行了一次网络判断】

        top_scores, top_idx = sample_scores[:, 1].topk(5)
        top_idx = top_idx.cpu()
        target_score = top_scores.mean()
        target_bbox = samples[top_idx]
        if top_idx.shape[0] > 1:
            target_bbox = target_bbox.mean(axis=0)  #多个target_bbox进行均值优化
        success = target_score > 0

        # Expand search area at failure:在【跟踪的同时】,根据跟踪情况,采用不同的sample生成参数
        if success:
            sample_generator.set_trans(opts['trans'])
        else:
            sample_generator.expand_trans(opts['trans_limit'])

        # Bbox regression :利用【当前:只是对当前帧图像进行回归,没有考虑前几帧】几名对应的box,最为regression的输入,来产生一个更好的Bbox
        if success:
            bbreg_samples = samples[top_idx]
            if top_idx.shape[0] == 1:
                bbreg_samples = bbreg_samples[None, :]
            bbreg_feats = forward_samples(model, image, bbreg_samples)
            bbreg_samples = bbreg.predict(bbreg_feats, bbreg_samples)
            bbreg_bbox = bbreg_samples.mean(axis=0)
        else:
            bbreg_bbox = target_bbox

        # Save result
        result[i] = target_bbox
        result_bb[i] = bbreg_bbox

        # Data collect
        if success:
            pos_examples = pos_generator(target_bbox, opts['n_pos_update'],
                                         opts['overlap_pos_update'])
            pos_feats = forward_samples(model, image, pos_examples)
            pos_feats_all.append(pos_feats)
            if len(pos_feats_all) > opts['n_frames_long']:
                del pos_feats_all[0]

            neg_examples = neg_generator(target_bbox, opts['n_neg_update'],
                                         opts['overlap_neg_update'])
            neg_feats = forward_samples(model, image, neg_examples)
            neg_feats_all.append(neg_feats)
            if len(neg_feats_all) > opts['n_frames_short']:
                del neg_feats_all[0]

        # Short term update :每张图片都要更新一次模型,这次的输入数据就是【前几帧累计的正样本和负样本】
        if not success:
            nframes = min(opts['n_frames_short'], len(pos_feats_all))
            pos_data = torch.cat(pos_feats_all[-nframes:], 0)
            neg_data = torch.cat(neg_feats_all, 0)
            train(model, criterion, update_optimizer, pos_data, neg_data,
                  opts['maxiter_update'])

        # Long term update
        elif i % opts['long_interval'] == 0:
            pos_data = torch.cat(pos_feats_all, 0)
            neg_data = torch.cat(neg_feats_all, 0)
            train(model, criterion, update_optimizer, pos_data, neg_data,
                  opts['maxiter_update'])

        torch.cuda.empty_cache()
        spf = time.time() - tic
        spf_total += spf

        # Display
        if display or savefig:
            im.set_data(image)

            if gt is not None:
                gt_rect.set_xy(gt[i, :2])
                gt_rect.set_width(gt[i, 2])
                gt_rect.set_height(gt[i, 3])

            rect.set_xy(result_bb[i, :2])
            rect.set_width(result_bb[i, 2])
            rect.set_height(result_bb[i, 3])

            if display:
                plt.pause(.01)
                plt.draw()
            if savefig:
                fig.savefig(os.path.join(savefig_dir, '{:04d}.jpg'.format(i)),
                            dpi=dpi)

        if gt is None:
            print('Frame {:d}/{:d}, Score {:.3f}, Time {:.3f}'.format(
                i, len(img_list), target_score, spf))
        else:
            overlap[i] = overlap_ratio(gt[i], result_bb[i])[0]
            print('Frame {:d}/{:d}, Overlap {:.3f}, Score {:.3f}, Time {:.3f}'.
                  format(i, len(img_list), overlap[i], target_score, spf))

    if gt is not None:
        print('meanIOU: {:.3f}'.format(overlap.mean()))
    fps = len(img_list) / spf_total
    return result, result_bb, fps
Пример #3
0
def run_mdnet(img_list,
              init_bbox,
              gt=None,
              savefig_dir='',
              display=False,
              model_path='models/model001.pth'):

    # Init bbox
    target_bbox = np.array(init_bbox)
    result = np.zeros((len(img_list), 4))
    result_bb = np.zeros((len(img_list), 4))
    result[0] = target_bbox
    result_bb[0] = target_bbox

    if gt is not None:
        overlap = np.zeros(len(img_list))
        overlap[0] = 1

    # Init model
    opts['model_path'] = model_path

    print('********')
    print('model:', opts['model_path'])
    print('********')

    assert (model_path == 'models/model000.pth'
            or model_path == 'models/model001.pth')

    if model_path == 'models/model000.pth': model = MDNet0(opts['model_path'])
    else: model = MDNet1(opts['model_path'])

    if opts['use_gpu']:
        model = model.cuda()

    # Init criterion and optimizer
    criterion = BCELoss()
    model.set_learnable_params(opts['ft_layers'])
    init_optimizer = set_optimizer(model, opts['lr_init'], opts['lr_mult'])
    update_optimizer = set_optimizer(model, opts['lr_update'], opts['lr_mult'])

    tic = time.time()
    # Load first image
    image = Image.open(img_list[0]).convert('RGB')

    # Draw pos/neg samples
    pos_examples = SampleGenerator('gaussian', image.size, opts['trans_pos'],
                                   opts['scale_pos'])(target_bbox,
                                                      opts['n_pos_init'],
                                                      opts['overlap_pos_init'])

    neg_examples = np.concatenate([
        SampleGenerator('uniform', image.size, opts['trans_neg_init'],
                        opts['scale_neg_init'])(target_bbox,
                                                int(opts['n_neg_init'] * 0.5),
                                                opts['overlap_neg_init']),
        SampleGenerator('whole', image.size)(target_bbox,
                                             int(opts['n_neg_init'] * 0.5),
                                             opts['overlap_neg_init'])
    ])
    neg_examples = np.random.permutation(neg_examples)

    # Extract pos/neg features
    pos_feats = forward_samples(model, image, pos_examples)
    print(pos_feats)
    neg_feats = forward_samples(model, image, neg_examples)
    print(neg_feats)

    # Initial training
    train(model, criterion, init_optimizer, pos_feats, neg_feats,
          opts['maxiter_init'])
    del init_optimizer, neg_feats
    torch.cuda.empty_cache()

    # Train bbox regressor
    bbreg_examples = SampleGenerator(
        'uniform', image.size, opts['trans_bbreg'], opts['scale_bbreg'],
        opts['aspect_bbreg'])(target_bbox, opts['n_bbreg'],
                              opts['overlap_bbreg'])
    bbreg_feats = forward_samples(model, image, bbreg_examples)
    bbreg = BBRegressor(image.size)
    bbreg.train(bbreg_feats, bbreg_examples, target_bbox)
    del bbreg_feats
    torch.cuda.empty_cache()

    # Init sample generators for update
    sample_generator = SampleGenerator('gaussian', image.size, opts['trans'],
                                       opts['scale'])
    pos_generator = SampleGenerator('gaussian', image.size, opts['trans_pos'],
                                    opts['scale_pos'])
    neg_generator = SampleGenerator('uniform', image.size, opts['trans_neg'],
                                    opts['scale_neg'])

    # Init pos/neg features for update
    neg_examples = neg_generator(target_bbox, opts['n_neg_update'],
                                 opts['overlap_neg_init'])
    neg_feats = forward_samples(model, image, neg_examples)
    pos_feats_all = [pos_feats]
    neg_feats_all = [neg_feats]

    spf_total = time.time() - tic

    # Display
    savefig = savefig_dir != ''
    if display or savefig:
        dpi = 80.0
        figsize = (image.size[0] / dpi, image.size[1] / dpi)

        fig = plt.figure(frameon=False, figsize=figsize, dpi=dpi)
        ax = plt.Axes(fig, [0., 0., 1., 1.])
        ax.set_axis_off()
        fig.add_axes(ax)
        im = ax.imshow(image, aspect='auto')

        if gt is not None:
            gt_rect = plt.Rectangle(tuple(gt[0, :2]),
                                    gt[0, 2],
                                    gt[0, 3],
                                    linewidth=3,
                                    edgecolor="#00ff00",
                                    zorder=1,
                                    fill=False)
            ax.add_patch(gt_rect)

        rect = plt.Rectangle(tuple(result_bb[0, :2]),
                             result_bb[0, 2],
                             result_bb[0, 3],
                             linewidth=3,
                             edgecolor="#ff0000",
                             zorder=1,
                             fill=False)
        ax.add_patch(rect)

        if display:
            plt.pause(.01)
            plt.draw()
        if savefig:
            fig.savefig(os.path.join(savefig_dir, '0000.jpg'), dpi=dpi)

    # Main loop
    for i in range(1, len(img_list)):

        tic = time.time()
        # Load image
        image = Image.open(img_list[i]).convert('RGB')

        # Estimate target bbox
        samples = sample_generator(target_bbox, opts['n_samples'])
        sample_scores = forward_samples(model, image, samples, out_layer='fc6')

        top_scores, top_idx = sample_scores[:, 1].topk(5)

        # for top 5 samples, maximize score using hill-climbing algorithm
        for j in range(5):
            sample_ = samples[top_idx[j]]
            last_top_score = None

            # hill-climbing search
            while True:
                sample_left_p = [
                    sample_[0] + 1, sample_[1], sample_[2] - 1, sample_[3]
                ]
                sample_left_n = [
                    sample_[0] - 1, sample_[1], sample_[2] + 1, sample_[3]
                ]
                sample_up_p = [
                    sample_[0], sample_[1] + 1, sample_[2], sample_[3] - 1
                ]
                sample_up_n = [
                    sample_[0], sample_[1] - 1, sample_[2], sample_[3] + 1
                ]
                sample_right_p = [
                    sample_[0], sample_[1], sample_[2] + 1, sample_[3]
                ]
                sample_right_n = [
                    sample_[0], sample_[1], sample_[2] - 1, sample_[3]
                ]
                sample_bottom_p = [
                    sample_[0], sample_[1], sample_[2], sample_[3] + 1
                ]
                sample_bottom_n = [
                    sample_[0], sample_[1], sample_[2], sample_[3] - 1
                ]

                all_samples = [
                    sample_left_p, sample_left_n, sample_up_p, sample_up_n,
                    sample_right_p, sample_right_n, sample_bottom_p,
                    sample_bottom_n
                ]

                hillClimbingSS = forward_samples(model,
                                                 image,
                                                 np.array(all_samples),
                                                 out_layer='fc6')
                top_score, top_index = hillClimbingSS[:, 1].topk(1)
                top_score_float = top_score.cpu().numpy()[0]

                # End of hill climbing: this is THE BEST!
                if last_top_score != None:
                    if top_score_float < last_top_score: break

                sample_ = all_samples[top_index]
                samples[top_idx[j]] = all_samples[top_index]
                last_top_score = top_score_float

        # modify sample scores array
        sample_scores = forward_samples(model, image, samples, out_layer='fc6')
        top_scores, top_idx = sample_scores[:, 1].topk(5)

        sampleStore = []
        for j in range(len(samples)):
            temp = []
            for k in range(4):
                temp.append(samples[j][k])
            sampleStore.append(temp)

        # if mean score of bbox < 0, find everywhere
        target_score = top_scores.mean()

        if target_score < 0:
            # print('')
            # print('last bbox:')
            # print(result[i-1])
            last_left = result[i - 1][0]
            last_top = result[i - 1][1]

            # print('')
            # for j in range(len(samples)): print(j, samples[j], sample_scores[j])
            # print('')
            # print('sample top scores (before):')
            # print(top_scores)
            # print(top_idx)

            cnt = 0
            rl = [32, 16]

            for _ in range(len(rl)):
                everywhere_sample = []

                # find everywhere (near the last bbox)
                meanWidth = 0.0
                meanHeight = 0.0
                for j in range(len(samples)):
                    meanWidth += samples[j][2]
                    meanHeight += samples[j][3]
                meanWidth /= len(samples)
                meanHeight /= len(samples)

                width = image.size[0]
                height = image.size[1]

                for j in range(32):
                    for k in range(32):
                        jk = [
                            last_left + (31 - 2 * j) * meanWidth / rl[_],
                            last_top + (31 - 2 * k) * meanHeight / rl[_],
                            meanWidth, meanHeight
                        ]
                        # print(j, k, jk)
                        everywhere_sample.append(jk)

                everywhere_scores = forward_samples(
                    model, image, np.array(everywhere_sample), out_layer='fc6')
                everywhere_top_scores, everywhere_top_idx = everywhere_scores[:,
                                                                              1].topk(
                                                                                  5
                                                                              )

                # print('')
                # print('everywhere_sample:')
                # for j in range(len(everywhere_sample)): print(j, everywhere_sample[j], everywhere_scores[j])

                # print('')
                # print('everywhere top scores (before):')
                # print(everywhere_top_scores)
                # print(everywhere_top_idx)
                # for j in range(5): print(everywhere_sample[everywhere_top_idx[j]])

                # for top 5 samples in everywhere_sample, maximize score using hill-climbing algorithm
                for j in range(5):
                    # print('')
                    sample_ = everywhere_sample[everywhere_top_idx[j]]
                    last_top_score = None

                    # hill-climbing search
                    while True:
                        sample_left_p = [
                            sample_[0] + 1, sample_[1], sample_[2] - 1,
                            sample_[3]
                        ]
                        sample_left_n = [
                            sample_[0] - 1, sample_[1], sample_[2] + 1,
                            sample_[3]
                        ]
                        sample_up_p = [
                            sample_[0], sample_[1] + 1, sample_[2],
                            sample_[3] - 1
                        ]
                        sample_up_n = [
                            sample_[0], sample_[1] - 1, sample_[2],
                            sample_[3] + 1
                        ]
                        sample_right_p = [
                            sample_[0], sample_[1], sample_[2] + 1, sample_[3]
                        ]
                        sample_right_n = [
                            sample_[0], sample_[1], sample_[2] - 1, sample_[3]
                        ]
                        sample_bottom_p = [
                            sample_[0], sample_[1], sample_[2], sample_[3] + 1
                        ]
                        sample_bottom_n = [
                            sample_[0], sample_[1], sample_[2], sample_[3] - 1
                        ]

                        all_samples = [
                            sample_left_p, sample_left_n, sample_up_p,
                            sample_up_n, sample_right_p, sample_right_n,
                            sample_bottom_p, sample_bottom_n
                        ]

                        hillClimbingSS = forward_samples(model,
                                                         image,
                                                         np.array(all_samples),
                                                         out_layer='fc6')
                        top_score, top_index = hillClimbingSS[:, 1].topk(1)
                        top_score_float = top_score.cpu().numpy()[0]

                        # End of hill climbing: this is THE BEST!
                        if last_top_score != None:
                            # print(last_top_score)
                            if top_score_float < last_top_score: break

                        sample_ = all_samples[top_index]
                        everywhere_sample[
                            everywhere_top_idx[j]] = all_samples[top_index]
                        last_top_score = top_score_float

                everywhere_scores = forward_samples(
                    model, image, np.array(everywhere_sample), out_layer='fc6')
                everywhere_top_scores, everywhere_top_idx = everywhere_scores[:,
                                                                              1].topk(
                                                                                  5
                                                                              )

                # print('')
                # print('everywhere top scores (after):')
                # print(everywhere_top_scores)
                # print(everywhere_top_idx)
                # for j in range(5): print(everywhere_sample[everywhere_top_idx[j]])

                # merge 'samples' with everywhere samples
                everywhere_top5 = []
                for j in range(5):
                    everywhere_top5.append(
                        everywhere_sample[everywhere_top_idx[j]])
                samples = np.concatenate((samples, np.array(everywhere_top5)))

                sample_scores = forward_samples(model,
                                                image,
                                                samples,
                                                out_layer='fc6')
                top_scores, top_idx = sample_scores[:, 1].topk(5)

                if top_scores.mean() > 0:
                    # print('')
                    # for j in range(len(samples)): print(j, samples[j], sample_scores[j])
                    # print('')
                    # print('sample top scores (after):')
                    # print(top_scores)
                    # print(top_idx)
                    break
                cnt += 1

            # failure -> recover original samples
            if cnt == 2:
                # print('recovered')
                samples = np.array(sampleStore)
                sample_scores = forward_samples(model,
                                                image,
                                                samples,
                                                out_layer='fc6')
                top_scores, top_idx = sample_scores[:, 1].topk(5)

        # finally modify sample scores array
        sample_scores = forward_samples(model, image, samples, out_layer='fc6')
        top_scores, top_idx = sample_scores[:, 1].topk(5)

        top_idx = top_idx.cpu()
        target_score = top_scores.mean()
        target_bbox = samples[top_idx]
        if top_idx.shape[0] > 1:
            target_bbox = target_bbox.mean(axis=0)
        success = target_score > 0

        # Expand search area at failure
        if success:
            sample_generator.set_trans(opts['trans'])
        else:
            sample_generator.expand_trans(opts['trans_limit'])

        # Bbox regression
        if success:
            bbreg_samples = samples[top_idx]
            if top_idx.shape[0] == 1:
                bbreg_samples = bbreg_samples[None, :]
            bbreg_feats = forward_samples(model, image, bbreg_samples)
            bbreg_samples = bbreg.predict(bbreg_feats, bbreg_samples)
            bbreg_bbox = bbreg_samples.mean(axis=0)
        else:
            bbreg_bbox = target_bbox

        # Save result
        result[i] = target_bbox
        result_bb[i] = bbreg_bbox

        # Data collect
        if success:
            pos_examples = pos_generator(target_bbox, opts['n_pos_update'],
                                         opts['overlap_pos_update'])
            pos_feats = forward_samples(model, image, pos_examples)
            pos_feats_all.append(pos_feats)
            if len(pos_feats_all) > opts['n_frames_long']:
                del pos_feats_all[0]

            neg_examples = neg_generator(target_bbox, opts['n_neg_update'],
                                         opts['overlap_neg_update'])
            neg_feats = forward_samples(model, image, neg_examples)
            neg_feats_all.append(neg_feats)
            if len(neg_feats_all) > opts['n_frames_short']:
                del neg_feats_all[0]

        # Short term update
        if not success:
            nframes = min(opts['n_frames_short'], len(pos_feats_all))
            pos_data = torch.cat(pos_feats_all[-nframes:], 0)
            neg_data = torch.cat(neg_feats_all, 0)
            train(model, criterion, update_optimizer, pos_data, neg_data,
                  opts['maxiter_update'])

        # Long term update
        elif i % opts['long_interval'] == 0:
            pos_data = torch.cat(pos_feats_all, 0)
            neg_data = torch.cat(neg_feats_all, 0)
            train(model, criterion, update_optimizer, pos_data, neg_data,
                  opts['maxiter_update'])

        torch.cuda.empty_cache()
        spf = time.time() - tic
        spf_total += spf

        # Display
        if display or savefig:
            im.set_data(image)

            if gt is not None:
                gt_rect.set_xy(gt[i, :2])
                gt_rect.set_width(gt[i, 2])
                gt_rect.set_height(gt[i, 3])

            rect.set_xy(result_bb[i, :2])
            rect.set_width(result_bb[i, 2])
            rect.set_height(result_bb[i, 3])

            if display:
                plt.pause(.01)
                plt.draw()
            if savefig:
                fig.savefig(os.path.join(
                    savefig_dir,
                    ('M' + model_path[14] + 'T3_' + '{:04d}.jpg'.format(i))),
                            dpi=dpi)

        if gt is None:
            print('Frame {:d}/{:d}, Score {:.3f}, Time {:.3f}'.format(
                i, len(img_list), target_score, spf))
        else:
            overlap[i] = overlap_ratio(gt[i], result_bb[i])[0]
            print('Frame {:d}/{:d}, Overlap {:.3f}, Score {:.3f}, Time {:.3f}'.
                  format(i, len(img_list), overlap[i], target_score, spf))

    if gt is not None:
        print('meanIOU: {:.3f}'.format(overlap.mean()))
    fps = len(img_list) / spf_total
    plt.close('all')
    return result, result_bb, fps, overlap
Пример #4
0
    def search_track(self, track_num, frame_idx, init_bbox, previous_num,
                     init_conf):
        # img_list, init_bbox

        # Init bbox

        # result = np.zeros((len(img_list), 4))
        # result_bb = np.zeros((len(img_list), 4))
        # result[0] = target_bbox
        # result_bb[0] = target_bbox

        # superorange params
        track_list = [(-1, -1)] * (self.TRACKLET_NUM + 1)
        bbox_list = [(frame_idx, init_bbox, init_conf)]
        IOU_count = 0
        # last_bbox = -1
        # next_frame = -1
        #        frameA_path = os.path.join(self.FILE_PATH, str(frame_idx) + "."+ext)

        frameA_path = os.path.join(self.FILE_PATH,
                                   str(frame_idx).zfill(5) + ".png")

        target_bbox = np.array(init_bbox)
        # Init model
        model = MDNet(opts['model_path'])
        if opts['use_gpu']:
            model = model.cuda()

        # Init criterion and optimizer
        criterion = BCELoss()
        model.set_learnable_params(opts['ft_layers'])
        init_optimizer = set_optimizer(model, opts['lr_init'], opts['lr_mult'])
        update_optimizer = set_optimizer(model, opts['lr_update'],
                                         opts['lr_mult'])

        tic = time.time()
        # Load first image
        image = Image.open(frameA_path).convert('RGB')

        # Draw pos/neg samples
        pos_examples = SampleGenerator('gaussian', image.size,
                                       opts['trans_pos'], opts['scale_pos'])(
                                           target_bbox, opts['n_pos_init'],
                                           opts['overlap_pos_init'])

        neg_examples = np.concatenate([
            SampleGenerator('uniform', image.size, opts['trans_neg_init'],
                            opts['scale_neg_init'])(target_bbox,
                                                    int(opts['n_neg_init'] *
                                                        0.5),
                                                    opts['overlap_neg_init']),
            SampleGenerator('whole', image.size)(target_bbox,
                                                 int(opts['n_neg_init'] * 0.5),
                                                 opts['overlap_neg_init'])
        ])
        neg_examples = np.random.permutation(neg_examples)

        # Extract pos/neg features
        if len(pos_examples) == 0 or len(neg_examples) == 0:
            print("!!pos_examples=0 skip!!")
            return
        pos_feats = self.forward_samples(model, image, pos_examples)
        neg_feats = self.forward_samples(model, image, neg_examples)

        # Initial training
        self.train(model, criterion, init_optimizer, pos_feats, neg_feats,
                   opts['maxiter_init'])
        del init_optimizer, neg_feats
        torch.cuda.empty_cache()

        # Train bbox regressor
        bbreg_examples = SampleGenerator(
            'uniform', image.size, opts['trans_bbreg'], opts['scale_bbreg'],
            opts['aspect_bbreg'])(target_bbox, opts['n_bbreg'],
                                  opts['overlap_bbreg'])
        bbreg_feats = self.forward_samples(model, image, bbreg_examples)
        bbreg = BBRegressor(image.size)
        bbreg.train(bbreg_feats, bbreg_examples, target_bbox)
        del bbreg_feats
        torch.cuda.empty_cache()

        # Init sample generators for update
        sample_generator = SampleGenerator('gaussian', image.size,
                                           opts['trans'], opts['scale'])
        pos_generator = SampleGenerator('gaussian', image.size,
                                        opts['trans_pos'], opts['scale_pos'])
        neg_generator = SampleGenerator('uniform', image.size,
                                        opts['trans_neg'], opts['scale_neg'])

        # Init pos/neg features for update
        neg_examples = neg_generator(target_bbox, opts['n_neg_update'],
                                     opts['overlap_neg_init'])
        neg_feats = self.forward_samples(model, image, neg_examples)
        pos_feats_all = [pos_feats]
        neg_feats_all = [neg_feats]

        spf_total = time.time() - tic

        # fps = len(img_list) / spf_total
        # return result, result_bb, fps

        # Main loop
        for i in range(0, self.TRACKLET_NUM):  # next 10 frame
            frameB_idx = frame_idx + i + 1
            #    print("frameB_idx="+str(frameB_idx))
            if frameB_idx > self.frame_num:
                break
            else:

                frameB_path = os.path.join(self.FILE_PATH,
                                           str(frameB_idx).zfill(5) + ".png")

                # ------------track by MDNet------------
                # Load image
                image = Image.open(frameB_path).convert('RGB')

                # Estimate target bbox
                samples = sample_generator(target_bbox, opts['n_samples'])
                sample_scores = self.forward_samples(model,
                                                     image,
                                                     samples,
                                                     out_layer='fc6')

                top_scores, top_idx = sample_scores[:, 1].topk(5)
                top_idx = top_idx.cpu()
                target_score = top_scores.mean()
                target_bbox = samples[top_idx]
                if top_idx.shape[0] > 1:
                    target_bbox = target_bbox.mean(axis=0)
                success = target_score > 0

                # Expand search area at failure
                if success:
                    sample_generator.set_trans(opts['trans'])
                else:
                    sample_generator.expand_trans(opts['trans_limit'])

                # Bbox regression
                if success:
                    bbreg_samples = samples[top_idx]
                    if top_idx.shape[0] == 1:
                        bbreg_samples = bbreg_samples[None, :]
                    bbreg_feats = self.forward_samples(model, image,
                                                       bbreg_samples)
                    bbreg_samples = bbreg.predict(bbreg_feats, bbreg_samples)
                    bbreg_bbox = bbreg_samples.mean(axis=0)
                else:
                    bbreg_bbox = target_bbox

                # Save result
                # result[i] = target_bbox
                # result_bb[i] = bbreg_bbox

        #       print(target_bbox)

        # Data collect
                if success:
                    pos_examples = pos_generator(target_bbox,
                                                 opts['n_pos_update'],
                                                 opts['overlap_pos_update'])
                    pos_feats = self.forward_samples(model, image,
                                                     pos_examples)
                    pos_feats_all.append(pos_feats)
                    if len(pos_feats_all) > opts['n_frames_long']:
                        del pos_feats_all[0]

                    neg_examples = neg_generator(target_bbox,
                                                 opts['n_neg_update'],
                                                 opts['overlap_neg_update'])
                    neg_feats = self.forward_samples(model, image,
                                                     neg_examples)
                    neg_feats_all.append(neg_feats)
                    if len(neg_feats_all) > opts['n_frames_short']:
                        del neg_feats_all[0]

                # Short term update
                if not success:
                    nframes = min(opts['n_frames_short'], len(pos_feats_all))
                    pos_data = torch.cat(pos_feats_all[-nframes:], 0)
                    neg_data = torch.cat(neg_feats_all, 0)
                    self.train(model, criterion, update_optimizer, pos_data,
                               neg_data, opts['maxiter_update'])

                # Long term update
                elif i % opts['long_interval'] == 0:
                    pos_data = torch.cat(pos_feats_all, 0)
                    neg_data = torch.cat(neg_feats_all, 0)
                    self.train(model, criterion, update_optimizer, pos_data,
                               neg_data, opts['maxiter_update'])

                torch.cuda.empty_cache()

                bboxT = bbreg_bbox
                anyIOU = False

                for bbox_id, bboxD in enumerate(frameBBoxList[frameB_idx]):
                    if bboxD.match == False:
                        IOU_ratio = overlap_ratio(bboxD.getRec(), bboxT)
                        print("IOUratio={}".format(IOU_ratio))
                        #             print(IOU_ratio)
                        if IOU_ratio > 0.3:
                            #                print("overlap")
                            track_list[i] = (frameB_idx, bbox_id)
                            IOU_count = IOU_count + 1
                            bbox_list.append(
                                (frameB_idx, bboxD.getRec(), bboxD.confidence))
                            anyIOU = True
                            break
                if not anyIOU:
                    bbox_list.append((frameB_idx, bboxT, 0))

        print("bbox_list=====")
        print(bbox_list)

        print("track_list=====")
        print(track_list)

        ##debug   show track_list
        # for idx in range(0,10):
        #    print(track_list[idx])
        #    print(bbox_list[idx])

        if IOU_count >= self.IOU_count_th:  # add track
            for idx in range(0, self.TRACKLET_NUM + 1):
                bbox_id = track_list[idx][1]
                if bbox_id != -1:
                    # print(frame_idx+idx+1)
                    # print(bbox_id)
                    frameBBoxList[frame_idx + idx + 1][bbox_id].setMatch()

            next_track_frame = -1
            init_bbox_next = []
            bbox_length = len(bbox_list)
            #!!! track_list length=10, bbox_list length=11
            for ii in range(1, bbox_length):

                idx = bbox_length - ii
                print("{} {} {}".format(ii, idx, track_list[idx][0]))
                if track_list[idx - 1][0] != -1:
                    next_track_frame = (-1) * ii  # count from back
                    print("idx ={} ii={} next_track_frame={}".format(
                        idx, ii, next_track_frame))
                    init_bbox_next = bbox_list[idx]
                    break

            # remove overlap range
            start_rm = -1
            while start_rm >= previous_num:
                del self.GLOBAL_TRACK_LIST[track_num][-1]  # rm 1 per move
                # GLOBAL_TRACK_LIST[track_num].remove(-1)
                start_rm = start_rm - 1

            if previous_num == 0:
                self.GLOBAL_TRACK_LIST.append(bbox_list)
            else:
                self.GLOBAL_TRACK_LIST[track_num].extend(bbox_list)

            next_start_frame = frame_idx + bbox_length + next_track_frame

            init_bbox_next = bbox_list[bbox_length + next_track_frame][1]
            init_conf_next = bbox_list[bbox_length + next_track_frame][2]
            # print(GLOBAL_TRACK_LIST)
            print("( " + str(track_num) + " " + str(next_start_frame) + " " +
                  str(next_track_frame) + ")")
            print("init_bbox_next:" + str(init_bbox_next))
            self.search_track(
                track_num, next_start_frame, init_bbox_next, next_track_frame,
                init_conf_next
            )  # next_start_frame = 310  next_track_frame=-3(will be deleted)
Пример #5
0
def run_mdnet(img_list,
              init_bbox,
              gt=None,
              savefig_dir='',
              display=False,
              loss_index=1,
              model_path=opts['model_path'],
              seq_name=None):

    #def run_mdnet(k, img_list, init_bbox, gt=None, savefig_dir='', display=False,
    #              loss_index=1, model_path=opts['model_path'], seq_name=None):

    ############################
    if fewer_images:
        num_images = min(sequence_len_limit, len(img_list))
    else:
        num_images = len(img_list)
    ############################

    # Init bbox
    target_bbox = np.array(init_bbox)
    result = np.zeros((len(img_list), 4))
    result_bb = np.zeros((len(img_list), 4))
    result[0] = target_bbox
    result_bb[0] = target_bbox

    # Init iou and pred_iou
    iou_list = np.zeros((len(img_list), 1))  # shape: [113.1) ### list of ious
    iou_list[0] = 1.0  ### in first frame gt=result_bb (by definition)

    if gt is not None:
        overlap = np.zeros(len(img_list))
        overlap[0] = 1

    # Init model
    # model = MDNet(model_path=opts['model_path'],use_gpu=opts['use_gpu'])
    model = MDNet(model_path=model_path, use_gpu=opts['use_gpu'])
    if opts['use_gpu']:
        model = model.cuda()

    print('Init criterion and optimizer')
    criterion = BCELoss()
    model.set_learnable_params(opts['ft_layers'])
    init_optimizer = set_optimizer(model, opts['lr_init'], opts['lr_mult'])
    update_optimizer = set_optimizer(model, opts['lr_update'], opts['lr_mult'])

    tic = time.time()
    # Load first image
    image = Image.open(img_list[0]).convert('RGB')

    print('Draw pos/neg samples')
    # Draw pos/neg samples
    pos_examples = SampleGenerator('gaussian', image.size, opts['trans_pos'],
                                   opts['scale_pos'])(target_bbox,
                                                      opts['n_pos_init'],
                                                      opts['overlap_pos_init'])
    neg_examples = np.concatenate([
        SampleGenerator('uniform', image.size, opts['trans_neg_init'],
                        opts['scale_neg_init'])(target_bbox,
                                                int(opts['n_neg_init'] * 0.5),
                                                opts['overlap_neg_init']),
        SampleGenerator('whole', image.size)(target_bbox,
                                             int(opts['n_neg_init'] * 0.5),
                                             opts['overlap_neg_init'])
    ])
    neg_examples = np.random.permutation(neg_examples)

    print('Extract pos/neg features')
    # Extract pos/neg features
    if fewer_images:  # shorter run in general, less accurate
        pos_feats = forward_samples(model, image, pos_examples)
        neg_feats = forward_samples(model, image, neg_examples)
    else:
        pos_feats = forward_samples(model, image, pos_examples)
        neg_feats = forward_samples(model, image, neg_examples)

    print('Initial training')
    # Initial training
    train(model,
          criterion,
          init_optimizer,
          pos_feats,
          neg_feats,
          opts['maxiter_init'],
          loss_index=loss_index)  ### iou_pred_list
    del init_optimizer, neg_feats
    torch.cuda.empty_cache()

    print('Train bbox regressor')
    # Train bbox regressor
    bbreg_examples = SampleGenerator(
        'uniform', image.size, opts['trans_bbreg'], opts['scale_bbreg'],
        opts['aspect_bbreg'])(target_bbox, opts['n_bbreg'],
                              opts['overlap_bbreg'])
    bbreg_feats = forward_samples(
        model, image,
        bbreg_examples)  # calc features ### shape: [927, 4608] ###
    bbreg = BBRegressor(image.size)
    bbreg.train(bbreg_feats, bbreg_examples, target_bbox)
    del bbreg_feats
    torch.cuda.empty_cache()

    # Init sample generators for update
    sample_generator = SampleGenerator('gaussian', image.size, opts['trans'],
                                       opts['scale'])
    pos_generator = SampleGenerator('gaussian', image.size, opts['trans_pos'],
                                    opts['scale_pos'])
    neg_generator = SampleGenerator('uniform', image.size, opts['trans_neg'],
                                    opts['scale_neg'])

    print('Init pos/neg features for update')
    # Init pos/neg features for update
    neg_examples = neg_generator(target_bbox, opts['n_neg_update'],
                                 opts['overlap_neg_init'])
    neg_feats = forward_samples(model, image, neg_examples)
    pos_feats_all = [pos_feats]
    neg_feats_all = [neg_feats]

    spf_total = time.time() - tic

    # Display
    savefig = savefig_dir != ''
    if display or savefig:
        dpi = 80.0
        figsize = (image.size[0] / dpi, image.size[1] / dpi)

        fig = plt.figure(frameon=False, figsize=figsize, dpi=dpi)
        ax = plt.Axes(fig, [0., 0., 1., 1.])
        ax.set_axis_off()
        fig.add_axes(ax)
        im = ax.imshow(image, aspect='auto')

        if gt is not None:
            gt_rect = plt.Rectangle(tuple(gt[0, :2]),
                                    gt[0, 2],
                                    gt[0, 3],
                                    linewidth=3,
                                    edgecolor="#00ff00",
                                    zorder=1,
                                    fill=False)
            ax.add_patch(gt_rect)
            #################
            num_gts = np.minimum(gt.shape[0], num_images)
            # print('num_gts.shape: ', num_gts.shape)
            gt_centers = gt[:num_gts, :2] + gt[:num_gts, 2:] / 2
            result_centers = np.zeros_like(gt[:num_gts, :2])
            result_centers[0] = gt_centers[0]
            result_ious = np.zeros(num_gts, dtype='float64')
            result_ious[0] = 1.
            #################

        rect = plt.Rectangle(tuple(result_bb[0, :2]),
                             result_bb[0, 2],
                             result_bb[0, 3],
                             linewidth=3,
                             edgecolor="#ff0000",
                             zorder=1,
                             fill=False)
        ax.add_patch(rect)

        if display:
            plt.pause(.01)
            plt.draw()
        if savefig:
            fig.savefig(os.path.join(savefig_dir, '0000.jpg'), dpi=dpi)

    print('Main Loop')
    # Main loop
    spf_total = 0  # I don't want to take into account initialization
    for i in tqdm(range(1, num_images)):
        # for i in range(1, len(img_list)):
        #print('Frame: ', i)
        tic = time.time()
        # Load image
        image = Image.open(img_list[i]).convert('RGB')

        #print('Estimate target bbox (in run_mdnet)')
        # Estimate target bbox
        samples = sample_generator(target_bbox, opts['n_samples'])
        sample_scores = forward_samples(model, image, samples, out_layer='fc6')

        top_scores, top_idx = sample_scores[:, 1].topk(5)
        top_idx = top_idx.cpu()
        target_score = top_scores.mean()
        target_bbox = samples[top_idx]
        if top_idx.shape[0] > 1:
            target_bbox = target_bbox.mean(axis=0)
        success = target_score > 0

        # Expand search area at failure
        if success:
            sample_generator.set_trans(opts['trans'])
        else:
            sample_generator.expand_trans(opts['trans_limit'])

        #print('Bbox regression (in run_mdnet)')
        # Bbox regression
        if success:
            bbreg_samples = samples[top_idx]
            if top_idx.shape[0] == 1:
                bbreg_samples = bbreg_samples[None, :]
            bbreg_feats = forward_samples(model, image, bbreg_samples)
            bbreg_samples = bbreg.predict(bbreg_feats, bbreg_samples)
            bbreg_bbox = bbreg_samples.mean(axis=0)
        else:
            bbreg_bbox = target_bbox

        # Save result
        result[i] = target_bbox
        result_bb[i] = bbreg_bbox
        iou_list[i] = overlap_ratio(gt[i], result_bb[i])

        ###########################################
        # identify tracking failure and abort when in VOT mode

        IoU = overlap_ratio(result_bb[i], gt[i])[0]
        if (IoU == 0) and init_after_loss:
            print('    * lost track in frame %d since init*' % (i))
            result_distances = scipy.spatial.distance.cdist(
                result_centers[:i], gt_centers[:i],
                metric='euclidean').diagonal()
            num_images_tracked = i - 1  # we don't count frame 0 and current frame (lost track)

            im.set_data(image)
            if gt is not None:
                if i < gt.shape[0]:
                    gt_rect.set_xy(gt[i, :2])
                    gt_rect.set_width(gt[i, 2])
                    gt_rect.set_height(gt[i, 3])
                else:
                    gt_rect.set_xy(np.array([np.nan, np.nan]))
                    gt_rect.set_width(np.nan)
                    gt_rect.set_height(np.nan)

            rect.set_xy(result_bb[i, :2])
            rect.set_width(result_bb[i, 2])
            rect.set_height(result_bb[i, 3])

            plt.pause(.01)
            plt.draw()

            print(
                'Finished identify tracking failure and abort when in VOT mode'
            )
            return result[:
                          i], result_bb[:
                                        i], num_images_tracked, spf_total, result_distances, result_ious[:
                                                                                                         i], True
        ########################################

        # Data collect
        if success:
            pos_examples = pos_generator(target_bbox, opts['n_pos_update'],
                                         opts['overlap_pos_update'])
            pos_feats = forward_samples(model, image, pos_examples)
            pos_feats_all.append(pos_feats)
            if len(pos_feats_all) > opts['n_frames_long']:
                del pos_feats_all[0]

            neg_examples = neg_generator(target_bbox, opts['n_neg_update'],
                                         opts['overlap_neg_update'])
            neg_feats = forward_samples(model, image, neg_examples)
            neg_feats_all.append(neg_feats)
            if len(neg_feats_all) > opts['n_frames_short']:
                del neg_feats_all[0]

        # Short term update
        if not success:
            nframes = min(opts['n_frames_short'], len(pos_feats_all))
            pos_data = torch.cat(pos_feats_all[-nframes:], 0)
            neg_data = torch.cat(neg_feats_all, 0)
            train(model, criterion, update_optimizer, pos_data, neg_data,
                  opts['maxiter_update'])

        # Long term update
        elif i % opts['long_interval'] == 0:
            pos_data = torch.cat(pos_feats_all, 0)
            neg_data = torch.cat(neg_feats_all, 0)
            train(model, criterion, update_optimizer, pos_data, neg_data,
                  opts['maxiter_update'])

        torch.cuda.empty_cache()
        spf = time.time() - tic
        spf_total += spf
        #print('Time: ', spf)

        # Display
        if display or savefig:
            im.set_data(image)

            if gt is not None:
                gt_rect.set_xy(gt[i, :2])
                gt_rect.set_width(gt[i, 2])
                gt_rect.set_height(gt[i, 3])

                #################
                result_ious[i] = overlap_ratio(result_bb[i], gt[i])[0]
                result_centers[i] = result_bb[i, :2] + result_bb[i, 2:] / 2
                #################

            rect.set_xy(result_bb[i, :2])
            rect.set_width(result_bb[i, 2])
            rect.set_height(result_bb[i, 3])

            if display:
                plt.pause(.01)
                plt.draw()
            if savefig:
                fig.savefig(os.path.join(savefig_dir, '{:04d}.jpg'.format(i)),
                            dpi=dpi)

        ####################################
        if detailed_printing:
            if gt is None:
                print("      Frame %d/%d, Score %.3f, Time %.3f" % \
                      (i, num_images-1, target_score, spf))
            else:
                if i < gt.shape[0]:
                    print("      Frame %d/%d, Overlap %.3f, Score %.3f, Time %.3f" % \
                        (i, num_images-1, overlap_ratio(gt[i], result_bb[i])[0], target_score, spf))
                else:
                    print("      Frame %d/%d, Overlap %.3f, Score %.3f, Time %.3f" % \
                        (i, num_images-1, overlap_ratio(np.array([np.nan,np.nan,np.nan,np.nan]), result_bb[i])[0], target_score, spf))
        ####################################

        # if gt is None:
        #     print('Frame {:d}/{:d}, Score {:.3f}, Time {:.3f}'
        #         .format(i, len(img_list), target_score, spf))
        # else:
        #     overlap[i] = overlap_ratio(gt[i], result_bb[i])[0]
        #     print('Frame {:d}/{:d}, Overlap {:.3f}, Score {:.3f}, Time {:.3f}'
        #         .format(i, len(img_list), overlap[i], target_score, spf))

    ########################
    plt.close()
    result_distances = scipy.spatial.distance.cdist(
        result_centers, gt_centers, metric='euclidean').diagonal()
    num_images_tracked = num_images - 1  # I don't want to count initialization frame (i.e. frame 0)
    print('    main loop finished, %d frames' % (num_images))

    print('mean IoU: ', iou_list.mean())
    print('Finished run_mdnet()')

    return result, result_bb, num_images_tracked, spf_total, result_distances, result_ious, False
Пример #6
0
def run_vtaan(img_list, init_bbox, gt=None, savefig_dir='', display=False):

    # Init bbox
    target_bbox = np.array(init_bbox)
    result = np.zeros((len(img_list), 4))
    result_bb = np.zeros((len(img_list), 4))
    result[0] = target_bbox
    result_bb[0] = target_bbox

    if gt is not None:
        overlap = np.zeros(len(img_list))
        overlap[0] = 1

    # Init model
    model = MDNet(opts['model_path'])
    model_g = NetG()
    if opts['use_gpu']:
        model = model.cuda()
        model_g = model_g.cuda()
    GBP = guided_backprop.GuidedBackprop(model, 1)

    # Init criterion and optimizer
    criterion = BCELoss()
    criterion_g = torch.nn.MSELoss(reduction='sum')
    model.set_learnable_params(opts['ft_layers'])
    model_g.set_learnable_params(opts['ft_layers'])
    init_optimizer = set_optimizer(model, opts['lr_init'], opts['lr_mult'])
    update_optimizer = set_optimizer(model, opts['lr_update'], opts['lr_mult'])

    tic = time.time()
    # Load first image
    image = Image.open(img_list[0]).convert('RGB')

    # Draw pos/neg samples
    pos_examples = SampleGenerator('gaussian', image.size, opts['trans_pos'],
                                   opts['scale_pos'])(target_bbox,
                                                      opts['n_pos_init'],
                                                      opts['overlap_pos_init'])

    neg_examples = np.concatenate([
        SampleGenerator('uniform', image.size, opts['trans_neg_init'],
                        opts['scale_neg_init'])(target_bbox,
                                                int(opts['n_neg_init'] * 0.5),
                                                opts['overlap_neg_init']),
        SampleGenerator('whole', image.size)(target_bbox,
                                             int(opts['n_neg_init'] * 0.5),
                                             opts['overlap_neg_init'])
    ])
    neg_examples = np.random.permutation(neg_examples)

    # Extract pos/neg features
    pos_feats = forward_samples(model, image, pos_examples)
    neg_feats = forward_samples(model, image, neg_examples)

    pos_imgids = np.array([[0]] * pos_feats.size(0))
    neg_imgids = np.array([[0]] * neg_feats.size(0))

    feat_dim = pos_feats.size(-1)

    # Initial training
    train(model, None, criterion, init_optimizer, pos_feats, neg_feats,
          opts['maxiter_init'], pos_imgids, pos_examples, neg_imgids,
          neg_examples, img_list, GBP)
    del init_optimizer, neg_feats
    torch.cuda.empty_cache()
    g_pretrain(model, model_g, criterion_g, pos_feats)
    torch.cuda.empty_cache()

    # Train bbox regressor
    bbreg_examples = SampleGenerator(
        'uniform', image.size, opts['trans_bbreg'], opts['scale_bbreg'],
        opts['aspect_bbreg'])(target_bbox, opts['n_bbreg'],
                              opts['overlap_bbreg'])
    bbreg_feats = forward_samples(model, image, bbreg_examples)
    bbreg = BBRegressor(image.size)
    bbreg.train(bbreg_feats, bbreg_examples, target_bbox)
    del bbreg_feats
    torch.cuda.empty_cache()

    # Init sample generators for update
    sample_generator = SampleGenerator('gaussian', image.size, opts['trans'],
                                       opts['scale'])
    pos_generator = SampleGenerator('gaussian', image.size, opts['trans_pos'],
                                    opts['scale_pos'])
    neg_generator = SampleGenerator('uniform', image.size, opts['trans_neg'],
                                    opts['scale_neg'])

    # Init pos/neg features for update
    neg_examples = neg_generator(target_bbox, opts['n_neg_update'],
                                 opts['overlap_neg_init'])
    neg_feats = forward_samples(model, image, neg_examples)
    pos_feats_all = [pos_feats[:opts['n_pos_update']]]
    neg_feats_all = [neg_feats[:opts['n_neg_update']]]

    pos_examples_all = [pos_examples[:opts['n_pos_update']]]
    neg_examples_all = [neg_examples[:opts['n_neg_update']]]

    pos_imgids_all = [pos_imgids[:opts['n_pos_update']]]
    neg_imgids_all = [neg_imgids[:opts['n_neg_update']]]

    spf_total = time.time() - tic

    # Display
    savefig = savefig_dir != ''
    if display or savefig:
        dpi = 80.0
        figsize = (image.size[0] / dpi, image.size[1] / dpi)

        fig = plt.figure(frameon=False, figsize=figsize, dpi=dpi)
        ax = plt.Axes(fig, [0., 0., 1., 1.])
        ax.set_axis_off()
        fig.add_axes(ax)
        im = ax.imshow(image, aspect='auto')

        if gt is not None:
            gt_rect = plt.Rectangle(tuple(gt[0, :2]),
                                    gt[0, 2],
                                    gt[0, 3],
                                    linewidth=3,
                                    edgecolor="#00ff00",
                                    zorder=1,
                                    fill=False)
            ax.add_patch(gt_rect)

        rect = plt.Rectangle(tuple(result_bb[0, :2]),
                             result_bb[0, 2],
                             result_bb[0, 3],
                             linewidth=3,
                             edgecolor="#ff0000",
                             zorder=1,
                             fill=False)
        ax.add_patch(rect)

        if display:
            plt.pause(.01)
            plt.draw()
        if savefig:
            fig.savefig(os.path.join(savefig_dir, '0000.jpg'), dpi=dpi)

    # Main loop
    for i in range(1, len(img_list)):

        tic = time.time()
        # Load image
        image = Image.open(img_list[i]).convert('RGB')

        # Estimate target bbox
        samples = sample_generator(target_bbox, opts['n_samples'])
        sample_scores = forward_samples(model, image, samples, out_layer='fc6')

        top_scores, top_idx = sample_scores[:, 1].topk(5)
        top_idx = top_idx.cpu()
        target_score = top_scores.mean()
        target_bbox = samples[top_idx]
        if top_idx.shape[0] > 1:
            target_bbox = target_bbox.mean(axis=0)
        success = target_score > 0

        # Expand search area at failure
        if success:
            sample_generator.set_trans(opts['trans'])
        else:
            sample_generator.expand_trans(opts['trans_limit'])

        # Bbox regression
        if success:
            bbreg_samples = samples[top_idx]
            if top_idx.shape[0] == 1:
                bbreg_samples = bbreg_samples[None, :]
            bbreg_feats = forward_samples(model, image, bbreg_samples)
            bbreg_samples = bbreg.predict(bbreg_feats, bbreg_samples)
            bbreg_bbox = bbreg_samples.mean(axis=0)
        else:
            bbreg_bbox = target_bbox

        # Save result
        result[i] = target_bbox
        result_bb[i] = bbreg_bbox

        # Data collect
        if success:
            pos_examples = pos_generator(target_bbox, opts['n_pos_update'],
                                         opts['overlap_pos_update'])
            pos_feats = forward_samples(model, image, pos_examples)
            pos_feats_all.append(pos_feats)
            if len(pos_feats_all) > opts['n_frames_long']:
                del pos_feats_all[0]
                del pos_examples_all[0]
                del pos_imgids_all[0]

            neg_examples = neg_generator(target_bbox, opts['n_neg_update'],
                                         opts['overlap_neg_update'])
            neg_feats = forward_samples(model, image, neg_examples)
            neg_feats_all.append(neg_feats)

            pos_examples_all.append(pos_examples)
            neg_examples_all.append(neg_examples)

            pos_imgids_all.append(np.array([[i]] * pos_feats.size(0)))
            neg_imgids_all.append(np.array([[i]] * neg_feats.size(0)))

            if len(neg_feats_all) > opts['n_frames_short']:
                del neg_feats_all[0]
                del neg_examples_all[0]
                del neg_imgids_all[0]

        # Short term update
        if not success:
            nframes = min(opts['n_frames_short'], len(pos_feats_all))
            pos_data = torch.stack(pos_feats_all[-nframes:],
                                   0).view(-1, feat_dim)
            neg_data = torch.stack(neg_feats_all, 0).view(-1, feat_dim)

            pos_examples_data = torch.from_numpy(
                np.stack(pos_examples_all[-nframes:], 0)).view(-1, 4).numpy()
            neg_examples_data = torch.from_numpy(np.stack(neg_examples_all,
                                                          0)).view(-1,
                                                                   4).numpy()

            pos_imgids_data = torch.from_numpy(
                np.stack(pos_imgids_all[-nframes:], 0)).view(-1, 1).numpy()
            neg_imgids_data = torch.from_numpy(np.stack(neg_imgids_all,
                                                        0)).view(-1,
                                                                 1).numpy()

            train(model, None, criterion, update_optimizer, pos_data, neg_data,
                  opts['maxiter_update'], pos_imgids_data, pos_examples_data,
                  neg_imgids_data, neg_examples_data, img_list, GBP)

        # Long term update
        elif i % opts['long_interval'] == 0:
            pos_data = t.stack(pos_feats_all, 0).view(-1, feat_dim)
            neg_data = t.stack(neg_feats_all, 0).view(-1, feat_dim)

            pos_examples_data = torch.from_numpy(np.stack(pos_examples_all,
                                                          0)).view(-1,
                                                                   4).numpy()
            neg_examples_data = torch.from_numpy(np.stack(neg_examples_all,
                                                          0)).view(-1,
                                                                   4).numpy()

            pos_imgids_data = torch.from_numpy(np.stack(pos_imgids_all,
                                                        0)).view(-1,
                                                                 1).numpy()
            neg_imgids_data = torch.from_numpy(np.stack(neg_imgids_all,
                                                        0)).view(-1,
                                                                 1).numpy()

            # train(model, model_g, criterion, update_optimizer, pos_data, neg_data, opts['maxiter_update'],
            #       pos_imgids_data, pos_examples_data, neg_imgids_data, neg_examples_data, img_list, GBP)
            train(model, model_g, criterion, update_optimizer, pos_data,
                  neg_data, opts['maxiter_update'], None, None, None, None,
                  img_list, GBP)

        torch.cuda.empty_cache()
        spf = time.time() - tic
        spf_total += spf

        # Display
        if display or savefig:
            im.set_data(image)

            if gt is not None:
                gt_rect.set_xy(gt[i, :2])
                gt_rect.set_width(gt[i, 2])
                gt_rect.set_height(gt[i, 3])

            rect.set_xy(result_bb[i, :2])
            rect.set_width(result_bb[i, 2])
            rect.set_height(result_bb[i, 3])

            if display:
                plt.pause(.01)
                plt.draw()
            if savefig:
                fig.savefig(os.path.join(savefig_dir, '{:04d}.jpg'.format(i)),
                            dpi=dpi)

        if gt is None:
            print('Frame {:d}/{:d}, Score {:.3f}, Time {:.3f}'.format(
                i + 1, len(img_list), target_score, spf))
        else:
            overlap[i] = overlap_ratio(gt[i], result_bb[i])[0]
            print('Frame {:d}/{:d}, Overlap {:.3f}, Score {:.3f}, Time {:.3f}'.
                  format(i + 1, len(img_list), overlap[i], target_score, spf))

    if gt is not None:
        print('meanIOU: {:.3f}'.format(overlap.mean()))
    fps = len(img_list) / spf_total
    return result, result_bb, fps