def _getSuperPointKps(image, confidence_threshold):
    nparr = np.fromstring(image, np.uint8)
    image = cv2.imdecode(nparr, cv2.IMREAD_GRAYSCALE)
    parser = argparse.ArgumentParser()
    parser.add_argument('--config',
                        default="../models/superpoint_config.yaml",
                        type=str)
    parser.add_argument('--experiment_name',
                        default="superpoint_coco",
                        type=str)

    args = parser.parse_args()

    experiment_name = args.experiment_name

    with open(args.config, 'r') as f:
        config = yaml.load(f)
    assert 'eval_iter' in config

    checkpoint = Path(EXPER_PATH, experiment_name)
    image = np.expand_dims(image, 2)
    with experiment._init_graph(config, with_dataset=False) as (net):
        net.load(str(checkpoint))
        prob = net.predict({'image': image}, keys='prob_nms')
        pts = select_top_k(prob, thresh=confidence_threshold)
        return pts
Beispiel #2
0
    assert 'eval_iter' in config

    if args.config.find('gl3d') >= 0:
        output_dir = Path(DATA_PATH, 'gl3d', 'data')
    else:
        output_dir = Path(EXPER_PATH, 'outputs/{}/'.format(export_name))
        if not output_dir.exists():
            os.makedirs(output_dir)
    checkpoint = Path(EXPER_PATH, experiment_name)
    if 'checkpoint' in config:
        checkpoint = Path(checkpoint, config['checkpoint'])

    config['model']['pred_batch_size'] = batch_size
    batch_size *= experiment.get_num_gpus()

    with experiment._init_graph(config, with_dataset=True) as (net, dataset):
        if net.trainable:
            net.load(str(checkpoint))
        test_set = dataset.get_test_set()

        for _ in tqdm(range(config.get('skip', 0))):
            next(test_set)

        pbar = tqdm(
            total=config['eval_iter'] if config['eval_iter'] > 0 else None)
        i = 0
        while True:
            # Gather dataset
            data = []
            try:
                for _ in range(batch_size):