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
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):