def __init__(self):
self.model = 'models/r2d2_WASF_N16.pt'
self.tag = 'r2d2' # output file tag
self.top_k = 2000 # number of keypoints
self.scale_f = 2**0.25
self.min_size = 256
self.max_size = 1024
self.min_scale = 0
self.max_scale = 1
self.reliability_thr = 0.7
self.repeatability_thr = 0.7
self.gpu = [0]
self.iscuda = common.torch_set_gpu(self.gpu)
self.net = None
# load_network
checkpoint = torch.load(self.model)
print("\n>> Creating net = " + checkpoint['net'])
net = eval(checkpoint['net'])
nb_of_weights = common.model_size(net)
#print(f"( Model size: {nb_of_weights/1000:.0f}K parameters )")
# initialization
weights = checkpoint['state_dict']
net.load_state_dict(
{k.replace('module.', ''): v
for k, v in weights.items()})
self.net = net.eval()
def extract_keypoints(args):
iscuda = common.torch_set_gpu(args.gpu)
# load the network...
net = load_network(args.model)
if iscuda: net = net.cuda()
# create the non-maxima detector
detector = NonMaxSuppression(
rel_thr = args.reliability_thr,
rep_thr = args.repeatability_thr)
for seq_name in sorted(os.listdir(args.in_dir)):
seq_path = os.path.join(args.in_dir, seq_name)
for img_name in sorted(os.listdir(seq_path)):
if img_name[-4:] != '.ppm':
continue
img_path = os.path.join(seq_path, img_name)
output_path = os.path.join(seq_path, img_name + '.r2d2')
print(f"Extracting features for {img_path}")
img = Image.open(img_path).convert('RGB')
W, H = img.size
img = norm_RGB(img)[None]
if iscuda: img = img.cuda()
# extract keypoints/descriptors for a single image
xys, desc, scores = extract_multiscale(net, img, detector,
scale_f = args.scale_f,
min_scale = args.min_scale,
max_scale = args.max_scale,
min_size = args.min_size,
max_size = args.max_size,
verbose = False)
xys = xys.cpu().numpy()
desc = desc.cpu().numpy()
scores = scores.cpu().numpy()
idxs = scores.argsort()[-args.top_k or None:]
#outpath = img_path + '.' + args.tag
#print(f"Saving {len(idxs)} keypoints to {outpath}")
#np.savez(open(outpath,'wb'),
# imsize = (W,H),
# keypoints = xys[idxs],
# descriptors = desc[idxs],
# scores = scores[idxs])
keypoints = xys[idxs][:,:2]
descriptors = desc[idxs]
scores = scores[idxs]
#print(keypoints.shape)
with open(output_path, 'wb') as output_file:
np.savez(output_file, keypoints=keypoints, scores=scores,
descriptors=descriptors)
def extract_keypoints(args):
iscuda = common.torch_set_gpu(args.gpu)
# load the network...
net = load_network(args.model)
if iscuda: net = net.cuda()
# create the non-maxima detector
detector = NonMaxSuppression(
rel_thr = args.reliability_thr,
rep_thr = args.repeatability_thr)
# while args.images:
# img_path = args.images.pop(0)
# if img_path.endswith('.txt'):
# args.images = open(img_path).read().splitlines() + args.images
# continue
# print(f"\nExtracting features for {img_path}")
# img = Image.open(img_path).convert('RGB')
if True:
img_path='kitti06-12-color.png'
img = cv2.imread('../data/kitti06-12-color.png')
#W, H = img.size
H, W = img.shape[:2]
img = norm_RGB(img)[None]
if iscuda: img = img.cuda()
# extract keypoints/descriptors for a single image
xys, desc, scores = extract_multiscale(net, img, detector,
scale_f = args.scale_f,
min_scale = args.min_scale,
max_scale = args.max_scale,
min_size = args.min_size,
max_size = args.max_size,
verbose = True)
xys = xys.cpu().numpy()
desc = desc.cpu().numpy()
scores = scores.cpu().numpy()
idxs = scores.argsort()[-args.top_k or None:]
outpath = img_path + '.' + args.tag
print(f"Saving {len(idxs)} keypoints to {outpath}")
np.savez(open(outpath,'wb'),
imsize = (W,H),
keypoints = xys[idxs],
descriptors = desc[idxs],
scores = scores[idxs])
def extract_keypoints(args):
iscuda = common.torch_set_gpu(args.gpu)
# load net
net = load_network(args.model)
if iscuda:
net = net.cuda()
detector = NonMaxSuppression(rel_thr=args.reliability_thr,
rep_thr=args.repeatability_thr)
# 将输入的图片一张张的提取
while args.images:
img_path = args.images.pop(0)
# 这个就是输入的是图片名称数组
if img_path.endswith('.txt'):
args.images = open(img_path).read().splitlines() + args.images
continue
print(f"\nExtracting features for {img_path}")
img = Image.open(img_path).convert('RGB')
W, H = img.size
img = norm_RGB(img)[None]
if iscuda: img = img.cuda()
xys, desc, scores = extract_multiscale(net,
img,
detector,
scale_f=args.scale_f,
min_scale=args.min_scale,
max_scale=args.max_scale,
min_size=args.min_size,
max_size=args.max_size,
verbose=True)
xys = xys.cpu().numpy()
desc = desc.cpu().numpy()
scores = scores.cpu().numpy()
# 这里返回前n名的位置
idxs = scores.argsort()[-args.top_k or None:]
outpath = img_path + '.' + args.tag
print(f"Saving {len(idxs)} keypoints to {outpath}")
np.savez(open(outpath, 'wb'),
imsize=(W, H),
keypoints=xys[idxs],
descriptors=desc[idxs],
scores=scores[idxs])
def extract_keypoints(args, root_to_current_dataset, scale_f):
iscuda = common.torch_set_gpu(args.gpu)
output_folder = 'r2d2_features_' + str(scale_f)
rgb_folder = 'rgb.txt'
root_to_save_features = root_to_current_dataset + '/' + output_folder
if not os.path.exists(root_to_save_features):
os.makedirs(root_to_save_features)
# load the network...
net = load_network(args.model)
if iscuda: net = net.cuda()
# create the non-maxima detector
detector = NonMaxSuppression(rel_thr=args.reliability_thr,
rep_thr=args.repeatability_thr)
imgs_under_current_folder = np.loadtxt(Path(root_to_current_dataset,
rgb_folder),
dtype=str)
for img in imgs_under_current_folder:
img_path = img[-1]
img_name = img[-1].split('/')[-1]
# print(f"\nExtracting features for {img_path}")
img = Image.open(Path(root_to_current_dataset,
img_path)).convert('RGB')
img = norm_RGB(img)[None]
if iscuda: img = img.cuda()
# extract keypoints/descriptors for a single image
desc = extract_and_save_byscale(net, img, detector, scale_f)
desc = desc.cpu().numpy().squeeze()
outpath = Path(root_to_save_features, img_name[:-4])
print(f"Saving features to {outpath}")
np.save(outpath, desc)
parser.add_argument("--loss", type=str, default=default_loss, help="loss function")
parser.add_argument("--sampler", type=str, default=default_sampler, help="AP sampler")
parser.add_argument("--N", type=int, default=16, help="patch size for repeatability")
parser.add_argument("--epochs", type=int, default=25, help='number of training epochs')
parser.add_argument("--batch-size", "--bs", type=int, default=8, help="batch size")
parser.add_argument("--learning-rate", "--lr", type=str, default=1e-4)
parser.add_argument("--weight-decay", "--wd", type=float, default=5e-4)
parser.add_argument("--threads", type=int, default=8, help='number of worker threads')
parser.add_argument("--gpu", type=int, nargs='+', default=[0], help='-1 for CPU')
args = parser.parse_args()
iscuda = common.torch_set_gpu(args.gpu)
common.mkdir_for(args.save_path)
# Create data loader
from datasets import *
db = [data_sources[key] for key in args.train_data]
db = eval(args.data_loader.replace('`data`',','.join(db)).replace('\n',''))
print("Training image database =", db)
loader = threaded_loader(db, iscuda, args.threads, args.batch_size, shuffle=True)
# create network
print("\n>> Creating net = " + args.net)
net = eval(args.net)
print(f" ( Model size: {common.model_size(net)/1000:.0f}K parameters )")
# initialization
def extract_kapture_keypoints(kapture_root,
config,
output_dir='',
overwrite=False):
"""
Extract r2d2 keypoints and descritors to the kapture format directly
"""
print('extract_kapture_keypoints...')
kdata = kapture_from_dir(kapture_root, matches_pairsfile_path=None,
skip_list= [kapture.GlobalFeatures,
kapture.Matches,
kapture.Points3d,
kapture.Observations])
export_dir = output_dir if output_dir else kapture_root # root of output directory for features
os.makedirs(export_dir, exist_ok=True)
assert kdata.records_camera is not None
image_list = [filename for _, _, filename in kapture.flatten(kdata.records_camera)]
# resume extraction if some features exist
try:
# load existing features, if any
kdata.keypoints = keypoints_from_dir(export_dir, None)
kdata.descriptors = descriptors_from_dir(export_dir, None)
if kdata.keypoints is not None and kdata.descriptors is not None and not overwrite:
image_list = [name for name in image_list if name not in kdata.keypoints or name not in kdata.descriptors]
except FileNotFoundError:
pass
except:
logging.exception("Error with importing existing local features.")
# clear features first if overwriting
if overwrite: delete_existing_kapture_files(export_dir, True, only=[kapture.Descriptors, kapture.Keypoints])
if len(image_list) == 0:
print('All features were already extracted')
return
else:
print(f'Extracting r2d2 features for {len(image_list)} images')
iscuda = common.torch_set_gpu([torch.cuda.is_available()])
# load the network...
net = load_network(config['checkpoint'])
if iscuda: net = net.cuda()
# create the non-maxima detector
detector = NonMaxSuppression(
rel_thr = config['reliability_thr'],
rep_thr = config['repeatability_thr'])
keypoints_dtype = None if kdata.keypoints is None else kdata.keypoints.dtype
descriptors_dtype = None if kdata.descriptors is None else kdata.descriptors.dtype
keypoints_dsize = None if kdata.keypoints is None else kdata.keypoints.dsize
descriptors_dsize = None if kdata.descriptors is None else kdata.descriptors.dsize
for image_name in image_list:
img_path = get_image_fullpath(kapture_root, image_name)
if img_path.endswith('.txt'):
images = open(img_path).read().splitlines() + images
continue
print(f"\nExtracting features for {img_path}")
img = Image.open(img_path).convert('RGB')
W, H = img.size
img = norm_RGB(img)[None]
if iscuda: img = img.cuda()
# extract keypoints/descriptors for a single image
xys, desc, scores = extract_multiscale(net, img, detector,
scale_f = config['scale_f'],
min_scale = config['min_scale'],
max_scale = config['max_scale'],
min_size = config['min_size'],
max_size = config['max_size'],
verbose = True)
xys = xys.cpu().numpy()
desc = desc.cpu().numpy()
scores = scores.cpu().numpy()
idxs = scores.argsort()[-config['top_k'] or None:]
xys = xys[idxs]
desc = desc[idxs]
if keypoints_dtype is None or descriptors_dtype is None:
keypoints_dtype = xys.dtype
descriptors_dtype = desc.dtype
keypoints_dsize = xys.shape[1]
descriptors_dsize = desc.shape[1]
kdata.keypoints = kapture.Keypoints('r2d2', keypoints_dtype, keypoints_dsize)
kdata.descriptors = kapture.Descriptors('r2d2', descriptors_dtype, descriptors_dsize)
keypoints_config_absolute_path = get_csv_fullpath(kapture.Keypoints, export_dir)
descriptors_config_absolute_path = get_csv_fullpath(kapture.Descriptors, export_dir)
keypoints_to_file(keypoints_config_absolute_path, kdata.keypoints)
descriptors_to_file(descriptors_config_absolute_path, kdata.descriptors)
else:
assert kdata.keypoints.type_name == 'r2d2'
assert kdata.descriptors.type_name == 'r2d2'
assert kdata.keypoints.dtype == xys.dtype
assert kdata.descriptors.dtype == desc.dtype
assert kdata.keypoints.dsize == xys.shape[1]
assert kdata.descriptors.dsize == desc.shape[1]
keypoints_fullpath = get_keypoints_fullpath(export_dir, image_name)
print(f"Saving {xys.shape[0]} keypoints to {keypoints_fullpath}")
image_keypoints_to_file(keypoints_fullpath, xys)
kdata.keypoints.add(image_name)
descriptors_fullpath = get_descriptors_fullpath(export_dir, image_name)
print(f"Saving {desc.shape[0]} descriptors to {descriptors_fullpath}")
image_descriptors_to_file(descriptors_fullpath, desc)
kdata.descriptors.add(image_name)
if not keypoints_check_dir(kdata.keypoints, export_dir) or \
not descriptors_check_dir(kdata.descriptors, export_dir):
print('local feature extraction ended successfully but not all files were saved')
def extract_kapture_keypoints(args):
"""
Extract r2d2 keypoints and descritors to the kapture format directly
"""
print('extract_kapture_keypoints...')
kdata = kapture_from_dir(args.kapture_root,
matches_pairs_file_path=None,
skip_list=[
kapture.GlobalFeatures, kapture.Matches,
kapture.Points3d, kapture.Observations
])
assert kdata.records_camera is not None
image_list = [
filename for _, _, filename in kapture.flatten(kdata.records_camera)
]
if kdata.keypoints is not None and kdata.descriptors is not None:
image_list = [
name for name in image_list
if name not in kdata.keypoints or name not in kdata.descriptors
]
if len(image_list) == 0:
print('All features were already extracted')
return
else:
print(f'Extracting r2d2 features for {len(image_list)} images')
iscuda = common.torch_set_gpu(args.gpu)
# load the network...
net = load_network(args.model)
if iscuda: net = net.cuda()
# create the non-maxima detector
detector = NonMaxSuppression(rel_thr=args.reliability_thr,
rep_thr=args.repeatability_thr)
keypoints_dtype = None if kdata.keypoints is None else kdata.keypoints.dtype
descriptors_dtype = None if kdata.descriptors is None else kdata.descriptors.dtype
keypoints_dsize = None if kdata.keypoints is None else kdata.keypoints.dsize
descriptors_dsize = None if kdata.descriptors is None else kdata.descriptors.dsize
for image_name in image_list:
img_path = get_image_fullpath(args.kapture_root, image_name)
print(f"\nExtracting features for {img_path}")
img = Image.open(img_path).convert('RGB')
W, H = img.size
img = norm_RGB(img)[None]
if iscuda: img = img.cuda()
# extract keypoints/descriptors for a single image
xys, desc, scores = extract_multiscale(net,
img,
detector,
scale_f=args.scale_f,
min_scale=args.min_scale,
max_scale=args.max_scale,
min_size=args.min_size,
max_size=args.max_size,
verbose=True)
xys = xys.cpu().numpy()
desc = desc.cpu().numpy()
scores = scores.cpu().numpy()
idxs = scores.argsort()[-args.top_k or None:]
xys = xys[idxs]
desc = desc[idxs]
if keypoints_dtype is None or descriptors_dtype is None:
keypoints_dtype = xys.dtype
descriptors_dtype = desc.dtype
keypoints_dsize = xys.shape[1]
descriptors_dsize = desc.shape[1]
kdata.keypoints = kapture.Keypoints('r2d2', keypoints_dtype,
keypoints_dsize)
kdata.descriptors = kapture.Descriptors('r2d2', descriptors_dtype,
descriptors_dsize)
keypoints_config_absolute_path = get_csv_fullpath(
kapture.Keypoints, args.kapture_root)
descriptors_config_absolute_path = get_csv_fullpath(
kapture.Descriptors, args.kapture_root)
keypoints_to_file(keypoints_config_absolute_path, kdata.keypoints)
descriptors_to_file(descriptors_config_absolute_path,
kdata.descriptors)
else:
assert kdata.keypoints.type_name == 'r2d2'
assert kdata.descriptors.type_name == 'r2d2'
assert kdata.keypoints.dtype == xys.dtype
assert kdata.descriptors.dtype == desc.dtype
assert kdata.keypoints.dsize == xys.shape[1]
assert kdata.descriptors.dsize == desc.shape[1]
keypoints_fullpath = get_keypoints_fullpath(args.kapture_root,
image_name)
print(f"Saving {xys.shape[0]} keypoints to {keypoints_fullpath}")
image_keypoints_to_file(keypoints_fullpath, xys)
kdata.keypoints.add(image_name)
descriptors_fullpath = get_descriptors_fullpath(
args.kapture_root, image_name)
print(f"Saving {desc.shape[0]} descriptors to {descriptors_fullpath}")
image_descriptors_to_file(descriptors_fullpath, desc)
kdata.descriptors.add(image_name)
if not keypoints_check_dir(kdata.keypoints, args.kapture_root) or \
not descriptors_check_dir(kdata.descriptors, args.kapture_root):
print(
'local feature extraction ended successfully but not all files were saved'
)