Spaces:
Running
Running
| import argparse | |
| import numpy as np | |
| from PIL import Image | |
| import torch | |
| import math | |
| from tqdm import tqdm | |
| from os import path | |
| # Kapture is a pivot file format, based on text and binary files, used to describe SfM (Structure From Motion) and more generally sensor-acquired data | |
| # it can be installed with | |
| # pip install kapture | |
| # for more information check out https://github.com/naver/kapture | |
| import kapture | |
| from kapture.io.records import get_image_fullpath | |
| from kapture.io.csv import kapture_from_dir, get_all_tar_handlers | |
| from kapture.io.csv import get_feature_csv_fullpath, keypoints_to_file, descriptors_to_file | |
| from kapture.io.features import get_keypoints_fullpath, keypoints_check_dir, image_keypoints_to_file | |
| from kapture.io.features import get_descriptors_fullpath, descriptors_check_dir, image_descriptors_to_file | |
| from lib.model_test import D2Net | |
| from lib.utils import preprocess_image | |
| from lib.pyramid import process_multiscale | |
| # import imageio | |
| # CUDA | |
| use_cuda = torch.cuda.is_available() | |
| device = torch.device("cuda:0" if use_cuda else "cpu") | |
| # Argument parsing | |
| parser = argparse.ArgumentParser(description='Feature extraction script') | |
| parser.add_argument( | |
| '--kapture-root', type=str, required=True, | |
| help='path to kapture root directory' | |
| ) | |
| parser.add_argument( | |
| '--preprocessing', type=str, default='caffe', | |
| help='image preprocessing (caffe or torch)' | |
| ) | |
| parser.add_argument( | |
| '--model_file', type=str, default='models/d2_tf.pth', | |
| help='path to the full model' | |
| ) | |
| parser.add_argument( | |
| '--keypoints-type', type=str, default=None, | |
| help='keypoint type_name, default is filename of model' | |
| ) | |
| parser.add_argument( | |
| '--descriptors-type', type=str, default=None, | |
| help='descriptors type_name, default is filename of model' | |
| ) | |
| parser.add_argument( | |
| '--max_edge', type=int, default=1600, | |
| help='maximum image size at network input' | |
| ) | |
| parser.add_argument( | |
| '--max_sum_edges', type=int, default=2800, | |
| help='maximum sum of image sizes at network input' | |
| ) | |
| parser.add_argument( | |
| '--multiscale', dest='multiscale', action='store_true', | |
| help='extract multiscale features' | |
| ) | |
| parser.set_defaults(multiscale=False) | |
| parser.add_argument( | |
| '--no-relu', dest='use_relu', action='store_false', | |
| help='remove ReLU after the dense feature extraction module' | |
| ) | |
| parser.set_defaults(use_relu=True) | |
| parser.add_argument("--max-keypoints", type=int, default=float("+inf"), | |
| help='max number of keypoints save to disk') | |
| args = parser.parse_args() | |
| print(args) | |
| with get_all_tar_handlers(args.kapture_root, | |
| mode={kapture.Keypoints: 'a', | |
| kapture.Descriptors: 'a', | |
| kapture.GlobalFeatures: 'r', | |
| kapture.Matches: 'r'}) as tar_handlers: | |
| kdata = kapture_from_dir(args.kapture_root, | |
| skip_list=[kapture.GlobalFeatures, | |
| kapture.Matches, | |
| kapture.Points3d, | |
| kapture.Observations], | |
| tar_handlers=tar_handlers) | |
| if kdata.keypoints is None: | |
| kdata.keypoints = {} | |
| if kdata.descriptors is None: | |
| kdata.descriptors = {} | |
| assert kdata.records_camera is not None | |
| image_list = [filename for _, _, filename in kapture.flatten(kdata.records_camera)] | |
| if args.keypoints_type is None: | |
| args.keypoints_type = path.splitext(path.basename(args.model_file))[0] | |
| print(f'keypoints_type set to {args.keypoints_type}') | |
| if args.descriptors_type is None: | |
| args.descriptors_type = path.splitext(path.basename(args.model_file))[0] | |
| print(f'descriptors_type set to {args.descriptors_type}') | |
| if args.keypoints_type in kdata.keypoints and args.descriptors_type in kdata.descriptors: | |
| image_list = [name | |
| for name in image_list | |
| if name not in kdata.keypoints[args.keypoints_type] or | |
| name not in kdata.descriptors[args.descriptors_type]] | |
| if len(image_list) == 0: | |
| print('All features were already extracted') | |
| exit(0) | |
| else: | |
| print(f'Extracting d2net features for {len(image_list)} images') | |
| # Creating CNN model | |
| model = D2Net( | |
| model_file=args.model_file, | |
| use_relu=args.use_relu, | |
| use_cuda=use_cuda | |
| ) | |
| if args.keypoints_type not in kdata.keypoints: | |
| keypoints_dtype = None | |
| keypoints_dsize = None | |
| else: | |
| keypoints_dtype = kdata.keypoints[args.keypoints_type].dtype | |
| keypoints_dsize = kdata.keypoints[args.keypoints_type].dsize | |
| if args.descriptors_type not in kdata.descriptors: | |
| descriptors_dtype = None | |
| descriptors_dsize = None | |
| else: | |
| descriptors_dtype = kdata.descriptors[args.descriptors_type].dtype | |
| descriptors_dsize = kdata.descriptors[args.descriptors_type].dsize | |
| # Process the files | |
| for image_name in tqdm(image_list, total=len(image_list)): | |
| img_path = get_image_fullpath(args.kapture_root, image_name) | |
| image = Image.open(img_path).convert('RGB') | |
| width, height = image.size | |
| resized_image = image | |
| resized_width = width | |
| resized_height = height | |
| max_edge = args.max_edge | |
| max_sum_edges = args.max_sum_edges | |
| if max(resized_width, resized_height) > max_edge: | |
| scale_multiplier = max_edge / max(resized_width, resized_height) | |
| resized_width = math.floor(resized_width * scale_multiplier) | |
| resized_height = math.floor(resized_height * scale_multiplier) | |
| resized_image = image.resize((resized_width, resized_height)) | |
| if resized_width + resized_height > max_sum_edges: | |
| scale_multiplier = max_sum_edges / (resized_width + resized_height) | |
| resized_width = math.floor(resized_width * scale_multiplier) | |
| resized_height = math.floor(resized_height * scale_multiplier) | |
| resized_image = image.resize((resized_width, resized_height)) | |
| fact_i = width / resized_width | |
| fact_j = height / resized_height | |
| resized_image = np.array(resized_image).astype('float') | |
| input_image = preprocess_image( | |
| resized_image, | |
| preprocessing=args.preprocessing | |
| ) | |
| with torch.no_grad(): | |
| if args.multiscale: | |
| keypoints, scores, descriptors = process_multiscale( | |
| torch.tensor( | |
| input_image[np.newaxis, :, :, :].astype(np.float32), | |
| device=device | |
| ), | |
| model | |
| ) | |
| else: | |
| keypoints, scores, descriptors = process_multiscale( | |
| torch.tensor( | |
| input_image[np.newaxis, :, :, :].astype(np.float32), | |
| device=device | |
| ), | |
| model, | |
| scales=[1] | |
| ) | |
| # Input image coordinates | |
| keypoints[:, 0] *= fact_i | |
| keypoints[:, 1] *= fact_j | |
| # i, j -> u, v | |
| keypoints = keypoints[:, [1, 0, 2]] | |
| if args.max_keypoints != float("+inf"): | |
| # keep the last (the highest) indexes | |
| idx_keep = scores.argsort()[-min(len(keypoints), args.max_keypoints):] | |
| keypoints = keypoints[idx_keep] | |
| descriptors = descriptors[idx_keep] | |
| if keypoints_dtype is None or descriptors_dtype is None: | |
| keypoints_dtype = keypoints.dtype | |
| descriptors_dtype = descriptors.dtype | |
| keypoints_dsize = keypoints.shape[1] | |
| descriptors_dsize = descriptors.shape[1] | |
| kdata.keypoints[args.keypoints_type] = kapture.Keypoints('d2net', keypoints_dtype, keypoints_dsize) | |
| kdata.descriptors[args.descriptors_type] = kapture.Descriptors('d2net', descriptors_dtype, | |
| descriptors_dsize, | |
| args.keypoints_type, 'L2') | |
| keypoints_config_absolute_path = get_feature_csv_fullpath(kapture.Keypoints, | |
| args.keypoints_type, | |
| args.kapture_root) | |
| descriptors_config_absolute_path = get_feature_csv_fullpath(kapture.Descriptors, | |
| args.descriptors_type, | |
| args.kapture_root) | |
| keypoints_to_file(keypoints_config_absolute_path, kdata.keypoints[args.keypoints_type]) | |
| descriptors_to_file(descriptors_config_absolute_path, kdata.descriptors[args.descriptors_type]) | |
| else: | |
| assert kdata.keypoints[args.keypoints_type].dtype == keypoints.dtype | |
| assert kdata.descriptors[args.descriptors_type].dtype == descriptors.dtype | |
| assert kdata.keypoints[args.keypoints_type].dsize == keypoints.shape[1] | |
| assert kdata.descriptors[args.descriptors_type].dsize == descriptors.shape[1] | |
| assert kdata.descriptors[args.descriptors_type].keypoints_type == args.keypoints_type | |
| assert kdata.descriptors[args.descriptors_type].metric_type == 'L2' | |
| keypoints_fullpath = get_keypoints_fullpath(args.keypoints_type, args.kapture_root, | |
| image_name, tar_handlers) | |
| print(f"Saving {keypoints.shape[0]} keypoints to {keypoints_fullpath}") | |
| image_keypoints_to_file(keypoints_fullpath, keypoints) | |
| kdata.keypoints[args.keypoints_type].add(image_name) | |
| descriptors_fullpath = get_descriptors_fullpath(args.descriptors_type, args.kapture_root, | |
| image_name, tar_handlers) | |
| print(f"Saving {descriptors.shape[0]} descriptors to {descriptors_fullpath}") | |
| image_descriptors_to_file(descriptors_fullpath, descriptors) | |
| kdata.descriptors[args.descriptors_type].add(image_name) | |
| if not keypoints_check_dir(kdata.keypoints[args.keypoints_type], args.keypoints_type, | |
| args.kapture_root, tar_handlers) or \ | |
| not descriptors_check_dir(kdata.descriptors[args.descriptors_type], args.descriptors_type, | |
| args.kapture_root, tar_handlers): | |
| print('local feature extraction ended successfully but not all files were saved') | |