Function main

tutorials/ctracker/test_byte.py:108–149

(args=None)

    

def main(args=None):
    parser = argparse.ArgumentParser(description='Simple script for testing a CTracker network.')
    parser.add_argument('--dataset_path', default='/dockerdata/home/jeromepeng/data/MOT/MOT17/', type=str,
                        help='Dataset path, location of the images sequence.')
    parser.add_argument('--model_dir', default='./trained_model/', help='Path to model (.pt) file.')
    parser.add_argument('--model_path', default='./trained_model/model_final.pth', help='Path to model (.pt) file.')
    parser.add_argument('--seq_nums', default=0, type=int)

    parser = parser.parse_args(args)

    if not os.path.exists(os.path.join(parser.model_dir, 'results')):
        os.makedirs(os.path.join(parser.model_dir, 'results'))

    retinanet = model.resnet50(num_classes=1, pretrained=True)
    # 	retinanet_save = torch.load(os.path.join(parser.model_dir, 'model_final.pth'))
    retinanet_save = torch.load(os.path.join(parser.model_path))

    # rename moco pre-trained keys
    state_dict = retinanet_save.state_dict()
    for k in list(state_dict.keys()):
        # retain only encoder up to before the embedding layer
        if k.startswith('module.'):
            # remove prefix
            state_dict[k[len("module."):]] = state_dict[k]
        # delete renamed or unused k
        del state_dict[k]

    retinanet.load_state_dict(state_dict)

    use_gpu = True

    if use_gpu: retinanet = retinanet.cuda()

    retinanet.eval()
    seq_nums = []
    if parser.seq_nums > 0:
        seq_nums.append(parser.seq_nums)
    else:
        seq_nums = [2, 4, 5, 9, 10, 11, 13]

    for seq_num in seq_nums:
        run_each_dataset(parser.model_dir, retinanet, parser.dataset_path, 'train', 'MOT17-{:02d}'.format(seq_num))


# 	for seq_num in [1, 3, 6, 7, 8, 12, 14]: