├── .gitattributes ├── .gitignore ├── LICENSE ├── README.md ├── colmap-helpers ├── README.md ├── colmap_model_from_nvm.py ├── create_cmu_query_db.py ├── example_scripts │ ├── run_aachen.sh │ ├── run_aachen_sift.sh │ ├── run_cmu_slice.sh │ ├── run_cmu_slice_sift.sh │ ├── run_robotcar.sh │ └── run_robotcar_sift.sh ├── export_for_sfm.py ├── features_from_npz.py ├── internal │ ├── __init__.py │ ├── db_handling.py │ ├── db_matching_images.py │ ├── frame_matching.py │ └── nvm_to_colmap_helper.py ├── magic_cmu_to_db.py ├── match_features_with_db_prior.py ├── nvm_to_model.py ├── robotcar_db_png_to_jpg.py ├── update_db_cmu_with_intrinsics.py ├── update_db_intrinsics_from_another_db.py └── update_db_with_nvm_intrinsics.py ├── demo.ipynb ├── doc ├── assets │ ├── hfnet.jpg │ ├── pipeline.jpg │ └── teaser.jpg ├── datasets.md ├── demo │ ├── db1.jpg │ ├── db2.jpg │ ├── db3.jpg │ ├── db4.jpg │ ├── db5.jpg │ ├── query1.jpg │ ├── query2.jpg │ └── query3.jpg ├── local_evaluation.md └── training.md ├── hfnet ├── __init__.py ├── configs │ ├── doap_export_aachen.yaml │ ├── doap_export_sfm.yaml │ ├── hfnet_export_aachen_db.yaml │ ├── hfnet_export_aachen_queries.yaml │ ├── hfnet_export_cmu_db.yaml │ ├── hfnet_export_cmu_queries.yaml │ ├── hfnet_export_hpatches.yaml │ ├── hfnet_export_model.yaml │ ├── hfnet_export_robotcar_db.yaml │ ├── hfnet_export_robotcar_queries.yaml │ ├── hfnet_export_sfm.yaml │ ├── hfnet_train_distill.yaml │ ├── lfnet_export_hpatches.yaml │ ├── lfnet_export_sfm.yaml │ ├── mobilenetvlad_export_aachen.yaml │ ├── mobilenetvlad_export_hpatches.yaml │ ├── mobilenetvlad_export_model.yaml │ ├── mobilenetvlad_export_nclt.yaml │ ├── mobilenetvlad_export_sfm.yaml │ ├── mobilenetvlad_train_distill.yaml │ ├── netvlad_export_aachen.yaml │ ├── netvlad_export_cmu.yaml │ ├── netvlad_export_distill.yaml │ ├── netvlad_export_hpatches.yaml │ ├── netvlad_export_nclt.yaml │ ├── netvlad_export_robotcar.yaml │ ├── netvlad_export_sfm.yaml │ ├── superpoint_export_aachen_db.yaml │ ├── superpoint_export_aachen_queries.yaml │ ├── superpoint_export_cmu_db.yaml │ ├── superpoint_export_cmu_queries.yaml │ ├── superpoint_export_distill.yaml │ ├── superpoint_export_hpatches.yaml │ ├── superpoint_export_robotcar_db.yaml │ ├── superpoint_export_robotcar_queries.yaml │ └── superpoint_export_sfm.yaml ├── datasets │ ├── __init__.py │ ├── aachen.py │ ├── base_dataset.py │ ├── cmu.py │ ├── colmap_utils │ │ ├── __init__.py │ │ ├── read_dense.py │ │ └── read_model.py │ ├── distillation.py │ ├── hpatches.py │ ├── nclt.py │ ├── robotcar.py │ ├── sfm.py │ └── utils │ │ ├── __init__.py │ │ ├── homographies.py │ │ ├── photometric_augmentation.py │ │ └── pipeline.py ├── evaluate_aachen.py ├── evaluate_cmu.py ├── evaluate_robotcar.py ├── evaluation │ ├── __init__.py │ ├── cpp_localization.py │ ├── image_retrieval.py │ ├── keypoint_detectors.py │ ├── loaders.py │ ├── local_descriptors.py │ ├── localization.py │ ├── utils │ │ ├── __init__.py │ │ ├── db_management.py │ │ ├── descriptors.py │ │ ├── keypoints.py │ │ ├── localization.py │ │ ├── metrics.py │ │ └── misc.py │ └── visualize.py ├── export_model.py ├── export_predictions.py ├── models │ ├── __init__.py │ ├── backbones │ │ ├── mobilenet_v2.py │ │ └── utils │ │ │ ├── __init__.py │ │ │ ├── conv_blocks.py │ │ │ └── mobilenet.py │ ├── base_model.py │ ├── doap.py │ ├── hf_net.py │ ├── lf_net.py │ ├── lfnet_utils │ │ ├── __init__.py │ │ ├── det_tools.py │ │ ├── inference.py │ │ ├── spatial_transformer.py │ │ ├── tf_layer_utils.py │ │ └── tf_train_utils.py │ ├── mobilenetvlad.py │ ├── netvlad_original.py │ ├── super_point.py │ ├── super_point_pytorch.py │ └── utils │ │ ├── __init__.py │ │ ├── layers.py │ │ └── transformer.py ├── time_model.py ├── train.py └── utils │ ├── __init__.py │ ├── stdout_capturing.py │ └── tools.py ├── hloc-cpp ├── CMakeLists.txt ├── include │ └── hloc │ │ └── nanoflann.hpp ├── package.xml ├── python │ └── test.py ├── setup.py └── src │ └── hloc.cc ├── makefile ├── notebooks ├── __init__.py ├── evaluation_descriptors_hpatches.ipynb ├── evaluation_descriptors_sfm.ipynb ├── evaluation_detectors_hpatches.ipynb ├── evaluation_detectors_sfm.ipynb ├── evaluation_retrieval_nclt.ipynb ├── model_quality.ipynb ├── scale_annotation.ipynb ├── summarize_timings.ipynb ├── time_matching.ipynb ├── utils.py ├── visualize_distillation_dataset.ipynb ├── visualize_hpatches.ipynb ├── visualize_keypoints_hpatches.ipynb ├── visualize_keypoints_sfm.ipynb ├── visualize_localization_aachen.ipynb ├── visualize_localization_cmu.ipynb ├── visualize_localization_robotcar.ipynb ├── visualize_matches_hpatches.ipynb ├── visualize_matches_sfm.ipynb └── visualize_sfm.ipynb ├── setup.py ├── setup ├── requirements.txt ├── scripts │ ├── download_google_landmarks.py │ ├── generate_robotcar_query_list.py │ └── import_superpoint_weights.py └── setup.sh └── utils ├── env_cluster.sh ├── env_cluster_cpu.sh └── launch_cluster.sh /.gitattributes: -------------------------------------------------------------------------------- 1 | *.ipynb linguist-documentation 2 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | env/ 2 | __pycache__/ 3 | *.egg-info/ 4 | settings.py 5 | .ipynb_checkpoints/ 6 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 Autonomous Systems Lab 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /colmap-helpers/README.md: -------------------------------------------------------------------------------- 1 | # Building SfM models 2 | 3 | We provide here scripts to build SfM models for the Aachen, RobotCar, and CMU datasets using [COLMAP](https://colmap.github.io/) and any learned features. While we used SuperPoint keypoints and descriptors, this process should work for any reasonable feature predictor. 4 | 5 | ## Exporting dense features 6 | 7 | If not already done when evaluating the localization, we first export the network predictions, e.g. dense keypoint scores and descriptors: 8 | ```bash 9 | python3 hfnet/export_predictions.py \ 10 | hfnet/configs/hfnet_export_[aachen|cmu|robotcar]_db.yaml \ 11 | [aachen|cmu|robotcar] \ 12 | --exper_name hfnet \ 13 | --keys keypoints,scores,local_descriptor_map 14 | ``` 15 | 16 | This will create an `.npz` file in `$EXPER_PATH/exports/hfnet/[aachen|cmu|robotcar]/` for each database image. 17 | 18 | ## Extracting sparse features 19 | 20 | For increased flexibility, we only subsequently extract features using non-maximum suppression (NMS) and bilinear interpolation: 21 | ```bash 22 | python3 colmap-helpers/export_for_sfm.py \ 23 | [aachen|cmu|robotcar] \ # dataset name 24 | hfnet/[aachen|cmu|robotcar] \ # input directory 25 | hfnet/[aachen|cmu|robotcar]_sfm # output directory 26 | ``` 27 | This creates new `.npz` files in `$EXPER_PATH/exports/hfnet/[aachen|cmu|robotcar]_sfm/`. Parameters for extraction, such as the NMS radius or the number of keypoints, can be adjusted in `colmap-helpers/export_for_sfm.py`. Why this complicated process? We want to keep dense predictions accessible on disk so as to experiment with the extraction parameters when evaluating the localization (e.g. impact of the number of keypoints). 28 | 29 | ## Building the model 30 | 31 | Assuming we have reference models that contain accurate poses (e.g. from SIFT, as provided by the benchmark authors), we can match keypoints between images and triangulate the 3D points using COLMAP. This is much faster than optimizing all the poses again with bundle adjunstment, and ensures that the estimated query poses are consistent across models. 32 | 33 | The process goes as follows: 34 | - Export `.txt` files containing keypoint locations with `features_from_npz.py`. 35 | - Match features accross frames with `match_features_with_db_prior.py`. To speed up the process, we use the original SIFT database file to only match frames that are covisible in the SIFT model. 36 | - Generate a new database file with COLMAP commands `database_creator`, `feature_importer`, and `matches_importer`. 37 | - Import camera poses from a SIFT NVM model file by creating a dummy intermediate model with `colmap_model_from_nvm.py`. 38 | - Triangulate the new model with the COLMAP command `point_triangulator`. 39 | 40 | As there are some slight variations between the datasets (e.g. importing intrinsics from the SIFT database or from an NVM model), we provide **reference** scripts in `example_scripts/`. Paths might need to be adjusted to work in your workspace. 41 | -------------------------------------------------------------------------------- /colmap-helpers/colmap_model_from_nvm.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import numpy as np 3 | import os 4 | import sqlite3 5 | 6 | from internal.nvm_to_colmap_helper import convert_nvm_pose_to_colmap_p 7 | from internal.db_handling import blob_to_array 8 | 9 | 10 | def parse_args(): 11 | parser = argparse.ArgumentParser() 12 | parser.add_argument('--database_file', required=True) 13 | parser.add_argument('--nvm_file', required=True) 14 | parser.add_argument('--output_dir', required=True) 15 | args = parser.parse_args() 16 | return args 17 | 18 | 19 | def db_image_name_dict(db): 20 | db.execute('SELECT image_id, camera_id, name FROM images;') 21 | return {name: (im_id, cam_id) for im_id, cam_id, name in db} 22 | 23 | 24 | def export_cameras(db, output_dir): 25 | # In db file: 26 | # model INTEGER NOT NULL, 27 | # width INTEGER NOT NULL, 28 | # height INTEGER NOT NULL, 29 | # params BLOB, 30 | # prior_focal_length INTEGER NOT NULL) 31 | 32 | # Camera list with one line of data per camera: 33 | # CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[] 34 | 35 | db.execute('SELECT camera_id, model, width, height, params FROM cameras;') 36 | outfile = open(os.path.join(output_dir, 'cameras.txt'), 'w') 37 | for camera_id, model, width, height, params in db: 38 | if model != 2: # Wrong model, skip this camera. 39 | continue 40 | model = 'SIMPLE_RADIAL' 41 | params = blob_to_array(params, np.double) 42 | out = [camera_id, model, width, height] + params.tolist() 43 | outfile.write(' '.join(map(str, out)) + '\n') 44 | outfile.close() 45 | 46 | 47 | def convert_image_data(nvm_data, name_to_image_id): 48 | # colmap format 49 | # IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME 50 | 51 | # nvm format 52 | # 0 53 | name, focal, q1, q2, q3, q4, p1, p2, p3, dist, _ = nvm_data 54 | 55 | if name.endswith('.png'): 56 | name = name[:-3] + 'jpg' 57 | if name.startswith('./'): 58 | name = name[2:] 59 | image_id, camera_id = name_to_image_id[name] 60 | assert image_id > 0 61 | 62 | q_nvm = np.array(list(map(float, [q1, q2, q3, q4]))) 63 | p_nvm = np.array(list(map(float, [p1, p2, p3]))) 64 | p1, p2, p3 = convert_nvm_pose_to_colmap_p(q_nvm, p_nvm).tolist() 65 | 66 | # Write the corresponding camera id from the database file, as some 67 | # intrinsics could have been merged in the database (and not in NVM). 68 | return [image_id, q1, q2, q3, q4, p1, p2, p3, camera_id, name] 69 | 70 | 71 | def main(): 72 | args = parse_args() 73 | 74 | print('Reading DB') 75 | connection = sqlite3.connect(args.database_file) 76 | db = connection.cursor() 77 | name_to_image_id = db_image_name_dict(db) 78 | 79 | print('Exporting cameras') 80 | export_cameras(db, args.output_dir) 81 | 82 | print('Creating empty Points3D file') 83 | open(os.path.join(args.output_dir, 'points3D.txt'), 'w+').close() 84 | 85 | print('Reading NVM') 86 | with open(args.nvm_file) as f: 87 | line = f.readline() 88 | while line == '\n' or line.startswith('NVM_V3'): 89 | line = f.readline() 90 | total_num_images = int(line) 91 | print('Num images', total_num_images) 92 | 93 | outfile = open(os.path.join(args.output_dir, 'images.txt'), 'w') 94 | i = 0 95 | while i < total_num_images: 96 | line = f.readline() 97 | if line == '\n': 98 | continue 99 | data = convert_image_data(line.split(), name_to_image_id) 100 | outfile.write(' '.join(map(str, data)) + '\n\n') 101 | i += 1 102 | outfile.close() 103 | 104 | print('Done') 105 | db.close() 106 | connection.close() 107 | 108 | 109 | if __name__ == '__main__': 110 | main() 111 | -------------------------------------------------------------------------------- /colmap-helpers/create_cmu_query_db.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import numpy as np 3 | from os import path 4 | import struct 5 | 6 | from internal import db_handling 7 | 8 | 9 | def parse_args(): 10 | parser = argparse.ArgumentParser() 11 | parser.add_argument('--sift_feature_dir', required=True) 12 | parser.add_argument('--query_txt_file', required=True) 13 | parser.add_argument('--database_file', required=True) 14 | args = parser.parse_args() 15 | return args 16 | 17 | 18 | def main(): 19 | args = parse_args() 20 | db = db_handling.COLMAPDatabase.connect(args.database_file) 21 | db.create_tables() 22 | 23 | with open(args.query_txt_file) as f: 24 | for line in f: 25 | name, _, h, w, fx, fy, cx, cy = line.split(' ') 26 | 27 | params = np.array([float(fx), float(fy), float(cx), float(cy)]) 28 | camera_id = db.add_camera(1, int(h), int(w), params) 29 | image_id = db.add_image(path.join('images', name), camera_id) 30 | 31 | featurefile = path.join(args.sift_feature_dir, 32 | path.splitext(name)[0] + '.sift') 33 | with open(featurefile, 'rb') as f: 34 | data = f.read() 35 | 36 | header = struct.unpack_from('iiiii', data, 0) 37 | _, _, num_points, num_entries, desc_size = header 38 | assert num_entries == 5 and desc_size == 128 39 | offset = 20 40 | 41 | keypoints = np.zeros((num_points, 2)) 42 | for i in range(num_points): 43 | point = struct.unpack_from('fffff', data, offset) 44 | offset += 20 45 | keypoints[i, :] = np.array((point[1], point[0])) 46 | 47 | descriptors = np.zeros((num_points, desc_size)) 48 | for i in range(num_points): 49 | descriptor = struct.unpack_from('128B', data, offset) 50 | offset += desc_size 51 | descriptors[i, :] = np.asarray(descriptor) 52 | 53 | db.add_keypoints(image_id, keypoints) 54 | db.add_descriptors(image_id, descriptors) 55 | 56 | db.commit() 57 | 58 | 59 | if __name__ == '__main__': 60 | main() 61 | -------------------------------------------------------------------------------- /colmap-helpers/example_scripts/run_aachen.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | 3 | image_dir="images_upright" 4 | sift_db="aachen.db" 5 | new_db="aachen-new.db" 6 | npz_dir="aachen_npz_sfm" 7 | nvm_file="aachen_cvpr2018_db.nvm" 8 | 9 | match_ratio="87" 10 | match_file="matches${match_ratio}.txt" 11 | 12 | temp_model="aachen_temp_model" 13 | final_model="aachen_model" 14 | 15 | # Removove old feature and matches txt files 16 | rm ${image_dir}/db/*.txt 17 | rm ${match_file} 18 | 19 | # Export txt file for the features 20 | python3 features_from_npz.py \ 21 | --npz_dir ${npz_dir} \ 22 | --image_dir ${image_dir}/db/ 23 | 24 | # Match the new features using the original SIFT-based DB as a prior 25 | python3 match_features_with_db_prior.py \ 26 | --database_file ${sift_db} \ 27 | --image_prefix db \ 28 | --image_dir ${image_dir} \ 29 | --npz_dir ${npz_dir} \ 30 | --min_num_matches 15 \ 31 | --num_points_per_frame 3000 \ 32 | --use_ratio_test \ 33 | --ratio_test_values "0.${match_ratio}" 34 | 35 | # Create an empty Colmap DB 36 | colmap database_creator --database_path ${new_db} 37 | 38 | # Import the features 39 | colmap feature_importer \ 40 | --database_path ${new_db} \ 41 | --image_path ${image_dir} \ 42 | --import_path ${image_dir} 43 | 44 | # Update the intrinsics 45 | python3 update_db_with_nvm_intrinsics.py \ 46 | --database_file ${new_db} \ 47 | --nvm_file ${nvm_file} 48 | 49 | # Import matches as two-view geometries 50 | colmap matches_importer \ 51 | --database_path ${new_db} \ 52 | --match_list_path ${match_file} \ 53 | --match_type raw \ 54 | --SiftMatching.max_num_trials 20000 \ 55 | --SiftMatching.min_inlier_ratio 0.20 56 | 57 | # Build an initial model using the camera poses from the NVM file 58 | mkdir ${temp_model} 59 | python3 colmap_model_from_nvm.py \ 60 | --database_file ${new_db} \ 61 | --nvm_file ${nvm_file} \ 62 | --output_dir ${temp_model} 63 | 64 | # Triangulate the superpoint features using the previously prepared poses 65 | # and the features matches stored in the DB 66 | mkdir ${final_model} 67 | colmap point_triangulator \ 68 | --database_path ${new_db} \ 69 | --image_path ${image_dir} \ 70 | --input_path ${temp_model} \ 71 | --output_path ${final_model} 72 | -------------------------------------------------------------------------------- /colmap-helpers/example_scripts/run_aachen_sift.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | 3 | image_dir="images_upright" 4 | sift_db="aachen.db" 5 | nvm_file="aachen_cvpr2018_db.nvm" 6 | 7 | temp_model="aachen_temp_sift_model" 8 | final_model="aachen_sift_model" 9 | 10 | # Update the provided SIFT DB file with the correct intrinsics. 11 | python3 update_db_with_nvm_intrinsics.py \ 12 | --database_file ${sift_db} \ 13 | --nvm_file ${nvm_file} 14 | 15 | # Use the provided NVM file for ground-truth poses of all the cameras. 16 | mkdir ${temp_model} 17 | python3 colmap_model_from_nvm.py \ 18 | --database_file ${sift_db} \ 19 | --nvm_file ${nvm_file} \ 20 | --output_dir ${temp_model} 21 | 22 | # Triangulate the 3D model according to the matches provided in the SIFT DB. 23 | mkdir ${final_model} 24 | colmap point_triangulator \ 25 | --database_path ${sift_db} \ 26 | --image_path ${image_dir} \ 27 | --input_path ${temp_model} \ 28 | --output_path ${final_model} 29 | -------------------------------------------------------------------------------- /colmap-helpers/example_scripts/run_cmu_slice.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | 3 | slice_num=$1 4 | 5 | image_dir="images" 6 | sift_ref_db="slice${slice_num}.db" 7 | sift_query_db="query_slice${slice_num}.db" 8 | new_db="cmu-slice${slice_num}_new.db" 9 | npz_dir="cmu_npz_sfm" 10 | nvm_file="slice${slice_num}.nvm" 11 | 12 | sift_feature_dir="sift/sift_descriptors" 13 | query_txt_file="slice${slice_num}.queries_with_intrinsics.txt" 14 | 15 | match_ratio="85" 16 | match_file="matches${match_ratio}.txt" 17 | 18 | temp_model="cmu-slice${slice_num}_temp_model" 19 | final_model="cmu-slice${slice_num}_model" 20 | 21 | 22 | # Creates a tentative reference DB for SIFT as we need prior matches later on. 23 | python3 magic_cmu_to_db.py \ 24 | --sift_feature_dir ${sift_feature_dir} \ 25 | --nvm_file ${nvm_file}\ 26 | --database_file ${sift_ref_db} 27 | 28 | colmap exhaustive_matcher --database_path ${sift_ref_db} 29 | 30 | # Create a query SIFT DB for our localization evaluation algorithms. 31 | python3 create_cmu_query_db.py 32 | --sift_feature_dir ${sift_feature_dir} \ 33 | --query_txt_file ${query_txt_file}\ 34 | --database_file ${sift_query_db} 35 | 36 | # Match the new features using the original SIFT-based DB as a prior 37 | python3 features_from_npz.py \ 38 | --npz_dir ${npz_dir}/slice${slice_num}/database/ \ 39 | --image_dir ${image_dir}/database/ 40 | 41 | python3 match_features_with_db_prior.py \ 42 | --database_file ${sift_ref_db} \ 43 | --image_prefix "" \ 44 | --image_dir ${image_dir} \ 45 | --npz_dir ${npz_dir}/slice${slice_num}/ \ 46 | --min_num_matches 15 \ 47 | --num_points_per_frame 3000 \ 48 | --use_ratio_test \ 49 | --ratio_test_values "0.${match_ratio}" 50 | 51 | # Create an empty Colmap DB 52 | colmap database_creator --database_path ${new_db} 53 | 54 | # Import the features 55 | colmap feature_importer \ 56 | --database_path ${new_db} \ 57 | --image_path ${image_dir} \ 58 | --import_path ${image_dir} 59 | 60 | # Update the intrinsics using the ones stored in the NVM file. 61 | python3 update_db_with_nvm_intrinsics.py \ 62 | --database_file ${new_db} \ 63 | --nvm_file ${nvm_file} 64 | 65 | # Necessary as no principal points are stored in the NVM file. 66 | python3 update_db_cmu_with_intrinsics.py --database_file ${new_db} 67 | 68 | # Import matches as two-view geometries 69 | colmap matches_importer \ 70 | --database_path ${new_db} \ 71 | --match_list_path ${match_file} \ 72 | --match_type raw 73 | 74 | # Build an initial model using the camera poses from the NVM file 75 | mkdir ${temp_model} 76 | python3 colmap_model_from_nvm.py \ 77 | --database_file ${new_db} \ 78 | --nvm_file ${nvm_file} \ 79 | --output_dir ${temp_model} 80 | 81 | # Triangulate the superpoint features using the previously prepared poses 82 | # and the features matches stored in the DB 83 | mkdir ${final_model} 84 | colmap point_triangulator \ 85 | --database_path ${new_db} \ 86 | --image_path ${image_dir} \ 87 | --input_path ${temp_model} \ 88 | --output_path ${final_model} \ 89 | -------------------------------------------------------------------------------- /colmap-helpers/example_scripts/run_cmu_slice_sift.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | 3 | slice_num=$1 4 | 5 | sift_ref_db="slice${slice_num}.db" 6 | nvm_file="slice${slice_num}.nvm" 7 | final_model="cmu-slice${slice_num}_sift_model" 8 | 9 | # Use the provided NVM file and convert it directly to a 3D model. 10 | mkdir ${final_model} 11 | python3 nvm_to_model.py \ 12 | --database_file ${sift_ref_db} \ 13 | --nvm_file ${nvm_file} \ 14 | --output_dir ${final_model} \ 15 | 16 | # Convert the model from a text format to a binary format. 17 | colmap model_converter \ 18 | --input_path ${final_model} \ 19 | --output_path ${final_model} \ 20 | --output_type bin 21 | -------------------------------------------------------------------------------- /colmap-helpers/example_scripts/run_robotcar.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | 3 | image_dir="images" 4 | sift_db="overcast-reference.db" 5 | new_db="robotcar_new.db" 6 | npz_dir="robotcar_npz_sfm" 7 | nvm_file="all.nvm" 8 | 9 | match_ratio="85" 10 | match_file="matches${match_ratio}.txt" 11 | 12 | temp_model="robotcar_temp_model" 13 | final_model="robotcar_new_model" 14 | 15 | # Removove old feature and matches txt files 16 | rm ${image_dir}/overcast-reference/left/*.txt 17 | rm ${image_dir}/overcast-reference/right/*.txt 18 | rm ${image_dir}/overcast-reference/rear/*.txt 19 | rm ${match_file} 20 | 21 | # Export txt file for the features 22 | python3 features_from_npz.py \ 23 | --npz_dir ${npz_dir}/overcast-reference/left/ \ 24 | --image_dir ${image_dir}/overcast-reference/left/ 25 | 26 | python3 features_from_npz.py \ 27 | --npz_dir ${npz_dir}/overcast-reference/right/ \ 28 | --image_dir ${image_dir}/overcast-reference/right/ 29 | 30 | python3 features_from_npz.py \ 31 | --npz_dir ${npz_dir}/overcast-reference/rear/ \ 32 | --image_dir ${image_dir}/overcast-reference/rear/ 33 | 34 | # Match the new features using the original SIFT-based DB as a prior 35 | python3 match_features_with_db_prior.py \ 36 | --database_file ${sift_db} \ 37 | --image_prefix "" \ 38 | --image_dir ${image_dir} \ 39 | --npz_dir ${npz_dir} \ 40 | --min_num_matches 15 \ 41 | --num_points_per_frame 3000 \ 42 | --use_ratio_test \ 43 | --ratio_test_values "0.${match_ratio}" 44 | 45 | # Create an empty Colmap DB 46 | colmap database_creator --database_path ${new_db} 47 | 48 | # Import the features 49 | colmap feature_importer \ 50 | --database_path ${new_db} \ 51 | --image_path ${image_dir} \ 52 | --import_path ${image_dir} 53 | 54 | # Update the intrinsics 55 | python3 update_db_with_nvm_intrinsics.py \ 56 | --database_file ${new_db} \ 57 | --nvm_file ${nvm_file} 58 | 59 | python3 update_db_intrinsics_from_another_db.py \ 60 | --intrinsics_database_file ${sift_db} \ 61 | --database_file_to_modify ${new_db} 62 | 63 | # Import matches as two-view geometries 64 | colmap matches_importer \ 65 | --database_path ${new_db} \ 66 | --match_list_path ${match_file} \ 67 | --match_type raw 68 | 69 | # Build an initial model using the camera poses from the NVM file 70 | mkdir ${temp_model} 71 | python3 colmap_model_from_nvm.py \ 72 | --database_file ${new_db} \ 73 | --nvm_file ${nvm_file} \ 74 | --output_dir ${temp_model} 75 | 76 | # Triangulate the superpoint features using the previously prepared poses 77 | # and the features matches stored in the DB 78 | mkdir ${final_model} 79 | colmap point_triangulator \ 80 | --database_path ${new_db} \ 81 | --image_path ${image_dir} \ 82 | --input_path ${temp_model} \ 83 | --output_path ${final_model} 84 | -------------------------------------------------------------------------------- /colmap-helpers/example_scripts/run_robotcar_sift.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | 3 | image_dir="images" 4 | sift_db="overcast-reference.db" 5 | nvm_file="all.nvm" 6 | 7 | temp_model="robotcar_temp_sift_model" 8 | final_model="robotcar_sift_model" 9 | 10 | # Fix the inconsistency in the image file extension. 11 | # The database fill be updated to JPG file paths. 12 | python3 robotcar_db_png_to_jpg.py --db_file ${sift_db} 13 | 14 | # Update the provided SIFT DB file with the correct intrinsics. 15 | python3 update_db_with_nvm_intrinsics.py \ 16 | --database_file ${sift_db} \ 17 | --nvm_file ${nvm_file} 18 | 19 | # Use the provided NVM file for ground-truth poses of all the cameras. 20 | mkdir ${temp_model} 21 | python3 colmap_model_from_nvm.py \ 22 | --database_file ${sift_db} \ 23 | --nvm_file ${nvm_file} \ 24 | --output_dir ${temp_model} 25 | 26 | # Triangulate the 3D model according to the matches provided in the SIFT DB. 27 | mkdir ${final_model} 28 | colmap point_triangulator \ 29 | --database_path ${sift_db} \ 30 | --image_path ${image_dir} \ 31 | --input_path ${temp_model} \ 32 | --output_path ${final_model} 33 | -------------------------------------------------------------------------------- /colmap-helpers/export_for_sfm.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import argparse 3 | from pathlib import Path 4 | from tqdm import tqdm 5 | import numpy as np 6 | 7 | logging.basicConfig(format='[%(asctime)s %(levelname)s] %(message)s', 8 | datefmt='%m/%d/%Y %H:%M:%S', 9 | level=logging.INFO) 10 | from hfnet.datasets import get_dataset # noqa: E402 11 | from hfnet.evaluation.loaders import export_loader # noqa: E402 12 | from hfnet.settings import EXPER_PATH # noqa: E402 13 | 14 | 15 | def export_for_sfm(data_config, exper_config, export_name): 16 | export_dir = Path(EXPER_PATH, 'exports', export_name) 17 | export_dir.mkdir(exist_ok=True) 18 | 19 | dataset = get_dataset(data_config['name'])(**data_config) 20 | data_iter = dataset.get_test_set() 21 | 22 | for data in tqdm(data_iter): 23 | predictions = exper_config['predictor']( 24 | data['image'], data['name'], **exper_config) 25 | # Scale the keypoints to the original image size 26 | # and convert to Colmap convention (origin = corner of upper left pix) 27 | scale = ((np.array(data['original_size'][:2]) - 1) 28 | / (np.array(data['image'].shape[:2]) - 1)) 29 | export = { 30 | 'keypoints': scale[::-1] * predictions['keypoints'] + 0.5, 31 | 'scores': predictions['scores'], 32 | 'descriptors': predictions['descriptors'], 33 | 'image_size': data['image'].shape[:2][::-1] 34 | } 35 | name = data['name'].decode('utf-8') 36 | Path(export_dir, Path(name).parent).mkdir(parents=True, exist_ok=True) 37 | np.savez(Path(export_dir, f'{name}.npz'), **export) 38 | 39 | 40 | if __name__ == '__main__': 41 | parser = argparse.ArgumentParser() 42 | parser.add_argument('dataset', type=str) 43 | parser.add_argument('export_name', type=str) 44 | parser.add_argument('new_export_name', type=str) 45 | args = parser.parse_args() 46 | 47 | data_configs = { 48 | 'aachen': { 49 | 'name': 'aachen', 50 | 'load_db': True, 51 | 'load_queries': False, 52 | 'resize_max': 960, 53 | }, 54 | 'robotcar': { 55 | 'name': 'robotcar', 56 | 'sequences': ['overcast-reference'], 57 | 'resize_max': 960, 58 | }, 59 | 'cmu': { 60 | 'name': 'cmu', 61 | 'sequences': ['slice2', 'slice3', 'slice4', 'slice5', 'slice6', 62 | 'slice7', 'slice8', 'slice9', 'slice10', 'slice17'], 63 | 'load_db': True, 64 | 'load_queries': False, 65 | 'resize_max': 1024, 66 | } 67 | } 68 | 69 | exper_config = { 70 | 'experiment': args.export_name, 71 | 'predictor': export_loader, 72 | 'num_features': 3000, 73 | 'do_nms': True, 74 | 'nms_thresh': 4, 75 | } 76 | 77 | export_for_sfm( 78 | data_configs[args.dataset], exper_config, args.new_export_name) 79 | -------------------------------------------------------------------------------- /colmap-helpers/features_from_npz.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import numpy as np 3 | import os 4 | from tqdm import tqdm 5 | 6 | 7 | def parse_args(): 8 | parser = argparse.ArgumentParser() 9 | parser.add_argument('--npz_dir', required=True) 10 | # The images must be in the same directory as output. See: 11 | # https://colmap.github.io/tutorial.html#feature-detection-and-extraction 12 | parser.add_argument('--image_dir', required=True) 13 | args = parser.parse_args() 14 | return args 15 | 16 | 17 | def export_features_from_npz(filename, in_path, out_path): 18 | path_file = os.path.join(in_path, filename) 19 | frame1 = np.load(path_file) 20 | 21 | filename = os.path.splitext(os.path.basename(path_file))[0] 22 | 23 | out_path_and_name = os.path.join(out_path, filename) + '.jpg.txt' 24 | outfile = open(out_path_and_name, 'w+') 25 | 26 | SIFT_SIZE = 128 27 | kp1 = frame1['keypoints'] 28 | outfile.write(str(kp1.shape[0]) + ' ' + str(SIFT_SIZE) + '\n') 29 | 30 | for keypoint in kp1: 31 | # Generate some dummy SIFT values as we will anyway use external 32 | # from a matches.txt file. 33 | s = str(keypoint[0]) + ' ' + str(keypoint[1]) 34 | s += ' 1 1 ' + '1 '*SIFT_SIZE + '\n' 35 | outfile.write(s) 36 | 37 | outfile.close() 38 | 39 | 40 | def export_feature_detections(): 41 | args = parse_args() 42 | for filename in tqdm(os.listdir(args.npz_dir), unit='npz'): 43 | if filename.endswith('.npz'): 44 | export_features_from_npz(filename, args.npz_dir, args.image_dir) 45 | 46 | 47 | if __name__ == '__main__': 48 | export_feature_detections() 49 | -------------------------------------------------------------------------------- /colmap-helpers/internal/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethz-asl/hfnet/01577845583373470a0cf156a2acd972e9223bc7/colmap-helpers/internal/__init__.py -------------------------------------------------------------------------------- /colmap-helpers/internal/db_matching_images.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018, ETH Zurich and UNC Chapel Hill. 2 | # All rights reserved. 3 | # 4 | # Redistribution and use in source and binary forms, with or without 5 | # modification, are permitted provided that the following conditions are met: 6 | # 7 | # * Redistributions of source code must retain the above copyright 8 | # notice, this list of conditions and the following disclaimer. 9 | # 10 | # * Redistributions in binary form must reproduce the above copyright 11 | # notice, this list of conditions and the following disclaimer in the 12 | # documentation and/or other materials provided with the distribution. 13 | # 14 | # * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of 15 | # its contributors may be used to endorse or promote products derived 16 | # from this software without specific prior written permission. 17 | # 18 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 21 | # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE 22 | # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 23 | # CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 24 | # SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 25 | # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 26 | # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 27 | # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 28 | # POSSIBILITY OF SUCH DAMAGE. 29 | # 30 | # Author: Johannes L. Schoenberger (jsch-at-demuc-dot-de) 31 | 32 | 33 | from collections import defaultdict 34 | import numpy as np 35 | import sqlite3 36 | from tqdm import tqdm 37 | 38 | 39 | from .db_handling import blob_to_array 40 | 41 | 42 | def pair_id_to_image_ids(pair_id): 43 | image_id2 = pair_id % 2147483647 44 | image_id1 = (pair_id - image_id2) / 2147483647 45 | return image_id1, image_id2 46 | 47 | 48 | def get_matching_images(database_file, min_num_matches, filter_image_dir): 49 | connection = sqlite3.connect(database_file) 50 | cursor = connection.cursor() 51 | 52 | cursor.execute('SELECT image_id, name FROM images;') 53 | images = {image_id: name for image_id, name in cursor} 54 | 55 | two_way_matches = defaultdict(list) 56 | cursor.execute( 57 | 'SELECT pair_id, data FROM two_view_geometries WHERE rows>=?;', 58 | (min_num_matches,)) 59 | for pair_id, data in cursor: 60 | inlier_matches = blob_to_array(data, np.uint32, shape=(-1, 2)) 61 | 62 | image_id1, image_id2 = pair_id_to_image_ids(pair_id) 63 | image_name1 = images[image_id1] 64 | image_name2 = images[image_id2] 65 | num_matches = inlier_matches.shape[0] 66 | 67 | # Make sure the match comes from the desired directory. 68 | if (image_name1.startswith(filter_image_dir) 69 | and image_name2.startswith(filter_image_dir)): 70 | two_way_matches[image_id1].append((image_id2, num_matches)) 71 | two_way_matches[image_id2].append((image_id1, num_matches)) 72 | 73 | matching_image_pairs = [] 74 | for image_id, direct_matching_frames in tqdm(two_way_matches.items()): 75 | image_name = images[image_id] 76 | 77 | matching_frames = set() 78 | for matching_frame in direct_matching_frames: 79 | assert matching_frame[1] >= min_num_matches 80 | if matching_frame[0] > image_id: 81 | matching_frames.add(matching_frame[0]) 82 | 83 | # Insert the direct matching pairs. 84 | for match in matching_frames: 85 | assert match > image_id 86 | matching_image_pairs.append((image_name, images[match])) 87 | 88 | cursor.close() 89 | connection.close() 90 | return matching_image_pairs 91 | -------------------------------------------------------------------------------- /colmap-helpers/internal/frame_matching.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | from matplotlib import pyplot as plt 4 | 5 | 6 | def baseline_sift_matching(img1, img2): 7 | sift = cv2.xfeatures2d.SIFT_create() 8 | kp1, des1 = sift.detectAndCompute(img1, None) 9 | kp2, des2 = sift.detectAndCompute(img2, None) 10 | matches = cv2.BFMatcher().knnMatch(des1, des2, k=2) 11 | 12 | good = [[m] for m, n in matches if m.distance < 0.7*n.distance] 13 | img3 = cv2.drawMatchesKnn(img1, kp1, img2, kp2, good, None, 14 | matchColor=(0, 255, 0), matchesMask=None, 15 | singlePointColor=(255, 0, 0), flags=0) 16 | return img3 17 | 18 | 19 | def debug_matching(frame1, frame2, path_image1, path_image2, matches, 20 | matches_mask, num_points, use_ratio_test): 21 | img1 = cv2.imread(path_image1, 0) 22 | img2 = cv2.imread(path_image2, 0) 23 | 24 | kp1 = get_ocv_kpts_from_np(frame1['keypoints'][:num_points, :]) 25 | kp2 = get_ocv_kpts_from_np(frame2['keypoints'][:num_points, :]) 26 | 27 | if use_ratio_test: 28 | img = cv2.drawMatchesKnn(img1, kp1, img2, kp2, matches, None, 29 | matchColor=(0, 255, 0), 30 | matchesMask=matches_mask, 31 | singlePointColor=(255, 0, 0), flags=0) 32 | else: 33 | img = cv2.drawMatches(img1, kp1, img2, kp2, matches, None, 34 | matchColor=(0, 255, 0), 35 | singlePointColor=(255, 0, 0), flags=0) 36 | 37 | img_sift = baseline_sift_matching(img1, img2) 38 | 39 | fig = plt.figure(figsize=(2, 1)) 40 | fig.add_subplot(2, 1, 1) 41 | plt.imshow(img) 42 | plt.title('Custom features') 43 | fig.add_subplot(2, 1, 2) 44 | plt.imshow(img_sift) 45 | plt.title('SIFT') 46 | plt.show() 47 | 48 | 49 | def get_ocv_kpts_from_np(keypoints_np): 50 | return [cv2.KeyPoint(x=x, y=y, _size=1) for x, y in keypoints_np] 51 | 52 | 53 | def match_frames(path_npz1, path_npz2, path_image1, path_image2, num_points, 54 | use_ratio_test, ratio_test_values, debug): 55 | frame1 = np.load(path_npz1) 56 | frame2 = np.load(path_npz2) 57 | 58 | # Assert the keypoints are sorted according to the score. 59 | assert np.all(np.sort(frame1['scores'])[::-1] == frame1['scores']) 60 | 61 | # WARNING: scores are not taken into account as of now. 62 | des1 = frame1['descriptors'].astype('float32')[:num_points] 63 | des2 = frame2['descriptors'].astype('float32')[:num_points] 64 | 65 | if use_ratio_test: 66 | keypoint_matches = [[] for _ in ratio_test_values] 67 | matcher = cv2.BFMatcher(cv2.NORM_L2) 68 | matches = matcher.knnMatch(des1, des2, k=2) 69 | 70 | smallest_distances = [dict() for _ in ratio_test_values] 71 | 72 | matches_mask = [[0, 0] for _ in range(len(matches))] 73 | for i, (m, n) in enumerate(matches): 74 | for ratio_idx, ratio in enumerate(ratio_test_values): 75 | if m.distance < ratio * n.distance: 76 | if m.trainIdx not in smallest_distances[ratio_idx]: 77 | smallest_distances[ratio_idx][m.trainIdx] = ( 78 | m.distance, m.queryIdx) 79 | matches_mask[i] = [1, 0] 80 | keypoint_matches[ratio_idx].append( 81 | (m.queryIdx, m.trainIdx)) 82 | else: 83 | old_dist, old_queryIdx = smallest_distances[ 84 | ratio_idx][m.trainIdx] 85 | if m.distance < old_dist: 86 | old_distance, old_queryIdx = smallest_distances[ 87 | ratio_idx][m.trainIdx] 88 | smallest_distances[ratio_idx][m.trainIdx] = ( 89 | m.distance, m.queryIdx) 90 | matches_mask[i] = [1, 0] 91 | keypoint_matches[ratio_idx].remove( 92 | (old_queryIdx, m.trainIdx)) 93 | keypoint_matches[ratio_idx].append( 94 | (m.queryIdx, m.trainIdx)) 95 | else: 96 | keypoint_matches = [[]] 97 | matches_mask = [] 98 | matcher = cv2.BFMatcher(cv2.NORM_L2, crossCheck=True) 99 | matches = matcher.match(des1, des2) 100 | 101 | # Matches are already cross-checked. 102 | for match in matches: 103 | # match.trainIdx belongs to des2. 104 | keypoint_matches[0].append((match.queryIdx, match.trainIdx)) 105 | 106 | if debug: 107 | debug_matching(frame1, frame2, path_image1, path_image2, matches, 108 | matches_mask, num_points, use_ratio_test) 109 | 110 | return keypoint_matches 111 | -------------------------------------------------------------------------------- /colmap-helpers/internal/nvm_to_colmap_helper.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | # Code from http://nipy.org/nibabel/reference/nibabel.quaternions.html 5 | def quat2mat(q): 6 | w, x, y, z = q 7 | Nq = w*w + x*x + y*y + z*z 8 | if Nq < 1e-10: 9 | return np.eye(3) 10 | s = 2.0/Nq 11 | X = x*s 12 | Y = y*s 13 | Z = z*s 14 | wX = w*X; wY = w*Y; wZ = w*Z 15 | xX = x*X; xY = x*Y; xZ = x*Z 16 | yY = y*Y; yZ = y*Z; zZ = z*Z 17 | return np.array( 18 | [[ 1.0-(yY+zZ), xY-wZ, xZ+wY ], 19 | [ xY+wZ, 1.0-(xX+zZ), yZ-wX ], 20 | [ xZ-wY, yZ+wX, 1.0-(xX+yY) ]]) 21 | 22 | 23 | def convert_nvm_pose_to_colmap_p(q_nvm, p_nvm): 24 | R = quat2mat(q_nvm) 25 | p_colmap = R.dot(-p_nvm) 26 | return p_colmap 27 | 28 | 29 | def test(): 30 | # Colmap output 31 | # 3157 0.921457 -0.0279726 0.372923 -0.105178 -3.43436 -1.03257 -1.86432 3157 db/384.jpg 32 | q_colmap = np.array([0.921457, -0.0279726, 0.372923, -0.105178]) 33 | p_colmap = np.array([-3.43436, -1.03257, -1.86432]) 34 | 35 | # NVM 36 | # db/384.jpg 3186.77 0.921457 -0.0279726 0.372923 -0.105178 0.911128 1.3598 3.6956 -0.0538065 0 37 | q_nvm = np.array([0.921457, -0.0279726, 0.372923, -0.105178]) 38 | p_nvm = np.array([0.911128, 1.3598, 3.6956]) 39 | 40 | np.testing.assert_equal(q_colmap, q_nvm) 41 | 42 | p_colmap_computed = convert_nvm_pose_to_colmap_p(q_nvm, p_nvm) 43 | 44 | np.testing.assert_almost_equal(p_colmap_computed, p_colmap, decimal=4) 45 | 46 | 47 | if __name__ == "__main__": 48 | test() 49 | -------------------------------------------------------------------------------- /colmap-helpers/magic_cmu_to_db.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import numpy as np 3 | from os import path 4 | import struct 5 | 6 | from internal import db_handling 7 | 8 | 9 | def parse_args(): 10 | parser = argparse.ArgumentParser() 11 | parser.add_argument('--sift_feature_dir', required=True) 12 | parser.add_argument('--nvm_file', required=True) 13 | parser.add_argument('--database_file', required=True) 14 | args = parser.parse_args() 15 | return args 16 | 17 | 18 | def main(): 19 | args = parse_args() 20 | db = db_handling.COLMAPDatabase.connect(args.database_file) 21 | db.create_tables() 22 | 23 | camera_model = 2 24 | w = 1024 25 | h = 768 26 | 27 | with open(args.nvm_file) as f: 28 | f.readline() 29 | total_num_images = int(f.readline()) 30 | for _, line in zip(range(total_num_images), f): 31 | line = line.split(' ') 32 | name, focal, dist = line[0], line[1], line[9] 33 | 34 | # = 0 35 | params = np.array([float(focal), h/2, w/2, float(dist)]) 36 | camera_id = db.add_camera(camera_model, h, w, params) 37 | image_id = db.add_image(name, camera_id) 38 | 39 | featurefile = path.join(args.sift_feature_dir, 40 | path.splitext(name)[0] + '.sift') 41 | with open(featurefile, 'rb') as f: 42 | data = f.read() 43 | 44 | header = struct.unpack_from('iiiii', data, 0) 45 | _, _, num_points, num_entries, desc_size = header 46 | assert num_entries == 5 and desc_size == 128 47 | offset = 20 48 | 49 | keypoints = np.zeros((num_points, 2)) 50 | for i in range(num_points): 51 | point = struct.unpack_from('fffff', data, offset) 52 | offset += 20 53 | keypoints[i, :] = np.array((point[1], point[0])) 54 | 55 | descriptors = np.zeros((num_points, desc_size)) 56 | for i in range(num_points): 57 | descriptor = struct.unpack_from('128B', data, offset) 58 | offset += desc_size 59 | descriptors[i, :] = np.asarray(descriptor) 60 | 61 | db.add_keypoints(image_id, keypoints) 62 | db.add_descriptors(image_id, descriptors) 63 | 64 | db.commit() 65 | 66 | 67 | if __name__ == '__main__': 68 | main() 69 | -------------------------------------------------------------------------------- /colmap-helpers/match_features_with_db_prior.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | from tqdm import tqdm 4 | 5 | from internal import db_matching_images 6 | from internal import frame_matching 7 | 8 | 9 | def parse_args(): 10 | parser = argparse.ArgumentParser() 11 | parser.add_argument('--database_file', required=True) 12 | parser.add_argument('--min_num_matches', type=int, default=15) 13 | parser.add_argument('--num_points_per_frame', type=int, default=2500) 14 | parser.add_argument('--image_dir', required=True) 15 | parser.add_argument('--npz_dir', required=True) 16 | 17 | # This argument lets us only look at the matches from a certain folder. 18 | # We want to avoid adding matches from other folders, e.g. query. This 19 | # filters images according to the prefix as stored in the db file. 20 | parser.add_argument('--image_prefix', required=True) 21 | 22 | parser.add_argument('--use_ratio_test', action='store_true') 23 | parser.add_argument('--ratio_test_values', type=str, default='0.85') 24 | parser.add_argument('--debug', action='store_true') 25 | args = parser.parse_args() 26 | return args 27 | 28 | 29 | def main(): 30 | args = parse_args() 31 | 32 | ratio_test_values = [float(v) for v in args.ratio_test_values.split(',')] 33 | print('Ratio test values to use:', ratio_test_values) 34 | outfiles = [open('matches{}.txt'.format(x), 'w+') 35 | for x in [int(i * 100) for i in ratio_test_values]] 36 | 37 | print('Looking for matching image pairs...') 38 | matching_image_pairs = db_matching_images.get_matching_images( 39 | args.database_file, args.min_num_matches, args.image_prefix) 40 | print('Got', len(matching_image_pairs), 'matching image pairs.') 41 | 42 | num_missing_images = 0 43 | for (name1, name2) in tqdm(matching_image_pairs, unit='pairs'): 44 | # Get npz instead of image files. 45 | npz1 = os.path.join(args.npz_dir, os.path.splitext(name1)[0] + '.npz') 46 | npz2 = os.path.join(args.npz_dir, os.path.splitext(name2)[0] + '.npz') 47 | image1 = os.path.join(args.image_dir, name1) 48 | image2 = os.path.join(args.image_dir, name2) 49 | 50 | # Some images might be missing, e.g. in the Robotcar case. 51 | if not os.path.isfile(image1) or not os.path.isfile(image2): 52 | num_missing_images += 1 53 | continue 54 | 55 | assert os.path.isfile(npz1), npz1 56 | assert os.path.isfile(npz2), npz2 57 | 58 | num_points = args.num_points_per_frame 59 | matches_for_different_ratios = frame_matching.match_frames( 60 | npz1, npz2, image1, image2, num_points, 61 | args.use_ratio_test, ratio_test_values, args.debug) 62 | 63 | if(args.use_ratio_test): 64 | assert len(matches_for_different_ratios) == len(ratio_test_values) 65 | 66 | for i, keypoint_matches in enumerate(matches_for_different_ratios): 67 | if len(keypoint_matches) > args.min_num_matches: 68 | outfiles[i].write(name1 + ' ' + name2 + '\n') 69 | for (match1, match2) in keypoint_matches: 70 | outfiles[i].write(str(match1) + ' ' + str(match2) + '\n') 71 | outfiles[i].write('\n') 72 | 73 | for outfile in outfiles: 74 | outfile.close() 75 | 76 | print('Missing', num_missing_images, 'images skipped.') 77 | 78 | 79 | if __name__ == '__main__': 80 | main() 81 | -------------------------------------------------------------------------------- /colmap-helpers/robotcar_db_png_to_jpg.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import sqlite3 4 | 5 | 6 | def parse_args(): 7 | parser = argparse.ArgumentParser() 8 | parser.add_argument('--db_file', required=True) 9 | args = parser.parse_args() 10 | return args 11 | 12 | 13 | def main(): 14 | args = parse_args() 15 | connection = sqlite3.connect(args.db_file) 16 | cursor = connection.cursor() 17 | 18 | cursor.execute('SELECT image_id, name FROM images;') 19 | ids_and_names = [(int(image_id), name) for image_id, name in cursor] 20 | print('Got', len(ids_and_names), 'image ids.') 21 | 22 | for image_id, name in ids_and_names: 23 | jpg_name = os.path.splitext(name)[0] + '.jpg' 24 | cursor.execute('UPDATE images SET name = ? WHERE image_id = ?;', 25 | [jpg_name, image_id]) 26 | 27 | cursor.close() 28 | connection.commit() 29 | connection.close() 30 | 31 | 32 | if __name__ == '__main__': 33 | main() 34 | -------------------------------------------------------------------------------- /colmap-helpers/update_db_cmu_with_intrinsics.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import numpy as np 3 | import sqlite3 4 | 5 | from internal.db_handling import array_to_blob, blob_to_array 6 | 7 | # c0 8 | # 868.99 0 525.94 9 | # 0 866.06 420.04 10 | # 0 0 1 11 | 12 | # c1 13 | # 873.38 0 529.32 14 | # 0 876.49 397.27 15 | # 0 0 1 16 | 17 | cmu_intrinsics = { 18 | 'c0': { 19 | 'focal_length': (868.99 + 866.06) / 2, 20 | 'cx': 525.94, 21 | 'cy': 420.04, 22 | }, 23 | 'c1': { 24 | 'focal_length': (873.38 + 876.49) / 2, 25 | 'cx': 529.32, 26 | 'cy': 397.27, 27 | }, 28 | } 29 | 30 | 31 | def parse_args(): 32 | parser = argparse.ArgumentParser() 33 | parser.add_argument('--database_file', required=True) 34 | args = parser.parse_args() 35 | return args 36 | 37 | 38 | def update_db_intrinsics(db_connection, camera_id, focal_length, cx, cy): 39 | cursor = db_connection.cursor() 40 | cursor.execute('SELECT params FROM cameras WHERE camera_id=?;', 41 | (camera_id,)) 42 | data = cursor.fetchall() 43 | assert len(data) == 1 44 | intrinsics = blob_to_array(data[0][0], np.double) 45 | 46 | # Update intrinsic values. 47 | intrinsics[0] = focal_length 48 | intrinsics[1] = cx 49 | intrinsics[2] = cy 50 | 51 | cursor.execute('UPDATE cameras SET params = ? WHERE camera_id = ?;', 52 | [array_to_blob(intrinsics), camera_id]) 53 | cursor.close() 54 | 55 | 56 | def main(): 57 | args = parse_args() 58 | 59 | connection = sqlite3.connect(args.database_file) 60 | cursor = connection.cursor() 61 | 62 | # Get a mapping between image ids and image names and camera ids. 63 | cursor.execute('SELECT image_id, camera_id, name FROM images;') 64 | for image_id, camera_id, name in cursor: 65 | if 'c0' in name: 66 | assert 'c1' not in name 67 | update_db_intrinsics(connection, camera_id, **cmu_intrinsics['c0']) 68 | else: 69 | assert 'c0' not in name 70 | update_db_intrinsics(connection, camera_id, **cmu_intrinsics['c1']) 71 | 72 | cursor.close() 73 | connection.commit() 74 | connection.close() 75 | 76 | 77 | if __name__ == '__main__': 78 | main() 79 | -------------------------------------------------------------------------------- /colmap-helpers/update_db_intrinsics_from_another_db.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import numpy as np 3 | import sqlite3 4 | 5 | from internal.db_handling import array_to_blob, blob_to_array 6 | 7 | 8 | def parse_args(): 9 | parser = argparse.ArgumentParser() 10 | parser.add_argument('--intrinsics_database_file', required=True) 11 | parser.add_argument('--database_file_to_modify', required=True) 12 | args = parser.parse_args() 13 | return args 14 | 15 | 16 | def db_update_intrinsics_entry(db_connection, camera_id, intrinsics_to_write): 17 | cursor = db_connection.cursor() 18 | cursor.execute('UPDATE cameras SET params = ? WHERE camera_id = ?;', 19 | [array_to_blob(intrinsics_to_write), camera_id]) 20 | cursor.close() 21 | 22 | 23 | def update_target_db(database_file, image_name_to_id_and_camera_id, 24 | camera_intrinsics): 25 | connection = sqlite3.connect(database_file) 26 | cursor = connection.cursor() 27 | 28 | cursor.execute('SELECT camera_id, name FROM images;') 29 | for camera_id, name in cursor: 30 | assert name in image_name_to_id_and_camera_id 31 | src_image_id, src_camera_id = image_name_to_id_and_camera_id[name] 32 | assert(src_camera_id in camera_intrinsics) 33 | intrinsics_to_write = camera_intrinsics[src_camera_id] 34 | 35 | db_update_intrinsics_entry(connection, camera_id, intrinsics_to_write) 36 | 37 | cursor.close() 38 | connection.commit() 39 | connection.close() 40 | 41 | 42 | def parse_input_db(database_file): 43 | connection = sqlite3.connect(database_file) 44 | cursor = connection.cursor() 45 | 46 | # Get a mapping between image ids and image names and camera ids. 47 | cursor.execute('SELECT image_id, camera_id, name FROM images;') 48 | name_to_ids = {name: (im_id, cam_id) for im_id, cam_id, name in cursor} 49 | 50 | cursor.execute('SELECT camera_id, params FROM cameras;') 51 | intrinsics = {cam_id: blob_to_array(p, np.double) for cam_id, p in cursor} 52 | 53 | cursor.close() 54 | connection.close() 55 | return name_to_ids, intrinsics 56 | 57 | 58 | def main(): 59 | args = parse_args() 60 | 61 | print('Reading the input DB') 62 | image_name_to_id_and_camera_id, camera_intrinsics = parse_input_db( 63 | args.intrinsics_database_file) 64 | 65 | print('Updating the target DB') 66 | update_target_db( 67 | args.database_file_to_modify, image_name_to_id_and_camera_id, 68 | camera_intrinsics) 69 | 70 | print('Done!') 71 | 72 | 73 | if __name__ == '__main__': 74 | main() 75 | -------------------------------------------------------------------------------- /colmap-helpers/update_db_with_nvm_intrinsics.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import numpy as np 3 | import sqlite3 4 | import os 5 | 6 | from internal.db_handling import array_to_blob, blob_to_array 7 | 8 | 9 | def parse_args(): 10 | parser = argparse.ArgumentParser() 11 | parser.add_argument('--database_file', required=True) 12 | parser.add_argument('--nvm_file', required=True) 13 | args = parser.parse_args() 14 | return args 15 | 16 | 17 | def db_image_name_dict(connection): 18 | cursor = connection.cursor() 19 | cursor.execute('SELECT image_id, camera_id, name FROM images;') 20 | name_to_ids = {name: (im_id, cam_id) for im_id, cam_id, name in cursor} 21 | cursor.close() 22 | return name_to_ids 23 | 24 | 25 | def update_db_intrinsics(db_connection, camera_id, focal_length, radial_dist): 26 | cursor = db_connection.cursor() 27 | cursor.execute('SELECT params FROM cameras WHERE camera_id=?;', 28 | (camera_id,)) 29 | data = cursor.fetchall() 30 | assert len(data) == 1 31 | intrinsics = blob_to_array(data[0][0], np.double) 32 | 33 | # Update intrinsic values. 34 | intrinsics[0] = focal_length 35 | intrinsics[3] = radial_dist 36 | 37 | # Write intrinsics back. 38 | cursor.execute('UPDATE cameras SET params = ? WHERE camera_id = ?;', 39 | [array_to_blob(intrinsics), camera_id]) 40 | cursor.close() 41 | 42 | 43 | def process(nvm_data, image_name_to_id_and_camera_id, db_connection): 44 | name = nvm_data[0] 45 | focal_length = float(nvm_data[1]) 46 | radial_dist = -float(nvm_data[9]) 47 | 48 | name = os.path.splitext(nvm_data[0])[0] + '.jpg' 49 | if name.startswith('./'): 50 | name = name[2:] 51 | 52 | image_id, camera_id = image_name_to_id_and_camera_id[name] 53 | update_db_intrinsics(db_connection, camera_id, focal_length, radial_dist) 54 | 55 | 56 | def main(): 57 | args = parse_args() 58 | 59 | connection = sqlite3.connect(args.database_file) 60 | image_name_to_id_and_camera_id = db_image_name_dict(connection) 61 | 62 | with open(args.nvm_file) as f: 63 | line = f.readline() 64 | while line == '\n' or line.startswith('NVM_V3'): 65 | line = f.readline() 66 | num_images = int(line) 67 | 68 | i = 0 69 | while i < num_images: 70 | line = f.readline() 71 | if line == '\n': 72 | continue 73 | data = line.split(' ') 74 | process(data, image_name_to_id_and_camera_id, connection) 75 | i += 1 76 | 77 | print('Done parsing data for', num_images, 'images.') 78 | connection.commit() 79 | connection.close() 80 | 81 | 82 | if __name__ == '__main__': 83 | main() 84 | -------------------------------------------------------------------------------- /doc/assets/hfnet.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethz-asl/hfnet/01577845583373470a0cf156a2acd972e9223bc7/doc/assets/hfnet.jpg -------------------------------------------------------------------------------- /doc/assets/pipeline.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethz-asl/hfnet/01577845583373470a0cf156a2acd972e9223bc7/doc/assets/pipeline.jpg -------------------------------------------------------------------------------- /doc/assets/teaser.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethz-asl/hfnet/01577845583373470a0cf156a2acd972e9223bc7/doc/assets/teaser.jpg -------------------------------------------------------------------------------- /doc/datasets.md: -------------------------------------------------------------------------------- 1 | # Datasets 2 | 3 | All datasets should be downloaded in `$DATA_PATH`. We give below additional details as well as the expected directory structures. 4 | 5 | ## 6-DoF Localization 6 | 7 | These datasets are introduced in [Benchmarking 6DOF Outdoor Visual Localization in Changing Conditions](https://arxiv.org/abs/1707.09092) by Sattler et al., and can be downloaded from [the associated website](http://www.visuallocalization.net/). 8 | 9 | For each dataset, there are separate directories for the images, the SfM models, and the localization databases. 10 | 11 | #### Aachen Day-Night 12 | 13 | ``` 14 | aachen/ 15 | ├── aachen.db 16 | ├── day_time_queries_with_intrinsics.txt 17 | ├── night_time_queries_with_intrinsics.txt 18 | ├── databases/ 19 | ├── images_upright/ 20 | │   ├── db/ 21 | │   │   └── ... 22 | │   └── query/ 23 | │      └── ... 24 | └── models/ 25 |    └── hfnet_model/ 26 |       ├── cameras.bin 27 |       ├── images.bin 28 |       └── points3D.bin 29 | ``` 30 | 31 | #### RobotCar Seasons 32 | 33 | ``` 34 | robotcar/ 35 | ├── overcast-reference.db 36 | ├── query.db 37 | ├── images/ 38 | │   ├── overcast-reference/ 39 | │   ├── sun/ 40 | │   ├── dusk/ 41 | │   ├── night/ 42 | │   └── night-rain/ 43 | ├── intrinsics/ 44 | │   ├── left_intrinsics.txt 45 | │   ├── rear_intrinsics.txt 46 | │   └── right_intrinsics.txt 47 | ├── queries/ 48 | │   ├── dusk_queries_with_intrinsics.txt 49 | │   ├── night_queries_with_intrinsics.txt 50 | │   ├── night-rain_queries_with_intrinsics.txt 51 | │   └── sun_queries_with_intrinsics.txt 52 | └── models/ 53 |    └── hfnet_model/ 54 |       ├── cameras.bin 55 |       ├── images.bin 56 |       └── points3D.bin 57 | ``` 58 | The query lists generated with `setup/utils/generate_robotcar_query_list.py` are available [here](http://robotics.ethz.ch/~asl-datasets/2019_CVPR_hierarchical_localization/query_lists_robotcar.tar.gz). 59 | 60 | #### CMU Seasons 61 | 62 | ``` 63 | cmu/ 64 | ├── images/ 65 | │   ├── slice2/ 66 | │   │ ├── database/ 67 | │   │ └── query/ 68 | │   └── ... 69 | ├── slice2/ 70 | │   ├── sift_database.db 71 | │   ├── sift_queries.db 72 | │   ├── slice2.queries_with_intrinsics.txt 73 | │   └── models/ 74 | │   └── hfnet_model/ 75 | │      ├── cameras.bin 76 | │      ├── images.bin 77 | │      └── points3D.bin 78 | ├── slice3/ 79 | │   └── ... 80 | └── ... 81 | ``` 82 | 83 | ## Local features evaluation 84 | 85 | Local feature detectors and descriptors can be evaluated on the HPatches and SfM datasets, as reported in our paper. 86 | 87 | #### HPatches 88 | 89 | The dataset is described in the paper [HPatches: A benchmark and evaluation of handcrafted and learned local descriptors](https://arxiv.org/pdf/1704.05939.pdf) by Balntas et al., and can be downloaded [here](https://github.com/hpatches/hpatches-dataset). 90 | 91 | ``` 92 | hpatches/ 93 | ├── i_ajuntament/ 94 | └── ... 95 | ``` 96 | 97 | #### SfM 98 | 99 | The dataset is introduced by Ono et al. in their paper [LF-Net: Learning Local Features from Images and can be obtained [here](https://github.com/vcg-uvic/sfm_benchmark_release). 100 | 101 | ``` 102 | sfm/ 103 | ├── british_museum/ 104 | │   └── dense/ 105 | ├── florence_cathedral_side/ 106 | │   └── dense/ 107 | ├── lincoln_memorial_statue/ 108 | │   └── dense/ 109 | ├── london_bridge/ 110 | │   └── dense/ 111 | ├── milan_cathedral/ 112 | │   └── dense/ 113 | ├── mount_rushmore/ 114 | │   └── dense/ 115 | ├── piazza_san_marco/ 116 | │   └── dense/ 117 | ├── reichstag/ 118 | │   └── dense/ 119 | ├── sacre_coeur/ 120 | │   └── dense/ 121 | ├── sagrada_familia/ 122 | │   └── dense/ 123 | ├── st_pauls_cathedral/ 124 | │   └── dense/ 125 | ├── united_states_capitol/ 126 | │   └── dense/ 127 | ├── scales.txt 128 | └── exif 129 |    ├── brandenburg_gate 130 |    ├── british_museum 131 |    ├── buckingham 132 |    ├── colosseum_exterior 133 |    ├── florence_cathedral_side 134 |    ├── grand_place_brussels 135 |    ├── hagia_sophia_interior 136 |    ├── lincoln_memorial_statue 137 |    ├── london_bridge 138 |    ├── milan_cathedral 139 |    ├── mount_rushmore 140 |    ├── notre_dame_front_facade 141 |    ├── pantheon_exterior 142 |    ├── piazza_san_marco 143 |    ├── sagrada_familia 144 |    ├── st_pauls_cathedral 145 |    ├── st_peters_square 146 |    ├── taj_mahal 147 |    ├── temple_nara_japan 148 |    ├── trevi_fountain 149 |    ├── united_states_capitol 150 |    └── westminster_abbey 151 | ``` 152 | 153 | ## Multi-task distillation 154 | 155 | HF-Net is trained on the Google Landmarks and Berkeley Deep Drive datasets. For the former, first download the [index of images](https://github.com/ethz-asl/hierarchical_loc/releases/download/1.0/google_landmarks_index.csv) and then the dataset itself using the script `setup/scripts/download_google_landmarks.py`. The latter can be downloaded on the [dataset website](https://bdd-data.berkeley.edu/) (we used the night and dawn sequences). 156 | 157 | The labels are predictions of SuperPoint and NetVLAD. Their export is described in the [training documentation](doc/training.md). 158 | 159 | ``` 160 | google_landmarks/ 161 | ├── images/ 162 | ├── global_descriptors/ 163 | └── superpoint_predictions/ 164 | bdd/ 165 | ├── dawn_images_vga/ 166 | ├── night_images_vga/ 167 | ├── global_descriptors/ 168 | └── superpoint_predictions/ 169 | ``` 170 | -------------------------------------------------------------------------------- /doc/demo/db1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethz-asl/hfnet/01577845583373470a0cf156a2acd972e9223bc7/doc/demo/db1.jpg -------------------------------------------------------------------------------- /doc/demo/db2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethz-asl/hfnet/01577845583373470a0cf156a2acd972e9223bc7/doc/demo/db2.jpg -------------------------------------------------------------------------------- /doc/demo/db3.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethz-asl/hfnet/01577845583373470a0cf156a2acd972e9223bc7/doc/demo/db3.jpg -------------------------------------------------------------------------------- /doc/demo/db4.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethz-asl/hfnet/01577845583373470a0cf156a2acd972e9223bc7/doc/demo/db4.jpg -------------------------------------------------------------------------------- /doc/demo/db5.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethz-asl/hfnet/01577845583373470a0cf156a2acd972e9223bc7/doc/demo/db5.jpg -------------------------------------------------------------------------------- /doc/demo/query1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethz-asl/hfnet/01577845583373470a0cf156a2acd972e9223bc7/doc/demo/query1.jpg -------------------------------------------------------------------------------- /doc/demo/query2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethz-asl/hfnet/01577845583373470a0cf156a2acd972e9223bc7/doc/demo/query2.jpg -------------------------------------------------------------------------------- /doc/demo/query3.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethz-asl/hfnet/01577845583373470a0cf156a2acd972e9223bc7/doc/demo/query3.jpg -------------------------------------------------------------------------------- /doc/local_evaluation.md: -------------------------------------------------------------------------------- 1 | # Evaluation of local features 2 | 3 | As described in our paper, we evaluate several local feature detectors and descriptors on two dataset, [HPatches](https://github.com/hpatches/hpatches-dataset) and [SfM](https://github.com/vcg-uvic/sfm_benchmark_release), following standard metrics for matching and pose estimation. 4 | 5 | ## Exporting the predictions 6 | 7 | ``` 8 | python3 hfnet/export_predictions.py \ 9 | hfnet/configs/hfnet_export_[hfnet|superpoint|lfnet|doap]_[hpatches|sfm].yaml \ 10 | [hfnet|superpoint|lfnet|doap]/predictions_[hpatches|sfm] \ 11 | ``` 12 | 13 | This will export `.npz` files in `$EXPER_PATH/[hfnet|superpoint|lfnet|doap]/predictions_[hpatches|sfm]/`. 14 | 15 | ## Visualizing the predictions 16 | 17 | Head over to `notebooks/visualize_[keypoints|matches]_[hpatches|sfm].ipynb`. The names of the experiments might need to be adjusted. 18 | 19 | ## Running the evaluation 20 | 21 | Head over to `notebooks/evaluation_[detectors|descriptors]_[hpatches|sfm].ipynb`. 22 | 23 | -------------------------------------------------------------------------------- /doc/training.md: -------------------------------------------------------------------------------- 1 | # Training with multi-task distillation 2 | 3 | ## Exporting the target predictions 4 | 5 | We first export the predictions of NetVLAD (global descriptor) and SuperPoint (dense keypoint scores and descriptors), which will be the labels of the dataset. 6 | 7 | ```bash 8 | python3 hfnet/export_predictions.py 9 | hfnet/configs/netvlad_export_distill.yaml \ 10 | global_descriptors \ 11 | --keys global_descriptor \ 12 | --as_dataset 13 | python3 hfnet/export_predictions.py 14 | hfnet/configs/superpoint_export_distill.yaml \ 15 | superpoint_predictions \ 16 | --keys local_descriptor_map,dense_scores \ 17 | --as_dataset 18 | ``` 19 | 20 | ## Training HF-Net 21 | ```bash 22 | python3 hfnet/train.py hfnet/configs/hfnet_train_distill.yaml hfnet 23 | ``` 24 | 25 | The training can be interrupted at any time using `Ctrl+C` and can be monitored with Tensorboard summaries saved in `$EXPER_PATH/hfnet/`. The weights are also saved there. 26 | 27 | ## Exporting the model for deployment 28 | 29 | ```bash 30 | python3 hfnet/export_model.py config/hfnet_train_distill.yaml hfnet 31 | ``` 32 | will export the model to `$EXPER_PATH/saved_models/hfnet/`. 33 | -------------------------------------------------------------------------------- /hfnet/__init__.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | logging.basicConfig(format='[%(asctime)s %(levelname)s] %(message)s', 4 | datefmt='%m/%d/%Y %H:%M:%S', 5 | level=logging.INFO) 6 | -------------------------------------------------------------------------------- /hfnet/configs/doap_export_aachen.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'aachen' 3 | load_db: true 4 | load_queries: true 5 | resize_max: 960 6 | model: 7 | name: 'doap' 8 | weights: 'DOAP/HPatches_ST_LM_128d.mat' 9 | image_channels: 1 10 | weights: '' 11 | -------------------------------------------------------------------------------- /hfnet/configs/doap_export_sfm.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'sfm' 3 | sequences: [ 4 | 'british_museum', 'florence_cathedral_side', 5 | 'lincoln_memorial_statue', 'london_bridge', 6 | 'milan_cathedral', 'mount_rushmore', 7 | 'piazza_san_marco', 'reichstag', 'sagrada_familia', 8 | 'st_pauls_cathedral', 'united_states_capitol'] 9 | preprocessing: 10 | upright: true 11 | model: 12 | name: 'doap' 13 | weights: 'DOAP/HPatches_ST_LM_128d.mat' 14 | image_channels: 1 15 | weights: '' 16 | -------------------------------------------------------------------------------- /hfnet/configs/hfnet_export_aachen_db.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'aachen' 3 | load_db: true 4 | load_queries: false 5 | resize_max: 960 6 | grayscale: true 7 | model: 8 | name: 'hf_net' 9 | image_channels: 1 10 | local: 11 | detector_threshold: 0.005 12 | -------------------------------------------------------------------------------- /hfnet/configs/hfnet_export_aachen_queries.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'aachen' 3 | load_db: false 4 | load_queries: true 5 | resize_max: 960 6 | grayscale: true 7 | model: 8 | name: 'hf_net' 9 | image_channels: 1 10 | local: 11 | detector_threshold: 0.005 12 | nms_radius: 4 13 | num_keypoints: 2500 14 | -------------------------------------------------------------------------------- /hfnet/configs/hfnet_export_cmu_db.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'cmu' 3 | sequences: ['slice2', 'slice3', 'slice4', 'slice5', 'slice6', 'slice7', 4 | 'slice8', 'slice9', 'slice10', 'slice17'] 5 | load_db: true 6 | load_queries: false 7 | resize_max: 1024 8 | model: 9 | name: 'hf_net' 10 | image_channels: 1 11 | local: 12 | detector_threshold: 0.005 13 | -------------------------------------------------------------------------------- /hfnet/configs/hfnet_export_cmu_queries.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'cmu' 3 | sequences: ['slice2', 'slice3', 'slice4', 'slice5', 'slice6', 'slice7', 4 | 'slice8', 'slice9', 'slice10', 'slice17'] 5 | load_db: false 6 | load_queries: true 7 | resize_max: 1024 8 | model: 9 | name: 'hf_net' 10 | image_channels: 1 11 | local: 12 | detector_threshold: 0.005 13 | nms_radius: 4 14 | num_keypoints: 2500 15 | -------------------------------------------------------------------------------- /hfnet/configs/hfnet_export_hpatches.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'hpatches' 3 | alteration: 'all' 4 | model: 5 | name: 'hf_net' 6 | image_channels: 1 7 | local: 8 | detector_threshold: 0.005 9 | -------------------------------------------------------------------------------- /hfnet/configs/hfnet_export_model.yaml: -------------------------------------------------------------------------------- 1 | model: 2 | name: 'hf_net' 3 | image_channels: 1 4 | depth_multiplier: 0.75 5 | dropout_keep_prob: null 6 | global: 7 | intermediate_proj: 0 8 | dimensionality_reduction: 4096 9 | n_clusters: 32 10 | local: 11 | descriptor_dim: 256 12 | detector_threshold: 0.005 13 | nms_radius: 4 14 | num_keypoints: 1000 15 | -------------------------------------------------------------------------------- /hfnet/configs/hfnet_export_robotcar_db.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'robotcar' 3 | sequences: ['overcast-reference'] 4 | resize_max: 960 5 | model: 6 | name: 'hf_net' 7 | image_channels: 1 8 | local: 9 | detector_threshold: 0.005 10 | -------------------------------------------------------------------------------- /hfnet/configs/hfnet_export_robotcar_queries.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'robotcar' 3 | sequences: ['sun', 'dusk', 'night', 'night-rain'] 4 | resize_max: 960 5 | model: 6 | name: 'hf_net' 7 | image_channels: 1 8 | local: 9 | detector_threshold: 0.005 10 | nms_radius: 4 11 | num_keypoints: 2500 12 | -------------------------------------------------------------------------------- /hfnet/configs/hfnet_export_sfm.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'sfm' 3 | sequences: [ 4 | 'british_museum', 'florence_cathedral_side', 5 | 'lincoln_memorial_statue', 'london_bridge', 6 | 'milan_cathedral', 'mount_rushmore', 7 | 'piazza_san_marco', 'reichstag', 'sagrada_familia', 8 | 'st_pauls_cathedral', 'united_states_capitol'] 9 | preprocessing: 10 | grayscale: true 11 | model: 12 | name: 'hf_net' 13 | image_channels: 1 14 | local: 15 | detector_threshold: 0.005 16 | -------------------------------------------------------------------------------- /hfnet/configs/hfnet_train_distill.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'distillation' 3 | image_dirs: ['google_landmarks/images', 4 | 'bdd/dawn_images_vga', 'bdd/night_images_vga'] 5 | load_targets: True 6 | targets: 7 | - dir: 'global_descriptors' 8 | keys: ['global_descriptor'] 9 | - dir: 'superpoint_predictions' 10 | keys: ['local_descriptor_map', 'dense_scores'] 11 | validation_size: 192 12 | truncate: [185000, null, null] 13 | preprocessing: 14 | resize: [480, 640] 15 | grayscale: true 16 | augmentation: 17 | photometric: 18 | enable: true 19 | primitives: [ 20 | 'random_brightness', 'random_contrast', 21 | 'additive_gaussian_noise', 'additive_shade', 'motion_blur'] 22 | params: 23 | random_brightness: {delta_range: [-30., 40.]} 24 | random_contrast: {strength_range: [0.3, 1.2]} 25 | additive_gaussian_noise: {stddev_range: [0, 10]} 26 | additive_speckle_noise: {prob_range: [0, 0.0035]} 27 | additive_shade: 28 | transparency_range: [-0.5, 0.5] 29 | kernel_size_range: [100, 150] 30 | prob: 0.5 31 | motion_blur: {kernel_size: [13, 21], prob: 0.1} 32 | homographic: 33 | enable: false 34 | params: 35 | translation: true 36 | rotation: true 37 | scaling: true 38 | perspective: true 39 | scaling_amplitude: 0.2 40 | perspective_amplitude_x: 0.3 41 | perspective_amplitude_y: 0.3 42 | patch_ratio: 0.9 43 | max_angle: 0.3 44 | allow_artifacts: true 45 | valid_border_margin: 9 46 | model: 47 | name: 'hf_net' 48 | image_channels: 1 49 | depth_multiplier: 0.75 50 | dropout_keep_prob: null 51 | global: 52 | intermediate_proj: 0 53 | dimensionality_reduction: 4096 54 | n_clusters: 32 55 | local: 56 | descriptor_dim: 256 57 | #loss_weights: {local_desc: 1, global_desc: 1, detector: 1} 58 | loss_weights: 'uncertainties' 59 | train_backbone: true 60 | batch_size: 16 61 | eval_batch_size: 16 62 | learning_rate: [0.001, 0.0001, 0.00001] 63 | learning_rate_step: [60000, 80000] 64 | weights: 'mobilenet_v2_0.75_224/mobilenet_v2_0.75_224.ckpt' 65 | train_iter: 85000 66 | validation_interval: 500 67 | save_interval: 5000 68 | keep_checkpoints: 100 69 | -------------------------------------------------------------------------------- /hfnet/configs/lfnet_export_hpatches.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'hpatches' 3 | alteration: 'all' 4 | model: 5 | name: 'lf_net' 6 | top_k: 500 7 | #nms_ksize: 4 8 | image_channels: 1 9 | weights: 'lf-net/release/models/outdoor/models-latest-42000' 10 | -------------------------------------------------------------------------------- /hfnet/configs/lfnet_export_sfm.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'sfm' 3 | sequences: [ 4 | 'british_museum', 'florence_cathedral_side', 5 | 'lincoln_memorial_statue', 'london_bridge', 6 | 'milan_cathedral', 'mount_rushmore', 7 | 'piazza_san_marco', 'reichstag', 'sagrada_familia', 8 | 'st_pauls_cathedral', 'united_states_capitol'] 9 | preprocessing: 10 | upright: true 11 | model: 12 | name: 'lf_net' 13 | top_k: 1000 14 | config: 'lf-net/release/models/outdoor/config.pkl' 15 | #nms_ksize: 4 16 | image_channels: 1 17 | weights: 'lf-net/release/models/outdoor/models-latest-42000' 18 | -------------------------------------------------------------------------------- /hfnet/configs/mobilenetvlad_export_aachen.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'aachen' 3 | load_db: true 4 | load_queries: true 5 | resize_max: 960 6 | grayscale: true 7 | model: 8 | name: 'mobilenetvlad' 9 | local_descriptor_layer: 'layer_7' 10 | image_channels: 1 11 | -------------------------------------------------------------------------------- /hfnet/configs/mobilenetvlad_export_hpatches.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'hpatches' 3 | alteration: 'all' 4 | model: 5 | name: 'mobilenetvlad' 6 | local_descriptor_layer: 'layer_14' 7 | image_channels: 1 8 | -------------------------------------------------------------------------------- /hfnet/configs/mobilenetvlad_export_model.yaml: -------------------------------------------------------------------------------- 1 | model: 2 | name: 'mobilenetvlad' 3 | local_descriptor_layer: 'layer_7' 4 | image_channels: 1 5 | depth_multiplier: 0.35 6 | dimensionality_reduction: 4096 7 | dropout_keep_prob: null 8 | intermediate_proj: 0 9 | n_clusters: 32 10 | -------------------------------------------------------------------------------- /hfnet/configs/mobilenetvlad_export_nclt.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'nclt' 3 | test_sequences: ['2012-01-08', '2012-08-20', '2013-02-23'] 4 | camera: 4 5 | preprocessing: 6 | grayscale: false 7 | model: 8 | name: 'mobilenetvlad' 9 | image_channels: 3 10 | -------------------------------------------------------------------------------- /hfnet/configs/mobilenetvlad_export_sfm.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'sfm' 3 | sequences: [ 4 | 'british_museum', 'florence_cathedral_side', 5 | 'lincoln_memorial_statue', 'london_bridge', 6 | 'milan_cathedral', 'mount_rushmore', 7 | 'piazza_san_marco', 'reichstag', 'sagrada_familia', 8 | 'st_pauls_cathedral', 'united_states_capitol'] 9 | model: 10 | name: 'mobilenetvlad' 11 | local_descriptor_layer: 'layer_7' 12 | image_channels: 1 13 | -------------------------------------------------------------------------------- /hfnet/configs/mobilenetvlad_train_distill.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'distillation' 3 | image_dirs: ['google_landmarks/images', 'bdd/dawn_images_vga', 4 | 'bdd/night_images_vga'] 5 | load_targets: True 6 | targets: 7 | - dir: 'global_descriptors' 8 | keys: ['global_descriptor'] 9 | truncate: [185000, null, null] 10 | preprocessing: 11 | grayscale: true 12 | validation_size: 192 13 | cache_in_memory: false 14 | augmentation: 15 | photometric: 16 | enable: true 17 | primitives: [ 18 | 'random_brightness', 'random_contrast', 19 | 'additive_gaussian_noise', 'additive_shade', 'motion_blur'] 20 | params: 21 | random_brightness: {delta_range: [-75., 40.]} 22 | random_contrast: {strength_range: [0.3, 1.5]} 23 | additive_gaussian_noise: {stddev_range: [0, 10]} 24 | additive_speckle_noise: {prob_range: [0, 0.0035]} 25 | additive_shade: 26 | transparency_range: [-0.5, 0.5] 27 | kernel_size_range: [100, 150] 28 | prob: 0.5 29 | motion_blur: {kernel_size: [13, 21], prob: 0.1} 30 | model: 31 | name: 'mobilenetvlad' 32 | image_channels: 1 33 | dropout_keep_prob: null 34 | intermediate_proj: 0 35 | dimensionality_reduction: 4096 36 | depth_multiplier: 0.75 37 | n_clusters: 32 38 | train_backbone: true 39 | batch_size: 16 40 | eval_batch_size: 16 41 | learning_rate: [0.001, 0.0001, 0.00001, 0.000001] 42 | learning_rate_step: [40000, 60000, 70000] 43 | weights: 'mobilenet_v2_0.75_224/mobilenet_v2_0.75_224.ckpt' 44 | train_iter: 150000 45 | validation_interval: 500 46 | save_interval: 5000 47 | keep_checkpoints: 100 48 | -------------------------------------------------------------------------------- /hfnet/configs/netvlad_export_aachen.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'aachen' 3 | load_db: false 4 | load_queries: true 5 | resize_max: 960 6 | model: 7 | name: 'netvlad_original' 8 | local_descriptor_layer: 'conv3_3' 9 | image_channels: 1 10 | weights: 'vd16_pitts30k_conv5_3_vlad_preL2_intra_white/vd16_pitts30k_conv5_3_vlad_preL2_intra_white' 11 | -------------------------------------------------------------------------------- /hfnet/configs/netvlad_export_cmu.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'cmu' 3 | sequences: ['slice2', 'slice3', 'slice4', 'slice5', 'slice6', 'slice7', 4 | 'slice8', 'slice9', 'slice10', 'slice17'] 5 | grayscale: false 6 | load_db: true 7 | load_queries: true 8 | resize_max: 1024 9 | model: 10 | name: 'netvlad_original' 11 | local_descriptor_layer: 'conv3_3' 12 | image_channels: 3 13 | weights: 'vd16_pitts30k_conv5_3_vlad_preL2_intra_white/vd16_pitts30k_conv5_3_vlad_preL2_intra_white' 14 | -------------------------------------------------------------------------------- /hfnet/configs/netvlad_export_distill.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'distillation' 3 | image_dirs: ['google_landmarks/images', 4 | 'bdd/dawn_images_vga', 'bdd/night_images_vga'] 5 | shuffle: false 6 | preprocessing: 7 | resize: [480, 640] 8 | grayscale: false 9 | model: 10 | name: 'netvlad_original' 11 | image_channels: 3 12 | weights: 'vd16_pitts30k_conv5_3_vlad_preL2_intra_white/vd16_pitts30k_conv5_3_vlad_preL2_intra_white' 13 | -------------------------------------------------------------------------------- /hfnet/configs/netvlad_export_hpatches.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'hpatches' 3 | alteration: 'all' 4 | model: 5 | name: 'netvlad_original' 6 | local_descriptor_layer: 'conv3_3' 7 | image_channels: 1 8 | weights: 'vd16_pitts30k_conv5_3_vlad_preL2_intra_white/vd16_pitts30k_conv5_3_vlad_preL2_intra_white' 9 | -------------------------------------------------------------------------------- /hfnet/configs/netvlad_export_nclt.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'nclt' 3 | test_sequences: ['2012-01-08', '2012-08-20', '2013-02-23'] 4 | camera: 4 5 | model: 6 | name: 'netvlad_original' 7 | image_channels: 1 8 | weights: 'vd16_pitts30k_conv5_3_vlad_preL2_intra_white/vd16_pitts30k_conv5_3_vlad_preL2_intra_white' 9 | -------------------------------------------------------------------------------- /hfnet/configs/netvlad_export_robotcar.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'robotcar' 3 | sequences: ['overcast-reference', 'sun', 'dusk', 'night', 'night-rain'] 4 | grayscale: false 5 | resize_max: 960 6 | model: 7 | name: 'netvlad_original' 8 | local_descriptor_layer: 'conv3_3' 9 | image_channels: 3 10 | weights: 'vd16_pitts30k_conv5_3_vlad_preL2_intra_white/vd16_pitts30k_conv5_3_vlad_preL2_intra_white' 11 | -------------------------------------------------------------------------------- /hfnet/configs/netvlad_export_sfm.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'sfm' 3 | sequences: [ 4 | 'british_museum', 'florence_cathedral_side', 5 | 'lincoln_memorial_statue', 'london_bridge', 6 | 'milan_cathedral', 'mount_rushmore', 7 | 'piazza_san_marco', 'reichstag', 'sagrada_familia', 8 | 'st_pauls_cathedral', 'united_states_capitol'] 9 | preprocessing: 10 | upright: true 11 | model: 12 | name: 'netvlad_original' 13 | local_descriptor_layer: 'conv3_3' 14 | image_channels: 1 15 | weights: 'vd16_pitts30k_conv5_3_vlad_preL2_intra_white/vd16_pitts30k_conv5_3_vlad_preL2_intra_white' 16 | -------------------------------------------------------------------------------- /hfnet/configs/superpoint_export_aachen_db.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'aachen' 3 | load_db: true 4 | load_queries: false 5 | resize_max: 960 6 | model: 7 | name: 'super_point_pytorch' 8 | nms_radius: 0 9 | detector_threshold: 0.005 10 | image_channels: 1 11 | weights: 'superpoint_v1.pth' 12 | -------------------------------------------------------------------------------- /hfnet/configs/superpoint_export_aachen_queries.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'aachen' 3 | load_db: false 4 | load_queries: true 5 | resize_max: 960 6 | model: 7 | name: 'super_point_pytorch' 8 | nms_radius: 0 9 | detector_threshold: 0.005 10 | image_channels: 1 11 | nms_radius: 4 12 | num_keypoints: 2500 13 | weights: 'superpoint_v1.pth' 14 | -------------------------------------------------------------------------------- /hfnet/configs/superpoint_export_cmu_db.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'cmu' 3 | sequences: ['slice2', 'slice3', 'slice4', 'slice5', 'slice6', 'slice7', 4 | 'slice8', 'slice9', 'slice10', 'slice17'] 5 | load_db: true 6 | load_queries: false 7 | resize_max: 1024 8 | model: 9 | name: 'super_point_pytorch' 10 | nms_radius: 0 11 | detector_threshold: 0.005 12 | image_channels: 1 13 | weights: 'superpoint_v1.pth' 14 | -------------------------------------------------------------------------------- /hfnet/configs/superpoint_export_cmu_queries.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'cmu' 3 | sequences: ['slice2', 'slice3', 'slice4', 'slice5', 'slice6', 'slice7', 4 | 'slice8', 'slice9', 'slice10', 'slice17'] 5 | load_db: false 6 | load_queries: true 7 | resize_max: 1024 8 | model: 9 | name: 'super_point_pytorch' 10 | nms_radius: 0 11 | detector_threshold: 0.005 12 | image_channels: 1 13 | nms_radius: 4 14 | num_keypoints: 2500 15 | weights: 'superpoint_v1.pth' 16 | -------------------------------------------------------------------------------- /hfnet/configs/superpoint_export_distill.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'distillation' 3 | image_dirs: ['google_landmarks/images', 4 | 'bdd/dawn_images_vga', 'bdd/night_images_vga'] 5 | shuffle: false 6 | preprocessing: 7 | resize: [480, 640] 8 | grayscale: true 9 | model: 10 | name: 'super_point_pytorch' 11 | nms_radius: 0 12 | detector_threshold: 0.005 13 | image_channels: 1 14 | weights: 'superpoint_v1.pth' 15 | -------------------------------------------------------------------------------- /hfnet/configs/superpoint_export_hpatches.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'hpatches' 3 | alteration: 'all' 4 | model: 5 | name: 'super_point_pytorch' 6 | nms_radius: 0 7 | detector_threshold: 0.01 8 | image_channels: 1 9 | weights: 'superpoint_v1.pth' 10 | -------------------------------------------------------------------------------- /hfnet/configs/superpoint_export_robotcar_db.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'robotcar' 3 | sequences: ['overcast-reference'] 4 | resize_max: 960 5 | model: 6 | name: 'super_point_pytorch' 7 | nms_radius: 0 8 | detector_threshold: 0.005 9 | image_channels: 1 10 | weights: 'superpoint_v1.pth' 11 | -------------------------------------------------------------------------------- /hfnet/configs/superpoint_export_robotcar_queries.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'robotcar' 3 | sequences: ['sun', 'dusk', 'night', 'night-rain'] 4 | resize_max: 960 5 | model: 6 | name: 'super_point_pytorch' 7 | nms_radius: 0 8 | detector_threshold: 0.005 9 | image_channels: 1 10 | nms_radius: 4 11 | num_keypoints: 2500 12 | weights: 'superpoint_v1.pth' 13 | -------------------------------------------------------------------------------- /hfnet/configs/superpoint_export_sfm.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: 'sfm' 3 | sequences: [ 4 | 'british_museum', 'florence_cathedral_side', 5 | 'lincoln_memorial_statue', 'london_bridge', 6 | 'milan_cathedral', 'mount_rushmore', 7 | 'piazza_san_marco', 'reichstag', 'sagrada_familia', 8 | 'st_pauls_cathedral', 'united_states_capitol'] 9 | preprocessing: 10 | upright: true 11 | model: 12 | name: 'super_point_pytorch' 13 | nms_radius: 0 14 | detector_threshold: 0.005 15 | image_channels: 1 16 | weights: 'superpoint_v1.pth' 17 | -------------------------------------------------------------------------------- /hfnet/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | def get_dataset(name): 2 | mod = __import__('{}.{}'.format(__name__, name), fromlist=['']) 3 | return getattr(mod, _module_to_class(name)) 4 | 5 | 6 | def _module_to_class(name): 7 | return ''.join(n.capitalize() for n in name.split('_')) 8 | -------------------------------------------------------------------------------- /hfnet/datasets/aachen.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from pathlib import Path 3 | 4 | from .base_dataset import BaseDataset 5 | from hfnet.settings import DATA_PATH 6 | 7 | 8 | class Aachen(BaseDataset): 9 | default_config = { 10 | 'load_db': True, 11 | 'load_queries': True, 12 | 'image_names': None, 13 | 'grayscale': True, 14 | 'resize_max': 640, 15 | 'num_parallel_calls': 10, 16 | } 17 | dataset_folder = 'aachen/images_upright' 18 | 19 | def _init_dataset(self, **config): 20 | tf.data.Dataset.map_parallel = lambda self, fn: self.map( 21 | fn, num_parallel_calls=config['num_parallel_calls']) 22 | base_path = Path(DATA_PATH, self.dataset_folder) 23 | 24 | if config['image_names'] is not None: 25 | paths = [Path(base_path, n) for n in config['image_names']] 26 | else: 27 | search = [] 28 | if config['load_db']: 29 | search.append(Path(base_path, 'db')) 30 | if config['load_queries']: 31 | search.append(Path(base_path, 'query')) 32 | assert len(search) != 0 33 | paths = [p for s in search for p in s.glob('**/*.jpg')] 34 | 35 | data = {'image': [], 'name': []} 36 | for p in paths: 37 | data['image'].append(p.as_posix()) 38 | rel = p.relative_to(base_path) 39 | data['name'].append(Path(rel.parent, rel.stem).as_posix()) 40 | return data 41 | 42 | def _get_data(self, data, split_name, **config): 43 | def _read_image(path): 44 | image = tf.read_file(path) 45 | image = tf.image.decode_png(image, channels=3) 46 | return image 47 | 48 | def _resize_max(image, resize): 49 | target_size = tf.to_float(tf.convert_to_tensor(resize)) 50 | current_size = tf.to_float(tf.shape(image)[:2]) 51 | scale = target_size / tf.reduce_max(current_size) 52 | new_size = tf.to_int32(current_size * scale) 53 | return tf.image.resize_images( 54 | image, new_size, method=tf.image.ResizeMethod.BILINEAR) 55 | 56 | def _preprocess(data): 57 | image = data['image'] 58 | original_size = tf.shape(image) 59 | tf.Tensor.set_shape(image, [None, None, 3]) 60 | if config['grayscale']: 61 | image = tf.image.rgb_to_grayscale(image) 62 | if config['resize_max']: 63 | image = _resize_max(image, config['resize_max']) 64 | data['image'] = image 65 | data['original_size'] = original_size 66 | return data 67 | 68 | images = tf.data.Dataset.from_tensor_slices(data['image']) 69 | images = images.map_parallel(_read_image) 70 | names = tf.data.Dataset.from_tensor_slices(data['name']) 71 | dataset = tf.data.Dataset.zip({'image': images, 'name': names}) 72 | dataset = dataset.map_parallel(_preprocess) 73 | return dataset 74 | -------------------------------------------------------------------------------- /hfnet/datasets/base_dataset.py: -------------------------------------------------------------------------------- 1 | from abc import ABCMeta, abstractmethod 2 | import tensorflow as tf 3 | 4 | from hfnet.utils.tools import dict_update 5 | 6 | 7 | class BaseDataset(metaclass=ABCMeta): 8 | """Base model class. 9 | 10 | Arguments: 11 | config: A dictionary containing the configuration parameters. 12 | 13 | Datasets should inherit from this class and implement the following methods: 14 | `_init_dataset` and `_get_data`. 15 | Additionally, the following static attributes should be defined: 16 | default_config: A dictionary of potential default configuration values (e.g. the 17 | size of the validation set). 18 | """ 19 | default_split_names = ['training', 'validation', 'test'] 20 | 21 | @abstractmethod 22 | def _init_dataset(self, **config): 23 | """Prepare the dataset for reading. 24 | 25 | This method should configure the dataset for later fetching through `_get_data`, 26 | such as downloading the data if it is not stored locally, or reading the list of 27 | data files from disk. Ideally, especially in the case of large images, this 28 | method shoudl NOT read all the dataset into memory, but rather prepare for faster 29 | seubsequent fetching. 30 | 31 | Arguments: 32 | config: A configuration dictionary, given during the object instantiantion. 33 | 34 | Returns: 35 | An object subsequently passed to `_get_data`, e.g. a list of file paths and 36 | set splits. 37 | """ 38 | raise NotImplementedError 39 | 40 | @abstractmethod 41 | def _get_data(self, dataset, split_name, **config): 42 | """Reads the dataset splits using the Tensorflow `tf.data` API. 43 | 44 | This method should create a `tf.data.Dataset` object for the given data split, 45 | with named components defined through a dictionary mapping strings to tensors. 46 | 47 | It typically performs operations such as reading data from a file or from a 48 | Python generator, shuffling the elements or applying data augmentation to the 49 | training split. It should however NOT batch the dataset (left to the model). 50 | 51 | Arguments: 52 | dataset: An object returned by the `_init_dataset` method. 53 | split_name: A string, the name of the requested split, either `"training"`, 54 | `"validation"` or `"test"`. 55 | config: A configuration dictionary, given during the object instantiantion. 56 | 57 | Returns: 58 | An object of type `tf.data.Dataset` corresponding to the corresponding split. 59 | """ 60 | raise NotImplementedError 61 | 62 | def get_tf_datasets(self): 63 | """"Exposes data splits consistent with the Tensorflow `tf.data` API. 64 | 65 | Returns: 66 | A dictionary mapping split names (`str`, either `"training"`, `"validation"`, 67 | or `"test"`) to `tf.data.Dataset` objects. 68 | """ 69 | return self.tf_splits 70 | 71 | def get_training_set(self): 72 | """Processed training set. 73 | 74 | Returns: 75 | A generator of elements from the training set as dictionaries mapping 76 | component names to the corresponding data (e.g. Numpy array). 77 | """ 78 | return self._get_set_generator('training') 79 | 80 | def get_validation_set(self): 81 | """Processed validation set. 82 | 83 | Returns: 84 | A generator of elements from the training set as dictionaries mapping 85 | component names to the corresponding data (e.g. Numpy array). 86 | """ 87 | return self._get_set_generator('validation') 88 | 89 | def get_test_set(self): 90 | """Processed test set. 91 | 92 | Returns: 93 | A generator of elements from the training set as dictionaries mapping 94 | component names to the corresponding data (e.g. Numpy array). 95 | """ 96 | return self._get_set_generator('test') 97 | 98 | def __init__(self, **config): 99 | # Update config 100 | self.config = dict_update(getattr(self, 'default_config', {}), config) 101 | 102 | self.dataset = self._init_dataset(**self.config) 103 | self.split_names = getattr(self, 'split_names', self.default_split_names) 104 | 105 | self.tf_splits = {} 106 | self.tf_it = {} 107 | self.tf_next = {} 108 | with tf.device('/cpu:0'): 109 | for n in self.split_names: 110 | self.tf_splits[n] = self._get_data(self.dataset, n, **self.config) 111 | prefetched = self.tf_splits[n].prefetch( 112 | self.config.get('prefetch', 1)) 113 | self.tf_it[n] = prefetched.make_initializable_iterator() 114 | self.tf_next[n] = self.tf_it[n].get_next() 115 | self.end_set = tf.errors.OutOfRangeError 116 | 117 | sess_config = tf.ConfigProto() 118 | sess_config.gpu_options.allow_growth = True 119 | self.sess = tf.Session(config=sess_config) 120 | 121 | def _get_set_generator(self, set_name): 122 | self.sess.run(self.tf_it[set_name].initializer) 123 | while True: 124 | try: 125 | yield self.sess.run(self.tf_next[set_name]) 126 | except self.end_set: 127 | return 128 | -------------------------------------------------------------------------------- /hfnet/datasets/cmu.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import random 3 | from pathlib import Path 4 | 5 | from .base_dataset import BaseDataset 6 | from hfnet.settings import DATA_PATH 7 | 8 | 9 | class Cmu(BaseDataset): 10 | default_config = { 11 | 'prefix': '', 12 | 'sequences': [], 13 | 'load_db': True, 14 | 'load_queries': True, 15 | 'image_names': None, 16 | 'resize_max': 640, 17 | 'shuffle': False, 18 | 'grayscale': True, 19 | 'num_parallel_calls': 10, 20 | } 21 | dataset_folder = 'cmu/images' 22 | 23 | def _init_dataset(self, **config): 24 | tf.data.Dataset.map_parallel = lambda self, fn: self.map( 25 | fn, num_parallel_calls=config['num_parallel_calls']) 26 | base_path = Path(DATA_PATH, self.dataset_folder) 27 | 28 | if config['image_names'] is not None: 29 | paths = [Path(base_path, config['prefix'], n) 30 | for n in config['image_names']] 31 | base_path = Path(base_path, config['prefix']) 32 | else: 33 | search = [] 34 | for seq in config['sequences']: 35 | p = Path(base_path, seq) 36 | if config['load_db']: 37 | search.append(Path(p, 'database')) 38 | if config['load_queries']: 39 | search.append(Path(p, 'query')) 40 | assert len(search) > 0 41 | paths = [p for s in search for p in s.glob('**/*.jpg')] 42 | 43 | data = {'image': [], 'name': []} 44 | for p in paths: 45 | data['image'].append(p.as_posix()) 46 | rel = p.relative_to(base_path) 47 | data['name'].append(Path(rel.parent, rel.stem).as_posix()) 48 | if config['shuffle']: 49 | data_list = [dict(zip(data, d)) for d in zip(*data.values())] 50 | random.Random(0).shuffle(data_list) 51 | data = {k: [dic[k] for dic in data_list] for k in data_list[0]} 52 | return data 53 | 54 | def _get_data(self, data, split_name, **config): 55 | def _read_image(path): 56 | image = tf.read_file(path) 57 | image = tf.image.decode_png(image, channels=3) 58 | return image 59 | 60 | def _resize_max(image, resize): 61 | target_size = tf.to_float(tf.convert_to_tensor(resize)) 62 | current_size = tf.to_float(tf.shape(image)[:2]) 63 | scale = target_size / tf.reduce_max(current_size) 64 | new_size = tf.to_int32(current_size * scale) 65 | return tf.image.resize_images( 66 | image, new_size, method=tf.image.ResizeMethod.BILINEAR) 67 | 68 | def _preprocess(data): 69 | image = data['image'] 70 | original_size = tf.shape(image) 71 | tf.Tensor.set_shape(image, [None, None, 3]) 72 | if config['grayscale']: 73 | image = tf.image.rgb_to_grayscale(image) 74 | if config['resize_max']: 75 | image = _resize_max(image, config['resize_max']) 76 | data['image'] = image 77 | data['original_size'] = original_size 78 | return data 79 | 80 | images = tf.data.Dataset.from_tensor_slices(data['image']) 81 | images = images.map_parallel(_read_image) 82 | names = tf.data.Dataset.from_tensor_slices(data['name']) 83 | dataset = tf.data.Dataset.zip({'image': images, 'name': names}) 84 | dataset = dataset.map_parallel(_preprocess) 85 | return dataset 86 | -------------------------------------------------------------------------------- /hfnet/datasets/colmap_utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethz-asl/hfnet/01577845583373470a0cf156a2acd972e9223bc7/hfnet/datasets/colmap_utils/__init__.py -------------------------------------------------------------------------------- /hfnet/datasets/colmap_utils/read_dense.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018, ETH Zurich and UNC Chapel Hill. 2 | # All rights reserved. 3 | # 4 | # Redistribution and use in source and binary forms, with or without 5 | # modification, are permitted provided that the following conditions are met: 6 | # 7 | # * Redistributions of source code must retain the above copyright 8 | # notice, this list of conditions and the following disclaimer. 9 | # 10 | # * Redistributions in binary form must reproduce the above copyright 11 | # notice, this list of conditions and the following disclaimer in the 12 | # documentation and/or other materials provided with the distribution. 13 | # 14 | # * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of 15 | # its contributors may be used to endorse or promote products derived 16 | # from this software without specific prior written permission. 17 | # 18 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 21 | # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE 22 | # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 23 | # CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 24 | # SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 25 | # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 26 | # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 27 | # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 28 | # POSSIBILITY OF SUCH DAMAGE. 29 | # 30 | # Author: Johannes L. Schoenberger (jsch-at-demuc-dot-de) 31 | 32 | import numpy as np 33 | import pylab as plt 34 | 35 | 36 | def read_array(path): 37 | with open(path, "rb") as fid: 38 | width, height, channels = np.genfromtxt(fid, delimiter="&", max_rows=1, 39 | usecols=(0, 1, 2), dtype=int) 40 | fid.seek(0) 41 | num_delimiter = 0 42 | byte = fid.read(1) 43 | while True: 44 | if byte == b"&": 45 | num_delimiter += 1 46 | if num_delimiter >= 3: 47 | break 48 | byte = fid.read(1) 49 | array = np.fromfile(fid, np.float32) 50 | array = array.reshape((width, height, channels), order="F") 51 | return np.transpose(array, (1, 0, 2)).squeeze() 52 | 53 | 54 | def main(): 55 | # Read depth and normal maps corresponding to the same image. 56 | depth_map = read_array( 57 | "path/to/dense/stereo/depth_maps/image1.jpg.photometric.bin") 58 | normal_map = read_array( 59 | "path/to/dense/stereo/normal_maps/image1.jpg.photometric.bin") 60 | 61 | # Visualize the depth map. 62 | min_depth, max_depth = np.percentile(depth_map, [5, 95]) 63 | depth_map[depth_map < min_depth] = min_depth 64 | depth_map[depth_map > max_depth] = max_depth 65 | plt.imshow(depth_map) 66 | plt.show() 67 | 68 | 69 | if __name__ == "__main__": 70 | main() 71 | -------------------------------------------------------------------------------- /hfnet/datasets/hpatches.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import tensorflow as tf 3 | import cv2 4 | from pathlib import Path 5 | 6 | from .base_dataset import BaseDataset 7 | from hfnet.settings import DATA_PATH 8 | 9 | 10 | class Hpatches(BaseDataset): 11 | default_config = { 12 | 'alteration': 'all', # 'all', 'i' for illumination, 'v' for viewpoint 13 | 'truncate': None, 14 | 'make_pairs': False, 15 | 'hard': False, 16 | 'shuffle': False, 17 | 'random_seed': 0, 18 | 'preprocessing': { 19 | 'resize_max': 640, 20 | }, 21 | 'num_parallel_calls': 10, 22 | } 23 | dataset_folder = 'hpatches' 24 | num_images = 6 25 | image_ext = '.ppm' 26 | 27 | def _init_dataset(self, **config): 28 | tf.data.Dataset.map_parallel = lambda self, fn: self.map( 29 | fn, num_parallel_calls=config['num_parallel_calls']) 30 | base_path = Path(DATA_PATH, self.dataset_folder) 31 | scene_paths = sorted([x for x in base_path.iterdir() if x.is_dir()]) 32 | 33 | data = {'image': [], 'name': []} 34 | if config['make_pairs']: 35 | data = {**data, **{k: [] for k in 36 | ['image2', 'name2', 'homography']}} 37 | for path in scene_paths: 38 | if config['alteration'] == 'i' and path.stem[0] != 'i': 39 | continue 40 | if config['alteration'] == 'v' and path.stem[0] != 'v': 41 | continue 42 | for i in range(1, 1 + self.num_images): 43 | if config['make_pairs']: 44 | if i == 1: 45 | path2 = str(Path(path, '1' + self.image_ext)) 46 | name2 = path.stem + '/1' 47 | continue 48 | if config['hard'] and i < self.num_images: 49 | continue 50 | data['image2'].append(path2) 51 | data['name2'].append(name2) 52 | data['homography'].append( 53 | np.loadtxt(str(Path(path, 'H_1_' + str(i))))) 54 | data['image'].append(str(Path(path, str(i) + self.image_ext))) 55 | data['name'].append(path.stem + '/' + str(i)) 56 | 57 | if config['shuffle']: 58 | perm = np.random.RandomState( 59 | config['random_seed']).permutation(len(data['name'])) 60 | data = {k: [v[i] for i in perm] for k, v in data.items()} 61 | if config['truncate']: 62 | data = {k: v[:config['truncate']] for k, v in data.items()} 63 | return data 64 | 65 | def _get_data(self, data, split_name, **config): 66 | def _read_image(path): 67 | return cv2.imread(path.decode('utf-8')) 68 | 69 | def _resize_max(image, resize): 70 | target_size = tf.to_float(tf.convert_to_tensor(resize)) 71 | current_size = tf.to_float(tf.shape(image)[:2]) 72 | scale = target_size / tf.reduce_max(current_size) 73 | new_size = tf.to_int32(current_size * scale) 74 | return tf.image.resize_images( 75 | image, new_size, method=tf.image.ResizeMethod.BILINEAR) 76 | 77 | def _preprocess(image): 78 | tf.Tensor.set_shape(image, [None, None, 3]) 79 | image = tf.image.rgb_to_grayscale(image) 80 | original_size = tf.shape(image)[:2] 81 | if config['preprocessing']['resize_max']: 82 | image = _resize_max( 83 | image, config['preprocessing']['resize_max']) 84 | return tf.to_float(image), original_size 85 | 86 | def _adapt_homography_to_preprocessing(H, data): 87 | size1 = tf.to_float(tf.shape(data['image'])[:2]) 88 | size2 = tf.to_float(tf.shape(data['image2'])[:2]) 89 | s1 = size1 / tf.to_float(data['original_size']) 90 | s2 = size2 / tf.to_float(data['original_size2']) 91 | mult1 = tf.diag(tf.concat([s1, [1.]], 0)) 92 | mult2 = tf.diag(tf.concat([1/s2, [1.]], 0)) 93 | H = tf.matmul(mult1, tf.matmul(tf.to_float(H), mult2)) 94 | return H 95 | 96 | images = tf.data.Dataset.from_tensor_slices(data['image']) 97 | images = images.map_parallel( 98 | lambda path: tf.py_func(_read_image, [path], tf.uint8)) 99 | images = images.map_parallel(_preprocess) 100 | names = tf.data.Dataset.from_tensor_slices(data['name']) 101 | dataset = tf.data.Dataset.zip( 102 | (images, names)).map(lambda i, n: { 103 | 'image': i[0], 'original_size': i[1], 'name': n}) 104 | 105 | if config['make_pairs']: 106 | images2 = tf.data.Dataset.from_tensor_slices(data['image2']) 107 | images2 = images2.map_parallel( 108 | lambda path: tf.py_func(_read_image, [path], tf.uint8)) 109 | images2 = images2.map_parallel(_preprocess) 110 | names2 = tf.data.Dataset.from_tensor_slices(data['name2']) 111 | dataset = tf.data.Dataset.zip( 112 | (dataset, images2, names2)).map( 113 | lambda d, i, n: { 114 | 'image2': i[0], 'name2': n, 115 | 'original_size2': i[1], **d}) 116 | 117 | homographies = tf.data.Dataset.from_tensor_slices( 118 | np.array(data['homography'])) 119 | homographies = tf.data.Dataset.zip((homographies, dataset)) 120 | homographies = homographies.map_parallel( 121 | _adapt_homography_to_preprocessing) 122 | dataset = tf.data.Dataset.zip( 123 | (dataset, homographies)).map( 124 | lambda d, h: {'homography': h, **d}) 125 | 126 | return dataset 127 | -------------------------------------------------------------------------------- /hfnet/datasets/robotcar.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import random 3 | from pathlib import Path 4 | 5 | from .base_dataset import BaseDataset 6 | from hfnet.settings import DATA_PATH 7 | 8 | 9 | class Robotcar(BaseDataset): 10 | default_config = { 11 | 'sequences': [], 12 | 'image_names': None, 13 | 'resize_max': 640, 14 | 'shuffle': False, 15 | 'grayscale': True, 16 | 'num_parallel_calls': 10, 17 | } 18 | dataset_folder = 'robotcar/images' 19 | 20 | def _init_dataset(self, **config): 21 | tf.data.Dataset.map_parallel = lambda self, fn: self.map( 22 | fn, num_parallel_calls=config['num_parallel_calls']) 23 | base_path = Path(DATA_PATH, self.dataset_folder) 24 | 25 | if config['image_names'] is not None: 26 | paths = [Path(base_path, n) for n in config['image_names']] 27 | else: 28 | search = [Path(base_path, s) for s in config['sequences']] 29 | assert len(search) > 0 30 | paths = [p for s in search for p in s.glob('**/*.jpg')] 31 | 32 | data = {'image': [], 'name': []} 33 | for p in paths: 34 | data['image'].append(p.as_posix()) 35 | rel = p.relative_to(base_path) 36 | data['name'].append(Path(rel.parent, rel.stem).as_posix()) 37 | if config['shuffle']: 38 | data_list = [dict(zip(data, d)) for d in zip(*data.values())] 39 | random.Random(0).shuffle(data_list) 40 | data = {k: [dic[k] for dic in data_list] for k in data_list[0]} 41 | return data 42 | 43 | def _get_data(self, data, split_name, **config): 44 | def _read_image(path): 45 | image = tf.read_file(path) 46 | image = tf.image.decode_png(image, channels=3) 47 | return image 48 | 49 | def _resize_max(image, resize): 50 | target_size = tf.to_float(tf.convert_to_tensor(resize)) 51 | current_size = tf.to_float(tf.shape(image)[:2]) 52 | scale = target_size / tf.reduce_max(current_size) 53 | new_size = tf.to_int32(current_size * scale) 54 | return tf.image.resize_images( 55 | image, new_size, method=tf.image.ResizeMethod.BILINEAR) 56 | 57 | def _preprocess(data): 58 | image = data['image'] 59 | original_size = tf.shape(image) 60 | tf.Tensor.set_shape(image, [None, None, 3]) 61 | if config['grayscale']: 62 | image = tf.image.rgb_to_grayscale(image) 63 | if config['resize_max']: 64 | image = _resize_max(image, config['resize_max']) 65 | data['image'] = image 66 | data['original_size'] = original_size 67 | return data 68 | 69 | images = tf.data.Dataset.from_tensor_slices(data['image']) 70 | images = images.map_parallel(_read_image) 71 | names = tf.data.Dataset.from_tensor_slices(data['name']) 72 | dataset = tf.data.Dataset.zip({'image': images, 'name': names}) 73 | dataset = dataset.map_parallel(_preprocess) 74 | return dataset 75 | -------------------------------------------------------------------------------- /hfnet/datasets/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethz-asl/hfnet/01577845583373470a0cf156a2acd972e9223bc7/hfnet/datasets/utils/__init__.py -------------------------------------------------------------------------------- /hfnet/datasets/utils/photometric_augmentation.py: -------------------------------------------------------------------------------- 1 | import cv2 as cv 2 | import numpy as np 3 | import tensorflow as tf 4 | 5 | 6 | augmentations = [ 7 | 'additive_gaussian_noise', 8 | 'additive_speckle_noise', 9 | 'random_brightness', 10 | 'random_contrast', 11 | 'additive_shade', 12 | 'motion_blur' 13 | ] 14 | 15 | 16 | def additive_gaussian_noise(image, stddev_range=[5, 95]): 17 | stddev = tf.random_uniform((), *stddev_range) 18 | noise = tf.random_normal(tf.shape(image), stddev=stddev) 19 | noisy_image = tf.clip_by_value(image + noise, 0, 255) 20 | return noisy_image 21 | 22 | 23 | def additive_speckle_noise(image, prob_range=[0.0, 0.005]): 24 | prob = tf.random_uniform((), *prob_range) 25 | sample = tf.random_uniform(tf.shape(image)) 26 | noisy_image = tf.where(sample <= prob, tf.zeros_like(image), image) 27 | noisy_image = tf.where(sample >= (1. - prob), 255.*tf.ones_like(image), noisy_image) 28 | return noisy_image 29 | 30 | 31 | def random_brightness(image, delta_range=[-50, 50]): 32 | delta = tf.random_uniform((), *delta_range) 33 | return tf.clip_by_value(tf.image.adjust_brightness(image, delta), 0, 255) 34 | 35 | 36 | def random_contrast(image, strength_range=[0.5, 1.5]): 37 | return tf.clip_by_value(tf.image.random_contrast(image, *strength_range), 0, 255) 38 | 39 | 40 | def additive_shade(image, nb_ellipses=20, transparency_range=[-0.5, 0.8], 41 | kernel_size_range=[250, 350], prob=0.5): 42 | 43 | def _py_additive_shade(img): 44 | min_dim = min(img.shape[:2]) / 4 45 | mask = np.zeros(img.shape[:2], np.uint8) 46 | for i in range(nb_ellipses): 47 | ax = int(max(np.random.rand() * min_dim, min_dim / 5)) 48 | ay = int(max(np.random.rand() * min_dim, min_dim / 5)) 49 | max_rad = max(ax, ay) 50 | x = np.random.randint(max_rad, img.shape[1] - max_rad) # center 51 | y = np.random.randint(max_rad, img.shape[0] - max_rad) 52 | angle = np.random.rand() * 90 53 | cv.ellipse(mask, (x, y), (ax, ay), angle, 0, 360, 255, -1) 54 | 55 | transparency = np.random.uniform(*transparency_range) 56 | kernel_size = np.random.randint(*kernel_size_range) 57 | if (kernel_size % 2) == 0: # kernel_size has to be odd 58 | kernel_size += 1 59 | mask = cv.GaussianBlur(mask.astype(np.float32), (kernel_size, kernel_size), 0) 60 | shaded = img * (1 - transparency * mask[..., np.newaxis]/255.) 61 | return np.clip(shaded, 0, 255) 62 | 63 | f = lambda: tf.reshape( 64 | tf.py_func(_py_additive_shade, [image], tf.float32), tf.shape(image)) 65 | rand = tf.random_uniform(()) 66 | return tf.cond(rand < prob, f, lambda: tf.identity(image)) 67 | 68 | 69 | def motion_blur(image, kernel_size=[0, 10], prob=0.5): 70 | 71 | def _py_motion_blur(img): 72 | # Either vertial, hozirontal or diagonal blur 73 | mode = np.random.choice(['h', 'v', 'diag_down', 'diag_up']) 74 | # ksize = np.random.randint(0, (max_kernel_size+1)/2)*2 + 1 # make sure is odd 75 | ksize = np.random.randint(kernel_size[0]//2, (kernel_size[1]+1)//2)*2 + 1 # make sure is odd 76 | center = int((ksize-1)/2) 77 | kernel = np.zeros((ksize, ksize)) 78 | if mode == 'h': 79 | kernel[center, :] = 1. 80 | elif mode == 'v': 81 | kernel[:, center] = 1. 82 | elif mode == 'diag_down': 83 | kernel = np.eye(ksize) 84 | elif mode == 'diag_up': 85 | kernel = np.flip(np.eye(ksize), 0) 86 | var = ksize * ksize / 16. 87 | grid = np.repeat(np.arange(ksize)[:, np.newaxis], ksize, axis=-1) 88 | gaussian = np.exp(-(np.square(grid-center)+np.square(grid.T-center))/(2.*var)) 89 | kernel *= gaussian 90 | kernel /= np.sum(kernel) 91 | img = cv.filter2D(img, -1, kernel) 92 | return img 93 | 94 | f = lambda: tf.reshape( 95 | tf.py_func(_py_motion_blur, [image], tf.float32), tf.shape(image)) 96 | rand = tf.random_uniform(()) 97 | return tf.cond(rand < prob, f, lambda: tf.identity(image)) 98 | -------------------------------------------------------------------------------- /hfnet/datasets/utils/pipeline.py: -------------------------------------------------------------------------------- 1 | """ Augmentation pipeline by Paul-Edouard Sarlin and Rémi Pautrat (ETH Zurich) 2 | for the re-implementation of SuperPoint (DeTone et al.). 3 | Code: github.com/rpautrat/SuperPoint 4 | """ 5 | 6 | import tensorflow as tf 7 | import cv2 as cv 8 | import numpy as np 9 | 10 | from . import photometric_augmentation as photaug 11 | from .homographies import ( 12 | sample_homography, compute_valid_mask, warp_points, filter_points, 13 | adapt_homography_to_resizing) 14 | 15 | 16 | def parse_primitives(names, all_primitives): 17 | p = all_primitives if (names == 'all') \ 18 | else (names if isinstance(names, list) else [names]) 19 | assert set(p) <= set(all_primitives) 20 | return p 21 | 22 | 23 | def photometric_augmentation(data, **config): 24 | with tf.name_scope('photometric_augmentation'): 25 | primitives = parse_primitives( 26 | config['primitives'], photaug.augmentations) 27 | prim_configs = [config['params'].get( 28 | p, {}) for p in primitives] 29 | 30 | indices = tf.range(len(primitives)) 31 | if config['random_order']: 32 | indices = tf.random_shuffle(indices) 33 | 34 | def step(i, image): 35 | fn_pairs = [] 36 | for j, (p, c) in enumerate(zip(primitives, prim_configs)): 37 | fn_pairs.append( 38 | (tf.equal(indices[i], j), 39 | lambda p=p, c=c: getattr(photaug, p)(image, **c))) 40 | image = tf.case(fn_pairs) 41 | return i + 1, image 42 | 43 | _, image = tf.while_loop( 44 | lambda i, image: tf.less(i, len(primitives)), 45 | step, [0, data['image']], parallel_iterations=1) 46 | 47 | return {**data, 'image': image} 48 | 49 | 50 | def homographic_augmentation(data, add_homography=False, **config): 51 | with tf.name_scope('homographic_augmentation'): 52 | image_shape = tf.shape(data['image'])[:2] 53 | homography = sample_homography(image_shape, **config['params'])[0] 54 | warped_image = tf.contrib.image.transform( 55 | data['image'], homography, interpolation='BILINEAR') 56 | valid_mask = compute_valid_mask(image_shape, homography, 57 | config['valid_border_margin']) 58 | 59 | ret = {**data, 'image': warped_image, 'valid_mask': valid_mask} 60 | if 'keypoints' in data: 61 | warped_points = warp_points(data['keypoints'], homography) 62 | warped_points = filter_points(warped_points, image_shape) 63 | ret['keypoints'] = warped_points 64 | for k in ['local_descriptor_map', 'dense_scores']: 65 | if k in data: 66 | data[k].set_shape([None, None, None]) 67 | rescaled_homography = adapt_homography_to_resizing( 68 | homography, image_shape, tf.shape(data[k])) 69 | ret[k] = tf.contrib.image.transform( 70 | data[k], rescaled_homography, interpolation='BILINEAR') 71 | ret[k+'_valid_mask'] = compute_valid_mask( 72 | tf.shape(data[k])[:2], rescaled_homography, 0) 73 | if add_homography: 74 | ret['homography'] = homography 75 | return ret 76 | 77 | 78 | def add_dummy_valid_mask(data): 79 | with tf.name_scope('dummy_valid_mask'): 80 | valid_mask = tf.ones(tf.shape(data['image'])[:2], dtype=tf.int32) 81 | return {**data, 'valid_mask': valid_mask} 82 | 83 | 84 | def add_keypoint_map(data): 85 | with tf.name_scope('add_keypoint_map'): 86 | image_shape = tf.shape(data['image'])[:2] 87 | kp = tf.minimum(tf.to_int32(tf.round(data['keypoints'])), image_shape-1) 88 | kmap = tf.scatter_nd( 89 | kp, tf.ones([tf.shape(kp)[0]], dtype=tf.int32), image_shape) 90 | return {**data, 'keypoint_map': kmap} 91 | 92 | 93 | def downsample(image, coordinates, **config): 94 | with tf.name_scope('gaussian_blur'): 95 | k_size = config['blur_size'] 96 | kernel = cv.getGaussianKernel(k_size, 0)[:, 0] 97 | kernel = np.outer(kernel, kernel).astype(np.float32) 98 | kernel = tf.reshape(tf.convert_to_tensor(kernel), [k_size]*2+[1, 1]) 99 | pad_size = int(k_size/2) 100 | image = tf.pad(image, [[pad_size]*2, [pad_size]*2, [0, 0]], 'REFLECT') 101 | image = tf.expand_dims(image, axis=0) # add batch dim 102 | image = tf.nn.depthwise_conv2d(image, kernel, [1, 1, 1, 1], 'VALID')[0] 103 | 104 | with tf.name_scope('downsample'): 105 | ratio = tf.divide(tf.convert_to_tensor(config['resize']), tf.shape(image)[0:2]) 106 | coordinates = coordinates * tf.cast(ratio, tf.float32) 107 | image = tf.image.resize_images(image, config['resize'], 108 | method=tf.image.ResizeMethod.BILINEAR) 109 | 110 | return image, coordinates 111 | -------------------------------------------------------------------------------- /hfnet/evaluate_aachen.py: -------------------------------------------------------------------------------- 1 | import logging 2 | from pathlib import Path 3 | import argparse 4 | from pprint import pformat 5 | import yaml 6 | import numpy as np 7 | from pyquaternion import Quaternion 8 | 9 | from hfnet.evaluation.localization import Localization, evaluate 10 | from hfnet.evaluation.loaders import export_loader 11 | from hfnet.settings import EXPER_PATH 12 | 13 | 14 | configs_global = { 15 | 'netvlad': { 16 | 'db_name': 'globaldb_netvlad.pkl', 17 | 'experiment': 'netvlad/aachen', 18 | 'predictor': export_loader, 19 | 'has_keypoints': False, 20 | 'has_descriptors': False, 21 | 'pca_dim': 1024, 22 | 'num_prior': 10, 23 | }, 24 | 'hfnet': { 25 | 'db_name': 'globaldb_hfnet.pkl', 26 | 'experiment': 'hfnet/aachen', 27 | 'predictor': export_loader, 28 | 'has_keypoints': False, 29 | 'has_descriptors': False, 30 | 'pca_dim': 1024, 31 | 'num_prior': 10, 32 | }, 33 | } 34 | 35 | configs_local = { 36 | 'superpoint': { 37 | 'db_name': 'localdb_superpoint.pkl', 38 | 'experiment': 'superpoint/aachen', 39 | 'predictor': export_loader, 40 | 'has_keypoints': True, 41 | 'has_descriptors': True, 42 | 'binarize': False, 43 | # 'do_nms': True, 44 | # 'nms_thresh': 4, 45 | 'num_features': 2000, 46 | 'ratio_thresh': 0.9, 47 | }, 48 | 'hfnet': { 49 | 'db_name': 'localdb_hfnet.pkl', 50 | 'experiment': 'hfnet/aachen', 51 | 'predictor': export_loader, 52 | 'has_keypoints': True, 53 | 'has_descriptors': True, 54 | # 'do_nms': True, 55 | # 'nms_thresh': 4, 56 | 'num_features': 2000, 57 | 'ratio_thresh': 0.9, 58 | }, 59 | 'sift': { 60 | 'db_name': 'localdb_sift.pkl', 61 | 'colmap_db': 'aachen.db', 62 | 'root': False, 63 | 'ratio_thresh': 0.7, 64 | 'fast_matching': False, 65 | }, 66 | 'doap': { 67 | 'db_name': 'localdb_doap_kpts-sp-nms4.pkl', 68 | 'experiment': 'doap/aachen_resize-960', 69 | 'predictor': export_loader, 70 | 'keypoint_predictor': export_loader, 71 | 'keypoint_config': { 72 | 'experiment': 'superpoint/aachen', 73 | # 'do_nms': True, 74 | # 'nms_thresh': 4, 75 | }, 76 | 'num_features': 2000, 77 | 'ratio_thresh': 0.9, 78 | }, 79 | 'netvlad': { 80 | 'db_name': 'localdb_netvlad_kpts-sp-nms4.pkl', 81 | 'experiment': 'netvlad/aachen', 82 | 'predictor': export_loader, 83 | 'keypoint_predictor': export_loader, 84 | 'keypoint_config': { 85 | 'experiment': 'superpoint/aachen', 86 | # 'do_nms': True, 87 | # 'nms_thresh': 4, 88 | }, 89 | 'num_features': 1000, 90 | 'ratio_thresh': 0.7, 91 | } 92 | } 93 | 94 | config_pose = { 95 | 'reproj_error': 10, 96 | 'min_inliers': 12, 97 | } 98 | 99 | config_aachen = { 100 | 'resize_max': 960, 101 | } 102 | 103 | 104 | if __name__ == '__main__': 105 | parser = argparse.ArgumentParser() 106 | parser.add_argument('model', type=str) 107 | parser.add_argument('eval_name', type=str) 108 | parser.add_argument('--local_method', type=str) 109 | parser.add_argument('--global_method', type=str) 110 | parser.add_argument('--build_db', action='store_true') 111 | parser.add_argument('--queries', type=str, default='day_time') 112 | parser.add_argument('--max_iter', type=int) 113 | parser.add_argument('--export_poses', action='store_true') 114 | parser.add_argument('--cpp_backend', action='store_true') 115 | args = parser.parse_args() 116 | 117 | config = { 118 | 'global': configs_global[args.global_method], 119 | 'local': configs_local[args.local_method], 120 | 'aachen': config_aachen, 121 | 'pose': config_pose, 122 | 'model': args.model, 123 | 'max_iter': args.max_iter, 124 | 'queries': args.queries, 125 | 'use_cpp': args.cpp_backend, 126 | } 127 | logging.info(f'Configuration: \n'+pformat(config)) 128 | 129 | logging.info('Evaluating Aachen with configuration: ') 130 | loc = Localization('aachen', args.model, config, build_db=args.build_db) 131 | 132 | query_file = f'{args.queries}_queries_with_intrinsics.txt' 133 | queries, query_dataset = loc.init_queries(query_file, config_aachen) 134 | 135 | logging.info('Starting evaluation') 136 | metrics, results = evaluate( 137 | loc, queries, query_dataset, max_iter=args.max_iter) 138 | logging.info('Evaluation metrics: \n'+pformat(metrics)) 139 | 140 | output = {'config': config, 'metrics': metrics} 141 | output_dir = Path(EXPER_PATH, 'eval/aachen') 142 | output_dir.mkdir(exist_ok=True, parents=True) 143 | eval_path = Path(output_dir, f'{args.eval_name}.yaml') 144 | with open(eval_path, 'w') as f: 145 | yaml.dump(output, f, default_flow_style=False) 146 | 147 | if args.export_poses: 148 | poses_path = Path(output_dir, f'{args.eval_name}_poses.txt') 149 | with open(poses_path, 'w') as f: 150 | for query, result in zip(queries, results): 151 | query_T_w = np.linalg.inv(result.T) 152 | qvec_nvm = list(Quaternion(matrix=query_T_w)) 153 | pos_nvm = query_T_w[:3, 3].tolist() 154 | line = f'{Path(query.name).name}' 155 | line += ' ' + ' '.join(map(str, qvec_nvm+pos_nvm)) 156 | f.write(line+'\n') 157 | -------------------------------------------------------------------------------- /hfnet/evaluate_cmu.py: -------------------------------------------------------------------------------- 1 | import logging 2 | from pathlib import Path 3 | import argparse 4 | from pprint import pformat 5 | import yaml 6 | import numpy as np 7 | from pyquaternion import Quaternion 8 | 9 | from hfnet.evaluation.localization import Localization, evaluate 10 | from hfnet.evaluation.loaders import export_loader 11 | from hfnet.settings import EXPER_PATH 12 | 13 | 14 | configs_global = { 15 | 'netvlad': { 16 | 'db_name': 'globaldb_netvlad.pkl', 17 | 'experiment': 'netvlad/cmu', 18 | 'predictor': export_loader, 19 | 'has_keypoints': False, 20 | 'has_descriptors': False, 21 | 'pca_dim': 1024, 22 | 'num_prior': 10, 23 | }, 24 | 'hfnet': { 25 | 'db_name': 'globaldb_hfnet.pkl', 26 | 'experiment': 'hfnet/cmu', 27 | 'predictor': export_loader, 28 | 'has_keypoints': False, 29 | 'has_descriptors': False, 30 | 'pca_dim': 1024, 31 | 'num_prior': 10, 32 | }, 33 | } 34 | 35 | configs_local = { 36 | 'superpoint': { 37 | 'db_name': 'localdb_superpoint.pkl', 38 | 'experiment': 'superpoint/cmu', 39 | 'predictor': export_loader, 40 | 'has_keypoints': True, 41 | 'has_descriptors': True, 42 | 'binarize': False, 43 | # 'do_nms': True, 44 | # 'nms_thresh': 4, 45 | 'num_features': 2000, 46 | 'ratio_thresh': 0.9, 47 | }, 48 | 'sift': { 49 | 'db_name': 'localdb_sift.pkl', 50 | 'colmap_db': 'colmapdb_sift_database.db', 51 | 'colmap_db_queries': 'colmapdb_sift_queries.db', 52 | 'broken_paths': True, 53 | 'root': False, 54 | 'ratio_thresh': 0.7, 55 | 'fast_matching': False, 56 | }, 57 | 'hfnet': { 58 | 'db_name': 'localdb_hfnet.pkl', 59 | 'experiment': 'hfnet/cmu', 60 | 'predictor': export_loader, 61 | 'has_keypoints': True, 62 | 'has_descriptors': True, 63 | # 'do_nms': True, 64 | # 'nms_thresh': 4, 65 | 'num_features': 2000, 66 | 'ratio_thresh': 0.9, 67 | }, 68 | } 69 | 70 | config_pose = { 71 | 'reproj_error': 5, 72 | 'min_inliers': 12, 73 | } 74 | 75 | config_cmu = { 76 | 'name': 'cmu', 77 | 'resize_max': 1024, 78 | } 79 | 80 | 81 | if __name__ == '__main__': 82 | parser = argparse.ArgumentParser() 83 | parser.add_argument('model', type=str) 84 | parser.add_argument('eval_name', type=str) 85 | parser.add_argument('--local_method', type=str, required=True) 86 | parser.add_argument('--global_method', type=str, required=True) 87 | parser.add_argument('--build_db', action='store_true') 88 | parser.add_argument('--slice', type=str, required=True) 89 | parser.add_argument('--max_iter', type=int) 90 | parser.add_argument('--export_poses', action='store_true') 91 | parser.add_argument('--cpp_backend', action='store_true') 92 | args = parser.parse_args() 93 | 94 | config = { 95 | 'global': configs_global[args.global_method], 96 | 'local': configs_local[args.local_method], 97 | 'cmu': {**config_cmu, 'prefix': args.slice}, 98 | 'pose': config_pose, 99 | 'model': args.model, 100 | 'max_iter': args.max_iter, 101 | 'slice': args.slice, 102 | 'use_cpp': args.cpp_backend, 103 | } 104 | for i in ['local', 'global']: 105 | if 'experiment' in config[i]: 106 | config[i]['experiment'] += '/' + args.slice 107 | 108 | name = f'cmu/{args.slice}' 109 | logging.info(f'Evaluating {name} with configuration: \n'+pformat(config)) 110 | loc = Localization(name, config['model'], config, build_db=args.build_db) 111 | 112 | query_file = f'{args.slice}.queries_with_intrinsics.txt' 113 | queries, query_dataset = loc.init_queries(query_file, config['cmu']) 114 | 115 | logging.info('Starting evaluation') 116 | metrics, results = evaluate( 117 | loc, queries, query_dataset, max_iter=args.max_iter) 118 | logging.info('Evaluation metrics: \n'+pformat(metrics)) 119 | 120 | output = {'config': config, 'metrics': metrics} 121 | output_dir = Path(EXPER_PATH, 'eval/cmu') 122 | output_dir.mkdir(exist_ok=True, parents=True) 123 | eval_filename = f'{args.eval_name}_{args.slice}' 124 | eval_path = Path(output_dir, f'{eval_filename}.yaml') 125 | with open(eval_path, 'w') as f: 126 | yaml.dump(output, f, default_flow_style=False) 127 | 128 | if args.export_poses: 129 | poses_path = Path(output_dir, f'{eval_filename}_poses.txt') 130 | with open(poses_path, 'w') as f: 131 | for query, result in zip(queries, results): 132 | query_T_w = np.linalg.inv(result.T) 133 | qvec_nvm = list(Quaternion(matrix=query_T_w)) 134 | pos_nvm = query_T_w[:3, 3].tolist() 135 | name = '/'.join(query.name.split('/')[-2:]) 136 | line = name + ' ' + ' '.join(map(str, qvec_nvm+pos_nvm)) 137 | f.write(line+'\n') 138 | -------------------------------------------------------------------------------- /hfnet/evaluate_robotcar.py: -------------------------------------------------------------------------------- 1 | import logging 2 | from pathlib import Path 3 | import argparse 4 | from pprint import pformat 5 | import yaml 6 | import numpy as np 7 | from pyquaternion import Quaternion 8 | 9 | from hfnet.evaluation.localization import Localization, evaluate 10 | from hfnet.evaluation.loaders import export_loader 11 | from hfnet.settings import EXPER_PATH 12 | 13 | 14 | configs_global = { 15 | 'netvlad': { 16 | 'db_name': 'globaldb_netvlad.pkl', 17 | 'experiment': 'netvlad/robotcar', 18 | 'predictor': export_loader, 19 | 'has_keypoints': False, 20 | 'has_descriptors': False, 21 | 'pca_dim': 1024, 22 | 'num_prior': 10, 23 | }, 24 | 'hfnet': { 25 | 'db_name': 'globaldb_hfnet.pkl', 26 | 'experiment': 'hfnet/robotcar', 27 | 'predictor': export_loader, 28 | 'has_keypoints': False, 29 | 'has_descriptors': False, 30 | 'pca_dim': 1024, 31 | 'num_prior': 10, 32 | }, 33 | } 34 | 35 | configs_local = { 36 | 'superpoint': { 37 | 'db_name': 'localdb_superpoint.pkl', 38 | 'experiment': 'superpoint/robotcar', 39 | 'predictor': export_loader, 40 | 'has_keypoints': True, 41 | 'has_descriptors': True, 42 | 'binarize': False, 43 | # 'do_nms': True, 44 | # 'nms_thresh': 4, 45 | 'num_features': 2000, 46 | 'ratio_thresh': 0.9, 47 | }, 48 | 'sift': { 49 | 'db_name': 'localdb_sift.pkl', 50 | 'colmap_db': 'overcast-reference.db', 51 | 'colmap_db_queries': 'query.db', 52 | 'broken_db': True, 53 | 'root': False, 54 | 'ratio_thresh': 0.7, 55 | 'fast_matching': False, 56 | }, 57 | 'hfnet': { 58 | 'db_name': 'localdb_hfnet.pkl', 59 | 'experiment': 'hfnet/robotcar', 60 | 'predictor': export_loader, 61 | 'has_keypoints': True, 62 | 'has_descriptors': True, 63 | # 'do_nms': True, 64 | # 'nms_thresh': 4, 65 | 'num_features': 2000, 66 | 'ratio_thresh': 0.9, 67 | }, 68 | } 69 | 70 | config_pose = { 71 | 'reproj_error': 12, 72 | 'min_inliers': 15, 73 | } 74 | 75 | config_robotcar = { 76 | 'resize_max': 960, 77 | } 78 | 79 | 80 | if __name__ == '__main__': 81 | parser = argparse.ArgumentParser() 82 | parser.add_argument('model', type=str) 83 | parser.add_argument('eval_name', type=str) 84 | parser.add_argument('--local_method', type=str) 85 | parser.add_argument('--global_method', type=str) 86 | parser.add_argument('--build_db', action='store_true') 87 | parser.add_argument('--queries', type=str, default='dusk_left') 88 | parser.add_argument('--max_iter', type=int) 89 | parser.add_argument('--export_poses', action='store_true') 90 | parser.add_argument('--cpp_backend', action='store_true') 91 | args = parser.parse_args() 92 | 93 | config = { 94 | 'global': configs_global[args.global_method], 95 | 'local': configs_local[args.local_method], 96 | 'robotcar': config_robotcar, 97 | 'pose': config_pose, 98 | 'model': args.model, 99 | 'max_iter': args.max_iter, 100 | 'queries': args.queries, 101 | 'use_cpp': args.cpp_backend, 102 | } 103 | logging.info('Evaluating Robotcar with configuration: \n'+pformat(config)) 104 | loc = Localization('robotcar', args.model, config, build_db=args.build_db) 105 | 106 | query_file = f'queries/{args.queries}_queries_with_intrinsics.txt' 107 | queries, query_dataset = loc.init_queries(query_file, config_robotcar) 108 | 109 | logging.info('Starting evaluation') 110 | metrics, results = evaluate( 111 | loc, queries, query_dataset, max_iter=args.max_iter) 112 | logging.info('Evaluation metrics: \n'+pformat(metrics)) 113 | 114 | output = {'config': config, 'metrics': metrics} 115 | output_dir = Path(EXPER_PATH, 'eval/robotcar') 116 | output_dir.mkdir(exist_ok=True, parents=True) 117 | eval_filename = f'{args.eval_name}_{args.queries}' 118 | eval_path = Path(output_dir, f'{eval_filename}.yaml') 119 | with open(eval_path, 'w') as f: 120 | yaml.dump(output, f, default_flow_style=False) 121 | 122 | if args.export_poses: 123 | poses_path = Path(output_dir, f'{eval_filename}_poses.txt') 124 | with open(poses_path, 'w') as f: 125 | for query, result in zip(queries, results): 126 | query_T_w = np.linalg.inv(result.T) 127 | qvec_nvm = list(Quaternion(matrix=query_T_w)) 128 | pos_nvm = query_T_w[:3, 3].tolist() 129 | name = '/'.join(query.name.split('/')[-2:]) 130 | line = name + ' ' + ' '.join(map(str, qvec_nvm+pos_nvm)) 131 | f.write(line+'\n') 132 | -------------------------------------------------------------------------------- /hfnet/evaluation/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethz-asl/hfnet/01577845583373470a0cf156a2acd972e9223bc7/hfnet/evaluation/__init__.py -------------------------------------------------------------------------------- /hfnet/evaluation/cpp_localization.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | import logging 4 | 5 | from .utils.localization import LocResult 6 | 7 | 8 | class CppLocalization: 9 | def __init__(self, db_ids, local_db, global_descriptors, images, points): 10 | import _hloc_cpp 11 | self.hloc = _hloc_cpp.HLoc() 12 | 13 | id_to_idx = {} 14 | old_to_new_kpt = {} 15 | for idx, i in enumerate(db_ids): 16 | keypoints = local_db[i].keypoints.T.astype(np.float32).copy() 17 | local_desc = local_db[i].descriptors.T.astype(np.float32).copy() 18 | global_desc = global_descriptors[idx].astype(np.float32).copy() 19 | # keypoints are NOT undistorted or nomalized 20 | idx = self.hloc.addImage(global_desc, keypoints, local_desc) 21 | 22 | id_to_idx[i] = idx 23 | old_to_new_kpt[i] = { 24 | k: j for j, k 25 | in enumerate(np.where(images[i].point3D_ids >= 0)[0])} 26 | 27 | for i, pt in points.items(): 28 | observations = np.array( 29 | [[id_to_idx[im_id], old_to_new_kpt[im_id][kpt_id]] 30 | for im_id, kpt_id in zip(pt.image_ids, pt.point2D_idxs)], 31 | dtype=np.int32) 32 | self.hloc.add3dPoint( 33 | pt.xyz.astype(np.float32).copy(), observations.copy()) 34 | self.hloc.buildIndex() 35 | 36 | def localize(self, query_info, query_item, global_transf, local_transf): 37 | global_desc = global_transf(query_item.global_desc[np.newaxis])[0] 38 | local_desc = local_transf(query_item.local_desc) 39 | keypoints = cv2.undistortPoints( 40 | query_item.keypoints[np.newaxis], query_info.K, 41 | np.array([query_info.dist, 0, 0, 0]))[0] 42 | 43 | logging.info('Localizing image %s', query_info.name) 44 | ret = self.cpp_backend.localize( 45 | global_desc.astype(np.float32), 46 | keypoints.astype(np.float32).T.copy(), 47 | local_desc.astype(np.float32).T.copy()) 48 | (success, num_components_total, num_components_tested, 49 | last_component_size, num_db_landmarks, num_matches, 50 | num_inliers, num_iters, global_ms, covis_ms, local_ms, pnp_ms) = ret 51 | 52 | result = LocResult(success, num_inliers, 0, np.eye(4)) 53 | stats = { 54 | 'success': success, 55 | 'num_components_total': num_components_total, 56 | 'num_components_tested': num_components_tested, 57 | 'last_component_size': last_component_size, 58 | 'num_db_landmarks': num_db_landmarks, 59 | 'num_matches': num_matches, 60 | 'num_inliers': num_inliers, 61 | 'num_ransac_iters': num_iters, 62 | 'timings': { 63 | 'global': global_ms, 64 | 'covis': covis_ms, 65 | 'local': local_ms, 66 | 'pnp': pnp_ms, 67 | } 68 | } 69 | return (result, stats) 70 | -------------------------------------------------------------------------------- /hfnet/evaluation/image_retrieval.py: -------------------------------------------------------------------------------- 1 | from scipy.spatial import cKDTree 2 | from sklearn.decomposition import PCA 3 | import numpy as np 4 | 5 | from .utils.descriptors import normalize 6 | 7 | 8 | def is_gt_match_3D(query_poses, ref_poses, distance_thresh, angle_thresh): 9 | distances = np.linalg.norm(np.expand_dims(query_poses['pos'], axis=1) 10 | - np.expand_dims(ref_poses['pos'], axis=0), axis=-1) 11 | angle_errors = np.arccos( 12 | (np.trace( 13 | np.matmul(np.expand_dims(np.linalg.inv(query_poses['rot']), axis=1), 14 | np.expand_dims(ref_poses['rot'], axis=0)), 15 | axis1=2, axis2=3) - 1)/2) 16 | return np.logical_and(distances < distance_thresh, angle_errors < angle_thresh) 17 | 18 | 19 | def is_gt_match_2D(query_poses, ref_poses, distance_thresh, angle_thresh): 20 | distances = np.linalg.norm( 21 | np.expand_dims([query_poses['x'], query_poses['y']], axis=2) 22 | - np.expand_dims([ref_poses['x'], ref_poses['y']], axis=1), axis=0) 23 | angle_errors = np.abs(np.mod(np.expand_dims(query_poses['angle'], axis=1) 24 | - np.expand_dims(ref_poses['angle'], axis=0) + np.pi, 25 | 2*np.pi) - np.pi) # bring it in [-pi,+pi] 26 | return np.logical_and(distances < distance_thresh, angle_errors < angle_thresh) 27 | 28 | 29 | def retrieval(ref_descriptors, query_descriptors, max_num_nn, pca_dim=0): 30 | if pca_dim != 0: 31 | pca = PCA(n_components=pca_dim) 32 | ref_descriptors = normalize(pca.fit_transform(normalize(ref_descriptors))) 33 | query_descriptors = normalize(pca.transform(normalize(query_descriptors))) 34 | 35 | ref_tree = cKDTree(ref_descriptors) 36 | _, indices = ref_tree.query(query_descriptors, k=max_num_nn) 37 | return indices 38 | 39 | 40 | def compute_tp_fp(ref_descriptors, query_descriptors, 41 | gt_matches, *arg, **kwarg): 42 | indices = retrieval(ref_descriptors, query_descriptors, *arg, **kwarg) 43 | tp = gt_matches[np.expand_dims(np.arange(len(indices)), axis=1), indices] 44 | fp = np.logical_not(tp) 45 | tp_cum = np.cumsum(tp, axis=1) 46 | fp_cum = np.cumsum(fp, axis=1) 47 | valid = np.any(gt_matches, axis=1) 48 | return tp_cum, fp_cum, valid 49 | 50 | 51 | def compute_recall(*arg, **kwarg): 52 | tp, fp, valid = compute_tp_fp(*arg, **kwarg) 53 | return np.mean(tp[valid] > 0, axis=0) 54 | -------------------------------------------------------------------------------- /hfnet/evaluation/keypoint_detectors.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from tqdm import tqdm 3 | 4 | from .utils.keypoints import keypoints_warp_2D, keypoints_warp_3D 5 | from .utils.misc import div0 6 | from .utils.metrics import compute_pr, compute_average_precision 7 | 8 | 9 | def compute_correctness(kpts1, kpts2, kpts1_w, kpts2_w, thresh, mutual=True): 10 | 11 | def compute_correctness_single(kpts, kpts_w): 12 | dist = np.linalg.norm(kpts_w[:, np.newaxis] 13 | - kpts[np.newaxis], axis=-1) 14 | min_dist = np.min(dist, axis=1) 15 | correct = min_dist <= thresh 16 | if mutual: 17 | idx = np.argmin(dist, axis=1) 18 | idx_w = np.argmin(dist, axis=0) 19 | correct &= np.equal(np.arange(len(kpts_w)), idx_w[idx]) 20 | return min_dist, correct 21 | 22 | dist1, correct1 = compute_correctness_single(kpts2, kpts1_w) 23 | dist2, correct2 = compute_correctness_single(kpts1, kpts2_w) 24 | return correct1, correct2, dist1, dist2 25 | 26 | 27 | def evaluate(data_iter, config, is_2d=True): 28 | num_kpts = [] 29 | loc_error = [] 30 | repeatability = [] 31 | all_tp = [] 32 | all_num_gt = 0 33 | all_scores = [] 34 | 35 | for data in tqdm(data_iter): 36 | shape1 = data['image'].shape[:2][::-1] 37 | shape2 = data['image2'].shape[:2][::-1] 38 | pred1 = config['predictor']( 39 | data['image'], data['name'], **config) 40 | pred2 = config['predictor']( 41 | data['image2'], data['name2'], **config) 42 | 43 | num_kpts.extend([len(pred1['keypoints']), len(pred2['keypoints'])]) 44 | if len(pred1['keypoints']) == 0 or len(pred2['keypoints']) == 0: 45 | repeatability.append(0) 46 | continue 47 | 48 | if is_2d: 49 | H = data['homography'] 50 | kpts1_w, vis1 = keypoints_warp_2D( 51 | pred1['keypoints'], np.linalg.inv(H), shape2) 52 | kpts2_w, vis2 = keypoints_warp_2D( 53 | pred2['keypoints'], H, shape1) 54 | else: 55 | kpts1_w, vis1, _ = keypoints_warp_3D( 56 | pred1['keypoints'], data['depth'], data['K'], 57 | data['K2'], np.linalg.inv(data['1_T_2']), shape2, 58 | depth2=data['depth2'], consistency_check=True) 59 | kpts2_w, vis2, _ = keypoints_warp_3D( 60 | pred2['keypoints'], data['depth2'], data['K2'], 61 | data['K'], data['1_T_2'], shape1, 62 | depth2=data['depth'], consistency_check=True) 63 | 64 | correct1, correct2, dist1, dist2 = compute_correctness( 65 | pred1['keypoints'], pred2['keypoints'], kpts1_w, kpts2_w, 66 | config['correct_match_thresh']) 67 | 68 | error1 = dist1[vis1 & correct1] 69 | if len(error1) > 0: 70 | loc_error.append(error1.mean()) 71 | error2 = dist2[vis2 & correct2] 72 | if len(error2) > 0: 73 | loc_error.append(error2.mean()) 74 | 75 | repeat = div0(correct1[vis1].sum() + correct2[vis2].sum(), 76 | vis1.sum() + vis2.sum()) 77 | repeatability.append(repeat) 78 | 79 | all_tp.extend([correct1[vis1], correct2[vis2]]) 80 | all_scores.extend([pred1['scores'][vis1], pred2['scores'][vis2]]) 81 | all_num_gt += vis2.sum() + vis1.sum() 82 | 83 | precision, recall, scores = compute_pr( 84 | np.concatenate(all_tp, 0), np.concatenate(all_scores, 0), all_num_gt, 85 | reverse=True) # confidence is in decreasing order 86 | mAP = compute_average_precision(precision, recall) 87 | 88 | metrics = { 89 | 'average_num_keypoints': np.mean(num_kpts), 90 | 'localization_error': np.mean(loc_error), 91 | 'repeatability': np.mean(repeatability), 92 | 'mAP': mAP, 93 | } 94 | return metrics, precision, recall, scores 95 | -------------------------------------------------------------------------------- /hfnet/evaluation/loaders.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | from pathlib import Path 4 | 5 | from .utils.keypoints import ( 6 | keypoints_filter_borders, nms_fast, keypoints_cv2np) 7 | from .utils.descriptors import sample_descriptors, root_descriptors 8 | from hfnet.settings import EXPER_PATH 9 | 10 | 11 | def sift_loader(image, name, **config): 12 | num_features = config.get('num_features', 0) 13 | do_nms = config.get('do_nms', False) 14 | nms_thresh = config.get('nms_thresh', 4) 15 | do_root = config.get('root', False) 16 | 17 | sift = cv2.xfeatures2d.SIFT_create(contrastThreshold=1e-5) 18 | kpts, desc = sift.detectAndCompute(image.astype(np.uint8), None) 19 | kpts, scores = keypoints_cv2np(kpts) 20 | if do_nms: 21 | keep = nms_fast(kpts, scores, image.shape[:2], nms_thresh) 22 | kpts, scores, desc = kpts[keep], scores[keep], desc[keep] 23 | if num_features: 24 | keep_indices = np.argsort(scores)[::-1][:num_features] 25 | kpts, desc, scores = [i[keep_indices] for i in [kpts, desc, scores]] 26 | if do_root: 27 | desc = root_descriptors(desc) 28 | return {'keypoints': kpts, 'descriptors': desc, 'scores': scores} 29 | 30 | 31 | def fast_loader(image, name, **config): 32 | num_features = config.get('num_features', 0) 33 | do_nms = config.get('do_nms', False) 34 | nms_thresh = config.get('nms_thresh', 4) 35 | 36 | fast = cv2.FastFeatureDetector_create() 37 | kpts = fast.detect(image.astype(np.uint8), None) 38 | kpts, scores = keypoints_cv2np(kpts) 39 | if do_nms: 40 | keep = nms_fast(kpts, scores, image.shape[:2], nms_thresh) 41 | kpts, scores = kpts[keep], scores[keep] 42 | if num_features: 43 | keep_indices = np.argsort(scores)[::-1][:num_features] 44 | kpts, scores = [i[keep_indices] for i in [kpts, scores]] 45 | return {'keypoints': kpts, 'scores': scores} 46 | 47 | 48 | def harris_loader(image, name, **config): 49 | num_features = config.get('num_features', 0) 50 | do_nms = config.get('do_nms', False) 51 | nms_thresh = config.get('nms_thresh', 4) 52 | 53 | detect_map = cv2.cornerHarris(image.astype(np.uint8), 4, 3, 0.04) 54 | kpts = np.where(detect_map > 1e-6) 55 | scores = detect_map[kpts] 56 | kpts = np.stack([kpts[1], kpts[0]], axis=-1) 57 | if do_nms: 58 | keep = nms_fast(kpts, scores, image.shape[:2], nms_thresh) 59 | kpts, scores = kpts[keep], scores[keep] 60 | if num_features: 61 | keep_indices = np.argsort(scores)[::-1][:num_features] 62 | kpts, scores = [i[keep_indices] for i in [kpts, scores]] 63 | return {'keypoints': kpts, 'scores': scores} 64 | 65 | 66 | def export_loader(image, name, experiment, **config): 67 | has_keypoints = config.get('has_keypoints', True) 68 | has_descriptors = config.get('has_descriptors', True) 69 | 70 | num_features = config.get('num_features', 0) 71 | remove_borders = config.get('remove_borders', 0) 72 | keypoint_predictor = config.get('keypoint_predictor', None) 73 | do_nms = config.get('do_nms', False) 74 | nms_thresh = config.get('nms_thresh', 4) 75 | keypoint_refinement = config.get('keypoint_refinement', False) 76 | binarize = config.get('binarize', False) 77 | entries = ['keypoints', 'scores', 'descriptors', 'local_descriptors'] 78 | 79 | name = name.decode('utf-8') if isinstance(name, bytes) else name 80 | path = Path(EXPER_PATH, 'exports', experiment, name+'.npz') 81 | with np.load(path) as p: 82 | pred = {k: v.copy() for k, v in p.items()} 83 | image_shape = image.shape[:2] 84 | if keypoint_predictor: 85 | keypoint_config = config.get('keypoint_config', config) 86 | keypoint_config['keypoint_predictor'] = None 87 | pred_detector = keypoint_predictor( 88 | image, name, **{'experiment': experiment, **keypoint_config}) 89 | pred['keypoints'] = pred_detector['keypoints'] 90 | pred['scores'] = pred_detector['scores'] 91 | elif has_keypoints: 92 | assert 'keypoints' in pred 93 | if remove_borders: 94 | mask = keypoints_filter_borders( 95 | pred['keypoints'], image_shape, remove_borders) 96 | pred = {**pred, 97 | **{k: v[mask] for k, v in pred.items() if k in entries}} 98 | if do_nms: 99 | keep = nms_fast( 100 | pred['keypoints'], pred['scores'], image_shape, nms_thresh) 101 | pred = {**pred, 102 | **{k: v[keep] for k, v in pred.items() if k in entries}} 103 | if keypoint_refinement: 104 | criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 105 | 30, 0.001) 106 | pred['keypoints'] = cv2.cornerSubPix( 107 | image, np.float32(pred['keypoints']), 108 | (3, 3), (-1, -1), criteria) 109 | if num_features: 110 | keep = np.argsort(pred['scores'])[::-1][:num_features] 111 | pred = {**pred, 112 | **{k: v[keep] for k, v in pred.items() if k in entries}} 113 | if has_descriptors: 114 | if 'descriptors' in pred: 115 | pass 116 | elif 'local_descriptors' in pred: 117 | pred['descriptors'] = pred['local_descriptors'] 118 | else: 119 | assert 'local_descriptor_map' in pred 120 | pred['descriptors'] = sample_descriptors( 121 | pred['local_descriptor_map'], pred['keypoints'], image_shape, 122 | input_shape=pred['input_shape'][:2] if 'input_shape' in pred 123 | else None) 124 | if binarize: 125 | pred['descriptors'] = pred['descriptors'] > 0 126 | return pred 127 | -------------------------------------------------------------------------------- /hfnet/evaluation/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethz-asl/hfnet/01577845583373470a0cf156a2acd972e9223bc7/hfnet/evaluation/utils/__init__.py -------------------------------------------------------------------------------- /hfnet/evaluation/utils/descriptors.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | 4 | 5 | def normalize(l, axis=-1): 6 | return np.array(l) / np.linalg.norm(l, axis=axis, keepdims=True) 7 | 8 | 9 | def root_descriptors(d, axis=-1): 10 | return np.sqrt(d / np.sum(d, axis=axis, keepdims=True)) 11 | 12 | 13 | def sample_bilinear(data, points): 14 | # Pad the input data with zeros 15 | data = np.lib.pad( 16 | data, ((1, 1), (1, 1), (0, 0)), "constant", constant_values=0) 17 | points = np.asarray(points) + 1 18 | 19 | x, y = points.T 20 | x0, y0 = points.T.astype(int) 21 | x1, y1 = x0 + 1, y0 + 1 22 | 23 | x0 = np.clip(x0, 0, data.shape[1]-1) 24 | x1 = np.clip(x1, 0, data.shape[1]-1) 25 | y0 = np.clip(y0, 0, data.shape[0]-1) 26 | y1 = np.clip(y1, 0, data.shape[0]-1) 27 | 28 | Ia = data[y0, x0] 29 | Ib = data[y1, x0] 30 | Ic = data[y0, x1] 31 | Id = data[y1, x1] 32 | 33 | wa = (x1-x) * (y1-y) 34 | wb = (x1-x) * (y-y0) 35 | wc = (x-x0) * (y1-y) 36 | wd = (x-x0) * (y-y0) 37 | return (Ia.T*wa).T + (Ib.T*wb).T + (Ic.T*wc).T + (Id.T*wd).T 38 | 39 | 40 | def sample_descriptors(descriptor_map, keypoints, image_shape, 41 | input_shape=None, do_round=True): 42 | '''In some cases, the deep network computing the dense descriptors requires 43 | the input to be divisible by the downsampling rate, and crops the 44 | remaining pixels (PyTorch) or pad the input (Tensorflow). We assume the 45 | PyTorch behavior and round the factor between the network input 46 | (input shape) and the output (desc_shape). The keypoints are assumed to 47 | be at the scale of image_shape. 48 | ''' 49 | fix = np.round if do_round else lambda x: x 50 | image_shape = np.array(image_shape) 51 | desc_shape = np.array(descriptor_map.shape[:-1]) 52 | 53 | if input_shape is not None: 54 | input_shape = np.array(input_shape) 55 | factor = image_shape / input_shape 56 | effective_input_shape = desc_shape * fix(input_shape / desc_shape) 57 | factor = factor * (effective_input_shape - 1) / (desc_shape - 1) 58 | else: 59 | factor = (image_shape - 1) / (desc_shape - 1) 60 | 61 | desc = sample_bilinear(descriptor_map, keypoints/factor[::-1]) 62 | desc = normalize(desc, axis=1) 63 | assert np.all(np.isfinite(desc)) 64 | return desc 65 | 66 | 67 | def matching(desc1, desc2, do_ratio_test=False, cross_check=True): 68 | if desc1.dtype == np.bool and desc2.dtype == np.bool: 69 | desc1, desc2 = np.packbits(desc1, axis=1), np.packbits(desc2, axis=1) 70 | norm = cv2.NORM_HAMMING 71 | else: 72 | desc1, desc2 = np.float32(desc1), np.float32(desc2) 73 | norm = cv2.NORM_L2 74 | 75 | if do_ratio_test: 76 | matches = [] 77 | matcher = cv2.BFMatcher(norm) 78 | for m, n in matcher.knnMatch(desc1, desc2, k=2): 79 | m.distance = 1.0 if (n.distance == 0) else m.distance / n.distance 80 | matches.append(m) 81 | else: 82 | matcher = cv2.BFMatcher(norm, crossCheck=cross_check) 83 | matches = matcher.match(desc1, desc2) 84 | return matches_cv2np(matches) 85 | 86 | 87 | def fast_matching(desc1, desc2, ratio_thresh, labels=None): 88 | '''A fast matching method that matches multiple descriptors simultaneously. 89 | Assumes that descriptors are normalized and can run on GPU if available. 90 | Performs the landmark-aware ratio test if labels are provided. 91 | ''' 92 | import torch 93 | cuda = torch.cuda.is_available() 94 | 95 | desc1, desc2 = torch.from_numpy(desc1), torch.from_numpy(desc2) 96 | if cuda: 97 | desc1, desc2 = desc1.cuda(), desc2.cuda() 98 | 99 | with torch.no_grad(): 100 | dist = 2*(1 - desc1 @ desc2.t()) 101 | dist_nn, ind = dist.topk(2, dim=-1, largest=False) 102 | match_ok = (dist_nn[:, 0] <= (ratio_thresh**2)*dist_nn[:, 1]) 103 | 104 | if labels is not None: 105 | labels = torch.from_numpy(labels) 106 | if cuda: 107 | labels = labels.cuda() 108 | labels_nn = labels[ind] 109 | match_ok |= (labels_nn[:, 0] == labels_nn[:, 1]) 110 | 111 | if match_ok.any(): 112 | matches = torch.stack( 113 | [torch.nonzero(match_ok)[:, 0], ind[match_ok][:, 0]], dim=-1) 114 | else: 115 | matches = ind.new_empty((0, 2)) 116 | 117 | return matches.cpu().numpy() 118 | 119 | 120 | def topk_matching(query, database, k): 121 | '''Retrieve top k matches from a database (shape N x dim) with a single 122 | query. In order to reduce any overhead, use numpy instead of PyTorch 123 | ''' 124 | dist = 2 * (1 - database @ query) 125 | ind = np.argpartition(dist, k)[:k] 126 | ind = ind[np.argsort(dist[ind])] 127 | return ind 128 | 129 | 130 | def matches_cv2np(matches_cv): 131 | matches_np = np.int32([[m.queryIdx, m.trainIdx] for m in matches_cv]) 132 | distances = np.float32([m.distance for m in matches_cv]) 133 | return matches_np.reshape(-1, 2), distances 134 | -------------------------------------------------------------------------------- /hfnet/evaluation/utils/keypoints.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from .misc import from_homogeneous, to_homogeneous 4 | 5 | 6 | def nms_fast(kpts, scores, shape, dist_thresh=4): 7 | grid = np.zeros(shape, dtype=int) 8 | inds = np.zeros(shape, dtype=int) 9 | 10 | inds1 = np.argsort(-scores) # Sort by confidence 11 | kpts_sorted = kpts[inds1] 12 | kpts_sorted = kpts_sorted.round().astype(int) # Rounded corners. 13 | 14 | # Check for edge case of 0 or 1 corners. 15 | if kpts_sorted.shape[0] == 0: 16 | return np.zeros(0, dtype=int) 17 | if kpts_sorted.shape[0] == 1: 18 | return np.zeros((1), dtype=int) 19 | 20 | grid[kpts_sorted[:, 1], kpts_sorted[:, 0]] = 1 21 | inds[kpts_sorted[:, 1], kpts_sorted[:, 0]] = np.arange(len(kpts_sorted)) 22 | pad = dist_thresh 23 | grid = np.pad(grid, [[pad]*2]*2, mode='constant') 24 | 25 | # Iterate through points, highest to lowest conf, suppress neighborhood. 26 | for i, k in enumerate(kpts_sorted): 27 | pt = (k[0]+pad, k[1]+pad) 28 | if grid[pt[1], pt[0]] == 1: 29 | grid[pt[1]-pad:pt[1]+pad+1, pt[0]-pad:pt[0]+pad+1] = 0 30 | grid[pt[1], pt[0]] = -1 31 | 32 | keepy, keepx = np.where(grid == -1) 33 | keepy, keepx = keepy - pad, keepx - pad 34 | inds_keep = inds[keepy, keepx] 35 | scores_keep = scores[inds1][inds_keep] 36 | 37 | inds2 = np.argsort(-scores_keep) 38 | out_inds = inds1[inds_keep[inds2]] 39 | return out_inds 40 | 41 | 42 | def keypoints_cv2np(kpts_cv): 43 | kpts_np = np.array([k.pt for k in kpts_cv]) 44 | scores = np.array([k.response for k in kpts_cv]) 45 | return kpts_np, scores 46 | 47 | 48 | def keypoints_filter_borders(kpts, shape, border): 49 | good = np.all(np.logical_and( 50 | kpts < (np.asarray(shape[::-1]) - border), 51 | kpts >= border), -1) 52 | return good 53 | 54 | 55 | def keypoints_warp_2D(kpts, H, shape): 56 | kpts_w = from_homogeneous(np.dot(to_homogeneous(kpts), np.transpose(H))) 57 | vis = np.all((kpts_w >= 0) & (kpts_w <= (np.array(shape)-1)), axis=-1) 58 | return kpts_w, vis 59 | 60 | 61 | def keypoints_warp_3D(kpts1, depth1, K1, K2, T_1to2, shape2, 62 | consistency_check=False, depth2=None, thresh=0.1): 63 | kpts1_int = np.round(kpts1).astype(int) 64 | depth1_kpts = depth1[kpts1_int[:, 1], kpts1_int[:, 0]] 65 | kpts1_3d_1 = np.dot(to_homogeneous(kpts1), np.linalg.inv(K1).T) 66 | kpts1_3d_1 = depth1_kpts[:, np.newaxis]*kpts1_3d_1 67 | kpts1_3d_2 = from_homogeneous(np.dot(to_homogeneous(kpts1_3d_1), T_1to2.T)) 68 | kpts1_w = from_homogeneous(np.dot(kpts1_3d_2, K2.T)) 69 | 70 | vis = np.all((kpts1_w >= 0) & (kpts1_w <= (np.array(shape2)-1)), axis=-1) 71 | vis &= (depth1_kpts > 0) # visible in SfM 72 | vis &= (kpts1_3d_2[:, -1] > 0) # point in front of the camera 73 | 74 | if consistency_check: 75 | assert depth2 is not None 76 | kpts1_w_int = np.round(kpts1_w[vis]).astype(int) 77 | depth2_kpts = depth2[kpts1_w_int[:, 1], kpts1_w_int[:, 0]] 78 | kpt1_w_z = kpts1_3d_2[vis, -1] 79 | # Consistency of the two depth values for each point 80 | error = np.abs(kpt1_w_z - depth2_kpts) / np.maximum(depth2_kpts, 1e-4) 81 | vis[vis] &= (error < thresh) & (depth2_kpts > 0) 82 | 83 | return kpts1_w, vis, kpts1_3d_1 84 | -------------------------------------------------------------------------------- /hfnet/evaluation/utils/metrics.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from .misc import div0 4 | 5 | 6 | def compute_pr(tp, distances, num_gt, reverse=False): 7 | sort_idx = np.argsort(distances) 8 | if reverse: 9 | sort_idx = sort_idx[::-1] 10 | tp = tp[sort_idx] 11 | distances = distances[sort_idx] 12 | fp = np.logical_not(tp) 13 | 14 | tp_cum = np.cumsum(tp) 15 | fp_cum = np.cumsum(fp) 16 | recall = div0(tp_cum, num_gt) 17 | precision = div0(tp_cum, tp_cum + fp_cum) 18 | precision = np.maximum.accumulate(precision[::-1])[::-1] 19 | return precision, recall, distances 20 | 21 | 22 | def compute_average_precision(precision, recall): 23 | return np.sum(precision[1:] * (recall[1:] - recall[:-1])) 24 | -------------------------------------------------------------------------------- /hfnet/evaluation/utils/misc.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def to_homogeneous(points): 5 | return np.concatenate( 6 | [points, np.ones((points.shape[0], 1), dtype=points.dtype)], axis=-1) 7 | 8 | 9 | def from_homogeneous(points): 10 | return points[:, :-1] / points[:, -1:] 11 | 12 | 13 | def angle_error(R1, R2): 14 | cos = (np.trace(np.dot(np.linalg.inv(R1), R2)) - 1) / 2 15 | return np.rad2deg(np.abs(np.arccos(cos))) 16 | 17 | 18 | def div0(a, b): 19 | with np.errstate(divide='ignore', invalid='ignore'): 20 | c = np.true_divide(a, b) 21 | if np.isscalar(c): 22 | c = c if np.isfinite(c) else (1 if a == 0 else 0) 23 | else: 24 | idx = ~np.isfinite(c) 25 | c[idx] = np.where(a[idx] == 0, 1, 0) # -inf inf NaN 26 | return c 27 | -------------------------------------------------------------------------------- /hfnet/evaluation/visualize.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | from sklearn.decomposition import PCA 4 | 5 | 6 | def draw_keypoints(img, kpts, color=(0, 255, 0), radius=4, s=3): 7 | img = np.uint8(img) 8 | if s != 1: 9 | img = cv2.resize(img, None, fx=s, fy=s) 10 | if len(img.shape) == 2: 11 | img = img[..., np.newaxis] 12 | if img.shape[-1] == 1: 13 | img = np.repeat(img, 3, -1) 14 | for k in np.int32(kpts): 15 | cv2.circle(img, tuple(s*k), radius, color, 16 | thickness=-1, lineType=cv2.LINE_AA) 17 | return img 18 | 19 | 20 | def draw_matches(img1, kp1, img2, kp2, matches, color=None, kp_radius=5, 21 | thickness=2, margin=20): 22 | # Create frame 23 | if len(img1.shape) == 2: 24 | img1 = img1[..., np.newaxis] 25 | if len(img2.shape) == 2: 26 | img2 = img2[..., np.newaxis] 27 | if img1.shape[-1] == 1: 28 | img1 = np.repeat(img1, 3, -1) 29 | if img2.shape[-1] == 1: 30 | img2 = np.repeat(img2, 3, -1) 31 | new_shape = (max(img1.shape[0], img2.shape[0]), 32 | img1.shape[1]+img2.shape[1]+margin, 33 | img1.shape[2]) 34 | new_img = np.ones(new_shape, type(img1.flat[0]))*255 35 | 36 | # Place original images 37 | new_img[0:img1.shape[0], 0:img1.shape[1]] = img1 38 | new_img[0:img2.shape[0], 39 | img1.shape[1]+margin:img1.shape[1]+img2.shape[1]+margin] = img2 40 | 41 | # Draw lines between matches 42 | if not isinstance(color, list): 43 | color = [color]*len(matches) 44 | for m, c in zip(matches, color): 45 | # Generate random color for RGB/BGR and grayscale images as needed. 46 | if not c: 47 | if len(img1.shape) == 3: 48 | c = np.random.randint(0, 256, 3) 49 | else: 50 | c = np.random.randint(0, 256) 51 | c = (int(c[0]), int(c[1]), int(c[2])) 52 | 53 | end1 = tuple(np.round(kp1[m[0]]).astype(int)) 54 | end2 = tuple(np.round(kp2[m[1]]).astype(int) 55 | + np.array([img1.shape[1]+margin, 0])) 56 | cv2.line(new_img, end1, end2, c, thickness, lineType=cv2.LINE_AA) 57 | cv2.circle( 58 | new_img, end1, kp_radius, c, thickness, lineType=cv2.LINE_AA) 59 | cv2.circle( 60 | new_img, end2, kp_radius, c, thickness, lineType=cv2.LINE_AA) 61 | return new_img 62 | 63 | 64 | def draw_dense_descriptors(*maps): 65 | pca = PCA(n_components=3, svd_solver='full') 66 | dims = [m.shape[-1] for m in maps] 67 | assert len(np.unique(dims)) == 1 68 | dim = dims[0] 69 | all_maps = np.concatenate([m.reshape(-1, dim) for m in maps]) 70 | pca.fit(all_maps) 71 | projected = [pca.transform(m.reshape(-1, dim)).reshape(m.shape[:2]+(3,)) 72 | for m in maps] 73 | _min = np.min([np.min(p) for p in projected]) 74 | _max = np.max([np.max(p) for p in projected]) 75 | return [(p - _min) / (_max - _min) for p in projected] 76 | -------------------------------------------------------------------------------- /hfnet/export_model.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import yaml 3 | import argparse 4 | from pathlib import Path 5 | from pprint import pformat 6 | 7 | from hfnet.models import get_model 8 | from hfnet.utils import tools # noqa: E402 9 | from hfnet.settings import EXPER_PATH, DATA_PATH 10 | import tensorflow as tf 11 | 12 | 13 | if __name__ == '__main__': 14 | parser = argparse.ArgumentParser() 15 | parser.add_argument('config', type=str) 16 | parser.add_argument('export_name', type=str) 17 | parser.add_argument('--exper_name', type=str) 18 | args = parser.parse_args() 19 | 20 | export_name = args.export_name 21 | exper_name = args.exper_name 22 | 23 | with open(args.config, 'r') as f: 24 | config = yaml.load(f) 25 | 26 | export_dir = Path(EXPER_PATH, 'saved_models', export_name) 27 | 28 | if exper_name: 29 | assert Path(EXPER_PATH, exper_name).exists() 30 | with open(Path(EXPER_PATH, exper_name, 'config.yml'), 'r') as f: 31 | config['model'] = tools.dict_update( 32 | yaml.load(f)['model'], config.get('model', {})) 33 | checkpoint_path = Path(EXPER_PATH, exper_name) 34 | if config.get('weights', None): 35 | checkpoint_path = Path(checkpoint_path, config['weights']) 36 | else: 37 | checkpoint_path = Path(DATA_PATH, 'weights', config['weights']) 38 | logging.info(f'Exporting model with configuration:\n{pformat(config)}') 39 | 40 | with get_model(config['model']['name'])( 41 | data_shape={'image': [None, None, None, 42 | config['model']['image_channels']]}, 43 | **config['model']) as net: 44 | 45 | net.load(str(checkpoint_path)) 46 | 47 | tf.saved_model.simple_save( 48 | net.sess, 49 | str(export_dir), 50 | inputs=net.pred_in, 51 | outputs=net.pred_out) 52 | tf.train.write_graph(net.graph, str(export_dir), 'graph.pbtxt') 53 | -------------------------------------------------------------------------------- /hfnet/export_predictions.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import argparse 3 | import yaml 4 | import logging 5 | from pathlib import Path 6 | from tqdm import tqdm 7 | from pprint import pformat 8 | 9 | logging.basicConfig(format='[%(asctime)s %(levelname)s] %(message)s', 10 | datefmt='%m/%d/%Y %H:%M:%S', 11 | level=logging.INFO) 12 | from hfnet.models import get_model # noqa: E402 13 | from hfnet.datasets import get_dataset # noqa: E402 14 | from hfnet.utils import tools # noqa: E402 15 | from hfnet.settings import EXPER_PATH, DATA_PATH # noqa: E402 16 | 17 | 18 | if __name__ == '__main__': 19 | parser = argparse.ArgumentParser() 20 | parser.add_argument('config', type=str) 21 | parser.add_argument('export_name', type=str) 22 | parser.add_argument('--keys', type=str, default='*') 23 | parser.add_argument('--exper_name', type=str) 24 | parser.add_argument('--as_dataset', action='store_true') 25 | args = parser.parse_args() 26 | 27 | export_name = args.export_name 28 | exper_name = args.exper_name 29 | with open(args.config, 'r') as f: 30 | config = yaml.load(f) 31 | keys = '*' if args.keys == '*' else args.keys.split(',') 32 | 33 | if args.as_dataset: 34 | base_dir = Path(DATA_PATH, export_name) 35 | else: 36 | base_dir = Path(EXPER_PATH, 'exports') 37 | base_dir = Path(base_dir, ((exper_name+'/') if exper_name else '') + export_name) 38 | base_dir.mkdir(parents=True, exist_ok=True) 39 | 40 | if exper_name: 41 | # Update only the model config (not the dataset) 42 | with open(Path(EXPER_PATH, exper_name, 'config.yaml'), 'r') as f: 43 | config['model'] = tools.dict_update( 44 | yaml.load(f)['model'], config.get('model', {})) 45 | checkpoint_path = Path(EXPER_PATH, exper_name) 46 | if config.get('weights', None): 47 | checkpoint_path = Path(checkpoint_path, config['weights']) 48 | else: 49 | if config.get('weights', None): 50 | checkpoint_path = Path(DATA_PATH, 'weights', config['weights']) 51 | else: 52 | checkpoint_path = None 53 | logging.info('No weights provided.') 54 | logging.info(f'Starting export with configuration:\n{pformat(config)}') 55 | 56 | with get_model(config['model']['name'])( 57 | data_shape={'image': [None, None, None, config['model']['image_channels']]}, 58 | **config['model']) as net: 59 | if checkpoint_path is not None: 60 | net.load(str(checkpoint_path)) 61 | dataset = get_dataset(config['data']['name'])(**config['data']) 62 | test_set = dataset.get_test_set() 63 | 64 | for data in tqdm(test_set): 65 | predictions = net.predict(data, keys=keys) 66 | predictions['input_shape'] = data['image'].shape 67 | name = data['name'].decode('utf-8') 68 | Path(base_dir, Path(name).parent).mkdir(parents=True, exist_ok=True) 69 | np.savez(Path(base_dir, '{}.npz'.format(name)), **predictions) 70 | -------------------------------------------------------------------------------- /hfnet/models/__init__.py: -------------------------------------------------------------------------------- 1 | def get_model(name): 2 | mod = __import__('{}.{}'.format(__name__, name), fromlist=['']) 3 | return getattr(mod, _module_to_class(name)) 4 | 5 | 6 | def _module_to_class(name): 7 | return ''.join(n.capitalize() for n in name.split('_')) 8 | -------------------------------------------------------------------------------- /hfnet/models/backbones/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethz-asl/hfnet/01577845583373470a0cf156a2acd972e9223bc7/hfnet/models/backbones/utils/__init__.py -------------------------------------------------------------------------------- /hfnet/models/lfnet_utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethz-asl/hfnet/01577845583373470a0cf156a2acd972e9223bc7/hfnet/models/lfnet_utils/__init__.py -------------------------------------------------------------------------------- /hfnet/models/mobilenetvlad.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from tensorflow.contrib import slim 3 | 4 | from .base_model import BaseModel, Mode 5 | from .backbones import mobilenet_v2 as mobilenet 6 | from .utils.layers import vlad, dimensionality_reduction, image_normalization 7 | 8 | 9 | class Mobilenetvlad(BaseModel): 10 | input_spec = { 11 | 'image': {'shape': [None, None, None, 1], 'type': tf.float32}, 12 | } 13 | required_config_keys = [] 14 | default_config = { 15 | 'depth_multiplier': 1.0, 16 | 'tile_image_if_grayscale': False, 17 | 'dropout_keep_prob': None, 18 | 'encoder_endpoint': 'layer_18', 19 | 'intermediate_proj': None, 20 | 'n_clusters': 64, 21 | 'dimensionality_reduction': None, 22 | 'proj_regularizer': 0., 23 | 'train_backbone': True, 24 | 'train_vlad': True, 25 | 'local_descriptor_layer': None, 26 | } 27 | 28 | def _model(self, inputs, mode, **config): 29 | image = image_normalization(inputs['image']) 30 | if config['tile_image_if_grayscale']: 31 | if image.shape[-1] == 1: 32 | image = tf.tile(image, [1, 1, 1, 3]) 33 | 34 | is_training = config['train_backbone'] and (mode == Mode.TRAIN) 35 | with slim.arg_scope(mobilenet.training_scope( 36 | is_training=is_training, dropout_keep_prob=config['dropout_keep_prob'])): 37 | _, encoder = mobilenet.mobilenet(image, num_classes=None, base_only=True, 38 | depth_multiplier=config['depth_multiplier'], 39 | final_endpoint=config['encoder_endpoint']) 40 | feature_map = encoder[config['encoder_endpoint']] 41 | descriptor = vlad(feature_map, config, mode == Mode.TRAIN) 42 | if config['dimensionality_reduction']: 43 | descriptor = dimensionality_reduction(descriptor, config) 44 | 45 | ret = {'global_descriptor': descriptor} 46 | if config['local_descriptor_layer']: 47 | desc = encoder[config['local_descriptor_layer']] 48 | ret['local_descriptor_map'] = tf.nn.l2_normalize(desc, axis=-1) 49 | return ret 50 | 51 | def _descriptor_l2_error(self, inputs, outputs): 52 | dist = tf.square(inputs['global_descriptor'] 53 | - outputs['global_descriptor']) 54 | return tf.reduce_sum(dist, axis=-1)/2 55 | 56 | def _loss(self, outputs, inputs, **config): 57 | return tf.reduce_mean(self._descriptor_l2_error(inputs, outputs)) 58 | 59 | def _metrics(self, outputs, inputs, **config): 60 | return {'l2_error': self._descriptor_l2_error(inputs, outputs)} 61 | -------------------------------------------------------------------------------- /hfnet/models/netvlad_original.py: -------------------------------------------------------------------------------- 1 | """ Based on the TensorFlow re-implementation of NetVLAD by Titus Cieslewski. 2 | Paper: https://arxiv.org/abs/1511.07247 3 | TensorFlow: github.com/uzh-rpg/netvlad_tf_open 4 | Matlab: github.com/Relja/netvlad 5 | """ 6 | 7 | import tensorflow as tf 8 | import numpy as np 9 | 10 | from .base_model import BaseModel 11 | 12 | 13 | def vlad_legacy(inputs, num_clusters, assign_weight_initializer=None, 14 | cluster_initializer=None, skip_postnorm=False): 15 | K = num_clusters 16 | D = inputs.get_shape()[-1] 17 | 18 | # soft-assignment. 19 | s = tf.layers.conv2d(inputs, K, 1, use_bias=False, 20 | kernel_initializer=assign_weight_initializer, 21 | name='assignment') 22 | a = tf.nn.softmax(s) 23 | 24 | # Dims used hereafter: batch, H, W, desc_coeff, cluster 25 | # Move cluster assignment to corresponding dimension. 26 | a = tf.expand_dims(a, -2) 27 | 28 | # VLAD core. 29 | C = tf.get_variable('cluster_centers', [1, 1, 1, D, K], 30 | initializer=cluster_initializer, 31 | dtype=inputs.dtype) 32 | 33 | v = tf.expand_dims(inputs, -1) + C 34 | v = a * v 35 | v = tf.reduce_sum(v, axis=[1, 2]) 36 | v = tf.transpose(v, perm=[0, 2, 1]) 37 | 38 | if not skip_postnorm: 39 | # Result seems to be very sensitive to the normalization method 40 | # details, so sticking to matconvnet-style normalization here. 41 | v = matconvnetNormalize(v, 1e-12) 42 | v = tf.transpose(v, perm=[0, 2, 1]) 43 | v = matconvnetNormalize(tf.layers.flatten(v), 1e-12) 44 | 45 | return v 46 | 47 | 48 | def matconvnetNormalize(inputs, epsilon): 49 | return inputs / tf.sqrt(tf.reduce_sum(inputs ** 2, axis=-1, keep_dims=True) 50 | + epsilon) 51 | 52 | 53 | class NetvladOriginal(BaseModel): 54 | input_spec = { 55 | 'image': {'shape': [None, None, None, 3], 'type': tf.float32}, 56 | } 57 | required_config_keys = [] 58 | default_config = { 59 | 'num_clusters': 64, 60 | 'pca_dimension': 4096, 61 | 'local_descriptor_layer': None, 62 | } 63 | 64 | def _model(self, inputs, mode, **config): 65 | image_batch = inputs['image'] 66 | 67 | with tf.variable_scope('vgg16_netvlad_pca'): 68 | # Gray to color if necessary. 69 | if image_batch.shape[3] == 1: 70 | x = tf.nn.conv2d(image_batch, np.ones((1, 1, 1, 3)), 71 | np.ones(4).tolist(), 'VALID') 72 | else: 73 | assert image_batch.shape[3] == 3 74 | x = image_batch 75 | 76 | # Subtract trained average image. 77 | average_rgb = tf.get_variable( 78 | 'average_rgb', 3, dtype=image_batch.dtype) 79 | x = x - average_rgb 80 | endpoints = {} 81 | 82 | # VGG16 83 | def vggConv(inputs, numbers, out_dim, with_relu): 84 | activation = tf.nn.relu if with_relu else None 85 | return tf.layers.conv2d( 86 | inputs, out_dim, [3, 3], 1, padding='same', 87 | activation=activation, name='conv%s' % numbers) 88 | 89 | def vggPool(inputs): 90 | return tf.layers.max_pooling2d(inputs, 2, 2) 91 | 92 | x = vggConv(x, '1_1', 64, True) 93 | x = vggConv(x, '1_2', 64, False) 94 | endpoints['conv1_2'] = x 95 | x = vggPool(x) 96 | x = tf.nn.relu(x) 97 | 98 | x = vggConv(x, '2_1', 128, True) 99 | x = vggConv(x, '2_2', 128, False) 100 | endpoints['conv2_2'] = x 101 | x = vggPool(x) 102 | x = tf.nn.relu(x) 103 | 104 | x = vggConv(x, '3_1', 256, True) 105 | x = vggConv(x, '3_2', 256, True) 106 | x = vggConv(x, '3_3', 256, False) 107 | endpoints['conv3_3'] = x 108 | x = vggPool(x) 109 | x = tf.nn.relu(x) 110 | 111 | x = vggConv(x, '4_1', 512, True) 112 | x = vggConv(x, '4_2', 512, True) 113 | x = vggConv(x, '4_3', 512, False) 114 | endpoints['conv4_3'] = x 115 | x = vggPool(x) 116 | x = tf.nn.relu(x) 117 | 118 | x = vggConv(x, '5_1', 512, True) 119 | x = vggConv(x, '5_2', 512, True) 120 | x = vggConv(x, '5_3', 512, False) 121 | endpoints['conv5_3'] = x 122 | 123 | # NetVLAD 124 | x = tf.nn.l2_normalize(x, dim=-1) 125 | x = vlad_legacy(x, config['num_clusters']) 126 | 127 | # PCA 128 | x = tf.layers.conv2d(tf.expand_dims(tf.expand_dims(x, 1), 1), 129 | config['pca_dimension'], 1, 1, name='WPCA') 130 | x = tf.nn.l2_normalize(tf.layers.flatten(x), dim=-1) 131 | 132 | ret = {'global_descriptor': x} 133 | if config['local_descriptor_layer']: 134 | desc = tf.nn.l2_normalize( 135 | endpoints[config['local_descriptor_layer']], axis=-1) 136 | ret['local_descriptor_map'] = desc 137 | return ret 138 | 139 | def _loss(self, outputs, inputs, **config): 140 | raise NotImplementedError 141 | 142 | def _metrics(self, outputs, inputs, **config): 143 | raise NotImplementedError 144 | -------------------------------------------------------------------------------- /hfnet/models/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethz-asl/hfnet/01577845583373470a0cf156a2acd972e9223bc7/hfnet/models/utils/__init__.py -------------------------------------------------------------------------------- /hfnet/models/utils/layers.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from tensorflow.contrib import slim 3 | from tensorflow.python.ops import gen_nn_ops 4 | 5 | 6 | def image_normalization(image, pixel_value_offset=128.0, pixel_value_scale=128.0): 7 | return tf.div(tf.subtract(image, pixel_value_offset), pixel_value_scale) 8 | 9 | 10 | def simple_nms(scores, radius, iterations=3): 11 | """Performs non maximum suppression (NMS) on the heatmap using max-pooling. 12 | This method does not suppress contiguous points that have the same score. 13 | It is an approximate of the standard NMS and uses iterative propagation. 14 | Arguments: 15 | scores: the score heatmap, with shape `[B, H, W]`. 16 | size: an interger scalar, the radius of the NMS window. 17 | """ 18 | with tf.name_scope('simple_nms'): 19 | radius = tf.constant(radius, name='radius') 20 | size = radius*2 + 1 21 | 22 | max_pool = lambda x: gen_nn_ops.max_pool_v2( # supports dynamic ksize 23 | x[..., None], ksize=[1, size, size, 1], 24 | strides=[1, 1, 1, 1], padding='SAME')[..., 0] 25 | zeros = tf.zeros_like(scores) 26 | max_mask = tf.equal(scores, max_pool(scores)) 27 | for _ in range(iterations-1): 28 | supp_mask = tf.cast(max_pool(tf.to_float(max_mask)), tf.bool) 29 | supp_scores = tf.where(supp_mask, zeros, scores) 30 | new_max_mask = tf.equal(supp_scores, max_pool(supp_scores)) 31 | max_mask = max_mask | (new_max_mask & tf.logical_not(supp_mask)) 32 | return tf.where(max_mask, scores, zeros) 33 | 34 | 35 | def delf_attention(feature_map, config, is_training, arg_scope=None): 36 | with tf.variable_scope('attonly/attention/compute'): 37 | with slim.arg_scope(arg_scope): 38 | is_training = config['train_attention'] and is_training 39 | with slim.arg_scope([slim.conv2d, slim.batch_norm], 40 | trainable=is_training): 41 | with slim.arg_scope([slim.batch_norm], is_training=is_training): 42 | attention = slim.conv2d( 43 | feature_map, 512, config['attention_kernel'], rate=1, 44 | activation_fn=tf.nn.relu, scope='conv1') 45 | attention = slim.conv2d( 46 | attention, 1, config['attention_kernel'], rate=1, 47 | activation_fn=None, normalizer_fn=None, scope='conv2') 48 | attention = tf.nn.softplus(attention) 49 | if config['normalize_feature_map']: 50 | feature_map = tf.nn.l2_normalize(feature_map, -1) 51 | descriptor = tf.reduce_sum(feature_map*attention, axis=[1, 2]) 52 | if config['normalize_average']: 53 | descriptor /= tf.reduce_sum(attention, axis=[1, 2]) 54 | return descriptor 55 | 56 | 57 | def vlad(feature_map, config, training, mask=None): 58 | with tf.variable_scope('vlad'): 59 | if config['intermediate_proj']: 60 | with slim.arg_scope([slim.conv2d, slim.batch_norm], trainable=training): 61 | with slim.arg_scope([slim.batch_norm], is_training=training): 62 | feature_map = slim.conv2d( 63 | feature_map, config['intermediate_proj'], 1, rate=1, 64 | activation_fn=None, normalizer_fn=slim.batch_norm, 65 | weights_initializer=slim.initializers.xavier_initializer(), 66 | trainable=training, scope='pre_proj') 67 | 68 | batch_size = tf.shape(feature_map)[0] 69 | feature_dim = feature_map.shape[-1] 70 | 71 | with slim.arg_scope([slim.batch_norm], trainable=training, is_training=training): 72 | memberships = slim.conv2d( 73 | feature_map, config['n_clusters'], 1, rate=1, 74 | activation_fn=None, normalizer_fn=slim.batch_norm, 75 | weights_initializer=slim.initializers.xavier_initializer(), 76 | trainable=training, scope='memberships') 77 | memberships = tf.nn.softmax(memberships, axis=-1) 78 | 79 | clusters = slim.model_variable( 80 | 'clusters', shape=[1, 1, 1, config['n_clusters'], feature_dim], 81 | initializer=slim.initializers.xavier_initializer(), trainable=training) 82 | residuals = clusters - tf.expand_dims(feature_map, axis=3) 83 | residuals *= tf.expand_dims(memberships, axis=-1) 84 | if mask is not None: 85 | residuals *= tf.to_float(mask)[..., tf.newaxis, tf.newaxis] 86 | descriptor = tf.reduce_sum(residuals, axis=[1, 2]) 87 | 88 | descriptor = tf.nn.l2_normalize(descriptor, axis=1) # intra-normalization 89 | descriptor = tf.reshape(descriptor, 90 | [batch_size, feature_dim*config['n_clusters']]) 91 | descriptor = tf.nn.l2_normalize(descriptor, axis=1) 92 | return descriptor 93 | 94 | 95 | def dimensionality_reduction(descriptor, config): 96 | descriptor = tf.nn.l2_normalize(descriptor, -1) 97 | reg = slim.l2_regularizer(config['proj_regularizer']) \ 98 | if config['proj_regularizer'] else None 99 | descriptor = slim.fully_connected( 100 | descriptor, 101 | config['dimensionality_reduction'], 102 | activation_fn=None, 103 | weights_initializer=slim.initializers.xavier_initializer(), 104 | trainable=True, 105 | weights_regularizer=reg, 106 | scope='dimensionality_reduction') 107 | descriptor = tf.nn.l2_normalize(descriptor, -1) 108 | return descriptor 109 | -------------------------------------------------------------------------------- /hfnet/time_model.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import time 4 | import tensorflow as tf 5 | import numpy as np 6 | import cv2 7 | from tensorflow.python.saved_model import tag_constants 8 | 9 | tf.contrib.resampler # import C++ op 10 | 11 | 12 | if __name__ == '__main__': 13 | parser = argparse.ArgumentParser() 14 | parser.add_argument('model', type=str) 15 | parser.add_argument('--image', type=str) 16 | parser.add_argument('--input_size', type=int, default=960) 17 | parser.add_argument('--global_head', action='store_true') 18 | parser.add_argument('--local_head', action='store_true') 19 | parser.add_argument('--iterations', type=int, default=1000) 20 | parser.add_argument('--use_cpu', action='store_true') 21 | args = parser.parse_args() 22 | 23 | b, h, w = 1, args.input_size, int(args.input_size*3/4) 24 | 25 | output_names = [] 26 | if args.global_head: 27 | output_names += ['global_descriptor'] 28 | if args.local_head: 29 | output_names += ['keypoints', 'scores', 'local_descriptors'] 30 | assert len(output_names) > 0 31 | 32 | measure_iter = args.iterations 33 | warmup_iter = 100 34 | 35 | if args.use_cpu: 36 | device = '/cpu:0' 37 | else: 38 | gpus = os.environ['CUDA_VISIBLE_DEVICES'].split(',') 39 | assert len(gpus) > 0 40 | device = '/device:GPU:{}'.format(gpus[0]) 41 | 42 | with tf.Session(graph=tf.Graph()) as sess: 43 | with tf.device(device): 44 | if args.image: 45 | image_np = cv2.resize(cv2.imread(args.image), (w, h)) 46 | image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2GRAY) 47 | image_np = image_np.astype(np.float32)[None, ..., None] 48 | image = tf.constant(image_np) 49 | else: 50 | image = tf.random_uniform([b, h, w, 1], dtype=tf.float32) 51 | tf.saved_model.loader.load( 52 | sess, [tag_constants.SERVING], args.model, 53 | clear_devices=False, input_map={'image:0': image}) 54 | 55 | graph = tf.get_default_graph() 56 | outputs = [graph.get_tensor_by_name(n+':0') for n in output_names] 57 | 58 | for _ in range(warmup_iter): 59 | out = sess.run(outputs) 60 | del out 61 | 62 | start_time = time.time() 63 | for _ in range(measure_iter): 64 | out = sess.run(outputs) 65 | del out 66 | 67 | duration = time.time() - start_time 68 | 69 | print(f'Total: {duration:.4f}, per batch: {duration/measure_iter:.4f}') 70 | -------------------------------------------------------------------------------- /hfnet/train.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import yaml 3 | import os 4 | import argparse 5 | import numpy as np 6 | from contextlib import contextmanager 7 | from json import dumps as pprint 8 | 9 | from hfnet.datasets import get_dataset 10 | from hfnet.models import get_model 11 | from hfnet.utils.stdout_capturing import capture_outputs 12 | from hfnet.settings import EXPER_PATH, DATA_PATH 13 | 14 | logging.basicConfig(format='[%(asctime)s %(levelname)s] %(message)s', 15 | datefmt='%m/%d/%Y %H:%M:%S', 16 | level=logging.INFO) 17 | import tensorflow as tf # noqa: E402 18 | 19 | 20 | def train(config, n_iter, output_dir, checkpoint_name='model.ckpt'): 21 | checkpoint_path = os.path.join(output_dir, checkpoint_name) 22 | with _init_graph(config) as net: 23 | if 'weights' in config: 24 | net.load(os.path.join(DATA_PATH, 'weights', config['weights'])) 25 | elif 'weights_exper' in config: 26 | net.load(os.path.join(EXPER_PATH, config['weights_exper'])) 27 | try: 28 | net.train(n_iter, output_dir=output_dir, 29 | validation_interval=config.get('validation_interval', 100), 30 | save_interval=config.get('save_interval', None), 31 | checkpoint_path=checkpoint_path, 32 | keep_checkpoints=config.get('keep_checkpoints', 1)) 33 | except KeyboardInterrupt: 34 | logging.info('Got Keyboard Interrupt, saving model and closing.') 35 | net.save(checkpoint_path) 36 | 37 | 38 | def set_seed(seed): 39 | tf.set_random_seed(seed) 40 | np.random.seed(seed) 41 | 42 | 43 | @contextmanager 44 | def _init_graph(config, with_dataset=False): 45 | set_seed(config.get('seed', int.from_bytes(os.urandom(4), byteorder='big'))) 46 | n_gpus = len(os.environ['CUDA_VISIBLE_DEVICES'].split(',')) 47 | logging.info('Number of GPUs detected: {}'.format(n_gpus)) 48 | 49 | dataset = get_dataset(config['data']['name'])(**config['data']) 50 | model = get_model(config['model']['name'])( 51 | data=dataset.get_tf_datasets(), n_gpus=n_gpus, **config['model']) 52 | model.__enter__() 53 | if with_dataset: 54 | yield model, dataset 55 | else: 56 | yield model 57 | model.__exit__() 58 | tf.reset_default_graph() 59 | 60 | 61 | def _cli_train(config, output_dir): 62 | assert 'train_iter' in config 63 | 64 | with open(os.path.join(output_dir, 'config.yaml'), 'w') as f: 65 | yaml.dump(config, f, default_flow_style=False) 66 | train(config, config['train_iter'], output_dir) 67 | 68 | 69 | if __name__ == '__main__': 70 | parser = argparse.ArgumentParser() 71 | parser.add_argument('config', type=str) 72 | parser.add_argument('exper_name', type=str) 73 | 74 | args = parser.parse_args() 75 | with open(args.config, 'r') as f: 76 | config = yaml.load(f) 77 | output_dir = os.path.join(EXPER_PATH, args.exper_name) 78 | if not os.path.exists(output_dir): 79 | os.mkdir(output_dir) 80 | 81 | with capture_outputs(os.path.join(output_dir, 'log')): 82 | _cli_train(config, output_dir) 83 | -------------------------------------------------------------------------------- /hfnet/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethz-asl/hfnet/01577845583373470a0cf156a2acd972e9223bc7/hfnet/utils/__init__.py -------------------------------------------------------------------------------- /hfnet/utils/stdout_capturing.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # coding=utf-8 3 | from __future__ import division, print_function, unicode_literals 4 | import os 5 | import sys 6 | import subprocess 7 | from threading import Timer 8 | from contextlib import contextmanager 9 | 10 | ''' 11 | Based on sacred/stdout_capturing.py in project Sacred 12 | https://github.com/IDSIA/sacred 13 | ''' 14 | 15 | 16 | def flush(): 17 | """Try to flush all stdio buffers, both from python and from C.""" 18 | try: 19 | sys.stdout.flush() 20 | sys.stderr.flush() 21 | except (AttributeError, ValueError, IOError): 22 | pass # unsupported 23 | 24 | 25 | # Duplicate stdout and stderr to a file. Inspired by: 26 | # http://eli.thegreenplace.net/2015/redirecting-all-kinds-of-stdout-in-python/ 27 | # http://stackoverflow.com/a/651718/1388435 28 | # http://stackoverflow.com/a/22434262/1388435 29 | @contextmanager 30 | def capture_outputs(filename): 31 | """Duplicate stdout and stderr to a file on the file descriptor level.""" 32 | # with NamedTemporaryFile(mode='w+') as target: 33 | with open(filename, 'a+') as target: 34 | original_stdout_fd = 1 35 | original_stderr_fd = 2 36 | target_fd = target.fileno() 37 | 38 | # Save a copy of the original stdout and stderr file descriptors 39 | saved_stdout_fd = os.dup(original_stdout_fd) 40 | saved_stderr_fd = os.dup(original_stderr_fd) 41 | 42 | tee_stdout = subprocess.Popen( 43 | ['tee', '-a', '/dev/stderr'], start_new_session=True, 44 | stdin=subprocess.PIPE, stderr=target_fd, stdout=1) 45 | tee_stderr = subprocess.Popen( 46 | ['tee', '-a', '/dev/stderr'], start_new_session=True, 47 | stdin=subprocess.PIPE, stderr=target_fd, stdout=2) 48 | 49 | flush() 50 | os.dup2(tee_stdout.stdin.fileno(), original_stdout_fd) 51 | os.dup2(tee_stderr.stdin.fileno(), original_stderr_fd) 52 | 53 | try: 54 | yield 55 | finally: 56 | flush() 57 | 58 | # then redirect stdout back to the saved fd 59 | tee_stdout.stdin.close() 60 | tee_stderr.stdin.close() 61 | 62 | # restore original fds 63 | os.dup2(saved_stdout_fd, original_stdout_fd) 64 | os.dup2(saved_stderr_fd, original_stderr_fd) 65 | 66 | # wait for completion of the tee processes with timeout 67 | # implemented using a timer because timeout support is py3 only 68 | def kill_tees(): 69 | tee_stdout.kill() 70 | tee_stderr.kill() 71 | 72 | tee_timer = Timer(1, kill_tees) 73 | try: 74 | tee_timer.start() 75 | tee_stdout.wait() 76 | tee_stderr.wait() 77 | finally: 78 | tee_timer.cancel() 79 | 80 | os.close(saved_stdout_fd) 81 | os.close(saved_stderr_fd) 82 | -------------------------------------------------------------------------------- /hfnet/utils/tools.py: -------------------------------------------------------------------------------- 1 | import collections 2 | import time 3 | 4 | 5 | def dict_update(d, u): 6 | """Improved update for nested dictionaries. 7 | 8 | Arguments: 9 | d: The dictionary to be updated. 10 | u: The update dictionary. 11 | 12 | Returns: 13 | The updated dictionary. 14 | """ 15 | d = d.copy() 16 | for k, v in u.items(): 17 | if isinstance(v, collections.Mapping): 18 | d[k] = dict_update(d.get(k, {}), v) 19 | else: 20 | d[k] = v 21 | return d 22 | 23 | 24 | class Timer(object): 25 | def __init__(self, name=None): 26 | self.name = name 27 | 28 | def __enter__(self): 29 | self.tstart = time.time() 30 | return self 31 | 32 | def __exit__(self, type, value, traceback): 33 | self.duration = time.time() - self.tstart 34 | if self.name is not None: 35 | print('[%s] Elapsed: %s' % (self.name, self.duration)) 36 | -------------------------------------------------------------------------------- /hloc-cpp/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 2.8.3) 2 | project(hloc_cpp) 3 | 4 | find_package(catkin_simple REQUIRED) 5 | catkin_simple(ALL_DEPS_REQUIRED) 6 | 7 | include_directories(${EIGEN3_INCLUDE_DIR}) 8 | 9 | add_compile_options(-std=c++11 -Ofast -ffast-math -march=native -mavx2) 10 | 11 | pybind_add_module(_hloc_cpp MODULE src/hloc.cc) 12 | 13 | catkin_python_setup() 14 | 15 | cs_install() 16 | cs_export() 17 | 18 | # For Python 3.6 use: 19 | # --cmake-args -DPYBIND_PYTHON_EXECUTABLE=/usr/bin/python3.6 20 | 21 | install(TARGETS _hloc_cpp LIBRARY DESTINATION ${CATKIN_GLOBAL_PYTHON_DESTINATION}) 22 | -------------------------------------------------------------------------------- /hloc-cpp/package.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | hloc_cpp 4 | 0.0.0 5 | 6 | Hierarchical Localization in C++ 7 | 8 | ASL 9 | Closed 10 | 11 | catkin_simple 12 | catkin 13 | 14 | eigen_catkin 15 | gflags_catkin 16 | glog_catkin 17 | pybind11_catkin 18 | opengv 19 | libnabo 20 | 21 | 22 | -------------------------------------------------------------------------------- /hloc-cpp/python/test.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import _hloc_cpp 4 | 5 | 6 | def main(): 7 | local_descriptor_size = 256 8 | global_descriptor_size = 1024 9 | 10 | print('test hloc') 11 | localdescr = np.random.rand(local_descriptor_size, 4).astype(np.float32) 12 | keypoints = np.array([[-1, 1, 0, 0], [0, 0, -1, 1]]).astype(np.float32) # Why other input and transpose doesn't work? 13 | 14 | globaldescr = np.random.rand(1, global_descriptor_size).astype(np.float32) 15 | 16 | hloc = _hloc_cpp.HLoc() 17 | image_idx = hloc.addImage(globaldescr, keypoints, localdescr) 18 | print('Added first image', image_idx) 19 | 20 | for i in range(0, 2): 21 | globaldescr2 = np.random.rand(1, global_descriptor_size).astype(np.float32) 22 | localdescr2 = np.random.rand(local_descriptor_size, 3).astype(np.float32) 23 | keypoints2 = np.random.rand(2, 3).astype(np.float32) 24 | image_idx2 = hloc.addImage(globaldescr2, keypoints2, localdescr2) 25 | print('Added image', image_idx2) 26 | 27 | # Add obs: image 0, kpt 2; image 1, kpt 5 etc. 28 | observing_keyframes = np.array([[0, 0]]).astype(np.int32) 29 | hloc.add3dPoint(np.array((-1, 0, 1)).astype(np.float32), observing_keyframes) 30 | observing_keyframes = np.array([[0, 1]]).astype(np.int32) 31 | hloc.add3dPoint(np.array((1, 0, 1)).astype(np.float32), observing_keyframes) 32 | observing_keyframes = np.array([[0, 2]]).astype(np.int32) 33 | hloc.add3dPoint(np.array((0, -1, 1)).astype(np.float32), observing_keyframes) 34 | observing_keyframes = np.array([[0, 3]]).astype(np.int32) 35 | hloc.add3dPoint(np.array((0, 1, 1)).astype(np.float32), observing_keyframes) 36 | 37 | point_xyz = np.array((1, 2, 3)).astype(np.float32) 38 | 39 | observing_keyframes = np.array([[1, 2], [1, 0], [2, 1]]).astype(np.int32) 40 | hloc.add3dPoint(point_xyz, observing_keyframes) 41 | observing_keyframes = np.array([[2, 2], [1, 1], [2, 0]]).astype(np.int32) 42 | hloc.add3dPoint(point_xyz, observing_keyframes) 43 | 44 | hloc.buildIndex() 45 | 46 | # Check retrieval 47 | ret = hloc.localize(globaldescr, keypoints, localdescr) 48 | (success, num_components_tested, num_inliers, 49 | num_iters, global_ms, covis_ms, local_ms, pnp_ms) = ret 50 | print(success, num_components_tested, num_inliers, num_iters) 51 | print('Timing: ', global_ms, covis_ms, local_ms, pnp_ms) 52 | 53 | 54 | if __name__ == "__main__": 55 | main() 56 | -------------------------------------------------------------------------------- /hloc-cpp/setup.py: -------------------------------------------------------------------------------- 1 | ## ! DO NOT MANUALLY INVOKE THIS setup.py, USE CATKIN INSTEAD 2 | 3 | from distutils.core import setup 4 | from catkin_pkg.python_setup import generate_distutils_setup 5 | 6 | # fetch values from package.xml 7 | setup_args = generate_distutils_setup( 8 | packages=['hloc_cpp'], 9 | package_dir={'': 'src'}) 10 | 11 | setup(**setup_args) 12 | -------------------------------------------------------------------------------- /makefile: -------------------------------------------------------------------------------- 1 | install: 2 | pip3 install -r setup/requirements.txt 3 | pip3 install -e . 4 | sh setup/setup.sh 5 | -------------------------------------------------------------------------------- /notebooks/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethz-asl/hfnet/01577845583373470a0cf156a2acd972e9223bc7/notebooks/__init__.py -------------------------------------------------------------------------------- /notebooks/evaluation_retrieval_nclt.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "import numpy as np\n", 10 | "from pathlib import Path\n", 11 | "from tqdm import tqdm_notebook as tqdm\n", 12 | "import matplotlib.pyplot as plt\n", 13 | "\n", 14 | "from hfnet.evaluation.image_retrieval import compute_recall, is_gt_match_2D\n", 15 | "from hfnet.datasets.nclt import Nclt\n", 16 | "from hfnet.settings import DATA_PATH, EXPER_PATH\n", 17 | "%load_ext autoreload\n", 18 | "%autoreload 2\n", 19 | "%matplotlib inline" 20 | ] 21 | }, 22 | { 23 | "cell_type": "code", 24 | "execution_count": null, 25 | "metadata": {}, 26 | "outputs": [], 27 | "source": [ 28 | "def get_data(seq, experiment):\n", 29 | " im_poses = Nclt.get_pose_file(seq)\n", 30 | " descriptors = []\n", 31 | " for t in im_poses['time']:\n", 32 | " with np.load(Path(EXPER_PATH, 'exports', experiment, seq, f'{t}.npz')) as npz:\n", 33 | " descriptors.append(npz['global_descriptor'].copy())\n", 34 | " return im_poses, np.array(descriptors)\n", 35 | "\n", 36 | "def nclt_recall(ref_seq, query_seqs, experiment, distance_thresh=10, angle_thresh=np.pi/2, *arg, **kwarg):\n", 37 | " ref_poses, ref_descriptors = get_data(ref_seq, experiment)\n", 38 | " query_poses, query_descriptors = [], []\n", 39 | " for s in query_seqs:\n", 40 | " poses, descriptors = get_data(s, experiment)\n", 41 | " query_poses.append(poses)\n", 42 | " query_descriptors.append(descriptors)\n", 43 | " query_poses = np.concatenate(query_poses, axis=0)\n", 44 | " query_descriptors = np.concatenate(query_descriptors, axis=0)\n", 45 | " gt_matches = is_gt_match_2D(query_poses, ref_poses, distance_thresh, angle_thresh)\n", 46 | " return compute_recall(ref_descriptors, query_descriptors, gt_matches, *arg, **kwarg)" 47 | ] 48 | }, 49 | { 50 | "cell_type": "code", 51 | "execution_count": null, 52 | "metadata": {}, 53 | "outputs": [], 54 | "source": [ 55 | "experiments = [\n", 56 | " 'netvlad/nclt', \n", 57 | " # experiments\n", 58 | "]\n", 59 | "ref_seq = '2012-01-08'\n", 60 | "query_seqs = ['2013-02-23', '2012-08-20']\n", 61 | "\n", 62 | "plt.figure(dpi=150)\n", 63 | "colors = plt.rcParams['axes.prop_cycle'].by_key()['color']\n", 64 | "for e, c in zip(experiments, colors):\n", 65 | " m = nclt_recall(ref_seq, query_seqs, e, max_num_nn=30, pca_dim=0)\n", 66 | " plt.plot(1+np.arange(len(m)), 100*m, label=e, color=c, linewidth=1);\n", 67 | " print(f'{e:<70} Recall@10: {m[9]:.3f}')\n", 68 | " \n", 69 | "# m = nclt_recall(ref_seq, query_seqs, e, max_num_nn=30, pca_dim=512)\n", 70 | "# plt.plot(1+np.arange(len(m)), 100*m, label=e+'_proj512', \n", 71 | "# color=c, linewidth=1.3, linestyle='--');\n", 72 | "# print(f'{e+\"_proj512\":<70} Recall@10: {m[9]:.3f}')\n", 73 | "\n", 74 | "plt.xticks([1]+np.arange(10, 31, step=10).tolist()); plt.grid(color=[0.85]*3);\n", 75 | "plt.legend(loc=9, bbox_to_anchor=(0.5, -0.2));\n", 76 | "plt.xlabel('Number of queried neighbors'), plt.ylabel('Recall@N (%)');" 77 | ] 78 | } 79 | ], 80 | "metadata": { 81 | "kernelspec": { 82 | "display_name": "Python 3", 83 | "language": "python", 84 | "name": "python3" 85 | }, 86 | "language_info": { 87 | "codemirror_mode": { 88 | "name": "ipython", 89 | "version": 3 90 | }, 91 | "file_extension": ".py", 92 | "mimetype": "text/x-python", 93 | "name": "python", 94 | "nbconvert_exporter": "python", 95 | "pygments_lexer": "ipython3", 96 | "version": "3.6.1" 97 | } 98 | }, 99 | "nbformat": 4, 100 | "nbformat_minor": 2 101 | } 102 | -------------------------------------------------------------------------------- /notebooks/summarize_timings.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": 2, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "import numpy as np\n", 10 | "from pathlib import Path\n", 11 | "import yaml\n", 12 | "import collections\n", 13 | "from pprint import pprint\n", 14 | "\n", 15 | "from hfnet.settings import EXPER_PATH" 16 | ] 17 | }, 18 | { 19 | "cell_type": "code", 20 | "execution_count": 21, 21 | "metadata": {}, 22 | "outputs": [], 23 | "source": [ 24 | "name = 'TODO'\n", 25 | "with open(Path(EXPER_PATH, 'eval', name), 'r') as f:\n", 26 | " raw_stats = yaml.load(f)['metrics']['all_stats']" 27 | ] 28 | }, 29 | { 30 | "cell_type": "code", 31 | "execution_count": 22, 32 | "metadata": {}, 33 | "outputs": [], 34 | "source": [ 35 | "timings = collections.defaultdict(list)\n", 36 | "warmup = 2\n", 37 | "for s in raw_stats[warmup:]:\n", 38 | " if s is not None:\n", 39 | " for k, v in s['timings'].items():\n", 40 | " timings[k].append(v)\n", 41 | "average = {k: (np.mean(v), np.median(v)) for k, v in timings.items()}\n", 42 | "print(name)\n", 43 | "pprint(average)" 44 | ] 45 | }, 46 | { 47 | "cell_type": "markdown", 48 | "metadata": {}, 49 | "source": [ 50 | "# RobotCar" 51 | ] 52 | }, 53 | { 54 | "cell_type": "code", 55 | "execution_count": 20, 56 | "metadata": {}, 57 | "outputs": [ 58 | { 59 | "name": "stdout", 60 | "output_type": "stream", 61 | "text": [ 62 | "robotcar/02-05/hf_shift-kpts_local-gpu-f32_global-np-f32_night.yaml\n", 63 | "{'covis': (0.01860559822582617, 0.018210291862487793),\n", 64 | " 'global': (0.013169680971924851, 0.013445377349853516),\n", 65 | " 'local': (0.005331631295564698, 0.005107879638671875),\n", 66 | " 'pnp': (0.17651329171366809, 0.16267991065979004)}\n" 67 | ] 68 | } 69 | ], 70 | "source": [] 71 | }, 72 | { 73 | "cell_type": "code", 74 | "execution_count": 24, 75 | "metadata": {}, 76 | "outputs": [ 77 | { 78 | "name": "stdout", 79 | "output_type": "stream", 80 | "text": [ 81 | "robotcar/02-05/hf_shift-kpts_local-gpu-f32_global-np-f32_dusk.yaml\n", 82 | "{'covis': (0.022095190670530673, 0.020515799522399902),\n", 83 | " 'global': (0.01575339062739227, 0.01618969440460205),\n", 84 | " 'local': (0.0041552727505312125, 0.003859877586364746),\n", 85 | " 'pnp': (0.02495658316854703, 0.005520820617675781)}\n" 86 | ] 87 | } 88 | ], 89 | "source": [] 90 | }, 91 | { 92 | "cell_type": "code", 93 | "execution_count": 16, 94 | "metadata": {}, 95 | "outputs": [ 96 | { 97 | "name": "stdout", 98 | "output_type": "stream", 99 | "text": [ 100 | "robotcar/02-05/sp+nv_shift-kpts_local-gpu-f32_global-np-f32_night.yaml\n", 101 | "{'covis': (0.012639381900066284, 0.008560895919799805),\n", 102 | " 'global': (0.013457701882211173, 0.013429522514343262),\n", 103 | " 'local': (0.005822661264640529, 0.00567173957824707),\n", 104 | " 'pnp': (0.21143789116929218, 0.20650744438171387)}\n" 105 | ] 106 | } 107 | ], 108 | "source": [] 109 | }, 110 | { 111 | "cell_type": "code", 112 | "execution_count": 12, 113 | "metadata": {}, 114 | "outputs": [ 115 | { 116 | "name": "stdout", 117 | "output_type": "stream", 118 | "text": [ 119 | "robotcar/02-05/sp+nv_shift-kpts_local-gpu-f32_global-np-f32_dusk.yaml\n", 120 | "{'covis': (0.022481102862600553, 0.02100074291229248),\n", 121 | " 'global': (0.015038933996426857, 0.01377570629119873),\n", 122 | " 'local': (0.0037523904089200293, 0.0035392045974731445),\n", 123 | " 'pnp': (0.015128081733897581, 0.0044841766357421875)}\n" 124 | ] 125 | } 126 | ], 127 | "source": [] 128 | }, 129 | { 130 | "cell_type": "markdown", 131 | "metadata": {}, 132 | "source": [ 133 | "# Aachen" 134 | ] 135 | }, 136 | { 137 | "cell_type": "code", 138 | "execution_count": 84, 139 | "metadata": {}, 140 | "outputs": [ 141 | { 142 | "name": "stdout", 143 | "output_type": "stream", 144 | "text": [ 145 | "aachen/02-05/hf_local-gpu-f32_global-np-f32_night.yaml\n", 146 | "{'covis': (0.06296504040559132, 0.05047762393951416),\n", 147 | " 'global': (0.0060204143325487776, 0.0061827898025512695),\n", 148 | " 'local': (0.00912059098482132, 0.009047865867614746),\n", 149 | " 'pnp': (0.1173677643140157, 0.10982358455657959)}\n" 150 | ] 151 | } 152 | ], 153 | "source": [] 154 | }, 155 | { 156 | "cell_type": "code", 157 | "execution_count": 100, 158 | "metadata": {}, 159 | "outputs": [ 160 | { 161 | "name": "stdout", 162 | "output_type": "stream", 163 | "text": [ 164 | "aachen/02-05/sp+nv_newdb_local-gpu-f32_global-np-f32_day.yaml\n", 165 | "{'covis': (0.06572581758754387, 0.05311107635498047),\n", 166 | " 'global': (0.0071457122654230345, 0.006694316864013672),\n", 167 | " 'local': (0.009518497181634833, 0.009199023246765137),\n", 168 | " 'pnp': (0.0496542801822189, 0.0076904296875)}\n" 169 | ] 170 | } 171 | ], 172 | "source": [] 173 | }, 174 | { 175 | "cell_type": "code", 176 | "execution_count": 96, 177 | "metadata": {}, 178 | "outputs": [ 179 | { 180 | "name": "stdout", 181 | "output_type": "stream", 182 | "text": [ 183 | "aachen/02-05/sp+nv_newdb_local-gpu-f32_global-np-f32_night.yaml\n", 184 | "{'covis': (0.0730412428577741, 0.06047165393829346),\n", 185 | " 'global': (0.007879391312599182, 0.007326364517211914),\n", 186 | " 'local': (0.009806322554747263, 0.009793996810913086),\n", 187 | " 'pnp': (0.08818015704552333, 0.035488009452819824)}\n" 188 | ] 189 | } 190 | ], 191 | "source": [] 192 | } 193 | ], 194 | "metadata": { 195 | "kernelspec": { 196 | "display_name": "Python 3", 197 | "language": "python", 198 | "name": "python3" 199 | }, 200 | "language_info": { 201 | "codemirror_mode": { 202 | "name": "ipython", 203 | "version": 3 204 | }, 205 | "file_extension": ".py", 206 | "mimetype": "text/x-python", 207 | "name": "python", 208 | "nbconvert_exporter": "python", 209 | "pygments_lexer": "ipython3", 210 | "version": "3.6.1" 211 | } 212 | }, 213 | "nbformat": 4, 214 | "nbformat_minor": 2 215 | } 216 | -------------------------------------------------------------------------------- /notebooks/utils.py: -------------------------------------------------------------------------------- 1 | import matplotlib.pyplot as plt 2 | import numpy as np 3 | 4 | 5 | def plot_images(imgs, titles=None, cmap='brg', ylabel='', normalize=True, 6 | r=(0, 1), dpi=100, titlefont=None, labelfont=None, title=None, 7 | keypoints=None, keypoint_colors=(0, 1, 0), keypoints_size=3): 8 | n = len(imgs) 9 | if not isinstance(cmap, list): 10 | cmap = [cmap]*n 11 | if keypoints is not None: 12 | assert len(keypoints) == n 13 | if not isinstance(keypoint_colors[0], (tuple, list)): 14 | keypoint_colors = [keypoint_colors]*n 15 | fig, ax = plt.subplots(1, n, figsize=(6*n, 6), dpi=dpi) 16 | if n == 1: 17 | ax = [ax] 18 | if title is not None: 19 | fig.suptitle(title) 20 | for i in range(n): 21 | if len(imgs[i].shape) == 3: 22 | if imgs[i].shape[-1] == 3: 23 | imgs[i] = imgs[i] 24 | elif imgs[i].shape[-1] == 1: 25 | imgs[i] = imgs[i][..., 0] 26 | if len(imgs[i].shape) == 2 and cmap[i] == 'brg': 27 | cmap[i] = 'gray' 28 | ax[i].imshow(imgs[i], cmap=plt.get_cmap(cmap[i]), 29 | vmin=None if normalize else r[0], 30 | vmax=None if normalize else r[1]) 31 | if keypoints is not None: 32 | ax[i].scatter(keypoints[i][:, 0], keypoints[i][:, 1], 33 | s=keypoints_size, c=keypoint_colors[i]) 34 | if titles: 35 | ax[i].set_title(titles[i], fontsize=titlefont) 36 | ax[i].get_yaxis().set_ticks([]) 37 | ax[i].get_xaxis().set_ticks([]) 38 | for spine in ax[i].spines.values(): # remove frame 39 | spine.set_visible(False) 40 | if ylabel: 41 | ax[0].set_ylabel(ylabel, fontsize=labelfont) 42 | plt.tight_layout() 43 | 44 | 45 | def plot_matches(img1, kp1, img2, kp2, matches, color=None, kp_size=4, 46 | thickness=2.0, margin=20, cmap='brg', ylabel='', 47 | normalize=True, r=(0, 1), dpi=100, title=None): 48 | # Create frame 49 | if len(img1.shape) == 2: 50 | img1 = img1[..., np.newaxis] 51 | if len(img2.shape) == 2: 52 | img2 = img2[..., np.newaxis] 53 | if img1.shape[-1] == 1: 54 | img1 = np.repeat(img1, 3, -1) 55 | if img2.shape[-1] == 1: 56 | img2 = np.repeat(img2, 3, -1) 57 | tile_shape = (max(img1.shape[0], img2.shape[0]), 58 | img1.shape[1]+img2.shape[1]+margin, 59 | img1.shape[2]) 60 | tile = np.ones(tile_shape, type(img1.flat[0])) 61 | if np.max(img1) > 1 or np.max(img2) > 1: 62 | tile *= 255 63 | 64 | # Place original images 65 | tile[0:img1.shape[0], 0:img1.shape[1]] = img1 66 | tile[0:img2.shape[0], 67 | img1.shape[1]+margin:img1.shape[1]+img2.shape[1]+margin] = img2 68 | 69 | # fig, ax = plt.subplots(1, 1, dpi=dpi, frameon=False, figsize=(12, 12)) 70 | fig = plt.figure(frameon=False, dpi=dpi) 71 | w, h, _ = tile.shape 72 | fig_size = 12 73 | fig.set_size_inches(fig_size, fig_size/h*w) 74 | ax = plt.Axes(fig, [0., 0., 1, 1]) 75 | ax.set_axis_off() 76 | fig.add_axes(ax) 77 | 78 | ax.imshow(tile, cmap=plt.get_cmap(cmap), 79 | vmin=None if normalize else r[0], 80 | vmax=None if normalize else r[1], aspect='auto') 81 | 82 | kp1 = kp1[matches[:, 0]] 83 | kp2 = kp2[matches[:, 1]] 84 | kp2[:, 0] += img1.shape[1]+margin 85 | xs = np.stack([kp1[:, 0], kp2[:, 0]], 0) 86 | ys = np.stack([kp1[:, 1], kp2[:, 1]], 0) 87 | 88 | if isinstance(color, list) and len(color) == len(matches): 89 | for i, c in enumerate(color): 90 | ax.plot( 91 | xs[:, i], ys[:, i], linestyle='-', linewidth=thickness, 92 | aa=True, marker='.', markersize=kp_size, color=c, 93 | ) 94 | else: 95 | ax.plot( 96 | xs, ys, linestyle='-', linewidth=thickness, 97 | aa=True, marker='.', markersize=kp_size, color=color, 98 | ) 99 | 100 | ax.get_yaxis().set_ticks([]) 101 | ax.get_xaxis().set_ticks([]) 102 | for spine in ax.spines.values(): # remove frame 103 | spine.set_visible(False) 104 | if title: 105 | ax.set_title(title, fontsize=None) 106 | if ylabel: 107 | ax.set_ylabel(ylabel, fontsize=None) 108 | 109 | 110 | def add_frame(img, c, b=40): 111 | """Add a colored frame around an image with color c and thickness b 112 | """ 113 | img = img.copy() 114 | img[:, :b] = img[:b, :] = img[:, -b:] = img[-b:, :] = c 115 | return img 116 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | setup(name='hfnet', 4 | version="0.0", 5 | packages=['hfnet']) 6 | -------------------------------------------------------------------------------- /setup/requirements.txt: -------------------------------------------------------------------------------- 1 | tensorflow-gpu==1.12 2 | torch==0.4.1 3 | numpy 4 | scipy 5 | opencv-python 6 | tqdm 7 | pyyaml 8 | flake8 9 | jupyter 10 | matplotlib 11 | protobuf 12 | sklearn 13 | pillow 14 | deepdish 15 | -------------------------------------------------------------------------------- /setup/scripts/download_google_landmarks.py: -------------------------------------------------------------------------------- 1 | import multiprocessing 2 | import csv 3 | import tqdm 4 | import argparse 5 | from pathlib import Path 6 | from urllib import request 7 | from PIL import Image 8 | from io import BytesIO 9 | 10 | from hfnet.settings import DATA_PATH 11 | 12 | 13 | def parse_data(data_file, num=None): 14 | csvfile = open(data_file, 'r') 15 | csvreader = csv.reader(csvfile) 16 | key_url_list = [line[:2] for line in csvreader] 17 | key_url_list = key_url_list[1:] # Chop off header 18 | if num is not None: 19 | key_url_list = key_url_list[:num] 20 | return key_url_list 21 | 22 | 23 | def download_image(key_url, output_dir): 24 | (key, url) = key_url 25 | filename = Path(output_dir, '{}.jpg'.format(key)) 26 | 27 | if filename.exists(): 28 | print('Image {} already exists. Skipping download.'.format(filename)) 29 | return 0 30 | 31 | try: 32 | response = request.urlopen(url) 33 | image_data = response.read() 34 | except: 35 | print('Warning: Could not download image {} from {}'.format(key, url)) 36 | return 1 37 | 38 | try: 39 | pil_image = Image.open(BytesIO(image_data)) 40 | except: 41 | print('Warning: Failed to parse image {}'.format(key)) 42 | return 1 43 | 44 | try: 45 | pil_image_rgb = pil_image.convert('RGB') 46 | except: 47 | print('Warning: Failed to convert image {} to RGB'.format(key)) 48 | return 1 49 | 50 | try: 51 | pil_image_rgb.save(str(filename), format='JPEG', quality=90) 52 | except: 53 | print('Warning: Failed to save image {}'.format(filename)) 54 | return 1 55 | 56 | return 0 57 | 58 | 59 | if __name__ == '__main__': 60 | parser = argparse.ArgumentParser() 61 | parser.add_argument('index_file', type=str, help='csv index file') 62 | parser.add_argument('--output_dir', action='store', type=str, 63 | default=str(Path(DATA_PATH, 'google_landmarks/images')), 64 | help='output directory') 65 | parser.add_argument('--truncate', action='store', default=None, type=int) 66 | parser.add_argument('--num_cpus', action='store', default=5, type=int) 67 | args = parser.parse_args() 68 | 69 | Path(args.output_dir).mkdir(parents=True, exist_ok=True) 70 | key_url_list = parse_data(args.index_file, args.truncate) 71 | 72 | def download_func(key_url): 73 | return download_image(key_url, args.output_dir) 74 | 75 | pool = multiprocessing.Pool(processes=args.num_cpus) 76 | failures = sum(tqdm.tqdm(pool.imap_unordered(download_func, key_url_list), 77 | total=len(key_url_list))) 78 | print('Total number of download failures: ', failures) 79 | pool.close() 80 | pool.terminate() 81 | -------------------------------------------------------------------------------- /setup/scripts/generate_robotcar_query_list.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | from tqdm import tqdm 3 | 4 | 5 | if __name__ == '__main__': 6 | images_dir = 'images' 7 | query_list_name = 'queries/{}_queries_with_intrinsics.txt' 8 | intrinsics_name = 'intrinsics/{}_intrinsics.txt' 9 | sequence = 'night-rain' 10 | h, w = 1024, 1024 11 | 12 | intrinsics = {} 13 | for side in ['left', 'right', 'rear']: 14 | with open(intrinsics_name.format(side), 'r') as f: 15 | fx = f.readline().split()[1] 16 | fy = f.readline().split()[1] 17 | cx = f.readline().split()[1] 18 | cy = f.readline().split()[1] 19 | assert fx == fy 20 | params = ['SIMPLE_RADIAL', w, h, fx, cx, cy, 0.0] 21 | intrinsics[side] = [str(p) for p in params] 22 | 23 | query_file = open(query_list_name.format(sequence), 'w') 24 | paths = sorted([p for p in Path(images_dir, sequence).glob('**/*.jpg')]) 25 | for p in tqdm(paths): 26 | name = str(Path(p).relative_to(images_dir)) 27 | side = Path(p).parent.name 28 | query_file.write(' '.join([name]+intrinsics[side])+'\n') 29 | query_file.close() 30 | -------------------------------------------------------------------------------- /setup/scripts/import_superpoint_weights.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from pathlib import Path 3 | import yaml 4 | import torch 5 | import tensorflow as tf 6 | 7 | from hfnet.settings import EXPER_PATH, DATA_PATH 8 | from hfnet.models.super_point import SuperPoint 9 | 10 | 11 | parser = argparse.ArgumentParser() 12 | parser.add_argument('--pytorch_weights', action='store', type=str, 13 | default=str(Path(DATA_PATH, 'weights/superpoint_v1.pth')), 14 | help='original Pytorch weights') 15 | parser.add_argument('--exper_name', action='store', type=str, 16 | default='super_point_original', 17 | help='name of the resulting experiment') 18 | args = parser.parse_args() 19 | print("Converting SuperPoint weights from Pytorch to Tensorflow") 20 | 21 | # Load weights from Torch 22 | weights = torch.load(args.pytorch_weights) 23 | 24 | # Initialize the Tensorflow model 25 | config = { 26 | 'descriptor_size': 256, 27 | 'detection_threshold': 0.015, 28 | 'nms': 4, 29 | } 30 | 31 | output_dir = Path(EXPER_PATH, args.exper_name) 32 | output_dir.mkdir(parents=True, exist_ok=True) 33 | with open(Path(output_dir, 'config.yml'), 'w') as f: 34 | yaml.dump(config, f, default_flow_style=False) 35 | checkpoint_path = Path(output_dir, 'model.ckpt').as_posix() 36 | 37 | net = SuperPoint(data={}, n_gpus=1, **config) 38 | 39 | # Assign the values to the weights 40 | with tf.variable_scope('', reuse=tf.AUTO_REUSE): 41 | # Shared encoder 42 | with tf.variable_scope('vgg', reuse=tf.AUTO_REUSE): 43 | conv1_1 = tf.get_variable('conv1_1/conv/kernel') 44 | conv1_1 = tf.assign(conv1_1, 45 | weights['conv1a.weight'].numpy().transpose([3, 2, 1, 0])) 46 | conv1_1_b = tf.get_variable('conv1_1/conv/bias') 47 | conv1_1_b = tf.assign(conv1_1_b, weights['conv1a.bias'].numpy()) 48 | 49 | conv1_2 = tf.get_variable('conv1_2/conv/kernel') 50 | conv1_2 = tf.assign(conv1_2, 51 | weights['conv1b.weight'].numpy().transpose([3, 2, 1, 0])) 52 | conv1_2_b = tf.get_variable('conv1_2/conv/bias') 53 | conv1_2_b = tf.assign(conv1_2_b, weights['conv1b.bias'].numpy()) 54 | 55 | conv2_1 = tf.get_variable('conv2_1/conv/kernel') 56 | conv2_1 = tf.assign(conv2_1, 57 | weights['conv2a.weight'].numpy().transpose([3, 2, 1, 0])) 58 | conv2_1_b = tf.get_variable('conv2_1/conv/bias') 59 | conv2_1_b = tf.assign(conv2_1_b, weights['conv2a.bias'].numpy()) 60 | 61 | conv2_2 = tf.get_variable('conv2_2/conv/kernel') 62 | conv2_2 = tf.assign(conv2_2, 63 | weights['conv2b.weight'].numpy().transpose([3, 2, 1, 0])) 64 | conv2_2_b = tf.get_variable('conv2_2/conv/bias') 65 | conv2_2_b = tf.assign(conv2_2_b, weights['conv2b.bias'].numpy()) 66 | 67 | conv3_1 = tf.get_variable('conv3_1/conv/kernel') 68 | conv3_1 = tf.assign(conv3_1, 69 | weights['conv3a.weight'].numpy().transpose([3, 2, 1, 0])) 70 | conv3_1_b = tf.get_variable('conv3_1/conv/bias') 71 | conv3_1_b = tf.assign(conv3_1_b, weights['conv3a.bias'].numpy()) 72 | 73 | conv3_2 = tf.get_variable('conv3_2/conv/kernel') 74 | conv3_2 = tf.assign(conv3_2, 75 | weights['conv3b.weight'].numpy().transpose([3, 2, 1, 0])) 76 | conv3_2_b = tf.get_variable('conv3_2/conv/bias') 77 | conv3_2_b = tf.assign(conv3_2_b, weights['conv3b.bias'].numpy()) 78 | 79 | conv4_1 = tf.get_variable('conv4_1/conv/kernel') 80 | conv4_1 = tf.assign(conv4_1, 81 | weights['conv4a.weight'].numpy().transpose([3, 2, 1, 0])) 82 | conv4_1_b = tf.get_variable('conv4_1/conv/bias') 83 | conv4_1_b = tf.assign(conv4_1_b, weights['conv4a.bias'].numpy()) 84 | 85 | conv4_2 = tf.get_variable('conv4_2/conv/kernel') 86 | conv4_2 = tf.assign(conv4_2, 87 | weights['conv4b.weight'].numpy().transpose([3, 2, 1, 0])) 88 | conv4_2_b = tf.get_variable('conv4_2/conv/bias') 89 | conv4_2_b = tf.assign(conv4_2_b, weights['conv4b.bias'].numpy()) 90 | 91 | # Detector head 92 | with tf.variable_scope('detector', reuse=tf.AUTO_REUSE): 93 | convP_1 = tf.get_variable('conv1/conv/kernel') 94 | convP_1 = tf.assign(convP_1, 95 | weights['convPa.weight'].numpy().transpose([3, 2, 1, 0])) 96 | convP_1_b = tf.get_variable('conv1/conv/bias') 97 | convP_1_b = tf.assign(convP_1_b, weights['convPa.bias'].numpy()) 98 | 99 | convP_2 = tf.get_variable('conv2/conv/kernel') 100 | convP_2 = tf.assign(convP_2, 101 | weights['convPb.weight'].numpy().transpose([3, 2, 1, 0])) 102 | convP_2_b = tf.get_variable('conv2/conv/bias') 103 | convP_2_b = tf.assign(convP_2_b, weights['convPb.bias'].numpy()) 104 | 105 | # Descriptor head 106 | with tf.variable_scope('descriptor', reuse=tf.AUTO_REUSE): 107 | convD_1 = tf.get_variable('conv1/conv/kernel') 108 | convD_1 = tf.assign(convD_1, 109 | weights['convDa.weight'].numpy().transpose([3, 2, 1, 0])) 110 | convD_1_b = tf.get_variable('conv1/conv/bias') 111 | convD_1_b = tf.assign(convD_1_b, weights['convDa.bias'].numpy()) 112 | 113 | convD_2 = tf.get_variable('conv2/conv/kernel') 114 | convD_2 = tf.assign(convD_2, 115 | weights['convDb.weight'].numpy().transpose([3, 2, 1, 0])) 116 | convD_2_b = tf.get_variable('conv2/conv/bias') 117 | convD_2_b = tf.assign(convD_2_b, weights['convDb.bias'].numpy()) 118 | 119 | # Run and save the tf model 120 | assign_op = [conv1_1, conv1_1_b, conv2_1, conv2_1_b, conv3_1, conv3_1_b, 121 | conv4_1, conv4_1_b, convP_1, convP_1_b, convD_1, convD_1_b, 122 | conv1_2, conv1_2_b, conv2_2, conv2_2_b, conv3_2, conv3_2_b, 123 | conv4_2, conv4_2_b, convP_2, convP_2_b, convD_2, convD_2_b] 124 | 125 | sess = tf.Session() 126 | sess.run(tf.initialize_all_variables()) 127 | sess.run(assign_op) 128 | saver = tf.train.Saver(save_relative_paths=True) 129 | saver.save(sess, checkpoint_path, write_meta_graph=False) 130 | print("Weights saved to " + checkpoint_path) 131 | -------------------------------------------------------------------------------- /setup/setup.sh: -------------------------------------------------------------------------------- 1 | read -p "Path of the directory where datasets are stored and read: " dir 2 | echo "DATA_PATH = '$dir'" >> ./hfnet/settings.py 3 | 4 | read -p "Path of the directory where experiments data (logs, checkpoints, configs) are written: " dir 5 | echo "EXPER_PATH = '$dir'" >> ./hfnet/settings.py 6 | -------------------------------------------------------------------------------- /utils/env_cluster.sh: -------------------------------------------------------------------------------- 1 | env_dir=./env 2 | python_version="3.6.1" 3 | 4 | module load eth_proxy python_gpu/$python_version cuda/9.0.176 cudnn/7.5 hdf5/1.10.1 5 | 6 | if [ ! -d "$env_dir" ]; then 7 | python -m pip install --user virtualenv 8 | python -m virtualenv \ 9 | --system-site-packages \ 10 | --python="/cluster/apps/python/$python_version/bin/python" \ 11 | "$env_dir" 12 | fi 13 | 14 | source "$env_dir/bin/activate" 15 | -------------------------------------------------------------------------------- /utils/env_cluster_cpu.sh: -------------------------------------------------------------------------------- 1 | env_dir=./env_cpu 2 | python_version="3.6.1" 3 | 4 | module load eth_proxy python_cpu/$python_version hdf5/1.10.1 5 | 6 | if [ ! -d "$env_dir" ]; then 7 | python -m pip install --user virtualenv 8 | python -m virtualenv \ 9 | --system-site-packages \ 10 | --python="/cluster/apps/python/$python_version/bin/python" \ 11 | "$env_dir" 12 | fi 13 | 14 | source "$env_dir/bin/activate" 15 | -------------------------------------------------------------------------------- /utils/launch_cluster.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | cores="15" 4 | memory="14000" # per core 5 | scratch="1000" 6 | gpus="1" 7 | clock="4:00" 8 | model="GeForceGTX1080Ti" 9 | warn="-wt 15 -wa INT" 10 | 11 | cmd="bsub 12 | -n $cores 13 | -W $clock $output 14 | $warn 15 | -R 'select[gpu_model0 == $model] rusage[mem=$memory,scratch=$scratch,ngpus_excl_p=$gpus]' 16 | $*" 17 | echo $cmd 18 | eval $cmd 19 | --------------------------------------------------------------------------------