├── .gitignore ├── .gitmodules ├── LICENSE ├── README.md ├── configs ├── aspanformer.yml ├── caps_sift.yml ├── caps_superpoint.yml ├── cotr.yml ├── d2net.yml ├── dogaffnethardnet.yml ├── loftr.yml ├── ncnet.yml ├── patch2pix.yml ├── patch2pix_superglue.yml ├── r2d2.yml ├── sift.yml ├── sparsencnet.yml ├── superglue.yml └── superpoint.yml ├── data └── download.sh ├── docs ├── evaluation.md ├── install.md └── patch2pix_example_matches.png ├── environment.yml ├── immatch ├── __init__.py ├── eval_aachen.py ├── eval_hpatches.py ├── eval_inloc.py ├── eval_relapose.py ├── eval_robotcar.py ├── modules │ ├── __init__.py │ ├── aspanformer.py │ ├── base.py │ ├── caps.py │ ├── cotr.py │ ├── d2net.py │ ├── dogaffnethardnet.py │ ├── loftr.py │ ├── nn_matching.py │ ├── patch2pix.py │ ├── r2d2.py │ ├── sift.py │ ├── sparsencnet.py │ ├── superglue.py │ └── superpoint.py └── utils │ ├── __init__.py │ ├── colmap │ ├── data_parsing.py │ ├── database.py │ └── read_write_model.py │ ├── data_io.py │ ├── hpatches_helper.py │ ├── localize_sfm_helper.py │ ├── metrics.py │ └── model_helper.py ├── notebooks ├── HPatches_MMA_Curves.ipynb └── visualize_matches_on_example_pairs.ipynb ├── outputs └── hpatches │ └── cache │ ├── CAPS_SIFT.npy │ ├── CAPS_SuperPoint_r4.npy │ ├── D2Net.npy │ ├── DoG1024-AffNet-HardNet.m0.95.npy │ ├── NCNet.im1024.m0.9.npy │ ├── Patch2Pix.im1024.m0.5.npy │ ├── Patch2Pix.im1024.m0.9.npy │ ├── R2D2.npy │ ├── SparseNCNet_N2000.im3200.npy │ ├── SuperGlue_r4.m0.2.npy │ ├── SuperGlue_r4.m0.5.npy │ ├── SuperGlue_r4.m0.9.npy │ ├── SuperPoint_r4.npy │ ├── aslfeat.npy │ ├── delf.npy │ ├── hesaff.npy │ └── hesaffnet.npy ├── pretrained └── download.sh └── setup.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Manually added ignores 2 | outputs/ 3 | data/ 4 | pretrained/ 5 | !pretrained/download.sh 6 | 7 | # Byte-compiled / optimized / DLL files 8 | __pycache__/ 9 | *.py[cod] 10 | *$py.class 11 | 12 | # C extensions 13 | *.so 14 | 15 | # Distribution / packaging 16 | .Python 17 | build/ 18 | develop-eggs/ 19 | dist/ 20 | downloads/ 21 | eggs/ 22 | .eggs/ 23 | lib/ 24 | lib64/ 25 | parts/ 26 | sdist/ 27 | var/ 28 | wheels/ 29 | pip-wheel-metadata/ 30 | share/python-wheels/ 31 | *.egg-info/ 32 | .installed.cfg 33 | *.egg 34 | MANIFEST 35 | 36 | # PyInstaller 37 | # Usually these files are written by a python script from a template 38 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 39 | *.manifest 40 | *.spec 41 | 42 | # Installer logs 43 | pip-log.txt 44 | pip-delete-this-directory.txt 45 | 46 | # Unit test / coverage reports 47 | htmlcov/ 48 | .tox/ 49 | .nox/ 50 | .coverage 51 | .coverage.* 52 | .cache 53 | nosetests.xml 54 | coverage.xml 55 | *.cover 56 | *.py,cover 57 | .hypothesis/ 58 | .pytest_cache/ 59 | 60 | # Translations 61 | *.mo 62 | *.pot 63 | 64 | # Django stuff: 65 | *.log 66 | local_settings.py 67 | db.sqlite3 68 | db.sqlite3-journal 69 | 70 | # Flask stuff: 71 | instance/ 72 | .webassets-cache 73 | 74 | # Scrapy stuff: 75 | .scrapy 76 | 77 | # Sphinx documentation 78 | docs/_build/ 79 | 80 | # PyBuilder 81 | target/ 82 | 83 | # Jupyter Notebook 84 | .ipynb_checkpoints 85 | 86 | # IPython 87 | profile_default/ 88 | ipython_config.py 89 | 90 | # pyenv 91 | .python-version 92 | 93 | # pipenv 94 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 95 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 96 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 97 | # install all needed dependencies. 98 | #Pipfile.lock 99 | 100 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 101 | __pypackages__/ 102 | 103 | # Celery stuff 104 | celerybeat-schedule 105 | celerybeat.pid 106 | 107 | # SageMath parsed files 108 | *.sage.py 109 | 110 | # Environments 111 | .env 112 | .venv 113 | env/ 114 | venv/ 115 | ENV/ 116 | env.bak/ 117 | venv.bak/ 118 | 119 | # Spyder project settings 120 | .spyderproject 121 | .spyproject 122 | 123 | # Rope project settings 124 | .ropeproject 125 | 126 | # mkdocs documentation 127 | /site 128 | 129 | # mypy 130 | .mypy_cache/ 131 | .dmypy.json 132 | dmypy.json 133 | 134 | # Pyre type checker 135 | .pyre/ 136 | -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "third_party/SuperGluePretrainedNetwork"] 2 | path = third_party/superglue 3 | url = https://github.com/magicleap/SuperGluePretrainedNetwork 4 | [submodule "third_party/caps"] 5 | path = third_party/caps 6 | url = https://github.com/GrumpyZhou/caps 7 | [submodule "third_party/d2-net"] 8 | path = third_party/d2net 9 | url = https://github.com/mihaidusmanu/d2-net 10 | [submodule "third_party/r2d2"] 11 | path = third_party/r2d2 12 | url = https://github.com/naver/r2d2 13 | ignore = untracked 14 | [submodule "third_party/sparse-ncnet"] 15 | path = third_party/sparsencnet 16 | url = https://github.com/ignacio-rocco/sparse-ncnet 17 | [submodule "third_party/hloc"] 18 | path = third_party/hloc 19 | url = https://github.com/GrumpyZhou/Hierarchical-Localization 20 | branch = extend_base_dev 21 | [submodule "third_party/patch2pix"] 22 | path = third_party/patch2pix 23 | url = https://github.com/GrumpyZhou/patch2pix 24 | [submodule "third_party/loftr"] 25 | path = third_party/loftr 26 | url = https://github.com/zju3dv/LoFTR 27 | [submodule "third_party/cotr"] 28 | path = third_party/cotr 29 | url = https://github.com/ubc-vision/COTR 30 | branch = master 31 | [submodule "third_party/aspanformer"] 32 | path = third_party/aspanformer 33 | url = https://github.com/apple/ml-aspanformer.git 34 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 Qunjie Zhou 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 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # A Toolbox for Image Feature Matching and Evaluations 2 | In this repository, we provide **easy interfaces** for several exisiting SotA methods to match image feature correspondences between image pairs. 3 | We provide **scripts to evaluate** their predicted correspondences on common benchmarks for the tasks of image matching, homography estimation and visual localization. 4 | 5 | ## TODOs & Updates 6 | - [x] Add LoFTR method (2021-7-8) 7 | - [x] Add simple match visualization (2021-7-8) 8 | - [x] Use ***immatch*** as a python lib under develop mode. Check [install.md](docs/install.md) for details. (2021-7-22) 9 | - [x] Add SIFT method (opencv version) (2021-7-25) 10 | - [x] Add script to eval on RobotCar using HLoc (2021-7-31) 11 | - [x] Add Dog-AffNet-Hardnet (Contributed by Dmytro Mishkin 👏, 2021-8-29) 12 | - [x] Add AUC metric and opencv solver for Homography estimation on HPatches (#20, 2022-1-12) 13 | - [x] Add COTR (A naive wrapper without tuning parameters, 2022-3-29) 14 | - [x] Add Aspanformer (2023-6-2) 15 | - [x] Add Megadepth relative pose estimation following LoFTR & Aspanformer (2023-6-2) 16 | - [x] Add ScanNet relative pose estimation following LoFTR & Aspanformer (2024-1-11) 17 | - [ ] Add support to eval on [Image Matching Challenge](https://www.cs.ubc.ca/research/image-matching-challenge/current/data) 18 | - [ ] Add scripts to eval on [SimLoc](https://github.com/simlocmatch/simlocmatch-benchmark) challenge. 19 | 20 | ***Comments from QJ***: Currently I am quite busy with my study & work. So it will take some time before I release the next two TODOs. 21 | 22 | ## Supported Methods & Evaluations 23 | **Sparse Keypoint-based Matching:** 24 | - Local Feature: 25 | [CAPS](https://arxiv.org/abs/2004.13324), [D2Net](https://arxiv.org/abs/1905.03561), [R2D2](https://arxiv.org/abs/1906.06195), [SuperPoint](https://arxiv.org/abs/1712.07629), [Dog-AffNet-HardNet](https://arxiv.org/abs/1711.06704) 26 | - Matcher: [SuperGlue](https://arxiv.org/abs/1911.11763) 27 | 28 | **Semi-dense Matching:** 29 | - Correspondence Network: [NCNet](https://arxiv.org/abs/1810.10510), [SparseNCNet](https://arxiv.org/pdf/2004.10566.pdf), 30 | - Transformer-based: [Aspanformer](https://aspanformer.github.io/),[LoFTR](https://zju3dv.github.io/loftr/), [COTR](https://github.com/ubc-vision/COTR) 31 | - Local Refinement: [Patch2Pix](https://arxiv.org/abs/2012.01909) 32 | 33 | **Supported Evaluations** : 34 | - Image feature matching on HPatches 35 | - Homography estimation on HPatches 36 | - Visual localization benchmarks: 37 | - InLoc 38 | - Aachen (original + v1.1) 39 | - RobotCar Seasons (v1 + v2) 40 | 41 | ## Repository Overview 42 | The repository is structured as follows: 43 | - **configs/**: Each method has its own yaml (.yml) file to configure its testing parameters. 44 | - **data/**: All datasets should be placed under this folder following our instructions described in **[Data Preparation](docs/evaluation.md#data-preparation)**. 45 | - **immatch/**: It contains implementations of method wrappers and evaluation interfaces. 46 | - **outputs/**: All evaluation results are supposed to be saved here. One folder per benchmark. 47 | - **pretrained/**: It contains the pretrained models of the supported methods. 48 | - **third_party/**: The real implementation of the supported methods from their original repositories, as git submodules. 49 | - **notebooks/**: It contains jupyter notebooks that show example codes to quickly access the methods implemented in this repo. 50 | - **docs/**: It contains separate documentation about installation and evaluation. To keep a clean face of this repo :). 51 | 52 | #### 👉Refer to [install.md](docs/install.md) for details about installation. 53 | #### 👉Refer to [evaluation.md](docs/evaluation.md) for details about evaluation on benchmarks. 54 | 55 | ## Example Code for Quick Testing 56 | To use a specific method to perform the matching task, you simply need to do: 57 | - **Initialize a matcher using its config file.** See examples of config yaml files under [configs](configs/) folder, eg., [patch2pix.yml](configs/patch2pix.yml). Each config file contains multiple sections, each section corresponds to one setting. Here, we use the setting (*tagged by 'example'*) for testing on example image pairs. 58 | - **Perform matching** 59 | ```python 60 | import immatch 61 | import yaml 62 | from immatch.utils import plot_matches 63 | 64 | # Initialize model 65 | with open('configs/patch2pix.yml', 'r') as f: 66 | args = yaml.load(f, Loader=yaml.FullLoader)['example'] 67 | model = immatch.__dict__[args['class']](args) 68 | matcher = lambda im1, im2: model.match_pairs(im1, im2) 69 | 70 | # Specify the image pair 71 | im1 = 'third_party/patch2pix/examples/images/pair_2/1.jpg' 72 | im2 = 'third_party/patch2pix/examples/images/pair_2/2.jpg' 73 | 74 | # Match and visualize 75 | matches, _, _, _ = matcher(im1, im2) 76 | plot_matches(im1, im2, matches, radius=2, lines=True) 77 | ``` 78 | ![example matches](docs/patch2pix_example_matches.png) 79 | 80 | #### 👉 Try out the code using [example notebook ](notebooks/visualize_matches_on_example_pairs.ipynb). 81 | 82 | ## Notice 83 | - This repository is expected to be actively maintained (at least before I graduate🤣🤣) and **gradually** (slowly) grow for new features of interest. 84 | - Suggestions regarding how to improve this repo, such as adding new **SotA** image matching methods or new benchmark evaluations, are welcome 👏. 85 | 86 | ### Regarding Patch2Pix 87 | With this reprository, one can **reproduce** the tables reported in our paper accepted at CVPR2021: Patch2Pix: Epipolar-Guided Pixel-Level Correspondences[[pdf]](https://arxiv.org/abs/2012.01909). Check [our patch2pix repository](https://github.com/GrumpyZhou/patch2pix) for its training code. 88 | 89 | ### Disclaimer 90 | - All of the supported methods and evaluations are **not implemented from scratch** by us. Instead, we modularize their original code to define unified interfaces. 91 | - If you are using the results of a method, **remember to cite the corresponding paper**. 92 | - All credits of the implemetation of those methods belong to their authors . 93 | -------------------------------------------------------------------------------- /configs/aspanformer.yml: -------------------------------------------------------------------------------- 1 | default: &default 2 | class: 'ASpanFormer' 3 | ckpt: 'pretrained/aspanformer/outdoor.ckpt' 4 | train_res: [832, 832] 5 | test_res: [832, 832] 6 | coarse_scale: 0.125 7 | coarsest_level: [26, 26] 8 | border_rm: 2 9 | imsize: 832 10 | online_resize: False 11 | match_threshold: 0.2 12 | no_match_upscale: False 13 | eval_coarse: False 14 | example: 15 | <<: *default 16 | hpatch: 17 | <<: *default 18 | eval_coarse: False 19 | test_res: [480, 480] 20 | coarsest_level: [15, 15] 21 | no_match_upscale: True 22 | megadepth: 23 | <<: *default 24 | train_res: [480, 480] 25 | coarsest_level: [36, 36] 26 | test_res: [1152, 1152] 27 | scannet: 28 | <<: *default 29 | ckpt: 'pretrained/aspanformer/indoor.ckpt' 30 | coarsest_level: [15, 20] 31 | train_res: [480, 640] 32 | test_res: [480, 640] 33 | no_match_upscale: True # Intrinsics were pre-processed to [480, 640] 34 | border_rm: 0 35 | inloc: 36 | <<: *default 37 | match_threshold: 0.2 38 | npts: 4096 39 | imsize: 1024 40 | pairs: 'pairs-query-netvlad40-temporal.txt' 41 | rthres: 48 42 | skip_matches: 20 43 | aachen: 44 | <<: *default 45 | match_threshold: 0.0 # Save all matches 46 | pairs: ['pairs-db-covis20.txt', 'pairs-query-netvlad50.txt'] 47 | npts: 4096 48 | imsize: 1024 49 | qt_dthres: 4 50 | qt_psize: 48 51 | qt_unique: True 52 | ransac_thres: [20] 53 | sc_thres: 0.2 # Filtering during quantization 54 | covis_cluster: True 55 | -------------------------------------------------------------------------------- /configs/caps_sift.yml: -------------------------------------------------------------------------------- 1 | default: &default 2 | class: 'CAPS' 3 | ckpt: 'pretrained/caps/caps-pretrained.pth' 4 | backbone: 'resnet50' 5 | pretrained: 1 6 | coarse_feat_dim: 128 7 | fine_feat_dim: 128 8 | prob_from: 'correlation' 9 | window_size: 0.125 10 | use_nn: 1 11 | detector: 'SIFT' 12 | npts: 1024 13 | hpatch: 14 | <<: *default 15 | match_threshold: 0.0 16 | imsize: -1 17 | inloc: 18 | <<: *default 19 | match_threshold: 0.75 20 | imsize: 1600 21 | npts: 4096 22 | pairs: 'pairs-query-netvlad40-temporal.txt' 23 | rthres: 48 24 | skip_matches: 20 25 | -------------------------------------------------------------------------------- /configs/caps_superpoint.yml: -------------------------------------------------------------------------------- 1 | default: &default 2 | class: 'CAPS' 3 | ckpt: 'pretrained/caps/caps-pretrained.pth' 4 | backbone: 'resnet50' 5 | pretrained: 1 6 | coarse_feat_dim: 128 7 | fine_feat_dim: 128 8 | prob_from: 'correlation' 9 | window_size: 0.125 10 | use_nn: 1 11 | detector: 'SuperPoint' 12 | nms_radius: 4 13 | example: 14 | <<: *default 15 | match_threshold: 0.5 16 | imsize: -1 17 | hpatch: 18 | <<: *default 19 | max_keypoints: -1 20 | match_threshold: 0.0 21 | imsize: -1 22 | inloc: 23 | <<: *default 24 | max_keypoints: 4096 25 | match_threshold: 0.75 26 | imsize: 1024 27 | pairs: 'pairs-query-netvlad40-temporal.txt' 28 | rthres: 48 29 | skip_matches: 20 30 | -------------------------------------------------------------------------------- /configs/cotr.yml: -------------------------------------------------------------------------------- 1 | default: &default 2 | class: 'COTR' 3 | ckpt: 'pretrained/cotr/cotr_default.pth.tar' 4 | backbone: 'resnet50' 5 | hidden_dim: 256 6 | dilation: False 7 | dropout: 0.1 8 | nheads: 8 9 | layer: 'layer3' 10 | backbone_layer_dims : { 11 | 'layer1': 256, 12 | 'layer2': 512, 13 | 'layer3': 1024, 14 | 'layer4': 2048, 15 | } 16 | enc_layers: 6 17 | dec_layers: 6 18 | position_embedding: 'lin_sine' 19 | max_corrs: 100 20 | match_threshold: 0.0 21 | imsize: -1 22 | batch_size: 32 23 | example: 24 | <<: *default 25 | hpatch: 26 | <<: *default 27 | -------------------------------------------------------------------------------- /configs/d2net.yml: -------------------------------------------------------------------------------- 1 | default: &default 2 | class: 'D2Net' 3 | ckpt: 'pretrained/d2net/d2_tf.pth' 4 | use_relu: True 5 | multiscale: False 6 | example: 7 | <<: *default 8 | match_threshold: 0.5 9 | imsize: -1 10 | hpatch: 11 | <<: *default 12 | match_threshold: 0.0 13 | imsize: -1 14 | inloc: 15 | <<: *default 16 | match_threshold: 0.0 17 | imsize: 1024 18 | pairs: 'pairs-query-netvlad40-temporal.txt' 19 | rthres: 48 20 | skip_matches: 20 21 | -------------------------------------------------------------------------------- /configs/dogaffnethardnet.yml: -------------------------------------------------------------------------------- 1 | default: &default 2 | class: 'DogAffNetHardNet' 3 | npts: 1024 4 | match_threshold: 0.95 5 | imsize: -1 6 | device: 'cpu' 7 | example: 8 | <<: *default 9 | hpatch: 10 | <<: *default 11 | -------------------------------------------------------------------------------- /configs/loftr.yml: -------------------------------------------------------------------------------- 1 | default: &default 2 | class: 'LoFTR' 3 | ckpt: 'pretrained/loftr/outdoor_ds.ckpt' 4 | match_threshold: 0.2 5 | imsize: -1 6 | no_match_upscale: False 7 | eval_coarse: False 8 | example: 9 | <<: *default 10 | match_threshold: 0.5 11 | imsize: -1 12 | hpatch: 13 | <<: *default 14 | imsize: 480 15 | no_match_upscale: True 16 | megadepth: 17 | <<: *default 18 | imsize: 1024 19 | inloc: 20 | <<: *default 21 | match_threshold: 0.5 22 | npts: 4096 23 | imsize: 1024 24 | pairs: 'pairs-query-netvlad40-temporal.txt' 25 | rthres: 48 26 | skip_matches: 20 27 | aachen: 28 | <<: *default 29 | match_threshold: 0.0 # Save all matches 30 | pairs: ['pairs-db-covis20.txt', 'pairs-query-netvlad50.txt'] 31 | npts: 4096 32 | imsize: 1024 33 | qt_dthres: 4 34 | qt_psize: 48 35 | qt_unique: True 36 | ransac_thres: [20] 37 | sc_thres: 0.2 # Filtering during quantization 38 | covis_cluster: True -------------------------------------------------------------------------------- /configs/ncnet.yml: -------------------------------------------------------------------------------- 1 | # We use the adapted NCNet in Patch2Pix paper, instead of the original NCNet model. 2 | hpatch: 3 | class: 'NCNet' 4 | ckpt: 'pretrained/patch2pix/ncn_ivd_5ep.pth' 5 | match_threshold: 0.9 6 | imsize: 1024 7 | ksize: 2 -------------------------------------------------------------------------------- /configs/patch2pix.yml: -------------------------------------------------------------------------------- 1 | default: &default 2 | class: 'Patch2Pix' 3 | ckpt: 'pretrained/patch2pix/patch2pix_pretrained.pth' 4 | ksize: 2 5 | imsize: 1024 6 | match_threshold: 0.25 7 | no_match_upscale: False 8 | example: 9 | <<: *default 10 | match_threshold: 0.5 11 | imsize: -1 12 | hpatch: 13 | <<: *default 14 | imsize: 1024 15 | no_match_upscale: True 16 | inloc: 17 | <<: *default 18 | pairs: 'pairs-query-netvlad40-temporal.txt' 19 | rthres: 48 20 | skip_matches: 20 21 | aachen: 22 | <<: *default 23 | match_threshold: 0.0 # Save all matches 24 | pairs: ['pairs-db-covis20.txt', 'pairs-query-netvlad30.txt'] 25 | qt_dthres: 4 26 | qt_psize: 48 27 | qt_unique: True 28 | ransac_thres: [12] 29 | sc_thres: 0.25 # Filtering during quantization 30 | covis_cluster: True 31 | aachen_v1.1: 32 | <<: *default 33 | match_threshold: 0.0 # Save all matches 34 | pairs: ['pairs-db-covis20.txt', 'pairs-query-netvlad50.txt'] 35 | qt_dthres: 4 36 | qt_psize: 48 37 | qt_unique: True 38 | ransac_thres: [20] 39 | sc_thres: 0.25 # Filtering during quantization 40 | covis_cluster: True 41 | aachen_night: 42 | <<: *default 43 | match_threshold: 0.0 # Save all matches 44 | pairs: ['pairs-db-night_benchmark.txt', 'pairs-query-night_benchmark.txt'] 45 | qt_dthres: 4 46 | qt_psize: 48 47 | qt_unique: True 48 | ransac_thres: [20] 49 | sc_thres: 0.25 # Filtering during quantization 50 | covis_cluster: False -------------------------------------------------------------------------------- /configs/patch2pix_superglue.yml: -------------------------------------------------------------------------------- 1 | default: &default 2 | class: 'Patch2PixRefined' 3 | ckpt: 'pretrained/patch2pix/patch2pix_pretrained.pth' 4 | imsize: 1024 5 | coarse_default: &coarse_default 6 | name: 'SuperGlue' 7 | weights: 'outdoor' 8 | sinkhorn_iterations: 50 9 | max_keypoints: 4096 10 | nms_radius: 4 11 | imsize: 1024 12 | inloc: 13 | <<: *default 14 | match_threshold: 0.1 15 | imsize: 1600 16 | pairs: 'pairs-query-netvlad40-temporal.txt' 17 | rthres: 48 18 | skip_matches: 20 19 | coarse: 20 | <<: *coarse_default 21 | match_threshold: 0.1 22 | imsize: 1600 23 | aachen: 24 | <<: *default 25 | match_threshold: 0.0 # Save all matches 26 | pairs: ['pairs-db-covis20.txt', 'pairs-query-netvlad50.txt'] 27 | qt_dthres: 4 28 | qt_psize: 48 29 | qt_unique: True 30 | ransac_thres: [20] 31 | sc_thres: 0.25 # Filtering during quantization 32 | covis_cluster: False 33 | coarse: 34 | <<: *coarse_default 35 | match_threshold: 0.1 36 | aachen_v1.1: 37 | <<: *default 38 | match_threshold: 0.0 # Save all matches 39 | pairs: ['pairs-db-covis20.txt', 'pairs-query-netvlad50.txt'] 40 | qt_dthres: 4 41 | qt_psize: 48 42 | qt_unique: True 43 | ransac_thres: [20] 44 | sc_thres: 0.25 # Filtering during quantization 45 | covis_cluster: True 46 | coarse: 47 | <<: *coarse_default 48 | match_threshold: 0.1 49 | -------------------------------------------------------------------------------- /configs/r2d2.yml: -------------------------------------------------------------------------------- 1 | # 'imsize' -1 is equivalent to original arg name 'max_size' 9999 in R2D2's repo. 2 | default: &default 3 | class: 'R2D2' 4 | ckpt: 'pretrained/r2d2/r2d2_WASF_N16.pt' 5 | top_k: 5000 6 | reliability_thr: 0.7 7 | repeatability_thr: 0.7 8 | min_scale: 0.3 9 | max_scale: 1 10 | min_size: 0 11 | example: 12 | <<: *default 13 | match_threshold: 0.5 14 | imsize: -1 15 | hpatch: 16 | <<: *default 17 | min_size: 0 18 | imsize: -1 19 | match_threshold: 0.0 20 | inloc: 21 | <<: *default 22 | imsize: 1600 23 | match_threshold: 0.0 24 | pairs: 'pairs-query-netvlad40-temporal.txt' 25 | rthres: 48 26 | skip_matches: 20 27 | -------------------------------------------------------------------------------- /configs/sift.yml: -------------------------------------------------------------------------------- 1 | default: &default 2 | class: 'SIFT' 3 | npts: 1024 4 | match_threshold: 0.0 5 | imsize: -1 6 | hpatch: 7 | <<: *default 8 | inloc: 9 | <<: *default 10 | match_threshold: 0.0 11 | npts: 4096 12 | imsize: 1600 13 | pairs: 'pairs-query-netvlad40-temporal.txt' 14 | rthres: 48 15 | skip_matches: 20 16 | aachen: 17 | <<: *default 18 | match_threshold: 0.0 # Save all matches 19 | pairs: ['pairs-db-covis20.txt', 'pairs-query-netvlad50.txt'] 20 | npts: 4096 21 | imsize: 1024 22 | qt_dthres: 0 23 | qt_psize: 0 24 | qt_unique: False 25 | ransac_thres: [12, 20, 25] 26 | sc_thres: 0.5 # Filtering during quantization 27 | covis_cluster: True 28 | -------------------------------------------------------------------------------- /configs/sparsencnet.yml: -------------------------------------------------------------------------------- 1 | default: &default 2 | class: 'SparseNCNet' 3 | ckpt: 'pretrained/sparsencnet/sparsencnet_k10.pth.tar' 4 | benchmark: 0 5 | no_ncnet: 0 6 | relocalize: 1 7 | reloc_type: 'hard_soft' 8 | reloc_hard_crop_size: 2 9 | change_stride: 1 10 | Npts: 2000 11 | hpatch: 12 | <<: *default 13 | match_threshold: -1 14 | imsize: 3200 15 | inloc: 16 | <<: *default 17 | imsize: 1600 18 | match_threshold: -1 19 | pairs: 'pairs-query-netvlad40-temporal.txt' 20 | rthres: 48 21 | skip_matches: 20 22 | -------------------------------------------------------------------------------- /configs/superglue.yml: -------------------------------------------------------------------------------- 1 | default: &default 2 | class: 'SuperGlue' 3 | weights: 'outdoor' 4 | sinkhorn_iterations: 20 5 | nms_radius: 4 6 | match_threshold: 0.2 7 | no_match_upscale: False 8 | example: 9 | <<: *default 10 | match_threshold: 0.5 11 | imsize: -1 12 | hpatch: 13 | <<: *default 14 | max_keypoints: 1024 15 | imsize: -1 16 | inloc: 17 | <<: *default 18 | sinkhorn_iterations: 50 19 | max_keypoints: 4096 20 | imsize: 1600 21 | pairs: 'pairs-query-netvlad40-temporal.txt' 22 | rthres: 48 23 | skip_matches: 20 24 | aachen: 25 | <<: *default 26 | sinkhorn_iterations: 50 27 | match_threshold: 0.0 28 | max_keypoints: 4096 29 | nms_radius: 3 30 | imsize: 1024 31 | pairs: ['pairs-db-covis20.txt', 'pairs-query-netvlad50.txt'] 32 | qt_dthres: 0 33 | qt_psize: 0 34 | qt_unique: False 35 | ransac_thres: [12] 36 | sc_thres: 0.2 37 | covis_cluster: False 38 | aachen_v1.1: 39 | <<: *default 40 | sinkhorn_iterations: 50 41 | match_threshold: 0.0 42 | max_keypoints: 4096 43 | nms_radius: 3 44 | imsize: 1024 45 | pairs: ['pairs-db-covis20.txt', 'pairs-query-netvlad50.txt'] 46 | qt_dthres: 0 47 | qt_psize: 0 48 | qt_unique: False 49 | ransac_thres: [12] 50 | sc_thres: 0.2 51 | covis_cluster: False 52 | robotcar: 53 | <<: *default 54 | sinkhorn_iterations: 50 55 | match_threshold: 0.0 56 | max_keypoints: 4096 57 | nms_radius: 3 58 | imsize: 1024 59 | pairs: ['pairs-db-covis20.txt', 'pairs-query-netvlad20-percam-perloc.txt'] 60 | qt_dthres: 0 61 | qt_psize: 0 62 | qt_unique: False 63 | ransac_thres: [12] 64 | sc_thres: 0.2 65 | covis_cluster: False -------------------------------------------------------------------------------- /configs/superpoint.yml: -------------------------------------------------------------------------------- 1 | default: &default 2 | class: 'SuperPoint' 3 | keypoint_threshold: 0.005 4 | nms_radius: 4 5 | example: 6 | <<: *default 7 | match_threshold: 0.5 8 | imsize: -1 9 | hpatch: 10 | <<: *default 11 | max_keypoints: -1 12 | match_threshold: 0.0 13 | imsize: -1 14 | inloc: 15 | <<: *default 16 | max_keypoints: 4096 17 | match_threshold: 0.75 18 | imsize: 1024 19 | pairs: 'pairs-query-netvlad40-temporal.txt' 20 | rthres: 48 21 | skip_matches: 20 22 | -------------------------------------------------------------------------------- /data/download.sh: -------------------------------------------------------------------------------- 1 | # Download HLoc Pairs 2 | gdown --id 1T2c3DHJlTdRpo_Y3Ad4AH8p4tS7cOyjP 3 | unzip pairs.zip 4 | rm pairs.zip 5 | 6 | https://drive.google.com/file/d/13vfAKPG5lylXvr86gVW4WpDWoSOgidZ3/view?usp=share_link 7 | 8 | # Download HLoc robot car query list with intrinsics 9 | # File name: robotcar_season_queries_list_with_intrinsics.zip 10 | # After download + unzip, place it under: repo_root/data/datasets/RobotCar/queries 11 | gdown --id 13vfAKPG5lylXvr86gVW4WpDWoSOgidZ3 12 | -------------------------------------------------------------------------------- /docs/evaluation.md: -------------------------------------------------------------------------------- 1 | ## Data Preparation 2 | ### Prepare Datasets 3 | One needs to prepare the target dataset(s) to evaluate on in advance. 4 | All datasets should be placed under **data/datasets**. Here, we provide brief data instructions of our supported benchmarks. 5 | 6 | - **HPatches** : We recommend to follow [D2Net's instructions](https://github.com/mihaidusmanu/d2-net/tree/master/hpatches_sequences) as we did. The dataset root folder needs to be named as **hpatches-sequences-release/** and contains 52 illumination sequences and 56 viewpoint sequences. 7 | 8 | - **InLoc**: Simply download the [InLoc](http://www.ok.sc.e.titech.ac.jp/INLOC/) dataset and place them under a folder named **InLoc/**. 9 | 10 | - **Aachen Day and Night**: We follow the instructions provided in [Local Feature Evaluation](https://github.com/tsattler/visuallocalizationbenchmark/tree/master/local_feature_evaluation) to prepare the data folder. The dataset folder is name as **AachenDayNight/**. If you want to work on **Aachen v1.1**, you also need to place it under the original version as described in [README_Aachen-Day-Night_v1_1.md](https://data.ciirc.cvut.cz/public/projects/2020VisualLocalization/Aachen-Day-Night/README_Aachen-Day-Night_v1_1.md). 11 | 12 | - **RobotCar Seasons**: Download the [RobotCar Seasons](https://data.ciirc.cvut.cz/public/projects/2020VisualLocalization/RobotCar-Seasons/) dataset contents and place them under a folder named **RobotCar/**. 13 | 14 | - **MegaDepth** and **ScanNet** for relative pose estimation: we followed [LoFTR](https://github.com/zju3dv/LoFTR?tab=readme-ov-file) to setup both datasets for testing, please refer to their repo for details. 15 | 16 | ### Prepare Image Pairs 17 | To evaluate on visual localization benchmarks, one needs to prepare image pairs that are required by HLoc pipeline in advance. 18 | For convenience, we cached the pairs that are extracted by [HLoc](https://github.com/cvg/Hierarchical-Localization) author Paul-Edouard Sarlin. 19 | You can download them by running from the repository root: 20 | ```bash 21 | cd data 22 | bash download.sh 23 | ``` 24 | Otherwise, feel free to use your own database pairs and query pairs. 25 | 26 | ### RobotCar Season Queries 27 | The above _download.sh_ also downloads the zip file containing RobotCar Season queries which was pre-extracted using [scripts from the original HLoc](https://github.com/cvg/Hierarchical-Localization/blob/robotcar/pipelines/RobotCar/robotcar_generate_query_list.py). 28 | You need to unzip it and place the folder under _data/RobotCar/queries_. 29 | 30 | Finally, the **data/** folder should contain look like: 31 | ``` 32 | data/ 33 | |-- datasets/ 34 | |-- hpatches-sequences-release/ 35 | |-- i_ajuntament/ 36 | |-- ... 37 | |-- AachenDayNight/ 38 | |-- 3D-models/ 39 | |-- aachen_v_1_1/ 40 | |-- images/ 41 | |-- queries/ 42 | |-- InLoc/ 43 | |-- dataset/ 44 | |-- alignments/ 45 | |-- cutouts/ 46 | |-- query/ 47 | |-- Robotcar/ 48 | |-- 3D-models/ 49 | |-- images/ 50 | |-- intrinsics/ 51 | |-- queries/ 52 | ... 53 | |-- pairs/ 54 | |-- aachen 55 | |-- aachen_v1.1 56 | |-- inloc 57 | |-- robotcar 58 | |-- readme.txt 59 | ``` 60 | 61 | ## Evaluations 62 | The following example commands are supposed to be executed under **the repository root**. 63 | 64 | ### HPatches 65 | Available tasks are image feature matching and homography estimation. One can specify to run on either or both at one time by setting `--task ` with one of ['matching' , 'homography', 'both']. 66 | 67 | For example, the following command evaluates **SuperPoint** and **NCNet** on **both** tasks using their settings defined under **configs/**: 68 | ```python 69 | python -m immatch.eval_hpatches --gpu 0 \ 70 | --config 'superpoint' 'ncnet' \ 71 | --task 'both' --save_npy \ 72 | --root_dir . 73 | ``` 74 | The following command evaluates **Patch2Pix** on under 3 matching thresholds: 75 | ```python 76 | python -m immatch.eval_hpatches --gpu 0 \ 77 | --config 'patch2pix' --match_thres 0.25 0.5 0.9 \ 78 | --task 'both' --save_npy \ 79 | --root_dir . 80 | ``` 81 | ### Relative Pose Estimation 82 | To reproduce [AspanFormer results](https://github.com/apple/ml-aspanformer/tree/main?tab=readme-ov-file#evaluation): 83 | ``` 84 | # MegaDepth 85 | 86 | python -m immatch.eval_relapose --config 'aspanformer' --benchmark 'megadepth' 87 | 88 | # ScanNet 89 | python -m immatch.eval_relapose --config 'aspanformer' --benchmark 'scannet' 90 | ``` 91 | 92 | 93 | ### Long-term Visual Localization 94 | We adopt the public implementation [Hierarchical Localization](https://github.com/cvg/Hierarchical-Localization) to evaluate matches on several long-term visual localization benchmarks, including: 95 | - InLoc 96 | - Aachen Day and Night (original + v1.1) 97 | - RobotCar Seasons (v1 + v2) 98 | 99 | Notice, our released scripts are tested on at least one of the following methods: 100 | - **Patch2Pix** 101 | - **Patch2Pix+SuperGlue** 102 | - **SuperGlue** 103 | 104 | to verify that they can reproduce their released results on the visual localization benchmark [leader board](https://www.visuallocalization.net/benchmark/). 105 | 106 | In the following, we give examples to evaluate **Patch2Pix** using its setting defined inside [patch2pix.yml](configs/patch2pix.yml). 107 | To use another method, simply replace `--config 'patch2pix'` with the name of its config file name, for example `--config 'patch2pix_superglue'` or `--config 'superglue'`. 108 | **Notice**, one needs to prepare datasets following the previous section before running the following evaluations. 109 | 110 | #### InLoc 111 | ```python 112 | python -m immatch.eval_inloc --gpu 0\ 113 | --config 'patch2pix' 114 | ``` 115 | #### Aachen Day and Night 116 | ```python 117 | # Original version 118 | python -m immatch.eval_aachen --gpu 0 \ 119 | --config 'patch2pix_superglue' \ 120 | --colmap $COLMAP_PATH \ 121 | --benchmark_name 'aachen' 122 | 123 | # Version 1.1 124 | python -m immatch.eval_aachen --gpu 0 \ 125 | --config 'patch2pix' \ 126 | --colmap $COLMAP_PATH \ 127 | --benchmark_name 'aachen_v1.1' 128 | ``` 129 | 130 | ### RobotCar Season 131 | By running the following commant, you will get qurey poses for both RobotCar v1 & v2. We output those results at the same time, since they only differ in the format rather than the pose results. 132 | Notice, running RobotCar takes rather long time (~days). 133 | For debugging, we suggest using only a subset of the provided pairs for both db & query pairs. 134 | 135 | ```python 136 | # Version 1 + 2 137 | python -m immatch.eval_robotcar --gpu 0 \ 138 | --config 'superglue' \ 139 | --colmap $COLMAP_PATH \ 140 | --benchmark_name 'robotcar' 141 | ``` 142 | -------------------------------------------------------------------------------- /docs/install.md: -------------------------------------------------------------------------------- 1 | # Installation 2 | 3 | To use our code, first download this repository and initialize the submodules: 4 | ```bash 5 | git clone https://github.com/GrumpyZhou/image-matching-toolbox.git 6 | 7 | # Install submodules non-recursively 8 | cd image-matching-toolbox/ 9 | git submodule update --init 10 | ``` 11 | 12 | Next, download the pretained models and place them to the correct places by running the followings: 13 | ```bash 14 | cd pretrained/ 15 | bash download.sh 16 | ``` 17 | 18 | ## Setup Running Environment 19 | Following the steps to setup the ready environment to run the matching toolbox. The code has been tested on Ubuntu 18.04 with Python 3.7 + Pytorch 1.7.0 + CUDA 10.2. 20 | ### 1. Create the immatch virtual environment 21 | ```bash 22 | conda env create -f environment.yml 23 | conda activate immatch 24 | ``` 25 | Notice, the **immatch** conda env allows to run all supported methods **expect for SparseNCNet**. In order to use it, please install its required dependencies according to its official [installation](https://github.com/ignacio-rocco/sparse-ncnet/blob/master/INSTALL.md), 26 | 27 | ### 2. Install the immatch toolbox as a python package 28 | ```bash 29 | # Install immatch toolbox 30 | cd image-matching-toolbox/ 31 | python setup.py develop 32 | ``` 33 | The developing mode allows you to change the code **without re-installing** it in the environment. You can also install the matching toolbox to any environment to use it **for your other projects**. 34 | To **uninstall** it from an environment: 35 | ``` 36 | pip uninstall immatch 37 | ``` 38 | 39 | ### 3. Install pycolmap 40 | This package is essential for evaluations on localization benchmarks. 41 | ```bash 42 | # Install pycolmap 43 | pip install git+https://github.com/mihaidusmanu/pycolmap 44 | ``` 45 | In case https link doesnt work, you can install it directly for Python 3.7 and Python 3.8 via pypi: 46 | ```bash 47 | pip install pycolmap 48 | ``` 49 | 50 | ### 4. Update immatch environment when needed 51 | Incase more packages are needed for new features, one can update your created immatch environment: 52 | #### Option 1: add new libs into [setup.py](../setup.py) (Recommended & Faster) 53 | ```bash 54 | # Update immatch toolbox 55 | cd image-matching-toolbox/ 56 | python setup.py develop 57 | ``` 58 | 59 | #### Option 2: add new libs into [environment.yml](../environment.yml) 60 | ``` 61 | conda activate immatch 62 | conda env update --file environment.yml --prune 63 | ``` 64 | -------------------------------------------------------------------------------- /docs/patch2pix_example_matches.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GrumpyZhou/image-matching-toolbox/470c9e1ba07bab28438afdff9f934c8589a92ed6/docs/patch2pix_example_matches.png -------------------------------------------------------------------------------- /environment.yml: -------------------------------------------------------------------------------- 1 | name: immatch 2 | channels: 3 | - pytorch 4 | - conda-forge 5 | - defaults 6 | dependencies: 7 | - cudatoolkit=10.2 8 | - python=3.7 9 | - pytorch==1.7.0 10 | - torchvision 11 | - jupyter 12 | - scipy 13 | - matplotlib 14 | - opencv==4.5.2 -------------------------------------------------------------------------------- /immatch/__init__.py: -------------------------------------------------------------------------------- 1 | from .modules.caps import CAPS 2 | from .modules.superpoint import SuperPoint 3 | from .modules.superglue import SuperGlue 4 | from .modules.d2net import D2Net 5 | from .modules.r2d2 import R2D2 6 | from .modules.patch2pix import Patch2Pix, NCNet, Patch2PixRefined 7 | from .modules.loftr import LoFTR # Cause warnings 8 | from .modules.sift import SIFT 9 | from .modules.dogaffnethardnet import DogAffNetHardNet 10 | from .modules.cotr import COTR 11 | 12 | try: 13 | from .modules.aspanformer import ASpanFormer 14 | except ImportError as e: 15 | print(f"Can not import ASpanFormer: {e}") 16 | pass 17 | 18 | try: 19 | import MinkowskiEngine 20 | import sys 21 | from pathlib import Path 22 | 23 | # To prevent naming conflict as D2Net also has module called lib 24 | d2net_path = Path(__file__).parent / "modules/../../third_party/d2net" 25 | sys.path.remove(str(d2net_path)) 26 | 27 | from .modules.sparsencnet import SparseNCNet 28 | 29 | use_sparsencnet = True 30 | except ImportError as e: 31 | print(f"Can not import sparsencnet: {e}") 32 | pass 33 | -------------------------------------------------------------------------------- /immatch/eval_aachen.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from argparse import Namespace 3 | import os 4 | from pathlib import Path 5 | 6 | from immatch.utils.model_helper import init_model 7 | from immatch.utils.localize_sfm_helper import * 8 | 9 | 10 | def generate_covis_db_pairs(args): 11 | from immatch.utils.colmap.data_parsing import ( 12 | covis_pairs_from_nvm, 13 | covis_pairs_from_reference_model, 14 | ) 15 | 16 | if args.benchmark_name in ["aachen"]: 17 | nvm_path = Path(args.dataset_dir) / "3D-models/aachen_cvpr2018_db.nvm" 18 | covis_pairs_from_nvm(nvm_path, args.pair_dir, topk=args.topk) 19 | 20 | if args.benchmark_name in ["aachen_v1.1"]: 21 | model_dir = Path(args.dataset_dir) / "3D-models/aachen_v_1_1" 22 | covis_pairs_from_reference_model(model_dir, args.pair_dir, topk=args.topk) 23 | 24 | 25 | def eval_aachen(args): 26 | # Initialize model 27 | model, model_conf = init_model(args.config, args.benchmark_name) 28 | matcher = lambda im1, im2: model.match_pairs(im1, im2) 29 | 30 | # Merge args 31 | args = Namespace(**vars(args), **model_conf) 32 | args.model_tag = model.name 33 | args.conf_tag = f"im{model.imsize}" 34 | if args.prefix: 35 | args.conf_tag += f".{args.prefix}" 36 | 37 | # Load db & query pairs 38 | pair_list = load_pairs(args) 39 | if not pair_list: 40 | return 41 | 42 | # Initialize experiment paths 43 | args.dataset_dir = Path(args.dataset_dir) 44 | args.im_dir = args.dataset_dir / "images/images_upright" 45 | args = init_paths(args) 46 | print(args) 47 | 48 | # Match pairs 49 | if not args.skip_match: 50 | match_pairs_exporth5(pair_list, matcher, args.im_dir, args.output_dir) 51 | 52 | # Extract keypoints 53 | process_matches_and_keypoints_exporth5( 54 | pair_list, 55 | args.output_dir, 56 | args.result_dir, 57 | qt_psize=args.qt_psize, 58 | qt_dthres=args.qt_dthres, 59 | sc_thres=args.sc_thres, 60 | ) 61 | 62 | # Localization 63 | init_empty_sfm(args) 64 | reconstruct_database_pairs(args) 65 | localize_queries(args) 66 | 67 | 68 | if __name__ == "__main__": 69 | parser = argparse.ArgumentParser(description="Localize Aachen Day-Night") 70 | parser.add_argument("--gpu", "-gpu", type=str, default="0") 71 | parser.add_argument("--config", type=str, default=None) 72 | parser.add_argument("--prefix", type=str, default=None) 73 | parser.add_argument("--colmap", type=str, required=True) 74 | parser.add_argument("--skip_match", action="store_true") 75 | parser.add_argument( 76 | "--dataset_dir", type=str, default="data/datasets/AachenDayNight" 77 | ) 78 | parser.add_argument("--pair_dir", type=str, default="data/pairs") 79 | parser.add_argument( 80 | "--benchmark_name", 81 | type=str, 82 | choices=["aachen", "aachen_v1.1"], 83 | default="aachen", 84 | ) 85 | 86 | # Turn on to only generate covis db pairs 87 | parser.add_argument("--generate_covis_db_pairs", action="store_true") 88 | parser.add_argument("--topk", type=int, default=20) 89 | 90 | args = parser.parse_args() 91 | os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu 92 | 93 | # Generate covis pairs 94 | if args.generate_covis_db_pairs: 95 | generate_covis_db_pairs(args) 96 | else: 97 | eval_aachen(args) 98 | -------------------------------------------------------------------------------- /immatch/eval_hpatches.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from argparse import Namespace 3 | import os 4 | 5 | import immatch 6 | from immatch.utils.data_io import lprint 7 | import immatch.utils.hpatches_helper as helper 8 | from immatch.utils.model_helper import parse_model_config 9 | 10 | 11 | def eval_hpatches( 12 | root_dir, 13 | config_list, 14 | task="both", 15 | benchmark="hpatch", 16 | h_solver="degensac", 17 | ransac_thres=2, 18 | match_thres=None, 19 | odir="outputs/hpatches", 20 | save_npy=False, 21 | print_out=False, 22 | debug=False, 23 | ): 24 | # Init paths 25 | data_root = os.path.join(root_dir, "data/datasets/hpatches-sequences-release") 26 | cache_dir = os.path.join(root_dir, odir, "cache") 27 | result_dir = os.path.join(root_dir, odir, "results", task) 28 | if not os.path.exists(cache_dir): 29 | os.makedirs(cache_dir) 30 | if not os.path.exists(result_dir): 31 | os.makedirs(result_dir) 32 | 33 | # Iterate over methods 34 | for config_name in config_list: 35 | # Load model 36 | args = parse_model_config(config_name, benchmark, root_dir) 37 | class_name = args["class"] 38 | 39 | # One log file per method 40 | log_file = os.path.join(result_dir, f"{class_name}.txt") 41 | log = open(log_file, "a") 42 | lprint_ = lambda ms: lprint(ms, log) 43 | 44 | # Iterate over matching thresholds 45 | thresholds = match_thres if match_thres else [args["match_threshold"]] 46 | lprint_( 47 | f"\n>>>> Method={class_name} Default config: {args} " f"Thres: {thresholds}" 48 | ) 49 | 50 | for thres in thresholds: 51 | args["match_threshold"] = thres # Set to target thresholds 52 | 53 | # Init model 54 | model = immatch.__dict__[class_name](args) 55 | matcher = lambda im1, im2: model.match_pairs(im1, im2) 56 | 57 | # Init result save path (for matching results) 58 | result_npy = None 59 | if save_npy: 60 | result_tag = model.name 61 | if args["imsize"] > 0: 62 | result_tag += f".im{args['imsize']}" 63 | if thres > 0: 64 | result_tag += f".m{thres}" 65 | result_npy = os.path.join(cache_dir, f"{result_tag}.npy") 66 | 67 | lprint_(f"Matching thres: {thres} Save to: {result_npy}") 68 | 69 | # Eval on the specified task(s) 70 | helper.eval_hpatches( 71 | matcher, 72 | data_root, 73 | model.name, 74 | task=task, 75 | scale_H=getattr(model, "no_match_upscale", False), 76 | h_solver=h_solver, 77 | ransac_thres=ransac_thres, 78 | lprint_=lprint_, 79 | print_out=print_out, 80 | save_npy=result_npy, 81 | debug=debug, 82 | ) 83 | log.close() 84 | 85 | 86 | if __name__ == "__main__": 87 | parser = argparse.ArgumentParser(description="Benchmark HPatches") 88 | parser.add_argument("--gpu", "-gpu", type=str, default="0") 89 | parser.add_argument("--root_dir", type=str, default=".") 90 | parser.add_argument("--odir", type=str, default="outputs/hpatches") 91 | parser.add_argument("--config", type=str, nargs="*", default=None) 92 | parser.add_argument("--match_thres", type=float, nargs="*", default=None) 93 | parser.add_argument( 94 | "--task", 95 | type=str, 96 | default="homography", 97 | choices=["matching", "homography", "both"], 98 | ) 99 | parser.add_argument( 100 | "--h_solver", type=str, default="degensac", choices=["degensac", "cv"] 101 | ) 102 | parser.add_argument("--ransac_thres", type=float, default=2) 103 | parser.add_argument("--save_npy", action="store_true") 104 | parser.add_argument("--print_out", action="store_true") 105 | 106 | args = parser.parse_args() 107 | os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu 108 | eval_hpatches( 109 | args.root_dir, 110 | args.config, 111 | args.task, 112 | h_solver=args.h_solver, 113 | ransac_thres=args.ransac_thres, 114 | match_thres=args.match_thres, 115 | odir=args.odir, 116 | save_npy=args.save_npy, 117 | print_out=args.print_out, 118 | ) 119 | -------------------------------------------------------------------------------- /immatch/eval_inloc.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from argparse import Namespace 3 | import os 4 | 5 | from third_party.hloc.hloc.localize_inloc import localize_with_matcher 6 | from immatch.utils.model_helper import init_model 7 | 8 | 9 | def eval_inloc(args): 10 | model, model_conf = init_model(args.config, args.benchmark_name) 11 | matcher = lambda im1, im2: model.match_pairs(im1, im2) 12 | args = Namespace(**vars(args), **model_conf) 13 | print(">>>>", args) 14 | 15 | # Setup output dir 16 | rthres = args.rthres 17 | mthres = args.match_threshold 18 | skip_matches = args.skip_matches 19 | retrieval_pairs = os.path.join(args.pair_dir, args.benchmark_name, args.pairs) 20 | pair_tag = retrieval_pairs.split("query-")[-1].replace(".txt", "") 21 | exp_name = f"{pair_tag}_sk{skip_matches}im{args.imsize}rth{rthres}mth{mthres}" 22 | if args.prefix: 23 | exp_name = f"{exp_name}.{args.prefix}" 24 | odir = os.path.join("outputs", args.benchmark_name, model.name, exp_name) 25 | method_tag = f"{model.name}_{exp_name}" 26 | print(">>>Method tag:", method_tag) 27 | 28 | # Localize InLoc queries 29 | localize_with_matcher( 30 | matcher, 31 | args.dataset_dir, 32 | retrieval_pairs, 33 | odir, 34 | method_tag, 35 | rthres=rthres, 36 | skip_matches=skip_matches, 37 | ) 38 | 39 | 40 | if __name__ == "__main__": 41 | parser = argparse.ArgumentParser(description="Localize Inloc") 42 | parser.add_argument("--gpu", "-gpu", type=str, default="0") 43 | parser.add_argument("--config", type=str, default=None) 44 | parser.add_argument("--prefix", type=str, default=None) 45 | parser.add_argument("--dataset_dir", type=str, default="data/datasets/InLoc") 46 | parser.add_argument("--pair_dir", type=str, default="data/pairs") 47 | parser.add_argument("--benchmark_name", type=str, default="inloc") 48 | 49 | args = parser.parse_args() 50 | os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu 51 | eval_inloc(args) 52 | -------------------------------------------------------------------------------- /immatch/eval_relapose.py: -------------------------------------------------------------------------------- 1 | from collections import defaultdict 2 | import os 3 | import numpy as np 4 | from tqdm import tqdm 5 | import argparse 6 | from argparse import Namespace 7 | import cv2 8 | from pathlib import Path 9 | 10 | from immatch.utils.model_helper import init_model 11 | import immatch.utils.metrics as M 12 | 13 | 14 | def compute_relapose_aspan(kpts0, kpts1, K0, K1, pix_thres=0.5, conf=0.99999): 15 | """Original code from ASpanFormer repo: 16 | https://github.com/apple/ml-aspanformer/blob/main/src/utils/metrics.py 17 | """ 18 | 19 | if len(kpts0) < 5: 20 | return None 21 | # normalize keypoints 22 | kpts0 = (kpts0 - K0[[0, 1], [2, 2]][None]) / K0[[0, 1], [0, 1]][None] 23 | kpts1 = (kpts1 - K1[[0, 1], [2, 2]][None]) / K1[[0, 1], [0, 1]][None] 24 | 25 | # normalize ransac threshold 26 | ransac_thr = pix_thres / np.mean([K0[0, 0], K0[1, 1], K1[0, 0], K1[1, 1]]) 27 | 28 | # compute pose with cv2 29 | E, mask = cv2.findEssentialMat( 30 | kpts0, kpts1, np.eye(3), threshold=ransac_thr, prob=conf, method=cv2.RANSAC 31 | ) 32 | if E is None: 33 | print("\nE is None while trying to recover pose.\n") 34 | return None 35 | 36 | # recover pose from E 37 | best_num_inliers = 0 38 | ret = None 39 | for _E in np.split(E, len(E) / 3): 40 | n, R, t, _ = cv2.recoverPose(_E, kpts0, kpts1, np.eye(3), 1e9, mask=mask) 41 | if n > best_num_inliers: 42 | ret = (R, t[:, 0], mask.ravel() > 0) 43 | best_num_inliers = n 44 | 45 | return ret 46 | 47 | 48 | def load_scannet_pairs_npz(intrinsic_path, npz_path, data_root): 49 | # Load intrinsics 50 | intrinsics = dict(np.load(intrinsic_path)) 51 | 52 | # Collect pairs 53 | pairs = [] 54 | scene_info = dict(np.load(npz_path)) 55 | for scene_id, scene_sub, stem_name_1, stem_name_2 in tqdm(scene_info["name"]): 56 | scene_name = f"scene{scene_id:04d}_{scene_sub:02d}" 57 | 58 | im1 = f"{scene_name}/color/{stem_name_1}.jpg" 59 | im2 = f"{scene_name}/color/{stem_name_2}.jpg" 60 | K1 = K2 = np.array(intrinsics[scene_name].copy(), dtype=np.float32).reshape( 61 | 3, 3 62 | ) 63 | 64 | # Compute relative pose 65 | T1 = np.linalg.inv( 66 | np.loadtxt( 67 | f"{data_root}/{scene_name}/pose/{stem_name_1}.txt", delimiter=" " 68 | ) 69 | ) 70 | T2 = np.linalg.inv( 71 | np.loadtxt( 72 | f"{data_root}/{scene_name}/pose/{stem_name_2}.txt", delimiter=" " 73 | ) 74 | ) 75 | T12 = np.matmul(T2, np.linalg.inv(T1)) 76 | pairs.append( 77 | Namespace(im1=im1, im2=im2, K1=K1, K2=K2, t=T12[:3, 3], R=T12[:3, :3]) 78 | ) 79 | print(f"Loaded {len(pairs)} pairs.") 80 | return pairs 81 | 82 | 83 | def load_megadepth_pairs_npz(npz_root, npz_list): 84 | with open(npz_list, "r") as f: 85 | npz_names = [name.split()[0] for name in f.readlines()] 86 | print(f"Parse {len(npz_names)} npz from {npz_list}.") 87 | 88 | pairs = [] 89 | for name in npz_names: 90 | scene_info = np.load(f"{npz_root}/{name}.npz", allow_pickle=True) 91 | 92 | # Collect pairs 93 | for pair_info in scene_info["pair_infos"]: 94 | (id1, id2), overlap, _ = pair_info 95 | im1 = scene_info["image_paths"][id1].replace("Undistorted_SfM/", "") 96 | im2 = scene_info["image_paths"][id2].replace("Undistorted_SfM/", "") 97 | K1 = scene_info["intrinsics"][id1].astype(np.float32) 98 | K2 = scene_info["intrinsics"][id2].astype(np.float32) 99 | 100 | # Compute relative pose 101 | T1 = scene_info["poses"][id1] 102 | T2 = scene_info["poses"][id2] 103 | T12 = np.matmul(T2, np.linalg.inv(T1)) 104 | pairs.append( 105 | Namespace( 106 | im1=im1, 107 | im2=im2, 108 | overlap=overlap, 109 | K1=K1, 110 | K2=K2, 111 | t=T12[:3, 3], 112 | R=T12[:3, :3], 113 | ) 114 | ) 115 | print(f"Loaded {len(pairs)} pairs.") 116 | return pairs 117 | 118 | 119 | def eval_relapose( 120 | matcher, 121 | data_root, 122 | pairs, 123 | method="", 124 | ransac_thres=0.5, 125 | thresholds=[1, 3, 5, 10, 20], 126 | print_out=False, 127 | debug=False, 128 | ): 129 | statis = defaultdict(list) 130 | np.set_printoptions(precision=2) 131 | 132 | # Eval on pairs 133 | print(f">>> Start eval on Megadepth: method={method} rthres={ransac_thres} ... \n") 134 | for i, pair in tqdm(enumerate(pairs), smoothing=0.1, total=len(pairs)): 135 | if debug and i > 10: 136 | break 137 | 138 | K1 = pair.K1 139 | K2 = pair.K2 140 | t_gt = pair.t 141 | R_gt = pair.R 142 | im1 = str(data_root / pair.im1) 143 | im2 = str(data_root / pair.im2) 144 | match_res = matcher(im1, im2) 145 | matches, pts1, pts2 = match_res[0:3] 146 | 147 | # Compute pose errors 148 | ret = compute_relapose_aspan(pts1, pts2, K1, K2, pix_thres=ransac_thres) 149 | 150 | if ret is None: 151 | statis["failed"].append(i) 152 | statis["R_errs"].append(np.inf) 153 | statis["t_errs"].append(np.inf) 154 | statis["inliers"].append(np.array([]).astype(np.bool_)) 155 | else: 156 | R, t, inliers = ret 157 | R_err, t_err = M.cal_relapose_error(R, R_gt, t, t_gt) 158 | statis["R_errs"].append(R_err) 159 | statis["t_errs"].append(t_err) 160 | statis["inliers"].append(inliers.sum() / len(pts1)) 161 | if print_out: 162 | print(f"#M={len(matches)} R={R_err:.3f}, t={t_err:.3f}") 163 | 164 | print(f"Total samples: {len(pairs)} Failed:{len(statis['failed'])}.") 165 | pose_auc = M.cal_relapose_auc(statis, thresholds=thresholds) 166 | return pose_auc 167 | 168 | 169 | def eval_megadepth_relapose( 170 | root_dir, 171 | method, 172 | benchmark="megadepth", 173 | ransac_thres=0.5, 174 | print_out=False, 175 | debug=False, 176 | ): 177 | 178 | # Init paths 179 | npz_root = ( 180 | root_dir / "third_party/aspanformer/assets/megadepth_test_1500_scene_info" 181 | ) 182 | npz_list = ( 183 | root_dir 184 | / "third_party/aspanformer/assets/megadepth_test_1500_scene_info/megadepth_test_1500.txt" 185 | ) 186 | data_root = root_dir / "data/datasets/MegaDepth_undistort" 187 | 188 | # Load pairs 189 | pairs = load_megadepth_pairs_npz(npz_root, npz_list) 190 | 191 | # Init model 192 | model, config = init_model(method, benchmark, root_dir=root_dir) 193 | matcher = lambda im1, im2: model.match_pairs(im1, im2) 194 | 195 | # Eval 196 | eval_relapose( 197 | matcher, 198 | data_root, 199 | pairs, 200 | model.name, 201 | print_out=print_out, 202 | debug=debug, 203 | ) 204 | 205 | 206 | def eval_scannet_relapose( 207 | root_dir, 208 | method, 209 | benchmark="scannet", 210 | ransac_thres=0.5, 211 | print_out=False, 212 | debug=False, 213 | ): 214 | 215 | # Init paths 216 | npz_path = root_dir / "third_party/aspanformer/assets/scannet_test_1500/test.npz" 217 | intrinsic_path = ( 218 | root_dir / "third_party/aspanformer/assets/scannet_test_1500/intrinsics.npz" 219 | ) 220 | data_root = root_dir / "data/datasets/scannet/test" 221 | 222 | # Load pairs 223 | pairs = load_scannet_pairs_npz(intrinsic_path, npz_path, data_root) 224 | 225 | # Init model 226 | model, config = init_model(method, benchmark, root_dir=root_dir) 227 | matcher = lambda im1, im2: model.match_pairs(im1, im2) 228 | 229 | # Eval 230 | eval_relapose( 231 | matcher, 232 | data_root, 233 | pairs, 234 | model.name, 235 | print_out=print_out, 236 | debug=debug, 237 | ) 238 | 239 | 240 | if __name__ == "__main__": 241 | parser = argparse.ArgumentParser(description="Benchmark Relative Pose") 242 | parser.add_argument("--gpu", "-gpu", type=str, default="0") 243 | parser.add_argument("--config", type=str, required=True) 244 | parser.add_argument( 245 | "--benchmark", type=str, required=True, choices=["scannet", "megadepth"] 246 | ) 247 | parser.add_argument("--root_dir", type=str, default=".") 248 | parser.add_argument("--ransac_thres", type=float, default=0.5) 249 | parser.add_argument("--print_out", action="store_true") 250 | parser.add_argument("--debug", action="store_true") 251 | 252 | args = parser.parse_args() 253 | os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu 254 | 255 | eval(f"eval_{args.benchmark}_relapose")( 256 | Path(args.root_dir), 257 | args.config, 258 | benchmark=args.benchmark, 259 | ransac_thres=args.ransac_thres, 260 | print_out=args.print_out, 261 | debug=args.debug, 262 | ) 263 | -------------------------------------------------------------------------------- /immatch/eval_robotcar.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from argparse import Namespace 3 | import os 4 | from pathlib import Path 5 | import yaml 6 | import immatch 7 | from immatch.utils.localize_sfm_helper import * 8 | 9 | 10 | def generate_covis_db_pairs(args): 11 | from immatch.utils.colmap.data_parsing import covis_pairs_from_nvm 12 | 13 | nvm_path = Path(args.dataset_dir) / "3D-models/all-merged/all.nvm" 14 | covis_pairs_from_nvm(nvm_path, args.pair_dir, topk=args.topk) 15 | 16 | 17 | def eval_robotcar(args): 18 | os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu 19 | 20 | # Initialize Model 21 | config_file = f"configs/{args.config}.yml" 22 | with open(config_file, "r") as f: 23 | model_conf = yaml.load(f, Loader=yaml.FullLoader)[args.benchmark_name] 24 | class_name = model_conf["class"] 25 | print(f"Method:{class_name} Conf: {model_conf}") 26 | model = immatch.__dict__[class_name](model_conf) 27 | matcher = lambda im1, im2: model.match_pairs(im1, im2) 28 | 29 | # Merge args 30 | args = Namespace(**vars(args), **model_conf) 31 | args.model_tag = model.name 32 | args.conf_tag = f"im{model.imsize}" 33 | if args.prefix: 34 | args.conf_tag += f".{args.prefix}" 35 | 36 | # Load db & query pairs 37 | pair_list = load_pairs(args) 38 | if not pair_list: 39 | print("No pairs for localizations! Please turn on --generate_covis_db_pairs!") 40 | return 41 | 42 | # Initialize experiment paths 43 | args.dataset_dir = Path(args.dataset_dir) 44 | args.im_dir = args.dataset_dir / "images" 45 | args = init_paths(args) 46 | print(args) 47 | 48 | # Match pairs 49 | if not args.skip_match: 50 | match_pairs_exporth5(pair_list, matcher, args.im_dir, args.output_dir) 51 | 52 | # Extract keypoints 53 | process_matches_and_keypoints_exporth5( 54 | pair_list, 55 | args.output_dir, 56 | args.result_dir, 57 | qt_psize=args.qt_psize, 58 | qt_dthres=args.qt_dthres, 59 | sc_thres=args.sc_thres, 60 | ) 61 | 62 | # Localization 63 | init_empty_sfm(args) 64 | reconstruct_database_pairs(args) 65 | localize_queries(args) 66 | 67 | 68 | if __name__ == "__main__": 69 | parser = argparse.ArgumentParser(description="Localize RobotCar Seasons") 70 | parser.add_argument("--gpu", "-gpu", type=str, default="0") 71 | parser.add_argument("--config", type=str, default=None) 72 | parser.add_argument("--prefix", type=str, default=None) 73 | parser.add_argument("--colmap", type=str, required=True) 74 | parser.add_argument("--skip_match", action="store_true") 75 | parser.add_argument("--dataset_dir", type=str, default="data/datasets/RobotCar") 76 | parser.add_argument("--pair_dir", type=str, default="data/pairs") 77 | parser.add_argument("--benchmark_name", type=str, default="robotcar") 78 | 79 | # Turn on to only generate covis db pairs 80 | parser.add_argument("--generate_covis_db_pairs", action="store_true") 81 | parser.add_argument("--topk", type=int, default=20) 82 | 83 | args = parser.parse_args() 84 | 85 | # Generate covis pairs 86 | if args.generate_covis_db_pairs: 87 | generate_covis_db_pairs(args) 88 | else: 89 | eval_robotcar(args) 90 | -------------------------------------------------------------------------------- /immatch/modules/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GrumpyZhou/image-matching-toolbox/470c9e1ba07bab28438afdff9f934c8589a92ed6/immatch/modules/__init__.py -------------------------------------------------------------------------------- /immatch/modules/aspanformer.py: -------------------------------------------------------------------------------- 1 | from argparse import Namespace 2 | import torch 3 | import numpy as np 4 | import cv2 5 | import torch.nn.functional as F 6 | 7 | from third_party.aspanformer.src.ASpanFormer.aspanformer import ( 8 | ASpanFormer as ASpanFormer_, 9 | ) 10 | from third_party.aspanformer.src.config.default import get_cfg_defaults 11 | from third_party.aspanformer.src.utils.misc import lower_config 12 | 13 | from .base import Matching 14 | from immatch.utils.data_io import load_gray_scale_tensor_cv 15 | from third_party.aspanformer.src.utils.dataset import read_megadepth_gray 16 | 17 | 18 | class ASpanFormer(Matching): 19 | def __init__(self, args): 20 | super().__init__() 21 | if type(args) == dict: 22 | args = Namespace(**args) 23 | 24 | self.imsize = args.imsize 25 | self.match_threshold = args.match_threshold 26 | self.no_match_upscale = args.no_match_upscale 27 | self.online_resize = args.online_resize 28 | self.im_padding = False 29 | self.coarse_scale = args.coarse_scale 30 | self.eval_coarse = args.eval_coarse 31 | 32 | # Load model 33 | config = get_cfg_defaults() 34 | conf = lower_config(config)["aspan"] 35 | conf["coarse"]["train_res"] = args.train_res 36 | conf["coarse"]["test_res"] = args.test_res 37 | conf["coarse"]["coarsest_level"] = args.coarsest_level 38 | conf["match_coarse"]["border_rm"] = args.border_rm 39 | conf["match_coarse"]["thr"] = args.match_threshold 40 | 41 | if args.test_res: 42 | self.imsize = args.test_res[::-1] 43 | self.im_padding = args.test_res[0] == args.test_res[1] 44 | self.model = ASpanFormer_(config=conf) 45 | ckpt_dict = torch.load(args.ckpt) 46 | self.model.load_state_dict(ckpt_dict["state_dict"], strict=False) 47 | self.model = self.model.eval().to(self.device) 48 | 49 | # Name the method 50 | self.ckpt_name = args.ckpt.split("/")[-1].split(".")[0] 51 | self.name = f"ASpanFormer_{self.ckpt_name}" 52 | if self.no_match_upscale: 53 | self.name += "_noms" 54 | print(f"Initialize {self.name} {args} ") 55 | 56 | def load_im(self, im_path): 57 | return load_gray_scale_tensor_cv( 58 | im_path, 59 | self.device, 60 | dfactor=8, 61 | imsize=self.imsize, 62 | value_to_scale=max, 63 | pad2sqr=self.im_padding, 64 | ) 65 | 66 | def match_inputs_(self, gray1, gray2, mask1=None, mask2=None): 67 | batch = {"image0": gray1, "image1": gray2} 68 | if mask1 is not None and mask2 is not None and self.coarse_scale: 69 | [ts_mask_1, ts_mask_2] = ( 70 | F.interpolate( 71 | torch.stack([mask1, mask2], dim=0)[None].float(), 72 | scale_factor=self.coarse_scale, 73 | mode="nearest", 74 | recompute_scale_factor=False, 75 | )[0] 76 | .bool() 77 | .to(self.device) 78 | ) 79 | batch.update( 80 | {"mask0": ts_mask_1.unsqueeze(0), "mask1": ts_mask_2.unsqueeze(0)} 81 | ) 82 | 83 | # Forward pass 84 | self.model(batch, online_resize=self.online_resize) 85 | 86 | # Output parsing 87 | if self.eval_coarse: 88 | kpts1 = batch["mkpts0_c"].cpu().numpy() 89 | kpts2 = batch["mkpts1_c"].cpu().numpy() 90 | else: 91 | kpts1 = batch["mkpts0_f"].cpu().numpy() 92 | kpts2 = batch["mkpts1_f"].cpu().numpy() 93 | scores = batch["mconf"].cpu().numpy() 94 | matches = np.concatenate([kpts1, kpts2], axis=1) 95 | return matches, kpts1, kpts2, scores 96 | 97 | def match_pairs(self, im1_path, im2_path): 98 | gray1, sc1, mask1 = self.load_im(im1_path) 99 | gray2, sc2, mask2 = self.load_im(im2_path) 100 | upscale = np.array([sc1 + sc2]) 101 | matches, kpts1, kpts2, scores = self.match_inputs_(gray1, gray2, mask1, mask2) 102 | 103 | if self.no_match_upscale: 104 | return matches, kpts1, kpts2, scores, upscale.squeeze(0) 105 | 106 | # Upscale matches & kpts 107 | matches = upscale * matches 108 | kpts1 = sc1 * kpts1 109 | kpts2 = sc2 * kpts2 110 | return matches, kpts1, kpts2, scores 111 | -------------------------------------------------------------------------------- /immatch/modules/base.py: -------------------------------------------------------------------------------- 1 | from abc import ABCMeta, abstractmethod 2 | import torch 3 | from .nn_matching import mutual_nn_matching 4 | 5 | 6 | class FeatureDetection(metaclass=ABCMeta): 7 | """An abstract class for local feature detection and description methods""" 8 | 9 | def __init__(self): 10 | super().__init__() 11 | self.device = torch.device( 12 | "cuda:{}".format(0) if torch.cuda.is_available() else "cpu" 13 | ) 14 | torch.set_grad_enabled(False) 15 | 16 | @abstractmethod 17 | def extract_features(self, im, **kwargs): 18 | """Given the processed input, the keypoints and descriptors are extracted by the model. 19 | Return: 20 | kpts : a Nx2 tensor, N is the number of keypoints. 21 | desc : a NxD tensor, N is the number of descriptors 22 | and D is dimension of each descriptor. 23 | """ 24 | 25 | @abstractmethod 26 | def load_and_extract(self, im_path, **kwargs): 27 | """Given an image path, the input image is firstly loaded and processed accordingly, 28 | the keypoints and descriptors are then extracted by the model. 29 | Return: 30 | kpts : a Nx2 tensor, N is the number of keypoints. 31 | desc : a NxD tensor, N is the number of descriptors 32 | and D is dimension of each descriptor. 33 | """ 34 | 35 | def describe(self, im, kpts, **kwargs): 36 | """Given the processed input and the pre-detected keypoint locations, 37 | feature descriptors are described by the model. 38 | Return: 39 | desc : a NxD tensor, N is the number of descriptors 40 | and D is dimension of each descriptor. 41 | """ 42 | 43 | def detect(self, im, **kwargs): 44 | """Given the processed input, the keypoints are detected by that method. 45 | Return: 46 | kpts : a Nx2 tensor, N is the number of keypoints. 47 | """ 48 | 49 | 50 | class Matching(metaclass=ABCMeta): 51 | """An abstract class for a method that perform matching from the input pairs""" 52 | 53 | def __init__(self): 54 | super().__init__() 55 | self.device = torch.device( 56 | "cuda:{}".format(0) if torch.cuda.is_available() else "cpu" 57 | ) 58 | torch.set_grad_enabled(False) 59 | 60 | @classmethod 61 | def mutual_nn_match(self, desc1, desc2, threshold=0.0): 62 | """The feature descriptors from the pair of images are matched 63 | using nearset neighbor search with mutual check and an optional 64 | outlier filtering. This is normally used by feature detection methods. 65 | Args: 66 | desc1, desc2: descriptors from the 1st and 2nd image of a pair. 67 | threshold: the cosine similarity threshold for the outlier filtering. 68 | Return: 69 | match_ids: the indices of the matched descriptors. 70 | """ 71 | return mutual_nn_matching(desc1, desc2, threshold) 72 | 73 | @abstractmethod 74 | def match_pairs(self, im1_path, im2_path, **kwargs): 75 | """The model detects correspondences from a pair of images. 76 | All steps that are required to estimate the correspondences by a method 77 | are implemented here. 78 | Input: 79 | im1_path, im2_path: the paths of the input image pair. 80 | other args depend on the model. 81 | 82 | Return: 83 | matches: the detected matches stored as numpy array with shape Nx4, 84 | N is the number of matches. 85 | kpts1, kpts2: the keypoints used for matching. For methods that don't 86 | explicitly define keypoints, e.g., SparseNCNet, 87 | the keypoints are the locations of points that get matched. 88 | scores: the matching score or confidence of each correspondence. 89 | Notices, matching scores are defined differently across methods. 90 | For NN matcher, they can be the cosine distance of descriptors; 91 | For SuperGlue, they are the probablities in the OT matrix; ..etc. 92 | """ 93 | -------------------------------------------------------------------------------- /immatch/modules/caps.py: -------------------------------------------------------------------------------- 1 | from argparse import Namespace 2 | import torch 3 | import numpy as np 4 | import cv2 5 | 6 | from third_party.caps.CAPS.network import CAPSNet 7 | from immatch.utils.data_io import load_im_tensor 8 | from .base import FeatureDetection, Matching 9 | from .superpoint import SuperPoint 10 | from .sift import SIFT 11 | 12 | 13 | class CAPS(FeatureDetection, Matching): 14 | def __init__(self, args): 15 | super().__init__() 16 | if type(args) == dict: 17 | args = Namespace(**args) 18 | self.imsize = args.imsize 19 | self.match_threshold = ( 20 | args.match_threshold if "match_threshold" in args else 0.0 21 | ) 22 | self.model = CAPSNet(args, self.device).eval() 23 | self.load_model(args.ckpt) 24 | 25 | # Initialize detector 26 | if args.detector.lower() == "sift": 27 | self.detector = SIFT(args) 28 | self.name = f"CAPS_SIFT" 29 | else: 30 | self.detector = SuperPoint(vars(args)) 31 | rad = self.detector.model.config["nms_radius"] 32 | self.name = f"CAPS_SuperPoint_r{rad}" 33 | print(f"Initialize {self.name}") 34 | 35 | def load_im(self, im_path): 36 | return load_im_tensor( 37 | im_path, 38 | self.device, 39 | imsize=self.imsize, 40 | with_gray=True, 41 | raw_gray=("SIFT" in self.name), 42 | ) 43 | 44 | def load_model(self, ckpt): 45 | print("Reloading from {}".format(ckpt)) 46 | ckpt_dict = torch.load(ckpt) 47 | self.model.load_state_dict(ckpt_dict["state_dict"]) 48 | 49 | def extract_features(self, im, gray): 50 | kpts = self.detector.detect(gray) 51 | if isinstance(kpts, np.ndarray): 52 | kpts = torch.from_numpy(kpts).float().to(self.device) 53 | desc = self.describe(im, kpts) 54 | return kpts, desc 55 | 56 | def describe(self, im, kpts): 57 | kpts = kpts.unsqueeze(0) 58 | feat_c, feat_f = self.model.extract_features(im, kpts) 59 | desc = torch.cat((feat_c, feat_f), -1).squeeze(0) 60 | return desc 61 | 62 | def load_and_extract(self, im_path): 63 | im, gray, scale = self.load_im(im_path) 64 | kpts, desc = self.extract_features(im, gray) 65 | kpts = kpts * torch.tensor(scale).to(kpts) # N, 2 66 | return kpts, desc 67 | 68 | def match_inputs_(self, im1, gray1, im2, gray2): 69 | kpts1, desc1 = self.extract_features(im1, gray1) 70 | kpts2, desc2 = self.extract_features(im2, gray2) 71 | kpts1 = kpts1.cpu().data.numpy() 72 | kpts2 = kpts2.cpu().data.numpy() 73 | 74 | # NN Match 75 | match_ids, scores = self.mutual_nn_match( 76 | desc1, desc2, threshold=self.match_threshold 77 | ) 78 | p1s = kpts1[match_ids[:, 0], :2] 79 | p2s = kpts2[match_ids[:, 1], :2] 80 | matches = np.concatenate([p1s, p2s], axis=1) 81 | return matches, kpts1, kpts2, scores 82 | 83 | def match_pairs(self, im1_path, im2_path): 84 | im1, gray1, sc1 = self.load_im(im1_path) 85 | im2, gray2, sc2 = self.load_im(im2_path) 86 | upscale = np.array([sc1 + sc2]) 87 | matches, kpts1, kpts2, scores = self.match_inputs_(im1, gray1, im2, gray2) 88 | matches = upscale * matches 89 | kpts1 = sc1 * kpts1 90 | kpts2 = sc2 * kpts2 91 | return matches, kpts1, kpts2, scores 92 | -------------------------------------------------------------------------------- /immatch/modules/cotr.py: -------------------------------------------------------------------------------- 1 | from argparse import Namespace 2 | import numpy as np 3 | import imageio 4 | import torch 5 | from pathlib import Path 6 | import sys 7 | 8 | cotr_path = Path(__file__).parent / "../../third_party/cotr" 9 | sys.path.append(str(cotr_path)) 10 | 11 | from COTR.models import build_model 12 | from .base import Matching 13 | from COTR.inference.sparse_engine import SparseEngine, FasterSparseEngine 14 | 15 | 16 | class COTR(Matching): 17 | def __init__(self, args): 18 | super().__init__() 19 | if type(args) == dict: 20 | args = Namespace(**args) 21 | self.imsize = args.imsize 22 | self.match_threshold = args.match_threshold 23 | self.batch_size = args.batch_size 24 | self.max_corrs = args.max_corrs 25 | args.dim_feedforward = args.backbone_layer_dims[args.layer] 26 | 27 | self.model = build_model(args) 28 | self.model.load_state_dict( 29 | torch.load(args.ckpt, map_location="cpu")["model_state_dict"] 30 | ) 31 | self.model = self.model.eval().to(self.device) 32 | self.name = "COTR" 33 | print(f"Initialize {self.name}") 34 | 35 | def match_pairs(self, im1_path, im2_path, queries_im1=None): 36 | im1 = imageio.imread(im1_path, pilmode="RGB") 37 | im2 = imageio.imread(im2_path, pilmode="RGB") 38 | engine = SparseEngine(self.model, self.batch_size, mode="tile") 39 | matches = engine.cotr_corr_multiscale( 40 | im1, 41 | im2, 42 | np.linspace(0.5, 0.0625, 4), 43 | 1, 44 | max_corrs=self.max_corrs, 45 | queries_a=queries_im1, 46 | force=True, 47 | ) 48 | 49 | # Fake scores as not output by the model 50 | scores = np.ones(len(matches)) 51 | kpts1 = matches[:, :2] 52 | kpts2 = matches[:, 2:4] 53 | return matches, kpts1, kpts2, scores 54 | -------------------------------------------------------------------------------- /immatch/modules/d2net.py: -------------------------------------------------------------------------------- 1 | from argparse import Namespace 2 | import torch 3 | import numpy as np 4 | import sys 5 | from pathlib import Path 6 | 7 | d2net_path = Path(__file__).parent / "../../third_party/d2net" 8 | sys.path.append(str(d2net_path)) 9 | 10 | from third_party.d2net.lib.model_test import D2Net as D2N 11 | from third_party.d2net.lib.utils import preprocess_image 12 | from third_party.d2net.lib.pyramid import process_multiscale 13 | from immatch.utils.data_io import read_im 14 | from .base import FeatureDetection, Matching 15 | 16 | 17 | class D2Net(FeatureDetection, Matching): 18 | def __init__(self, args=None): 19 | super().__init__() 20 | if type(args) == dict: 21 | args = Namespace(**args) 22 | self.model = D2N( 23 | model_file=args.ckpt, 24 | use_relu=args.use_relu, 25 | use_cuda=torch.cuda.is_available(), 26 | ) 27 | self.ms = args.multiscale 28 | self.imsize = args.imsize 29 | self.match_threshold = args.match_threshold 30 | self.name = "D2Net" 31 | print(f"Initialize {self.name}") 32 | 33 | def load_and_extract(self, im_path): 34 | im, scale = read_im(im_path, self.imsize) 35 | im = np.array(im) 36 | im = preprocess_image(im, preprocessing="caffe") 37 | kpts, desc = self.extract_features(im) 38 | kpts = kpts * scale # N, 2 39 | return kpts, desc 40 | 41 | def extract_features(self, im): 42 | if self.ms: 43 | keypoints, scores, descriptors = process_multiscale( 44 | torch.tensor( 45 | im[np.newaxis, :, :, :].astype(np.float32), device=self.device 46 | ), 47 | self.model, 48 | ) 49 | else: 50 | keypoints, scores, descriptors = process_multiscale( 51 | torch.tensor( 52 | im[np.newaxis, :, :, :].astype(np.float32), device=self.device 53 | ), 54 | self.model, 55 | scales=[1], 56 | ) 57 | 58 | kpts = keypoints[:, [1, 0]] # (x, y) and remove the scale 59 | desc = descriptors 60 | return kpts, desc 61 | 62 | def match_pairs(self, im1_path, im2_path): 63 | kpts1, desc1 = self.load_and_extract(im1_path) 64 | kpts2, desc2 = self.load_and_extract(im2_path) 65 | 66 | # NN Match 67 | match_ids, scores = self.mutual_nn_match( 68 | desc1, desc2, threshold=self.match_threshold 69 | ) 70 | p1s = kpts1[match_ids[:, 0], :2] 71 | p2s = kpts2[match_ids[:, 1], :2] 72 | matches = np.concatenate([p1s, p2s], axis=1) 73 | return matches, kpts1, kpts2, scores 74 | -------------------------------------------------------------------------------- /immatch/modules/dogaffnethardnet.py: -------------------------------------------------------------------------------- 1 | from argparse import Namespace 2 | import torch 3 | import numpy as np 4 | import cv2 5 | import kornia as K 6 | import kornia.feature as KF 7 | from kornia_moons.feature import laf_from_opencv_SIFT_kpts 8 | from .base import FeatureDetection, Matching 9 | from ..utils.data_io import read_im_gray 10 | 11 | 12 | class DogAffNetHardNet(FeatureDetection, Matching): 13 | def __init__(self, args=None): 14 | super().__init__() 15 | if type(args) == dict: 16 | args = Namespace(**args) 17 | self.imsize = args.imsize 18 | try: 19 | self.device = args.device 20 | except: 21 | self.device = torch.device("cpu") 22 | self.match_threshold = args.match_threshold 23 | self.det = cv2.SIFT_create( 24 | args.npts, contrastThreshold=-10000, edgeThreshold=-10000 25 | ) 26 | self.desc = KF.HardNet(True).eval().to(self.device) 27 | self.aff = KF.LAFAffNetShapeEstimator(True).eval().to(self.device) 28 | self.name = f"DoG{args.npts}-AffNet-HardNet" 29 | print(f"Initialize {self.name}") 30 | 31 | def load_im(self, im_path): 32 | im, scale = read_im_gray(im_path, self.imsize) 33 | im = np.array(im) 34 | return im, scale 35 | 36 | def load_and_extract(self, im_path): 37 | im, scale = self.load_im(im_path) 38 | kpts, desc = self.extract_features(im) 39 | kpts = kpts * scale 40 | return kpts, desc 41 | 42 | def extract_features(self, im): 43 | kpts = self.det.detect(im, None) 44 | # We will not train anything, so let's save time and memory by no_grad() 45 | with torch.no_grad(): 46 | timg = K.image_to_tensor(im, False).float() / 255.0 47 | timg = timg.to(self.device) 48 | if timg.shape[1] == 3: 49 | timg_gray = K.rgb_to_grayscale(timg) 50 | else: 51 | timg_gray = timg 52 | # kornia expects keypoints in the local affine frame format. 53 | # Luckily, kornia_moons has a conversion function 54 | lafs = laf_from_opencv_SIFT_kpts(kpts, device=self.device) 55 | lafs_new = self.aff(lafs, timg_gray) 56 | patches = KF.extract_patches_from_pyramid(timg_gray, lafs_new, 32) 57 | B, N, CH, H, W = patches.size() 58 | # Descriptor accepts standard tensor [B, CH, H, W], while patches are [B, N, CH, H, W] shape 59 | # So we need to reshape a bit :) 60 | descs = ( 61 | self.desc(patches.view(B * N, CH, H, W)) 62 | .view(B * N, -1) 63 | .detach() 64 | .cpu() 65 | .numpy() 66 | ) 67 | kpts = np.array([[kp.pt[0], kp.pt[1]] for kp in kpts]) 68 | return kpts, descs 69 | 70 | def load_and_detect(self, im_path): 71 | im, scale = self.load_im(im_path) 72 | kpts, desc = self.detect(im) 73 | kpts = kpts * scale 74 | return kpts 75 | 76 | def detect(self, im): 77 | kpts = self.det.detect(im) 78 | kpts = np.array([[kp.pt[0], kp.pt[1]] for kp in kpts]) 79 | return kpts 80 | 81 | def match_inputs_(self, im1, im2): 82 | kpts1, desc1 = self.extract_features(im1) 83 | kpts2, desc2 = self.extract_features(im2) 84 | 85 | # NN Match 86 | dists, match_ids = KF.match_smnn( 87 | torch.from_numpy(desc1), torch.from_numpy(desc2), self.match_threshold 88 | ) 89 | match_ids = match_ids.data.numpy() 90 | p1s = kpts1[match_ids[:, 0], :2] 91 | p2s = kpts2[match_ids[:, 1], :2] 92 | matches = np.concatenate([p1s, p2s], axis=1) 93 | scores = 1.0 - dists 94 | return matches, kpts1, kpts2, scores 95 | 96 | def match_pairs(self, im1_path, im2_path): 97 | im1, sc1 = self.load_im(im1_path) 98 | im2, sc2 = self.load_im(im2_path) 99 | 100 | upscale = np.array([sc1 + sc2]) 101 | matches, kpts1, kpts2, scores = self.match_inputs_(im1, im2) 102 | matches = upscale * matches 103 | kpts1 = sc1 * kpts1 104 | kpts2 = sc2 * kpts2 105 | return matches, kpts1, kpts2, scores 106 | -------------------------------------------------------------------------------- /immatch/modules/loftr.py: -------------------------------------------------------------------------------- 1 | from argparse import Namespace 2 | import torch 3 | import numpy as np 4 | import cv2 5 | 6 | from third_party.loftr.src.loftr import LoFTR as LoFTR_, default_cfg 7 | from .base import Matching 8 | from immatch.utils.data_io import load_gray_scale_tensor_cv 9 | 10 | 11 | class LoFTR(Matching): 12 | def __init__(self, args): 13 | super().__init__() 14 | if type(args) == dict: 15 | args = Namespace(**args) 16 | 17 | self.imsize = args.imsize 18 | self.match_threshold = args.match_threshold 19 | self.no_match_upscale = args.no_match_upscale 20 | self.eval_coarse = args.eval_coarse 21 | 22 | # Load model 23 | conf = dict(default_cfg) 24 | conf["match_coarse"]["thr"] = self.match_threshold 25 | self.model = LoFTR_(config=conf) 26 | ckpt_dict = torch.load(args.ckpt) 27 | self.model.load_state_dict(ckpt_dict["state_dict"]) 28 | self.model = self.model.eval().to(self.device) 29 | 30 | # Name the method 31 | self.ckpt_name = args.ckpt.split("/")[-1].split(".")[0] 32 | self.name = f"LoFTR_{self.ckpt_name}" 33 | if self.no_match_upscale: 34 | self.name += "_noms" 35 | print(f"Initialize {self.name}") 36 | 37 | def load_im(self, im_path): 38 | return load_gray_scale_tensor_cv( 39 | im_path, self.device, imsize=self.imsize, dfactor=8 40 | ) 41 | 42 | def match_inputs_(self, gray1, gray2): 43 | batch = {"image0": gray1, "image1": gray2} 44 | self.model(batch) 45 | 46 | if self.eval_coarse: 47 | kpts1 = batch["mkpts0_c"].cpu().numpy() 48 | kpts2 = batch["mkpts1_c"].cpu().numpy() 49 | else: 50 | kpts1 = batch["mkpts0_f"].cpu().numpy() 51 | kpts2 = batch["mkpts1_f"].cpu().numpy() 52 | scores = batch["mconf"].cpu().numpy() 53 | matches = np.concatenate([kpts1, kpts2], axis=1) 54 | return matches, kpts1, kpts2, scores 55 | 56 | def match_pairs(self, im1_path, im2_path): 57 | gray1, sc1, _ = self.load_im(im1_path) 58 | gray2, sc2, _ = self.load_im(im2_path) 59 | 60 | upscale = np.array([sc1 + sc2]) 61 | matches, kpts1, kpts2, scores = self.match_inputs_(gray1, gray2) 62 | 63 | if self.no_match_upscale: 64 | return matches, kpts1, kpts2, scores, upscale.squeeze(0) 65 | 66 | # Upscale matches & kpts 67 | matches = upscale * matches 68 | kpts1 = sc1 * kpts1 69 | kpts2 = sc2 * kpts2 70 | return matches, kpts1, kpts2, scores 71 | -------------------------------------------------------------------------------- /immatch/modules/nn_matching.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | 4 | 5 | def mutual_nn_matching_torch(desc1, desc2, threshold=None, eps=1e-9): 6 | if len(desc1) == 0 or len(desc2) == 0: 7 | return torch.empty((0, 2), dtype=torch.int64), torch.empty( 8 | (0, 2), dtype=torch.int64 9 | ) 10 | 11 | device = desc1.device 12 | desc1 = desc1 / (desc1.norm(dim=1, keepdim=True) + eps) 13 | desc2 = desc2 / (desc2.norm(dim=1, keepdim=True) + eps) 14 | similarity = torch.einsum("id, jd->ij", desc1, desc2) 15 | 16 | nn12 = similarity.max(dim=1)[1] 17 | nn21 = similarity.max(dim=0)[1] 18 | ids1 = torch.arange(0, similarity.shape[0], device=device) 19 | mask = ids1 == nn21[nn12] 20 | matches = torch.stack([ids1[mask], nn12[mask]]).t() 21 | scores = similarity.max(dim=1)[0][mask] 22 | if threshold: 23 | mask = scores > threshold 24 | matches = matches[mask] 25 | scores = scores[mask] 26 | return matches, scores 27 | 28 | 29 | def mutual_nn_matching(desc1, desc2, threshold=None): 30 | if isinstance(desc1, np.ndarray): 31 | desc1 = torch.from_numpy(desc1) 32 | desc2 = torch.from_numpy(desc2) 33 | matches, scores = mutual_nn_matching_torch(desc1, desc2, threshold=threshold) 34 | return matches.cpu().numpy(), scores.cpu().numpy() 35 | -------------------------------------------------------------------------------- /immatch/modules/patch2pix.py: -------------------------------------------------------------------------------- 1 | from argparse import Namespace 2 | import torch 3 | import numpy as np 4 | from pathlib import Path 5 | import sys 6 | 7 | patch2pix_path = Path(__file__).parent / "../../third_party/patch2pix" 8 | sys.path.append(str(patch2pix_path)) 9 | 10 | from third_party.patch2pix.utils.eval.model_helper import load_model, estimate_matches 11 | from third_party.patch2pix.utils.datasets.preprocess import load_im_flexible 12 | 13 | from .base import Matching 14 | import immatch 15 | from immatch.utils.data_io import load_im_tensor 16 | 17 | 18 | class Patch2Pix(Matching): 19 | def __init__(self, args): 20 | super().__init__() 21 | if type(args) == dict: 22 | args = Namespace(**args) 23 | self.imsize = args.imsize 24 | self.match_threshold = args.match_threshold 25 | self.no_match_upscale = args.no_match_upscale 26 | self.ksize = args.ksize 27 | self.model = load_model(args.ckpt, method="patch2pix") 28 | self.name = "Patch2Pix" 29 | print(f"Initialize {self.name}") 30 | 31 | def load_im(self, im_path): 32 | im, scale = load_im_flexible( 33 | im_path, 2, self.model.upsample, imsize=self.imsize 34 | ) 35 | im = im.unsqueeze(0).to(self.device) 36 | return im, scale 37 | 38 | def match_pairs(self, im1_path, im2_path): 39 | # Assume batch size is 1 40 | # Load images 41 | im1, sc1 = self.load_im(im1_path) 42 | im2, sc2 = self.load_im(im2_path) 43 | upscale = np.array([sc1 + sc2]) 44 | 45 | # Fine matches 46 | fine_matches, fine_scores, coarse_matches = self.model.predict_fine( 47 | im1, im2, ksize=self.ksize, ncn_thres=0.0, mutual=True 48 | ) 49 | coarse_matches = coarse_matches[0].cpu().data.numpy() 50 | fine_matches = fine_matches[0].cpu().data.numpy() 51 | fine_scores = fine_scores[0].cpu().data.numpy() 52 | 53 | # Inlier filtering 54 | pos_ids = np.where(fine_scores > self.match_threshold)[0] 55 | if len(pos_ids) > 0: 56 | coarse_matches = coarse_matches[pos_ids] 57 | matches = fine_matches[pos_ids] 58 | scores = fine_scores[pos_ids] 59 | else: 60 | # Simply take all matches for this case 61 | matches = fine_matches 62 | scores = fine_scores 63 | 64 | if self.no_match_upscale: 65 | return matches, matches[:, :2], matches[:, 2:4], scores, upscale.squeeze(0) 66 | 67 | matches = upscale * matches 68 | kpts1 = matches[:, :2] 69 | kpts2 = matches[:, 2:4] 70 | return matches, kpts1, kpts2, scores 71 | 72 | 73 | class NCNet(Matching): 74 | def __init__(self, args): 75 | super().__init__() 76 | if type(args) == dict: 77 | args = Namespace(**args) 78 | self.imsize = args.imsize 79 | self.match_threshold = args.match_threshold 80 | self.ksize = args.ksize 81 | self.model = load_model(args.ckpt, method="nc") 82 | self.name = "NCNet" 83 | print(f"Initialize {self.name}") 84 | 85 | def match_pairs(self, im1_path, im2_path): 86 | matches, scores, _ = estimate_matches( 87 | self.model, 88 | im1_path, 89 | im2_path, 90 | ksize=self.ksize, 91 | ncn_thres=self.match_threshold, 92 | eval_type="coarse", 93 | imsize=self.imsize, 94 | ) 95 | kpts1 = matches[:, :2] 96 | kpts2 = matches[:, 2:4] 97 | return matches, kpts1, kpts2, scores 98 | 99 | 100 | class Patch2PixRefined(Matching): 101 | def __init__(self, args): 102 | super().__init__() 103 | # Initialize coarse matcher 104 | cargs = args["coarse"] 105 | self.cname = cargs["name"] 106 | self.cmatcher = immatch.__dict__[self.cname](cargs) 107 | 108 | # Initialize patch2pix 109 | args = Namespace(**args) 110 | self.imsize = args.imsize 111 | self.match_threshold = args.match_threshold 112 | self.fmatcher = load_model(args.ckpt, method="patch2pix") 113 | self.name = f"Patch2Pix_{self.cmatcher.name}_m" + str(cargs["match_threshold"]) 114 | print(f"Initialize {self.name}") 115 | 116 | def match_pairs(self, im1_path, im2_path): 117 | im1, grey1, sc1 = load_im_tensor( 118 | im1_path, self.device, imsize=self.imsize, with_gray=True 119 | ) 120 | im2, grey2, sc2 = load_im_tensor( 121 | im2_path, self.device, imsize=self.imsize, with_gray=True 122 | ) 123 | if self.cname in ["SuperGlue"]: 124 | coarse_match_res = self.cmatcher.match_inputs_(grey1, grey2) 125 | elif self.cname == "CAPS": 126 | coarse_match_res = self.cmatcher.match_inputs_(im1, grey1, im2, grey2) 127 | coarse_matches = coarse_match_res[0] 128 | 129 | # Patch2Pix refinement 130 | refined_matches, scores, _ = self.fmatcher.refine_matches( 131 | im1, im2, coarse_matches, io_thres=self.match_threshold 132 | ) 133 | 134 | upscale = np.array([sc1 + sc2]) 135 | matches = upscale * refined_matches 136 | kpts1 = matches[:, :2] 137 | kpts2 = matches[:, 2:4] 138 | return matches, kpts1, kpts2, scores 139 | -------------------------------------------------------------------------------- /immatch/modules/r2d2.py: -------------------------------------------------------------------------------- 1 | from argparse import Namespace 2 | import torch 3 | import numpy as np 4 | import sys 5 | from pathlib import Path 6 | 7 | r2d2_path = Path(__file__).parent / "../../third_party/r2d2" 8 | sys.path.append(str(r2d2_path)) 9 | 10 | from third_party.r2d2.extract import NonMaxSuppression, extract_multiscale 11 | from third_party.r2d2.tools.dataloader import norm_RGB 12 | from third_party.r2d2.nets.patchnet import * 13 | from PIL import Image 14 | from .base import FeatureDetection, Matching 15 | 16 | 17 | class R2D2(FeatureDetection, Matching): 18 | def __init__(self, args=None): 19 | super().__init__() 20 | if type(args) == dict: 21 | args = Namespace(**args) 22 | self.args = args 23 | 24 | # Initialize model 25 | ckpt = torch.load(args.ckpt) 26 | self.model = eval(ckpt["net"]).to(self.device).eval() 27 | self.model.load_state_dict( 28 | {k.replace("module.", ""): v for k, v in ckpt["state_dict"].items()} 29 | ) 30 | self.name = "R2D2" 31 | print(f"Initialize {self.name}") 32 | 33 | # Init NMS Detector 34 | self.detector = NonMaxSuppression( 35 | rel_thr=args.reliability_thr, rep_thr=args.repeatability_thr 36 | ) 37 | 38 | def load_and_extract(self, im_path): 39 | # No image resize here 40 | im = Image.open(im_path).convert("RGB") 41 | im = norm_RGB(im)[None].to(self.device) 42 | kpts, desc = self.extract_features(im) 43 | return kpts, desc 44 | 45 | def extract_features(self, im): 46 | args = self.args 47 | max_size = 9999 if args.imsize < 0 else args.imsize 48 | xys, desc, scores = extract_multiscale( 49 | self.model, 50 | im, 51 | self.detector, 52 | min_scale=args.min_scale, 53 | max_scale=args.max_scale, 54 | min_size=args.min_size, 55 | max_size=max_size, 56 | verbose=False, 57 | ) 58 | idxs = scores.argsort()[-args.top_k or None :] 59 | kpts = xys[idxs] 60 | desc = desc[idxs] 61 | return kpts, desc 62 | 63 | def match_pairs(self, im1_path, im2_path): 64 | kpts1, desc1 = self.load_and_extract(im1_path) 65 | kpts2, desc2 = self.load_and_extract(im2_path) 66 | 67 | # NN Match 68 | match_ids, scores = self.mutual_nn_match( 69 | desc1, desc2, threshold=self.args.match_threshold 70 | ) 71 | p1s = kpts1[match_ids[:, 0], :2].cpu().numpy() 72 | p2s = kpts2[match_ids[:, 1], :2].cpu().numpy() 73 | matches = np.concatenate([p1s, p2s], axis=1) 74 | return matches, kpts1, kpts2, scores 75 | -------------------------------------------------------------------------------- /immatch/modules/sift.py: -------------------------------------------------------------------------------- 1 | from argparse import Namespace 2 | import torch 3 | import numpy as np 4 | import cv2 5 | 6 | from .base import FeatureDetection, Matching 7 | from ..utils.data_io import read_im_gray 8 | 9 | 10 | class SIFT(FeatureDetection, Matching): 11 | def __init__(self, args=None): 12 | super().__init__() 13 | if type(args) == dict: 14 | args = Namespace(**args) 15 | self.imsize = args.imsize 16 | self.match_threshold = args.match_threshold 17 | self.model = cv2.SIFT_create(args.npts) 18 | self.name = f"SIFT{args.npts}" 19 | print(f"Initialize {self.name}") 20 | 21 | def load_im(self, im_path): 22 | im, scale = read_im_gray(im_path, self.imsize) 23 | im = np.array(im) 24 | return im, scale 25 | 26 | def load_and_extract(self, im_path): 27 | im, scale = self.load_im(im_path) 28 | kpts, desc = self.extract_features(im) 29 | kpts = kpts * scale 30 | return kpts, desc 31 | 32 | def extract_features(self, im): 33 | kpts, desc = self.model.detectAndCompute(im, None) 34 | kpts = np.array([[kp.pt[0], kp.pt[1]] for kp in kpts]) 35 | return kpts, desc 36 | 37 | def load_and_detect(self, im_path): 38 | im, scale = self.load_im(im_path) 39 | kpts = self.detect(im) 40 | kpts = kpts * scale 41 | return kpts 42 | 43 | def detect(self, im): 44 | kpts = self.model.detect(im) 45 | kpts = np.array([[kp.pt[0], kp.pt[1]] for kp in kpts]) 46 | return kpts 47 | 48 | def match_inputs_(self, im1, im2): 49 | kpts1, desc1 = self.extract_features(im1) 50 | kpts2, desc2 = self.extract_features(im2) 51 | 52 | # NN Match 53 | match_ids, scores = self.mutual_nn_match( 54 | desc1, desc2, threshold=self.match_threshold 55 | ) 56 | p1s = kpts1[match_ids[:, 0], :2] 57 | p2s = kpts2[match_ids[:, 1], :2] 58 | matches = np.concatenate([p1s, p2s], axis=1) 59 | return matches, kpts1, kpts2, scores 60 | 61 | def match_pairs(self, im1_path, im2_path): 62 | im1, sc1 = self.load_im(im1_path) 63 | im2, sc2 = self.load_im(im2_path) 64 | 65 | upscale = np.array([sc1 + sc2]) 66 | matches, kpts1, kpts2, scores = self.match_inputs_(im1, im2) 67 | matches = upscale * matches 68 | kpts1 = sc1 * kpts1 69 | kpts2 = sc2 * kpts2 70 | return matches, kpts1, kpts2, scores 71 | -------------------------------------------------------------------------------- /immatch/modules/sparsencnet.py: -------------------------------------------------------------------------------- 1 | from argparse import Namespace 2 | import torch 3 | from pathlib import Path 4 | import sys 5 | import numpy as np 6 | 7 | sparsencnet_path = Path(__file__).parent / "../../third_party/sparsencnet" 8 | sys.path.append(str(sparsencnet_path)) 9 | 10 | from third_party.sparsencnet.lib.model import ImMatchNet 11 | from third_party.sparsencnet.lib.normalization import imreadth, resize, normalize 12 | from third_party.sparsencnet.lib.sparse import get_matches_both_dirs, unique 13 | from third_party.sparsencnet.lib.relocalize import ( 14 | relocalize, 15 | relocalize_soft, 16 | eval_model_reloc, 17 | ) 18 | from .base import Matching 19 | 20 | 21 | def load_im(im_path, scale_factor, imsize=None): 22 | im = imreadth(im_path) 23 | h, w = im.shape[-2:] 24 | if not imsize: 25 | imsize = max(h, w) 26 | else: 27 | imsize = imsize 28 | im = resize(normalize(im), imsize, scale_factor) 29 | return im, h, w 30 | 31 | 32 | class SparseNCNet(Matching): 33 | def __init__(self, args): 34 | super().__init__() 35 | if type(args) == dict: 36 | args = Namespace(**args) 37 | self.args = args 38 | self.match_threshold = args.match_threshold 39 | self.model = self.init_model(args) 40 | self.name = f"SparseNCNet" 41 | if args.Npts: 42 | self.name += f"_N{args.Npts}" 43 | print(f"Initialize {self.name}") 44 | 45 | def init_model(self, args): 46 | chp_args = torch.load(args.ckpt)["args"] 47 | model = ImMatchNet( 48 | use_cuda=torch.cuda.is_available(), 49 | checkpoint=args.ckpt, 50 | ncons_kernel_sizes=chp_args.ncons_kernel_sizes, 51 | ncons_channels=chp_args.ncons_channels, 52 | sparse=True, 53 | symmetric_mode=bool(chp_args.symmetric_mode), 54 | feature_extraction_cnn=chp_args.feature_extraction_cnn, 55 | bn=bool(chp_args.bn), 56 | k=chp_args.k, 57 | return_fs=True, 58 | change_stride=args.change_stride, 59 | ) 60 | 61 | scale_factor = 0.0625 62 | if args.relocalize == 1: 63 | scale_factor = scale_factor / 2 64 | if args.change_stride == 1: 65 | scale_factor = scale_factor * 2 66 | args.scale_factor = scale_factor 67 | return model 68 | 69 | def match_pairs(self, im1_path, im2_path): 70 | args = self.args 71 | 72 | im1, hA, wA = load_im(im1_path, args.scale_factor, args.imsize) 73 | im2, hB, wB = load_im(im2_path, args.scale_factor, args.imsize) 74 | # print('Ims', im1.shape, im2.shape) 75 | corr4d, feature_A_2x, feature_B_2x, fs1, fs2, fs3, fs4 = eval_model_reloc( 76 | self.model, {"source_image": im1, "target_image": im2}, args 77 | ) 78 | xA_, yA_, xB_, yB_, score_ = get_matches_both_dirs(corr4d, fs1, fs2, fs3, fs4) 79 | 80 | if args.Npts is not None: 81 | matches_idx_sorted = torch.argsort(-score_.view(-1)) 82 | N_matches = min(args.Npts, matches_idx_sorted.shape[0]) 83 | matches_idx_sorted = matches_idx_sorted[:N_matches] 84 | score_ = score_[:, matches_idx_sorted] 85 | xA_ = xA_[:, matches_idx_sorted] 86 | yA_ = yA_[:, matches_idx_sorted] 87 | xB_ = xB_[:, matches_idx_sorted] 88 | yB_ = yB_[:, matches_idx_sorted] 89 | 90 | if args.relocalize: 91 | fs1, fs2, fs3, fs4 = 2 * fs1, 2 * fs2, 2 * fs3, 2 * fs4 92 | # relocalization stage 1: 93 | if args.reloc_type.startswith("hard"): 94 | xA_, yA_, xB_, yB_, score_ = relocalize( 95 | xA_, 96 | yA_, 97 | xB_, 98 | yB_, 99 | score_, 100 | feature_A_2x, 101 | feature_B_2x, 102 | crop_size=args.reloc_hard_crop_size, 103 | ) 104 | if args.reloc_hard_crop_size == 3: 105 | _, uidx = unique( 106 | yA_.double() * fs2 * fs3 * fs4 107 | + xA_.double() * fs3 * fs4 108 | + yB_.double() * fs4 109 | + xB_.double(), 110 | return_index=True, 111 | ) 112 | xA_ = xA_[:, uidx] 113 | yA_ = yA_[:, uidx] 114 | xB_ = xB_[:, uidx] 115 | yB_ = yB_[:, uidx] 116 | score_ = score_[:, uidx] 117 | elif args.reloc_type == "soft": 118 | xA_, yA_, xB_, yB_, score_ = relocalize_soft( 119 | xA_, yA_, xB_, yB_, score_, feature_A_2x, feature_B_2x 120 | ) 121 | 122 | # relocalization stage 2: 123 | if args.reloc_type == "hard_soft": 124 | xA_, yA_, xB_, yB_, score_ = relocalize_soft( 125 | xA_, 126 | yA_, 127 | xB_, 128 | yB_, 129 | score_, 130 | feature_A_2x, 131 | feature_B_2x, 132 | upsample_positions=False, 133 | ) 134 | 135 | elif args.reloc_type == "hard_hard": 136 | xA_, yA_, xB_, yB_, score_ = relocalize( 137 | xA_, 138 | yA_, 139 | xB_, 140 | yB_, 141 | score_, 142 | feature_A_2x, 143 | feature_B_2x, 144 | upsample_positions=False, 145 | ) 146 | 147 | yA_ = (yA_ + 0.5) / (fs1) 148 | xA_ = (xA_ + 0.5) / (fs2) 149 | yB_ = (yB_ + 0.5) / (fs3) 150 | xB_ = (xB_ + 0.5) / (fs4) 151 | 152 | xA = xA_.view(-1).data.cpu().float().numpy() * wA 153 | yA = yA_.view(-1).data.cpu().float().numpy() * hA 154 | xB = xB_.view(-1).data.cpu().float().numpy() * wB 155 | yB = yB_.view(-1).data.cpu().float().numpy() * hB 156 | scores = score_.view(-1).data.cpu().float().numpy() 157 | 158 | matches = np.stack((xA, yA, xB, yB), axis=1) 159 | kpts1 = matches[:, :2] 160 | kpts2 = matches[:, 2:] 161 | return matches, kpts1, kpts2, scores 162 | -------------------------------------------------------------------------------- /immatch/modules/superglue.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from argparse import Namespace 3 | import numpy as np 4 | 5 | from third_party.superglue.models.superglue import SuperGlue as SG 6 | from .superpoint import SuperPoint 7 | from .base import Matching 8 | from immatch.utils.data_io import load_gray_scale_tensor_cv 9 | 10 | 11 | class SuperGlue(Matching): 12 | def __init__(self, args): 13 | super().__init__() 14 | self.imsize = args["imsize"] 15 | self.no_match_upscale = args.get("no_match_upscale", False) 16 | 17 | self.model = SG(args).eval().to(self.device) 18 | self.detector = SuperPoint(args) 19 | rad = self.detector.model.config["nms_radius"] 20 | self.name = f"SuperGlue_r{rad}" 21 | print(f"Initialize {self.name}") 22 | 23 | def load_im(self, im_path): 24 | return load_gray_scale_tensor_cv(im_path, self.device, imsize=self.imsize) 25 | 26 | def match_inputs_(self, gray1, gray2): 27 | # Detect SuperPoint features 28 | pred1 = self.detector.model({"image": gray1}) 29 | pred2 = self.detector.model({"image": gray2}) 30 | 31 | # Construct SuperGlue input 32 | data = {"image0": gray1, "image1": gray2} 33 | pred = {k + "0": v for k, v in pred1.items()} 34 | pred = {**pred, **{k + "1": v for k, v in pred2.items()}} 35 | data = {**data, **pred} 36 | for k in data: 37 | if isinstance(data[k], (list, tuple)): 38 | data[k] = torch.stack(data[k]) 39 | 40 | # SuperGlue matching 41 | pred = {**pred, **self.model(data)} 42 | pred = {k: v[0].cpu().numpy() for k, v in pred.items()} 43 | kpts1, kpts2 = pred["keypoints0"], pred["keypoints1"] 44 | matches, scores = pred["matches0"], pred["matching_scores0"] 45 | valid = matches > -1 46 | p1s = kpts1[valid] 47 | p2s = kpts2[matches[valid]] 48 | matches = np.concatenate([p1s, p2s], axis=1) 49 | scores = scores[valid] 50 | return matches, kpts1, kpts2, scores 51 | 52 | def match_pairs(self, im1_path, im2_path): 53 | gray1, sc1, _ = self.load_im(im1_path) 54 | gray2, sc2, _ = self.load_im(im2_path) 55 | upscale = np.array([sc1 + sc2]) 56 | matches, kpts1, kpts2, scores = self.match_inputs_(gray1, gray2) 57 | 58 | if self.no_match_upscale: 59 | return matches, kpts1, kpts2, scores, upscale.squeeze(0) 60 | 61 | # Upscale matches & kpts 62 | matches = upscale * matches 63 | kpts1 = sc1 * kpts1 64 | kpts2 = sc2 * kpts2 65 | return matches, kpts1, kpts2, scores 66 | -------------------------------------------------------------------------------- /immatch/modules/superpoint.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | 4 | from third_party.superglue.models.superpoint import SuperPoint as SP 5 | from .base import FeatureDetection, Matching 6 | from immatch.utils.data_io import load_gray_scale_tensor_cv 7 | 8 | 9 | class SuperPoint(FeatureDetection, Matching): 10 | def __init__(self, args=None): 11 | super().__init__() 12 | self.imsize = args["imsize"] 13 | self.match_threshold = ( 14 | args["match_threshold"] if "match_threshold" in args else 0.0 15 | ) 16 | self.model = SP(args).eval().to(self.device) 17 | rad = self.model.config["nms_radius"] 18 | self.name = f"SuperPoint_r{rad}" 19 | print(f"Initialize {self.name}") 20 | 21 | def load_im(self, im_path): 22 | return load_gray_scale_tensor_cv(im_path, self.device, imsize=self.imsize) 23 | 24 | def load_and_extract(self, im_path): 25 | gray, scale, _ = self.load_im(im_path) 26 | kpts, desc = self.extract_features(gray) 27 | kpts = kpts * torch.tensor(scale).to(kpts) # N, 2 28 | return kpts, desc 29 | 30 | def extract_features(self, gray): 31 | # SuperPoint outputs: {keypoints, scores, descriptors} 32 | pred = self.model({"image": gray}) 33 | kpts = pred["keypoints"][0] 34 | desc = pred["descriptors"][0].permute(1, 0) # N, D 35 | return kpts, desc 36 | 37 | def detect(self, gray): 38 | kpts, _ = self.extract_features(gray) 39 | return kpts 40 | 41 | def match_inputs_(self, gray1, gray2): 42 | kpts1, desc1 = self.extract_features(gray1) 43 | kpts2, desc2 = self.extract_features(gray2) 44 | kpts1 = kpts1.cpu().data.numpy() 45 | kpts2 = kpts2.cpu().data.numpy() 46 | 47 | # NN Match 48 | match_ids, scores = self.mutual_nn_match( 49 | desc1, desc2, threshold=self.match_threshold 50 | ) 51 | p1s = kpts1[match_ids[:, 0], :2] 52 | p2s = kpts2[match_ids[:, 1], :2] 53 | matches = np.concatenate([p1s, p2s], axis=1) 54 | return matches, kpts1, kpts2, scores 55 | 56 | def match_pairs(self, im1_path, im2_path): 57 | gray1, sc1, _ = self.load_im(im1_path) 58 | gray2, sc2, _ = self.load_im(im2_path) 59 | upscale = np.array([sc1 + sc2]) 60 | matches, kpts1, kpts2, scores = self.match_inputs_(gray1, gray2) 61 | matches = upscale * matches 62 | kpts1 = sc1 * kpts1 63 | kpts2 = sc2 * kpts2 64 | return matches, kpts1, kpts2, scores 65 | -------------------------------------------------------------------------------- /immatch/utils/__init__.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from third_party.patch2pix.utils.common.plotting import plot_matches 4 | 5 | 6 | def model_params(model): 7 | return sum(param.numel() for param in model.parameters()) / 1e6 8 | 9 | 10 | def dict_to_deivce(batch, device): 11 | for k, v in batch.items(): 12 | if isinstance(v, torch.Tensor): 13 | batch[k] = v.to(device) 14 | return batch 15 | -------------------------------------------------------------------------------- /immatch/utils/colmap/data_parsing.py: -------------------------------------------------------------------------------- 1 | from collections import defaultdict 2 | import sqlite3 3 | import numpy as np 4 | from argparse import Namespace 5 | from transforms3d.quaternions import quat2mat, mat2quat 6 | import shutil 7 | from pathlib import Path 8 | from .read_write_model import * 9 | from .database import COLMAPDatabase 10 | 11 | 12 | def load_ids_from_database(database_path): 13 | images = {} 14 | cameras = {} 15 | db = sqlite3.connect(str(database_path)) 16 | ret = db.execute("SELECT name, image_id, camera_id FROM images;") 17 | for name, image_id, camera_id in ret: 18 | images[name] = image_id 19 | cameras[name] = camera_id 20 | db.close() 21 | print(f"Found {len(images)} images and {len(cameras)} cameras in database.") 22 | return images, cameras 23 | 24 | 25 | def load_cameras_from_database(database_path): 26 | print(f"Parsing intrinsics from {database_path}...") 27 | 28 | cameras = {} 29 | db = sqlite3.connect(str(database_path)) 30 | ret = db.execute("SELECT camera_id, model, width, height, params FROM cameras;") 31 | for camera_id, camera_model, width, height, params in ret: 32 | params = np.frombuffer(params, dtype=np.double).reshape(-1) 33 | camera_model = CAMERA_MODEL_IDS[camera_model] 34 | camera = Camera( 35 | id=camera_id, 36 | model=camera_model.model_name, 37 | width=int(width), 38 | height=int(height), 39 | params=params, 40 | ) 41 | cameras[camera_id] = camera 42 | return cameras 43 | 44 | 45 | def load_cameras_from_intrinsics_and_ids(intrinsic_txt, camera_ids): 46 | print(f"Parsing intrinsics from {intrinsic_txt}...") 47 | 48 | cameras = {} 49 | with open(intrinsic_txt, "r") as f: 50 | for line in f: 51 | intrinsics = line.split() 52 | name, camera_model, width, height = intrinsics[:4] 53 | params = [float(p) for p in intrinsics[4:]] 54 | camera_model = CAMERA_MODEL_NAMES[camera_model] 55 | assert len(params) == camera_model.num_params 56 | camera_id = camera_ids[name] 57 | camera = Camera( 58 | id=camera_id, 59 | model=camera_model.model_name, 60 | width=int(width), 61 | height=int(height), 62 | params=params, 63 | ) 64 | cameras[camera_id] = camera 65 | return cameras 66 | 67 | 68 | def load_images_from_nvm(nvm_path): 69 | images = {} 70 | with open(nvm_path, "r") as f: 71 | # Skip headers 72 | line = next(f) 73 | while line == "\n" or line.startswith("NVM_V3"): 74 | line = next(f) 75 | 76 | # Parse images 77 | num_images = int(line) 78 | images_empty = dict() 79 | for i in range(num_images): 80 | data = next(f).split() 81 | im_name = data[0] 82 | qvec = np.array(data[2:6], float) 83 | c = np.array(data[6:9], float) 84 | 85 | # NVM -> COLMAP. 86 | R = quat2mat(qvec) 87 | tvec = -np.matmul(R, c) 88 | images[im_name] = dict(qvec=qvec, tvec=tvec) 89 | print(f"Loaded {len(images)} images from {nvm_path}.") 90 | return images 91 | 92 | 93 | def create_empty_model_from_reference_model(reference_model, empty_model): 94 | empty_model = Path(empty_model) 95 | if os.path.exists(empty_model): 96 | print(f"Empty sfm {empty_model} existed.") 97 | return 98 | os.makedirs(empty_model) 99 | print(f"Creating an empty sfm under {empty_model} from {reference_model}") 100 | 101 | # Construct images with fake points 102 | images = read_images_binary(str(reference_model / "images.bin")) 103 | print(f"Loaded {len(images)} images") 104 | images_empty = dict() 105 | for id_, image in images.items(): 106 | image = image._asdict() 107 | image["xys"] = np.zeros((0, 2), float) 108 | image["point3D_ids"] = np.full(0, -1, int) 109 | images_empty[id_] = Image(**image) 110 | write_images_binary(images_empty, empty_model / "images.bin") 111 | shutil.copy(reference_model / "cameras.bin", empty_model) 112 | write_points3d_binary(dict(), empty_model / "points3D.bin") 113 | 114 | 115 | def create_empty_model_from_nvm_and_database( 116 | nvm_path, database, empty_model, intrinsic_txt=None 117 | ): 118 | empty_model = Path(empty_model) 119 | if empty_model.exists(): 120 | print(f"Empty sfm {empty_model} existed.") 121 | return 122 | os.makedirs(empty_model) 123 | print(f"Creating an empty sfm under {empty_model} from {nvm_path}") 124 | 125 | # Construct images with fake points 126 | images_empty = dict() 127 | image_ids, camera_ids = load_ids_from_database(database) 128 | images = load_images_from_nvm(nvm_path) 129 | for name in images: 130 | qvec = images[name]["qvec"] 131 | tvec = images[name]["tvec"] 132 | name = name.lstrip("./") 133 | image_id = image_ids[name] 134 | image = Image( 135 | id=image_id, 136 | qvec=qvec, 137 | tvec=tvec, 138 | camera_id=camera_ids[name], 139 | name=name.replace("png", "jpg"), # Needed by RobotCar 140 | xys=np.zeros((0, 2), float), 141 | point3D_ids=np.full(0, -1, int), 142 | ) 143 | images_empty[image_id] = image 144 | write_images_binary(images_empty, empty_model / "images.bin") 145 | 146 | if intrinsic_txt and intrinsic_txt.exists(): 147 | cameras = load_cameras_from_intrinsics_and_ids( 148 | intrinsic_txt, camera_ids 149 | ) # For aachen v1 150 | else: 151 | cameras = load_cameras_from_database(database) 152 | write_cameras_binary(cameras, empty_model / "cameras.bin") 153 | write_points3d_binary(dict(), empty_model / "points3D.bin") 154 | 155 | 156 | def init_database_from_empty_model_binary(empty_model, database_path): 157 | if database_path.exists(): 158 | print("Database already exists.") 159 | 160 | cameras = read_cameras_binary(str(model / "cameras.bin")) 161 | images = read_images_binary(str(model / "images.bin")) 162 | 163 | db = COLMAPDatabase.connect(database_path) 164 | db.create_tables() 165 | 166 | for i, camera in cameras.items(): 167 | model_id = CAMERA_MODEL_NAMES[camera.model].model_id 168 | db.add_camera( 169 | model_id, 170 | camera.width, 171 | camera.height, 172 | camera.params, 173 | camera_id=i, 174 | prior_focal_length=True, 175 | ) 176 | 177 | for i, image in images.items(): 178 | db.add_image(image.name, image.camera_id, image_id=i) 179 | 180 | db.commit() 181 | db.close() 182 | return {image.name: i for i, image in images.items()} 183 | 184 | 185 | def covis_pairs_from_nvm(nvm_path, odir, topk=20): 186 | image_names = [] 187 | image_ids_to_point_ids = {} 188 | point_ids_to_image_ids = {} 189 | 190 | with open(nvm_path, "r") as f: 191 | # Skip headers 192 | line = next(f) 193 | while line == "\n" or line.startswith("NVM_V3"): 194 | line = next(f) 195 | 196 | # Load images 197 | num_images = int(line) 198 | for im_id in range(num_images): 199 | im_name = next(f).split()[0] 200 | im_name = im_name.lstrip("./").replace("png", "jpg") # For robotcar 201 | image_names.append(im_name) 202 | image_ids_to_point_ids[im_id] = [] 203 | 204 | # Load 3D points 205 | line = next(f) 206 | while line == "\n": 207 | line = next(f) 208 | num_points = int(line) 209 | for pid in range(num_points): 210 | line = next(f) 211 | data = line.split() 212 | num_measurements = int(data[6]) 213 | point_ids_to_image_ids[pid] = [] 214 | for j in range(num_measurements): 215 | im_id = int(data[7 + j * 4]) 216 | im_name = image_names[im_id] 217 | image_ids_to_point_ids[im_id].append(pid) 218 | point_ids_to_image_ids[pid].append(im_id) 219 | print(f"Loaded {num_images} images {num_points} points.") 220 | 221 | # Covisible pairs 222 | pairs = [] 223 | if not os.path.exists(odir): 224 | os.makedirs(odir) 225 | out_txt = os.path.join(odir, f"pairs-db-covis{topk}.txt") 226 | with open(out_txt, "w") as f: 227 | for im_id, im_name in enumerate(image_names): 228 | covis = defaultdict(int) 229 | visible_point_ids = image_ids_to_point_ids[im_id] 230 | for pid in visible_point_ids: 231 | covis_im_ids = point_ids_to_image_ids[pid] 232 | for cim_id in covis_im_ids: 233 | if cim_id != im_id: 234 | covis[cim_id] += 1 235 | 236 | if len(covis) == 0: 237 | print(f"Image {im_name} does not have any covisibility.") 238 | continue 239 | 240 | covis_ids = np.array(list(covis.keys())) 241 | covis_num = np.array([covis[i] for i in covis_ids]) 242 | covis_ids_sorted = covis_ids[np.argsort(-covis_num)] 243 | top_covis_ids = covis_ids_sorted[: min(topk, len(covis_ids))] 244 | for i in top_covis_ids: 245 | im1, im2 = im_name, image_names[i] 246 | pairs.append((im1, im2)) 247 | f.write(f"{im1} {im2}\n") 248 | print(f"Writing {len(pairs)} covis-{topk} pairs to {out_txt}") 249 | 250 | 251 | def covis_pairs_from_reference_model(reference_model, odir, topk=20): 252 | cameras, images, points3D = read_model(reference_model, ".bin") 253 | 254 | # Covisible pairs 255 | pairs = [] 256 | if not os.path.exists(odir): 257 | os.makedirs(odir) 258 | out_txt = os.path.join(odir, f"pairs-db-covis{topk}.txt") 259 | with open(out_txt, "w") as f: 260 | for im_id, image in images.items(): 261 | covis = defaultdict(int) 262 | visible_point_ids = image.point3D_ids[image.point3D_ids != -1] 263 | for pid in visible_point_ids: 264 | covis_im_ids = point_ids_to_image_ids[pid] 265 | for cim_id in points3D[pid].image_ids: 266 | if cim_id != im_id: 267 | covis[cim_id] += 1 268 | 269 | if len(covis) == 0: 270 | print(f"Image {image.name} does not have any covisibility.") 271 | continue 272 | 273 | covis_ids = np.array(list(covis.keys())) 274 | covis_num = np.array([covis[i] for i in covis_ids]) 275 | covis_ids_sorted = covis_ids[np.argsort(-covis_num)] 276 | top_covis_ids = covis_ids_sorted[: min(topk, len(covis_ids))] 277 | for i in top_covis_ids: 278 | im1, im2 = image.name, images[i].name 279 | pairs.append((im1, im2)) 280 | f.write(f"{im1} {im2}\n") 281 | print(f"Writing {len(pairs)} covis-{topk} pairs to {out_txt}") 282 | -------------------------------------------------------------------------------- /immatch/utils/colmap/database.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 | # This script is based on an original implementation by True Price. 33 | 34 | # Source code from: 35 | # https://github.com/colmap/colmap/blob/dev/scripts/python/database.py 36 | 37 | import sys 38 | import sqlite3 39 | import numpy as np 40 | 41 | 42 | IS_PYTHON3 = sys.version_info[0] >= 3 43 | 44 | MAX_IMAGE_ID = 2**31 - 1 45 | 46 | CREATE_CAMERAS_TABLE = """CREATE TABLE IF NOT EXISTS cameras ( 47 | camera_id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL, 48 | model INTEGER NOT NULL, 49 | width INTEGER NOT NULL, 50 | height INTEGER NOT NULL, 51 | params BLOB, 52 | prior_focal_length INTEGER NOT NULL)""" 53 | 54 | CREATE_DESCRIPTORS_TABLE = """CREATE TABLE IF NOT EXISTS descriptors ( 55 | image_id INTEGER PRIMARY KEY NOT NULL, 56 | rows INTEGER NOT NULL, 57 | cols INTEGER NOT NULL, 58 | data BLOB, 59 | FOREIGN KEY(image_id) REFERENCES images(image_id) ON DELETE CASCADE)""" 60 | 61 | CREATE_IMAGES_TABLE = """CREATE TABLE IF NOT EXISTS images ( 62 | image_id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL, 63 | name TEXT NOT NULL UNIQUE, 64 | camera_id INTEGER NOT NULL, 65 | prior_qw REAL, 66 | prior_qx REAL, 67 | prior_qy REAL, 68 | prior_qz REAL, 69 | prior_tx REAL, 70 | prior_ty REAL, 71 | prior_tz REAL, 72 | CONSTRAINT image_id_check CHECK(image_id >= 0 and image_id < {}), 73 | FOREIGN KEY(camera_id) REFERENCES cameras(camera_id)) 74 | """.format( 75 | MAX_IMAGE_ID 76 | ) 77 | 78 | CREATE_TWO_VIEW_GEOMETRIES_TABLE = """ 79 | CREATE TABLE IF NOT EXISTS two_view_geometries ( 80 | pair_id INTEGER PRIMARY KEY NOT NULL, 81 | rows INTEGER NOT NULL, 82 | cols INTEGER NOT NULL, 83 | data BLOB, 84 | config INTEGER NOT NULL, 85 | F BLOB, 86 | E BLOB, 87 | H BLOB, 88 | qvec BLOB, 89 | tvec BLOB) 90 | """ 91 | 92 | CREATE_KEYPOINTS_TABLE = """CREATE TABLE IF NOT EXISTS keypoints ( 93 | image_id INTEGER PRIMARY KEY NOT NULL, 94 | rows INTEGER NOT NULL, 95 | cols INTEGER NOT NULL, 96 | data BLOB, 97 | FOREIGN KEY(image_id) REFERENCES images(image_id) ON DELETE CASCADE) 98 | """ 99 | 100 | CREATE_MATCHES_TABLE = """CREATE TABLE IF NOT EXISTS matches ( 101 | pair_id INTEGER PRIMARY KEY NOT NULL, 102 | rows INTEGER NOT NULL, 103 | cols INTEGER NOT NULL, 104 | data BLOB)""" 105 | 106 | CREATE_NAME_INDEX = "CREATE UNIQUE INDEX IF NOT EXISTS index_name ON images(name)" 107 | 108 | CREATE_ALL = "; ".join( 109 | [ 110 | CREATE_CAMERAS_TABLE, 111 | CREATE_IMAGES_TABLE, 112 | CREATE_KEYPOINTS_TABLE, 113 | CREATE_DESCRIPTORS_TABLE, 114 | CREATE_MATCHES_TABLE, 115 | CREATE_TWO_VIEW_GEOMETRIES_TABLE, 116 | CREATE_NAME_INDEX, 117 | ] 118 | ) 119 | 120 | 121 | def image_ids_to_pair_id(image_id1, image_id2): 122 | if image_id1 > image_id2: 123 | image_id1, image_id2 = image_id2, image_id1 124 | return image_id1 * MAX_IMAGE_ID + image_id2 125 | 126 | 127 | def pair_id_to_image_ids(pair_id): 128 | image_id2 = pair_id % MAX_IMAGE_ID 129 | image_id1 = (pair_id - image_id2) / MAX_IMAGE_ID 130 | return image_id1, image_id2 131 | 132 | 133 | def array_to_blob(array): 134 | if IS_PYTHON3: 135 | return array.tobytes() 136 | else: 137 | return np.getbuffer(array) 138 | 139 | 140 | def blob_to_array(blob, dtype, shape=(-1,)): 141 | if IS_PYTHON3: 142 | return np.fromstring(blob, dtype=dtype).reshape(*shape) 143 | else: 144 | return np.frombuffer(blob, dtype=dtype).reshape(*shape) 145 | 146 | 147 | class COLMAPDatabase(sqlite3.Connection): 148 | 149 | @staticmethod 150 | def connect(database_path): 151 | return sqlite3.connect(database_path, factory=COLMAPDatabase) 152 | 153 | def __init__(self, *args, **kwargs): 154 | super(COLMAPDatabase, self).__init__(*args, **kwargs) 155 | 156 | self.create_tables = lambda: self.executescript(CREATE_ALL) 157 | self.create_cameras_table = lambda: self.executescript(CREATE_CAMERAS_TABLE) 158 | self.create_descriptors_table = lambda: self.executescript( 159 | CREATE_DESCRIPTORS_TABLE 160 | ) 161 | self.create_images_table = lambda: self.executescript(CREATE_IMAGES_TABLE) 162 | self.create_two_view_geometries_table = lambda: self.executescript( 163 | CREATE_TWO_VIEW_GEOMETRIES_TABLE 164 | ) 165 | self.create_keypoints_table = lambda: self.executescript(CREATE_KEYPOINTS_TABLE) 166 | self.create_matches_table = lambda: self.executescript(CREATE_MATCHES_TABLE) 167 | self.create_name_index = lambda: self.executescript(CREATE_NAME_INDEX) 168 | 169 | def add_camera( 170 | self, model, width, height, params, prior_focal_length=False, camera_id=None 171 | ): 172 | params = np.asarray(params, np.float64) 173 | cursor = self.execute( 174 | "INSERT INTO cameras VALUES (?, ?, ?, ?, ?, ?)", 175 | ( 176 | camera_id, 177 | model, 178 | width, 179 | height, 180 | array_to_blob(params), 181 | prior_focal_length, 182 | ), 183 | ) 184 | return cursor.lastrowid 185 | 186 | def add_image( 187 | self, 188 | name, 189 | camera_id, 190 | prior_q=np.full(4, np.NaN), 191 | prior_t=np.full(3, np.NaN), 192 | image_id=None, 193 | ): 194 | cursor = self.execute( 195 | "INSERT INTO images VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", 196 | ( 197 | image_id, 198 | name, 199 | camera_id, 200 | prior_q[0], 201 | prior_q[1], 202 | prior_q[2], 203 | prior_q[3], 204 | prior_t[0], 205 | prior_t[1], 206 | prior_t[2], 207 | ), 208 | ) 209 | return cursor.lastrowid 210 | 211 | def add_keypoints(self, image_id, keypoints): 212 | assert len(keypoints.shape) == 2 213 | assert keypoints.shape[1] in [2, 4, 6] 214 | 215 | keypoints = np.asarray(keypoints, np.float32) 216 | self.execute( 217 | "INSERT INTO keypoints VALUES (?, ?, ?, ?)", 218 | (image_id,) + keypoints.shape + (array_to_blob(keypoints),), 219 | ) 220 | 221 | def add_descriptors(self, image_id, descriptors): 222 | descriptors = np.ascontiguousarray(descriptors, np.uint8) 223 | self.execute( 224 | "INSERT INTO descriptors VALUES (?, ?, ?, ?)", 225 | (image_id,) + descriptors.shape + (array_to_blob(descriptors),), 226 | ) 227 | 228 | def add_matches(self, image_id1, image_id2, matches): 229 | assert len(matches.shape) == 2 230 | assert matches.shape[1] == 2 231 | 232 | if image_id1 > image_id2: 233 | matches = matches[:, ::-1] 234 | 235 | pair_id = image_ids_to_pair_id(image_id1, image_id2) 236 | matches = np.asarray(matches, np.uint32) 237 | self.execute( 238 | "INSERT INTO matches VALUES (?, ?, ?, ?)", 239 | (pair_id,) + matches.shape + (array_to_blob(matches),), 240 | ) 241 | 242 | def add_two_view_geometry( 243 | self, 244 | image_id1, 245 | image_id2, 246 | matches, 247 | F=np.eye(3), 248 | E=np.eye(3), 249 | H=np.eye(3), 250 | qvec=np.array([1.0, 0.0, 0.0, 0.0]), 251 | tvec=np.zeros(3), 252 | config=2, 253 | ): 254 | assert len(matches.shape) == 2 255 | assert matches.shape[1] == 2 256 | 257 | if image_id1 > image_id2: 258 | matches = matches[:, ::-1] 259 | 260 | pair_id = image_ids_to_pair_id(image_id1, image_id2) 261 | matches = np.asarray(matches, np.uint32) 262 | F = np.asarray(F, dtype=np.float64) 263 | E = np.asarray(E, dtype=np.float64) 264 | H = np.asarray(H, dtype=np.float64) 265 | qvec = np.asarray(qvec, dtype=np.float64) 266 | tvec = np.asarray(tvec, dtype=np.float64) 267 | self.execute( 268 | "INSERT INTO two_view_geometries VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", 269 | (pair_id,) 270 | + matches.shape 271 | + ( 272 | array_to_blob(matches), 273 | config, 274 | array_to_blob(F), 275 | array_to_blob(E), 276 | array_to_blob(H), 277 | array_to_blob(qvec), 278 | array_to_blob(tvec), 279 | ), 280 | ) 281 | 282 | 283 | def example_usage(): 284 | import os 285 | import argparse 286 | 287 | parser = argparse.ArgumentParser() 288 | parser.add_argument("--database_path", default="database.db") 289 | args = parser.parse_args() 290 | 291 | if os.path.exists(args.database_path): 292 | print("ERROR: database path already exists -- will not modify it.") 293 | return 294 | 295 | # Open the database. 296 | 297 | db = COLMAPDatabase.connect(args.database_path) 298 | 299 | # For convenience, try creating all the tables upfront. 300 | 301 | db.create_tables() 302 | 303 | # Create dummy cameras. 304 | 305 | model1, width1, height1, params1 = 0, 1024, 768, np.array((1024.0, 512.0, 384.0)) 306 | model2, width2, height2, params2 = ( 307 | 2, 308 | 1024, 309 | 768, 310 | np.array((1024.0, 512.0, 384.0, 0.1)), 311 | ) 312 | 313 | camera_id1 = db.add_camera(model1, width1, height1, params1) 314 | camera_id2 = db.add_camera(model2, width2, height2, params2) 315 | 316 | # Create dummy images. 317 | 318 | image_id1 = db.add_image("image1.png", camera_id1) 319 | image_id2 = db.add_image("image2.png", camera_id1) 320 | image_id3 = db.add_image("image3.png", camera_id2) 321 | image_id4 = db.add_image("image4.png", camera_id2) 322 | 323 | # Create dummy keypoints. 324 | # 325 | # Note that COLMAP supports: 326 | # - 2D keypoints: (x, y) 327 | # - 4D keypoints: (x, y, theta, scale) 328 | # - 6D affine keypoints: (x, y, a_11, a_12, a_21, a_22) 329 | 330 | num_keypoints = 1000 331 | keypoints1 = np.random.rand(num_keypoints, 2) * (width1, height1) 332 | keypoints2 = np.random.rand(num_keypoints, 2) * (width1, height1) 333 | keypoints3 = np.random.rand(num_keypoints, 2) * (width2, height2) 334 | keypoints4 = np.random.rand(num_keypoints, 2) * (width2, height2) 335 | 336 | db.add_keypoints(image_id1, keypoints1) 337 | db.add_keypoints(image_id2, keypoints2) 338 | db.add_keypoints(image_id3, keypoints3) 339 | db.add_keypoints(image_id4, keypoints4) 340 | 341 | # Create dummy matches. 342 | 343 | M = 50 344 | matches12 = np.random.randint(num_keypoints, size=(M, 2)) 345 | matches23 = np.random.randint(num_keypoints, size=(M, 2)) 346 | matches34 = np.random.randint(num_keypoints, size=(M, 2)) 347 | 348 | db.add_matches(image_id1, image_id2, matches12) 349 | db.add_matches(image_id2, image_id3, matches23) 350 | db.add_matches(image_id3, image_id4, matches34) 351 | 352 | # Commit the data to the file. 353 | 354 | db.commit() 355 | 356 | # Read and check cameras. 357 | 358 | rows = db.execute("SELECT * FROM cameras") 359 | 360 | camera_id, model, width, height, params, prior = next(rows) 361 | params = blob_to_array(params, np.float64) 362 | assert camera_id == camera_id1 363 | assert model == model1 and width == width1 and height == height1 364 | assert np.allclose(params, params1) 365 | 366 | camera_id, model, width, height, params, prior = next(rows) 367 | params = blob_to_array(params, np.float64) 368 | assert camera_id == camera_id2 369 | assert model == model2 and width == width2 and height == height2 370 | assert np.allclose(params, params2) 371 | 372 | # Read and check keypoints. 373 | 374 | keypoints = dict( 375 | (image_id, blob_to_array(data, np.float32, (-1, 2))) 376 | for image_id, data in db.execute("SELECT image_id, data FROM keypoints") 377 | ) 378 | 379 | assert np.allclose(keypoints[image_id1], keypoints1) 380 | assert np.allclose(keypoints[image_id2], keypoints2) 381 | assert np.allclose(keypoints[image_id3], keypoints3) 382 | assert np.allclose(keypoints[image_id4], keypoints4) 383 | 384 | # Read and check matches. 385 | 386 | pair_ids = [ 387 | image_ids_to_pair_id(*pair) 388 | for pair in ( 389 | (image_id1, image_id2), 390 | (image_id2, image_id3), 391 | (image_id3, image_id4), 392 | ) 393 | ] 394 | 395 | matches = dict( 396 | (pair_id_to_image_ids(pair_id), blob_to_array(data, np.uint32, (-1, 2))) 397 | for pair_id, data in db.execute("SELECT pair_id, data FROM matches") 398 | ) 399 | 400 | assert np.all(matches[(image_id1, image_id2)] == matches12) 401 | assert np.all(matches[(image_id2, image_id3)] == matches23) 402 | assert np.all(matches[(image_id3, image_id4)] == matches34) 403 | 404 | # Clean up. 405 | 406 | db.close() 407 | 408 | if os.path.exists(args.database_path): 409 | os.remove(args.database_path) 410 | 411 | 412 | if __name__ == "__main__": 413 | example_usage() 414 | -------------------------------------------------------------------------------- /immatch/utils/colmap/read_write_model.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 | # Source code from: 33 | # https://github.com/colmap/colmap/blob/dev/scripts/python/read_write_model.py 34 | 35 | import os 36 | import sys 37 | import collections 38 | import numpy as np 39 | import struct 40 | import argparse 41 | 42 | 43 | CameraModel = collections.namedtuple( 44 | "CameraModel", ["model_id", "model_name", "num_params"] 45 | ) 46 | Camera = collections.namedtuple("Camera", ["id", "model", "width", "height", "params"]) 47 | BaseImage = collections.namedtuple( 48 | "Image", ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"] 49 | ) 50 | Point3D = collections.namedtuple( 51 | "Point3D", ["id", "xyz", "rgb", "error", "image_ids", "point2D_idxs"] 52 | ) 53 | 54 | 55 | class Image(BaseImage): 56 | def qvec2rotmat(self): 57 | return qvec2rotmat(self.qvec) 58 | 59 | 60 | CAMERA_MODELS = { 61 | CameraModel(model_id=0, model_name="SIMPLE_PINHOLE", num_params=3), 62 | CameraModel(model_id=1, model_name="PINHOLE", num_params=4), 63 | CameraModel(model_id=2, model_name="SIMPLE_RADIAL", num_params=4), 64 | CameraModel(model_id=3, model_name="RADIAL", num_params=5), 65 | CameraModel(model_id=4, model_name="OPENCV", num_params=8), 66 | CameraModel(model_id=5, model_name="OPENCV_FISHEYE", num_params=8), 67 | CameraModel(model_id=6, model_name="FULL_OPENCV", num_params=12), 68 | CameraModel(model_id=7, model_name="FOV", num_params=5), 69 | CameraModel(model_id=8, model_name="SIMPLE_RADIAL_FISHEYE", num_params=4), 70 | CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5), 71 | CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12), 72 | } 73 | CAMERA_MODEL_IDS = dict( 74 | [(camera_model.model_id, camera_model) for camera_model in CAMERA_MODELS] 75 | ) 76 | CAMERA_MODEL_NAMES = dict( 77 | [(camera_model.model_name, camera_model) for camera_model in CAMERA_MODELS] 78 | ) 79 | 80 | 81 | def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"): 82 | """Read and unpack the next bytes from a binary file. 83 | :param fid: 84 | :param num_bytes: Sum of combination of {2, 4, 8}, e.g. 2, 6, 16, 30, etc. 85 | :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}. 86 | :param endian_character: Any of {@, =, <, >, !} 87 | :return: Tuple of read and unpacked values. 88 | """ 89 | data = fid.read(num_bytes) 90 | return struct.unpack(endian_character + format_char_sequence, data) 91 | 92 | 93 | def write_next_bytes(fid, data, format_char_sequence, endian_character="<"): 94 | """pack and write to a binary file. 95 | :param fid: 96 | :param data: data to send, if multiple elements are sent at the same time, 97 | they should be encapsuled either in a list or a tuple 98 | :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}. 99 | should be the same length as the data list or tuple 100 | :param endian_character: Any of {@, =, <, >, !} 101 | """ 102 | if isinstance(data, (list, tuple)): 103 | bytes = struct.pack(endian_character + format_char_sequence, *data) 104 | else: 105 | bytes = struct.pack(endian_character + format_char_sequence, data) 106 | fid.write(bytes) 107 | 108 | 109 | def read_cameras_text(path): 110 | """ 111 | see: src/base/reconstruction.cc 112 | void Reconstruction::WriteCamerasText(const std::string& path) 113 | void Reconstruction::ReadCamerasText(const std::string& path) 114 | """ 115 | cameras = {} 116 | with open(path, "r") as fid: 117 | while True: 118 | line = fid.readline() 119 | if not line: 120 | break 121 | line = line.strip() 122 | if len(line) > 0 and line[0] != "#": 123 | elems = line.split() 124 | camera_id = int(elems[0]) 125 | model = elems[1] 126 | width = int(elems[2]) 127 | height = int(elems[3]) 128 | params = np.array(tuple(map(float, elems[4:]))) 129 | cameras[camera_id] = Camera( 130 | id=camera_id, model=model, width=width, height=height, params=params 131 | ) 132 | return cameras 133 | 134 | 135 | def read_cameras_binary(path_to_model_file): 136 | """ 137 | see: src/base/reconstruction.cc 138 | void Reconstruction::WriteCamerasBinary(const std::string& path) 139 | void Reconstruction::ReadCamerasBinary(const std::string& path) 140 | """ 141 | cameras = {} 142 | with open(path_to_model_file, "rb") as fid: 143 | num_cameras = read_next_bytes(fid, 8, "Q")[0] 144 | for _ in range(num_cameras): 145 | camera_properties = read_next_bytes( 146 | fid, num_bytes=24, format_char_sequence="iiQQ" 147 | ) 148 | camera_id = camera_properties[0] 149 | model_id = camera_properties[1] 150 | model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name 151 | width = camera_properties[2] 152 | height = camera_properties[3] 153 | num_params = CAMERA_MODEL_IDS[model_id].num_params 154 | params = read_next_bytes( 155 | fid, num_bytes=8 * num_params, format_char_sequence="d" * num_params 156 | ) 157 | cameras[camera_id] = Camera( 158 | id=camera_id, 159 | model=model_name, 160 | width=width, 161 | height=height, 162 | params=np.array(params), 163 | ) 164 | assert len(cameras) == num_cameras 165 | return cameras 166 | 167 | 168 | def write_cameras_text(cameras, path): 169 | """ 170 | see: src/base/reconstruction.cc 171 | void Reconstruction::WriteCamerasText(const std::string& path) 172 | void Reconstruction::ReadCamerasText(const std::string& path) 173 | """ 174 | HEADER = "# Camera list with one line of data per camera:\n" 175 | "# CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n" 176 | "# Number of cameras: {}\n".format(len(cameras)) 177 | with open(path, "w") as fid: 178 | fid.write(HEADER) 179 | for _, cam in cameras.items(): 180 | to_write = [cam.id, cam.model, cam.width, cam.height, *cam.params] 181 | line = " ".join([str(elem) for elem in to_write]) 182 | fid.write(line + "\n") 183 | 184 | 185 | def write_cameras_binary(cameras, path_to_model_file): 186 | """ 187 | see: src/base/reconstruction.cc 188 | void Reconstruction::WriteCamerasBinary(const std::string& path) 189 | void Reconstruction::ReadCamerasBinary(const std::string& path) 190 | """ 191 | with open(path_to_model_file, "wb") as fid: 192 | write_next_bytes(fid, len(cameras), "Q") 193 | for _, cam in cameras.items(): 194 | model_id = CAMERA_MODEL_NAMES[cam.model].model_id 195 | camera_properties = [cam.id, model_id, cam.width, cam.height] 196 | write_next_bytes(fid, camera_properties, "iiQQ") 197 | for p in cam.params: 198 | write_next_bytes(fid, float(p), "d") 199 | return cameras 200 | 201 | 202 | def read_images_text(path): 203 | """ 204 | see: src/base/reconstruction.cc 205 | void Reconstruction::ReadImagesText(const std::string& path) 206 | void Reconstruction::WriteImagesText(const std::string& path) 207 | """ 208 | images = {} 209 | with open(path, "r") as fid: 210 | while True: 211 | line = fid.readline() 212 | if not line: 213 | break 214 | line = line.strip() 215 | if len(line) > 0 and line[0] != "#": 216 | elems = line.split() 217 | image_id = int(elems[0]) 218 | qvec = np.array(tuple(map(float, elems[1:5]))) 219 | tvec = np.array(tuple(map(float, elems[5:8]))) 220 | camera_id = int(elems[8]) 221 | image_name = elems[9] 222 | elems = fid.readline().split() 223 | xys = np.column_stack( 224 | [tuple(map(float, elems[0::3])), tuple(map(float, elems[1::3]))] 225 | ) 226 | point3D_ids = np.array(tuple(map(int, elems[2::3]))) 227 | images[image_id] = Image( 228 | id=image_id, 229 | qvec=qvec, 230 | tvec=tvec, 231 | camera_id=camera_id, 232 | name=image_name, 233 | xys=xys, 234 | point3D_ids=point3D_ids, 235 | ) 236 | return images 237 | 238 | 239 | def read_images_binary(path_to_model_file): 240 | """ 241 | see: src/base/reconstruction.cc 242 | void Reconstruction::ReadImagesBinary(const std::string& path) 243 | void Reconstruction::WriteImagesBinary(const std::string& path) 244 | """ 245 | images = {} 246 | with open(path_to_model_file, "rb") as fid: 247 | num_reg_images = read_next_bytes(fid, 8, "Q")[0] 248 | for _ in range(num_reg_images): 249 | binary_image_properties = read_next_bytes( 250 | fid, num_bytes=64, format_char_sequence="idddddddi" 251 | ) 252 | image_id = binary_image_properties[0] 253 | qvec = np.array(binary_image_properties[1:5]) 254 | tvec = np.array(binary_image_properties[5:8]) 255 | camera_id = binary_image_properties[8] 256 | image_name = "" 257 | current_char = read_next_bytes(fid, 1, "c")[0] 258 | while current_char != b"\x00": # look for the ASCII 0 entry 259 | image_name += current_char.decode("utf-8") 260 | current_char = read_next_bytes(fid, 1, "c")[0] 261 | num_points2D = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[ 262 | 0 263 | ] 264 | x_y_id_s = read_next_bytes( 265 | fid, 266 | num_bytes=24 * num_points2D, 267 | format_char_sequence="ddq" * num_points2D, 268 | ) 269 | xys = np.column_stack( 270 | [tuple(map(float, x_y_id_s[0::3])), tuple(map(float, x_y_id_s[1::3]))] 271 | ) 272 | point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3]))) 273 | images[image_id] = Image( 274 | id=image_id, 275 | qvec=qvec, 276 | tvec=tvec, 277 | camera_id=camera_id, 278 | name=image_name, 279 | xys=xys, 280 | point3D_ids=point3D_ids, 281 | ) 282 | return images 283 | 284 | 285 | def write_images_text(images, path): 286 | """ 287 | see: src/base/reconstruction.cc 288 | void Reconstruction::ReadImagesText(const std::string& path) 289 | void Reconstruction::WriteImagesText(const std::string& path) 290 | """ 291 | if len(images) == 0: 292 | mean_observations = 0 293 | else: 294 | mean_observations = sum( 295 | (len(img.point3D_ids) for _, img in images.items()) 296 | ) / len(images) 297 | HEADER = "# Image list with two lines of data per image:\n" 298 | "# IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n" 299 | "# POINTS2D[] as (X, Y, POINT3D_ID)\n" 300 | "# Number of images: {}, mean observations per image: {}\n".format( 301 | len(images), mean_observations 302 | ) 303 | 304 | with open(path, "w") as fid: 305 | fid.write(HEADER) 306 | for _, img in images.items(): 307 | image_header = [img.id, *img.qvec, *img.tvec, img.camera_id, img.name] 308 | first_line = " ".join(map(str, image_header)) 309 | fid.write(first_line + "\n") 310 | 311 | points_strings = [] 312 | for xy, point3D_id in zip(img.xys, img.point3D_ids): 313 | points_strings.append(" ".join(map(str, [*xy, point3D_id]))) 314 | fid.write(" ".join(points_strings) + "\n") 315 | 316 | 317 | def write_images_binary(images, path_to_model_file): 318 | """ 319 | see: src/base/reconstruction.cc 320 | void Reconstruction::ReadImagesBinary(const std::string& path) 321 | void Reconstruction::WriteImagesBinary(const std::string& path) 322 | """ 323 | with open(path_to_model_file, "wb") as fid: 324 | write_next_bytes(fid, len(images), "Q") 325 | for _, img in images.items(): 326 | write_next_bytes(fid, img.id, "i") 327 | write_next_bytes(fid, img.qvec.tolist(), "dddd") 328 | write_next_bytes(fid, img.tvec.tolist(), "ddd") 329 | write_next_bytes(fid, img.camera_id, "i") 330 | for char in img.name: 331 | write_next_bytes(fid, char.encode("utf-8"), "c") 332 | write_next_bytes(fid, b"\x00", "c") 333 | write_next_bytes(fid, len(img.point3D_ids), "Q") 334 | for xy, p3d_id in zip(img.xys, img.point3D_ids): 335 | write_next_bytes(fid, [*xy, p3d_id], "ddq") 336 | 337 | 338 | def read_points3D_text(path): 339 | """ 340 | see: src/base/reconstruction.cc 341 | void Reconstruction::ReadPoints3DText(const std::string& path) 342 | void Reconstruction::WritePoints3DText(const std::string& path) 343 | """ 344 | points3D = {} 345 | with open(path, "r") as fid: 346 | while True: 347 | line = fid.readline() 348 | if not line: 349 | break 350 | line = line.strip() 351 | if len(line) > 0 and line[0] != "#": 352 | elems = line.split() 353 | point3D_id = int(elems[0]) 354 | xyz = np.array(tuple(map(float, elems[1:4]))) 355 | rgb = np.array(tuple(map(int, elems[4:7]))) 356 | error = float(elems[7]) 357 | image_ids = np.array(tuple(map(int, elems[8::2]))) 358 | point2D_idxs = np.array(tuple(map(int, elems[9::2]))) 359 | points3D[point3D_id] = Point3D( 360 | id=point3D_id, 361 | xyz=xyz, 362 | rgb=rgb, 363 | error=error, 364 | image_ids=image_ids, 365 | point2D_idxs=point2D_idxs, 366 | ) 367 | return points3D 368 | 369 | 370 | def read_points3d_binary(path_to_model_file): 371 | """ 372 | see: src/base/reconstruction.cc 373 | void Reconstruction::ReadPoints3DBinary(const std::string& path) 374 | void Reconstruction::WritePoints3DBinary(const std::string& path) 375 | """ 376 | points3D = {} 377 | with open(path_to_model_file, "rb") as fid: 378 | num_points = read_next_bytes(fid, 8, "Q")[0] 379 | for _ in range(num_points): 380 | binary_point_line_properties = read_next_bytes( 381 | fid, num_bytes=43, format_char_sequence="QdddBBBd" 382 | ) 383 | point3D_id = binary_point_line_properties[0] 384 | xyz = np.array(binary_point_line_properties[1:4]) 385 | rgb = np.array(binary_point_line_properties[4:7]) 386 | error = np.array(binary_point_line_properties[7]) 387 | track_length = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[ 388 | 0 389 | ] 390 | track_elems = read_next_bytes( 391 | fid, 392 | num_bytes=8 * track_length, 393 | format_char_sequence="ii" * track_length, 394 | ) 395 | image_ids = np.array(tuple(map(int, track_elems[0::2]))) 396 | point2D_idxs = np.array(tuple(map(int, track_elems[1::2]))) 397 | points3D[point3D_id] = Point3D( 398 | id=point3D_id, 399 | xyz=xyz, 400 | rgb=rgb, 401 | error=error, 402 | image_ids=image_ids, 403 | point2D_idxs=point2D_idxs, 404 | ) 405 | return points3D 406 | 407 | 408 | def write_points3D_text(points3D, path): 409 | """ 410 | see: src/base/reconstruction.cc 411 | void Reconstruction::ReadPoints3DText(const std::string& path) 412 | void Reconstruction::WritePoints3DText(const std::string& path) 413 | """ 414 | if len(points3D) == 0: 415 | mean_track_length = 0 416 | else: 417 | mean_track_length = sum( 418 | (len(pt.image_ids) for _, pt in points3D.items()) 419 | ) / len(points3D) 420 | HEADER = "# 3D point list with one line of data per point:\n" 421 | "# POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n" 422 | "# Number of points: {}, mean track length: {}\n".format( 423 | len(points3D), mean_track_length 424 | ) 425 | 426 | with open(path, "w") as fid: 427 | fid.write(HEADER) 428 | for _, pt in points3D.items(): 429 | point_header = [pt.id, *pt.xyz, *pt.rgb, pt.error] 430 | fid.write(" ".join(map(str, point_header)) + " ") 431 | track_strings = [] 432 | for image_id, point2D in zip(pt.image_ids, pt.point2D_idxs): 433 | track_strings.append(" ".join(map(str, [image_id, point2D]))) 434 | fid.write(" ".join(track_strings) + "\n") 435 | 436 | 437 | def write_points3d_binary(points3D, path_to_model_file): 438 | """ 439 | see: src/base/reconstruction.cc 440 | void Reconstruction::ReadPoints3DBinary(const std::string& path) 441 | void Reconstruction::WritePoints3DBinary(const std::string& path) 442 | """ 443 | with open(path_to_model_file, "wb") as fid: 444 | write_next_bytes(fid, len(points3D), "Q") 445 | for _, pt in points3D.items(): 446 | write_next_bytes(fid, pt.id, "Q") 447 | write_next_bytes(fid, pt.xyz.tolist(), "ddd") 448 | write_next_bytes(fid, pt.rgb.tolist(), "BBB") 449 | write_next_bytes(fid, pt.error, "d") 450 | track_length = pt.image_ids.shape[0] 451 | write_next_bytes(fid, track_length, "Q") 452 | for image_id, point2D_id in zip(pt.image_ids, pt.point2D_idxs): 453 | write_next_bytes(fid, [image_id, point2D_id], "ii") 454 | 455 | 456 | def detect_model_format(path, ext): 457 | if ( 458 | os.path.isfile(os.path.join(path, "cameras" + ext)) 459 | and os.path.isfile(os.path.join(path, "images" + ext)) 460 | and os.path.isfile(os.path.join(path, "points3D" + ext)) 461 | ): 462 | print("Detected model format: '" + ext + "'") 463 | return True 464 | 465 | return False 466 | 467 | 468 | def read_model(path, ext=""): 469 | # try to detect the extension automatically 470 | if ext == "": 471 | if detect_model_format(path, ".bin"): 472 | ext = ".bin" 473 | elif detect_model_format(path, ".txt"): 474 | ext = ".txt" 475 | else: 476 | print("Provide model format: '.bin' or '.txt'") 477 | return 478 | 479 | if ext == ".txt": 480 | cameras = read_cameras_text(os.path.join(path, "cameras" + ext)) 481 | images = read_images_text(os.path.join(path, "images" + ext)) 482 | points3D = read_points3D_text(os.path.join(path, "points3D") + ext) 483 | else: 484 | cameras = read_cameras_binary(os.path.join(path, "cameras" + ext)) 485 | images = read_images_binary(os.path.join(path, "images" + ext)) 486 | points3D = read_points3d_binary(os.path.join(path, "points3D") + ext) 487 | return cameras, images, points3D 488 | 489 | 490 | def write_model(cameras, images, points3D, path, ext=".bin"): 491 | if ext == ".txt": 492 | write_cameras_text(cameras, os.path.join(path, "cameras" + ext)) 493 | write_images_text(images, os.path.join(path, "images" + ext)) 494 | write_points3D_text(points3D, os.path.join(path, "points3D") + ext) 495 | else: 496 | write_cameras_binary(cameras, os.path.join(path, "cameras" + ext)) 497 | write_images_binary(images, os.path.join(path, "images" + ext)) 498 | write_points3d_binary(points3D, os.path.join(path, "points3D") + ext) 499 | return cameras, images, points3D 500 | 501 | 502 | def qvec2rotmat(qvec): 503 | return np.array( 504 | [ 505 | [ 506 | 1 - 2 * qvec[2] ** 2 - 2 * qvec[3] ** 2, 507 | 2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3], 508 | 2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2], 509 | ], 510 | [ 511 | 2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3], 512 | 1 - 2 * qvec[1] ** 2 - 2 * qvec[3] ** 2, 513 | 2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1], 514 | ], 515 | [ 516 | 2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2], 517 | 2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1], 518 | 1 - 2 * qvec[1] ** 2 - 2 * qvec[2] ** 2, 519 | ], 520 | ] 521 | ) 522 | 523 | 524 | def rotmat2qvec(R): 525 | Rxx, Ryx, Rzx, Rxy, Ryy, Rzy, Rxz, Ryz, Rzz = R.flat 526 | K = ( 527 | np.array( 528 | [ 529 | [Rxx - Ryy - Rzz, 0, 0, 0], 530 | [Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0], 531 | [Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0], 532 | [Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz], 533 | ] 534 | ) 535 | / 3.0 536 | ) 537 | eigvals, eigvecs = np.linalg.eigh(K) 538 | qvec = eigvecs[[3, 0, 1, 2], np.argmax(eigvals)] 539 | if qvec[0] < 0: 540 | qvec *= -1 541 | return qvec 542 | 543 | 544 | def main(): 545 | parser = argparse.ArgumentParser( 546 | description="Read and write COLMAP binary and text models" 547 | ) 548 | parser.add_argument("--input_model", help="path to input model folder") 549 | parser.add_argument( 550 | "--input_format", 551 | choices=[".bin", ".txt"], 552 | help="input model format", 553 | default="", 554 | ) 555 | parser.add_argument("--output_model", help="path to output model folder") 556 | parser.add_argument( 557 | "--output_format", 558 | choices=[".bin", ".txt"], 559 | help="outut model format", 560 | default=".txt", 561 | ) 562 | args = parser.parse_args() 563 | 564 | cameras, images, points3D = read_model(path=args.input_model, ext=args.input_format) 565 | 566 | print("num_cameras:", len(cameras)) 567 | print("num_images:", len(images)) 568 | print("num_points3D:", len(points3D)) 569 | 570 | if args.output_model is not None: 571 | write_model( 572 | cameras, images, points3D, path=args.output_model, ext=args.output_format 573 | ) 574 | 575 | 576 | if __name__ == "__main__": 577 | main() 578 | -------------------------------------------------------------------------------- /immatch/utils/data_io.py: -------------------------------------------------------------------------------- 1 | from PIL import Image 2 | import cv2 3 | import numpy as np 4 | 5 | import torch 6 | import torchvision.transforms as transforms 7 | 8 | 9 | def lprint(ms, log=None): 10 | """Print message on console and in a log file""" 11 | print(ms) 12 | if log: 13 | log.write(ms + "\n") 14 | log.flush() 15 | 16 | 17 | def resize_im(wo, ho, imsize=None, dfactor=1, value_to_scale=max, enforce=False): 18 | if not isinstance(imsize, int) and len(imsize) == 2: 19 | # Resize to a fixed shape 20 | wt, ht = imsize 21 | scale = [wo / wt, ho / ht] 22 | return wt, ht, scale 23 | wt, ht = wo, ho 24 | 25 | # Resize only if the image is too big 26 | resize = imsize and value_to_scale(wo, ho) > imsize and imsize > 0 27 | if resize or enforce: 28 | scale = imsize / value_to_scale(wo, ho) 29 | ht, wt = int(round(ho * scale)), int(round(wo * scale)) 30 | 31 | # Make sure new sizes are divisible by the given factor 32 | wt, ht = map(lambda x: int(x // dfactor * dfactor), [wt, ht]) 33 | scale = [wo / wt, ho / ht] 34 | return wt, ht, scale 35 | 36 | 37 | def read_im(im_path, imsize=None, dfactor=1): 38 | im = Image.open(im_path) 39 | im = im.convert("RGB") 40 | 41 | # Resize 42 | wo, ho = im.width, im.height 43 | wt, ht, scale = resize_im(wo, ho, imsize=imsize, dfactor=dfactor) 44 | im = im.resize((wt, ht), Image.BICUBIC) 45 | return im, scale 46 | 47 | 48 | def read_im_gray(im_path, imsize=None): 49 | im, scale = read_im(im_path, imsize) 50 | return im.convert("L"), scale 51 | 52 | 53 | def load_gray_scale_tensor(im_path, device, imsize=None, dfactor=1): 54 | im_rgb, scale = read_im(im_path, imsize, dfactor=dfactor) 55 | gray = np.array(im_rgb.convert("L")) 56 | gray = transforms.functional.to_tensor(gray).unsqueeze(0).to(device) 57 | return gray, scale 58 | 59 | 60 | def load_gray_scale_tensor_cv( 61 | im_path, device, imsize=None, value_to_scale=min, dfactor=1, pad2sqr=False 62 | ): 63 | """Image loading function applicable for LoFTR & Aspanformer.""" 64 | 65 | im = cv2.imread(im_path, cv2.IMREAD_GRAYSCALE) 66 | ho, wo = im.shape 67 | wt, ht, scale = resize_im( 68 | wo, 69 | ho, 70 | imsize=imsize, 71 | dfactor=dfactor, 72 | value_to_scale=value_to_scale, 73 | enforce=pad2sqr, 74 | ) 75 | im = cv2.resize(im, (wt, ht)) 76 | mask = None 77 | if pad2sqr and (wt != ht): 78 | # Padding to square image 79 | im, mask = pad_bottom_right(im, max(wt, ht), ret_mask=True) 80 | mask = torch.from_numpy(mask).to(device) 81 | im = transforms.functional.to_tensor(im).unsqueeze(0).to(device) 82 | return im, scale, mask 83 | 84 | 85 | def pad_bottom_right(inp, pad_size, ret_mask=False): 86 | assert isinstance(pad_size, int) and pad_size >= max( 87 | inp.shape[-2:] 88 | ), f"{pad_size} < {max(inp.shape[-2:])}" 89 | mask = None 90 | if inp.ndim == 2: 91 | padded = np.zeros((pad_size, pad_size), dtype=inp.dtype) 92 | padded[: inp.shape[0], : inp.shape[1]] = inp 93 | if ret_mask: 94 | mask = np.zeros((pad_size, pad_size), dtype=bool) 95 | mask[: inp.shape[0], : inp.shape[1]] = True 96 | elif inp.ndim == 3: 97 | padded = np.zeros((inp.shape[0], pad_size, pad_size), dtype=inp.dtype) 98 | padded[:, : inp.shape[1], : inp.shape[2]] = inp 99 | if ret_mask: 100 | mask = np.zeros((inp.shape[0], pad_size, pad_size), dtype=bool) 101 | mask[:, : inp.shape[1], : inp.shape[2]] = True 102 | else: 103 | raise NotImplementedError() 104 | return padded, mask 105 | 106 | 107 | def load_im_tensor( 108 | im_path, 109 | device, 110 | imsize=None, 111 | normalize=True, 112 | with_gray=False, 113 | raw_gray=False, 114 | dfactor=1, 115 | ): 116 | im_rgb, scale = read_im(im_path, imsize, dfactor=dfactor) 117 | 118 | # RGB 119 | im = transforms.functional.to_tensor(im_rgb) 120 | if normalize: 121 | im = transforms.functional.normalize( 122 | im, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] 123 | ) 124 | im = im.unsqueeze(0).to(device) 125 | 126 | if with_gray: 127 | # Grey 128 | gray = np.array(im_rgb.convert("L")) 129 | if not raw_gray: 130 | gray = transforms.functional.to_tensor(gray).unsqueeze(0).to(device) 131 | return im, gray, scale 132 | return im, scale 133 | -------------------------------------------------------------------------------- /immatch/utils/hpatches_helper.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import glob 4 | import time 5 | import pydegensac 6 | import cv2 7 | from tqdm import tqdm 8 | 9 | from immatch.utils.metrics import cal_error_auc, check_data_hist, cal_reproj_dists_H 10 | 11 | 12 | def eval_summary_homography(dists_sa, dists_si, dists_sv, thres): 13 | correct_sa = np.mean( 14 | [[float(dist <= t) for t in thres] for dist in dists_sa], axis=0 15 | ) 16 | correct_si = np.mean( 17 | [[float(dist <= t) for t in thres] for dist in dists_si], axis=0 18 | ) 19 | correct_sv = np.mean( 20 | [[float(dist <= t) for t in thres] for dist in dists_sv], axis=0 21 | ) 22 | 23 | # Compute aucs 24 | auc_sa = cal_error_auc(dists_sa, thresholds=thres) 25 | auc_si = cal_error_auc(dists_si, thresholds=thres) 26 | auc_sv = cal_error_auc(dists_sv, thresholds=thres) 27 | 28 | # Generate summary 29 | summary = f"Hest Correct: a={correct_sa}\ni={correct_si}\nv={correct_sv}\n" 30 | summary += f"Hest AUC: a={auc_sa}\ni={auc_si}\nv={auc_sv}\n" 31 | return summary 32 | 33 | 34 | def eval_summary_matching(results, thres=[1, 3, 5, 10], save_npy=None): 35 | np.set_printoptions(precision=2) 36 | summary = "" 37 | n_i = 52 38 | n_v = 56 39 | i_err, v_err, stats = results 40 | seq_type, n_feats, n_matches = stats 41 | 42 | if save_npy: 43 | print(f"Save results to {save_npy}") 44 | np.save(save_npy, np.array(results, dtype=object)) 45 | 46 | summary += "#Features: mean={:.0f} min={:d} max={:d}\n".format( 47 | np.mean(n_feats), np.min(n_feats), np.max(n_feats) 48 | ) 49 | summary += "#(Old)Matches: a={:.0f}, i={:.0f}, v={:.0f}\n".format( 50 | np.sum(n_matches) / ((n_i + n_v) * 5), 51 | np.sum(n_matches[seq_type == "i"]) / (n_i * 5), 52 | np.sum(n_matches[seq_type == "v"]) / (n_v * 5), 53 | ) 54 | summary += "#Matches: a={:.0f}, i={:.0f}, v={:.0f}\n".format( 55 | np.mean(n_matches), 56 | np.mean(n_matches[seq_type == "i"]), 57 | np.mean(n_matches[seq_type == "v"]), 58 | ) 59 | 60 | thres = np.array(thres) 61 | ierr = np.array([i_err[th] / (n_i * 5) for th in thres]) 62 | verr = np.array([v_err[th] / (n_v * 5) for th in thres]) 63 | aerr = np.array([(i_err[th] + v_err[th]) / ((n_i + n_v) * 5) for th in thres]) 64 | summary += "MMA@{} px:\na={}\ni={}\nv={}\n".format(thres, aerr, ierr, verr) 65 | return summary 66 | 67 | 68 | def scale_homography(sw, sh): 69 | return np.array([[sw, 0, 0], [0, sh, 0], [0, 0, 1]]) 70 | 71 | 72 | def eval_hpatches( 73 | matcher, 74 | data_root, 75 | method="", 76 | task="both", 77 | scale_H=False, 78 | h_solver="degensac", 79 | ransac_thres=2, 80 | thres=[1, 3, 5, 10], 81 | lprint_=print, 82 | print_out=False, 83 | save_npy=None, 84 | debug=False, 85 | ): 86 | """Evaluate a matcher on HPatches sequences for image matching and homogray estimation. 87 | The matching metric is adopted from D2Net paper, i.e., the precentage of correctly matched 88 | keypoints at the given re-projection error thresholds. 89 | For homography estimation, the average distances between the corners transformed using 90 | the estimated and GT homographies are computed. Both percentage of the corner distance at 91 | the given thresholds and the area under the cumulative error curve (AUC) at those thresholds 92 | are reported. 93 | 94 | Args: 95 | - matcher: the matching function that inputs an image pair paths and 96 | outputs the matches and keypoints. 97 | - data_root: the folder directory of HPatches dataset. 98 | - method: the description of the evaluated method. 99 | - task: the target task, options = [matching|homography|both] 100 | - ransac_thres: the set of ransac thresholds used by the solver to estimate homographies. 101 | Results under each ransac threshold are printed per line. 102 | - thres: error thresholds in pixels to compute the metrics. 103 | - lprint: the printing function. If needed it can be implemented to outstream to a log file. 104 | - print_out: when set to True, per-pair results are printed during the evaluation. 105 | """ 106 | 107 | np.set_printoptions(precision=2) 108 | from PIL import Image 109 | 110 | if task == "both": 111 | task = "matching+homography" 112 | seq_dirs = sorted(glob.glob("{}/*".format(data_root))) 113 | lprint_( 114 | f"\n>>>>Eval hpatches: task={task} method={method} scale_H={scale_H} rthres={ransac_thres} thres={thres} " 115 | ) 116 | 117 | # Matching 118 | if "matching" in task: 119 | thres_range = np.arange(1, 16) 120 | i_err = {thr: 0 for thr in thres_range} 121 | v_err = {thr: 0 for thr in thres_range} 122 | n_feats = [] 123 | seq_type = [] 124 | 125 | # Homography 126 | if "homography" in task: 127 | inlier_ratio = [] 128 | h_failed = 0 129 | dists_sa = [] 130 | dists_si = [] 131 | dists_sv = [] 132 | 133 | match_failed = 0 134 | n_matches = [] 135 | match_time = [] 136 | start_time = time.time() 137 | match_errs = [] 138 | for seq_idx, seq_dir in tqdm( 139 | enumerate(seq_dirs[::-1]), total=len(seq_dirs), smoothing=0.5 140 | ): 141 | if debug and seq_idx > 10: 142 | break 143 | sname = seq_dir.split("/")[-1] 144 | im1_path = os.path.join(seq_dir, "1.ppm") 145 | 146 | # Eval on composed pairs within seq 147 | for im_idx in range(2, 7): 148 | im2_path = os.path.join(seq_dir, "{}.ppm".format(im_idx)) 149 | H_gt = np.loadtxt(os.path.join(seq_dir, "H_1_{}".format(im_idx))) 150 | scale = np.ones(4) 151 | 152 | # Predict matches 153 | try: 154 | t0 = time.time() 155 | match_res = matcher(im1_path, im2_path) 156 | match_time.append(time.time() - t0) 157 | matches, p1s, p2s = match_res[0:3] 158 | if scale_H: 159 | # scale = (wo / wt, ho / ht) for im1 & im2 160 | scale = match_res[4] 161 | 162 | # Scale gt homoragphies 163 | H_scale_im1 = scale_homography(scale[0], scale[1]) 164 | H_scale_im2 = scale_homography(scale[2], scale[3]) 165 | H_gt = np.linalg.inv(H_scale_im2) @ H_gt @ H_scale_im1 166 | except Exception as e: 167 | print(e) 168 | p1s = p2s = matches = [] 169 | match_failed += 1 170 | n_matches.append(len(matches)) 171 | 172 | if "matching" in task: 173 | n_feats.append(len(p1s)) 174 | n_feats.append(len(p2s)) 175 | seq_type.append(sname[0]) 176 | if len(matches) == 0: 177 | dist = np.array([float("inf")]) 178 | else: 179 | dist = cal_reproj_dists_H(matches[:, :2], matches[:, 2:], H_gt) 180 | match_errs.append(dist) 181 | for thr in thres_range: 182 | if sname[0] == "i": 183 | i_err[thr] += np.mean(dist <= thr) 184 | else: 185 | v_err[thr] += np.mean(dist <= thr) 186 | 187 | if "homography" in task: 188 | try: 189 | if h_solver == "cv": 190 | # By default ransac 191 | H_pred, inliers = cv2.findHomography( 192 | matches[:, :2], matches[:, 2:4], cv2.RANSAC, ransac_thres 193 | ) 194 | elif h_solver == "magsac": 195 | H_pred, inliers = cv2.findHomography( 196 | matches[:, :2], 197 | matches[:, 2:4], 198 | cv2.USAC_MAGSAC, 199 | ransac_thres, 200 | ) 201 | 202 | else: 203 | H_pred, inliers = pydegensac.findHomography( 204 | matches[:, :2], matches[:, 2:4], ransac_thres 205 | ) 206 | except: 207 | H_pred = None 208 | 209 | if H_pred is None: 210 | corner_dist = np.nan 211 | irat = 0 212 | h_failed += 1 213 | inliers = [] 214 | else: 215 | im = Image.open(im1_path) 216 | w, h = im.size 217 | w, h = w / scale[0], h / scale[1] 218 | corners = np.array( 219 | [[0, 0, 1], [0, h - 1, 1], [w - 1, 0, 1], [w - 1, h - 1, 1]] 220 | ) 221 | real_warped_corners = np.dot(corners, np.transpose(H_gt)) 222 | real_warped_corners = ( 223 | real_warped_corners[:, :2] / real_warped_corners[:, 2:] 224 | ) 225 | warped_corners = np.dot(corners, np.transpose(H_pred)) 226 | warped_corners = warped_corners[:, :2] / warped_corners[:, 2:] 227 | corner_dist = np.mean( 228 | np.linalg.norm(real_warped_corners - warped_corners, axis=1) 229 | ) 230 | irat = np.mean(inliers) 231 | inlier_ratio.append(irat) 232 | dists_sa.append(corner_dist) 233 | if sname[0] == "i": 234 | dists_si.append(corner_dist) 235 | if sname[0] == "v": 236 | dists_sv.append(corner_dist) 237 | 238 | if print_out: 239 | print(f"Scene {sname}, pair:1-{im_idx} matches:{len(matches)}") 240 | if "matching" in task: 241 | print( 242 | f"Median matching dist:{np.median(dist):.2f} <1px:{np.mean(dist <= 1):.3f}" 243 | ) 244 | if "homography" in task: 245 | print(f"Corner dist:{corner_dist:.2f} inliers:{np.sum(inliers)}") 246 | 247 | lprint_( 248 | f">>Finished, pairs={len(match_time)} match_failed={match_failed} matches={np.mean(n_matches):.1f} match_time={np.mean(match_time):.2f}s" 249 | ) 250 | 251 | if "matching" in task: 252 | results = ( 253 | i_err, 254 | v_err, 255 | [np.array(seq_type), np.array(n_feats), np.array(n_matches)], 256 | ) 257 | lprint_("==== Image Matching ====") 258 | lprint_(eval_summary_matching(results, thres, save_npy=save_npy)) 259 | if "homography" in task: 260 | lprint_("==== Homography Estimation ====") 261 | lprint_( 262 | f"Hest solver={h_solver} est_failed={h_failed} ransac_thres={ransac_thres} inlier_rate={np.mean(inlier_ratio):.2f}" 263 | ) 264 | lprint_(eval_summary_homography(dists_sa, dists_si, dists_sv, thres)) 265 | 266 | # Measure distributions 267 | print( 268 | check_data_hist( 269 | match_errs, bins=[0, 1, 3, 5, 10, 50, 100, 1000], tag="Matching " 270 | ) 271 | ) 272 | -------------------------------------------------------------------------------- /immatch/utils/localize_sfm_helper.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import time 3 | import logging 4 | from tqdm import tqdm 5 | import h5py 6 | import os 7 | from pathlib import Path 8 | 9 | 10 | def names_to_pair(name0, name1): 11 | return "_".join((name0.replace("/", "-"), name1.replace("/", "-"))) 12 | 13 | 14 | def load_pairs(args): 15 | args.pair_dir = Path(args.pair_dir, args.benchmark_name) 16 | args.db_pairs_path = args.pair_dir / args.pairs[0] 17 | args.query_pairs_path = args.pair_dir / args.pairs[1] 18 | if not args.db_pairs_path.exists(): 19 | print(f"{args.db_pairs_path} does not exist!!") 20 | return None 21 | if not args.db_pairs_path.exists(): 22 | print(f"{args.query_pairs_path} does not exist!!") 23 | return None 24 | 25 | # Load pairs 26 | print(f"Pair list: {args.pairs}") 27 | db_pairs = [] 28 | query_pairs = [] 29 | with open(args.db_pairs_path) as f: 30 | db_pairs += f.readlines() 31 | with open(args.query_pairs_path) as f: 32 | query_pairs += f.readlines() 33 | print(f"Loaded pairs db:{len(db_pairs)} query:{len(query_pairs)}") 34 | pair_list = db_pairs + query_pairs 35 | return pair_list 36 | 37 | 38 | def init_paths(args): 39 | # Define experiment tags 40 | if args.qt_psize > 0: 41 | args.qt_tag = f"qt{args.qt_psize}d{args.qt_dthres}sc{args.sc_thres}" 42 | if args.qt_unique: 43 | args.qt_tag += "uni" 44 | else: 45 | args.qt_tag = f"sc{args.sc_thres}" 46 | 47 | db_pair_tag = args.pairs[0].replace("db-pairs-", "").replace(".txt", "") 48 | query_pair_tag = args.pairs[1].replace("pairs-query-", "").replace(".txt", "") 49 | args.pair_tag = db_pair_tag + query_pair_tag 50 | 51 | # Output paths 52 | output_dir = Path("outputs", args.benchmark_name, args.model_tag, args.conf_tag) 53 | if not output_dir.exists(): 54 | os.makedirs(output_dir) 55 | args.output_dir = output_dir 56 | args.empty_sfm = Path("outputs", args.benchmark_name, "empty_sfm") 57 | 58 | # Result dir for sfm models and localization results 59 | result_dir = output_dir / f"{args.pair_tag}_{args.qt_tag}" 60 | if not result_dir.exists(): 61 | os.makedirs(result_dir) 62 | logging.info(f"Result folder: {result_dir}") 63 | args.result_dir = result_dir 64 | args.result_sfm = result_dir / "sfm_final" 65 | return args 66 | 67 | 68 | def init_empty_sfm(args): 69 | from immatch.utils.colmap.data_parsing import ( 70 | create_empty_model_from_nvm_and_database, 71 | create_empty_model_from_reference_model, 72 | ) 73 | 74 | if args.empty_sfm.exists(): 75 | logging.info("Empty sfm existed.") 76 | return 77 | dataset_dir = args.dataset_dir 78 | 79 | if args.benchmark_name == "robotcar": 80 | create_empty_model_from_nvm_and_database( 81 | dataset_dir / "3D-models/all-merged/all.nvm", 82 | dataset_dir / "3D-models/overcast-reference.db", 83 | args.empty_sfm, 84 | ) 85 | elif args.benchmark_name in ["aachen"]: 86 | # Original Aachen 87 | logging.info("Init empty sfm from nvm for aachen...") 88 | create_empty_model_from_nvm_and_database( 89 | dataset_dir / "3D-models/aachen_cvpr2018_db.nvm", 90 | dataset_dir / "database.db", 91 | args.empty_sfm, 92 | dataset_dir / "3D-models/database_intrinsics.txt", 93 | ) 94 | elif args.benchmark_name in ["aachen_v1.1"]: 95 | # Aachen v1.1 96 | logging.info("Init empty sfm from bins for aachenv1.1...") 97 | create_empty_model_from_reference_model( 98 | dataset_dir / "3D-models/aachen_v_1_1", args.empty_sfm 99 | ) 100 | else: 101 | logging.error(f"Invalid benchmark {args.benchmark_name}!!") 102 | 103 | 104 | def reconstruct_database_pairs(args): 105 | from third_party.hloc.hloc import triangulation 106 | 107 | # Reconstruct database pairs 108 | if (args.result_sfm / "model" / "images.bin").exists(): 109 | logging.info("Reconstruction existed.") 110 | return 111 | logging.info("\nReconstructing database images....") 112 | triangulation.main( 113 | args.result_sfm, 114 | args.empty_sfm, 115 | args.im_dir, 116 | args.db_pairs_path, 117 | args.result_dir / "keypoints.h5", 118 | args.result_dir / "matches.h5", 119 | colmap_path=args.colmap, 120 | ) 121 | logging.info("Finished reconstruction.") 122 | 123 | 124 | def localize_queries(args): 125 | from third_party.hloc.hloc import localize_sfm 126 | 127 | # Localize query pairs 128 | for ransac_thresh in args.ransac_thres: 129 | localize_txt = f"{args.model_tag}_{args.conf_tag}.{args.pair_tag}.{args.qt_tag}.r{ransac_thresh}.txt" 130 | if args.covis_cluster: 131 | localize_txt = localize_txt.replace("txt", "covis.txt") 132 | logging.info(f"Localize queries...") 133 | print(">>>>", localize_txt) 134 | localize_sfm.main( 135 | args.result_sfm / "model", 136 | args.dataset_dir / "queries/*_queries_with_intrinsics.txt", 137 | args.query_pairs_path, 138 | args.result_dir / "keypoints.h5", 139 | args.result_dir / "matches.h5", 140 | args.result_dir / localize_txt, 141 | ransac_thresh=ransac_thresh, 142 | covisibility_clustering=args.covis_cluster, 143 | changed_format=True, 144 | benchmark=args.benchmark_name, 145 | ) 146 | 147 | 148 | def get_grouped_ids(array): 149 | # Group array indices based on its values 150 | # all duplicates are grouped as a set 151 | idx_sort = np.argsort(array) 152 | sorted_array = array[idx_sort] 153 | vals, ids, count = np.unique(sorted_array, return_counts=True, return_index=True) 154 | res = np.split(idx_sort, ids[1:]) 155 | return res 156 | 157 | 158 | def get_unique_matches_ids(match_ids, scores): 159 | if len(match_ids.shape) == 1: 160 | return [0] 161 | 162 | k1s = match_ids[:, 0] 163 | k2s = match_ids[:, 1] 164 | isets1 = get_grouped_ids(k1s) 165 | isets2 = get_grouped_ids(k2s) 166 | 167 | uid1s = [] 168 | for ids in isets1: 169 | if len(ids) == 1: 170 | uid1s.append(ids[0]) # Unique 171 | else: 172 | uid1s.append(ids[scores[ids].argmax()]) 173 | uid2s = [] 174 | for ids in isets2: 175 | if len(ids) == 1: 176 | uid2s.append(ids[0]) # Unique 177 | else: 178 | uid2s.append(ids[scores[ids].argmax()]) 179 | uids = list(set(uid1s).intersection(uid2s)) 180 | return uids 181 | 182 | 183 | def quantize_keypoints(fpts, kp_data, psize=48, dthres=4): 184 | """Keypoints quantization algorithm. 185 | The image is divided into cells, where each cell represents a psize*psize local patch. 186 | Each input point has its linked coarse point (patch region). 187 | For all points inside a patch region, those points with distances smaller than the given 188 | threshold will be merged and represented by their mean pixel. 189 | 190 | Args: 191 | - fpts: the set of input keypoint coordinates, shape (N, 2) 192 | - kp_data: the keypoint data dict linked to an image 193 | - psize: the size to divide an image into multiple patch regions 194 | - dthres: the distance threshold (in pixels) for merging points within the same patch region 195 | Return: 196 | - fpt_ids: the keypoint ids of the input points 197 | """ 198 | 199 | # kp_data: {'kps':[], 'kp_means': kp_dict} 200 | fpt_ids = [] 201 | cpts = fpts // psize * psize # Point coordinates (x, y) 202 | for cpt, fpt in zip(cpts, fpts): 203 | cpt = tuple(cpt) 204 | kps = kp_data["kps"] 205 | kp_dict = kp_data["kp_means"] # {cpt : {'means':[], 'kids':[]}} 206 | if cpt not in kp_dict: 207 | kid = len(kps) 208 | kps.append(fpt) # Insert another keypoint 209 | kp_dict[cpt] = {"means": [fpt], "kids": [kid]} # Init 1st center 210 | else: 211 | kids = kp_dict[cpt]["kids"] 212 | centers = kp_dict[cpt]["means"] # N, 2 213 | dist = np.linalg.norm(fpt - np.array(centers), axis=1) 214 | cid = np.argmin(dist) 215 | if dist[cid] < dthres: 216 | centers[cid] = (centers[cid] + fpt) / 2 # Update center 217 | kid = kids[cid] 218 | kps[kid] = centers[cid] # Update key point value 219 | else: 220 | kid = len(kps) 221 | kps.append(fpt) # Insert another keypoint 222 | centers.append(fpt) # Insert as a new center 223 | kids.append(kid) 224 | fpt_ids.append(kid) 225 | return fpt_ids 226 | 227 | 228 | def compute_keypoints(pts, kp_data): 229 | kps = kp_data["kps"] 230 | kp_dict = kp_data["kpids"] 231 | pt_ids = [] 232 | for pt in pts: 233 | key = tuple(pt) 234 | if key not in kp_dict: 235 | kid = len(kps) # 0-based, the next inserting index 236 | kps.append(pt) 237 | kp_dict[key] = kid 238 | else: 239 | kid = kp_dict[key] 240 | pt_ids.append(kid) 241 | return pt_ids 242 | 243 | 244 | def match_pairs_with_keys_exporth5(matcher, pairs, pair_keys, match_file, debug=False): 245 | # Pairwise matching 246 | num_matches = [] 247 | match_times = [] 248 | start_time = time.time() 249 | 250 | with h5py.File(match_file, "a") as fmatch: 251 | matched = list(fmatch.keys()) 252 | print(f"\nLoad match file, existing matches {len(matched)}") 253 | print(f"Start matching, total {len(pairs)} pairs...") 254 | for pair, key in tqdm(zip(pairs, pair_keys), total=len(pairs), smoothing=0.1): 255 | im1_path, im2_path = pair 256 | if key in matched: 257 | num_matches.append(len(fmatch[key]["matches"])) 258 | continue 259 | 260 | try: 261 | t0 = time.time() 262 | match_res = matcher(im1_path, im2_path) 263 | match_times.append(time.time() - t0) 264 | except: 265 | print(f"##Failed matching on {key}") 266 | continue 267 | 268 | matches = match_res[0] 269 | scores = match_res[3] 270 | N = len(matches) 271 | num_matches.append(N) 272 | 273 | # Add print for easy debugging 274 | if debug: 275 | print(f"{pair} matches: {N}") 276 | 277 | # Save matches 278 | grp = fmatch.create_group(key) 279 | grp.create_dataset("matches", data=matches) 280 | grp.create_dataset("scores", data=scores) 281 | total_time = time.time() - start_time 282 | mean_time = np.mean(match_times) if len(match_times) > 0 else 0.0 283 | print( 284 | f"Finished matched pairs: {len(fmatch)} num_matches:{np.mean(num_matches):.2f} " 285 | f"match_time/pair:{mean_time:.2f}s time:{total_time:.2f}s." 286 | ) 287 | 288 | 289 | def match_pairs_exporth5(pair_list, matcher, im_dir, output_dir, debug=False): 290 | # Construct pairs and pair keys 291 | pairs = [] 292 | pair_keys = [] 293 | pair_keys_set = set() 294 | for pair_line in tqdm(pair_list, smoothing=0.1): 295 | name0, name1 = pair_line.split() 296 | key = names_to_pair(name0, name1) 297 | key_inv = names_to_pair(name1, name0) 298 | if key_inv in pair_keys_set: 299 | continue 300 | 301 | pair_keys.append(key) 302 | pair_keys_set.add(key) 303 | pair = (str(im_dir / name0), str(im_dir / name1)) 304 | pairs.append(pair) 305 | 306 | match_file = output_dir / "matches_raw.h5" 307 | match_pairs_with_keys_exporth5(matcher, pairs, pair_keys, match_file, debug=debug) 308 | 309 | 310 | def process_matches_and_keypoints_exporth5( 311 | pair_list, 312 | output_dir, 313 | result_dir, 314 | sc_thres=0.25, 315 | qt_dthres=4, 316 | qt_psize=48, 317 | qt_unique=True, 318 | ): 319 | """ 320 | This function should be executed after running match_pairs_exporth5(), which save the 321 | precomputed the raw matches and their scores in a hdf5. 322 | This function first filters out the less confident matches based on the given score threshold. 323 | Then the keypoints are computed accordingly from the filtered matches by finding the unique set 324 | of points that have been matched for each image. 325 | Afterwards, matches are represented by pairs of keypoint ids. 326 | We further provide an option to quantize keypoints, where keypoints are closer than 327 | a given distance will be merged into one. 328 | For the methods that directly obtain pixel-level matches without having an explicit keypoint 329 | detection stage, such quantization could be helpful or necessary to enable proper stability 330 | for the triangulation step of a localization pipeline. 331 | Notice, such quantization sacrifices the pixel-wise accuracy, so the methods with 332 | specific keypoint detection (SuperPoint, D2Net, ..etc) or the methods that produce 333 | patch-level matches (NCNet, SparseNCNet) should not use it. 334 | 335 | Args: 336 | - pair_list: list of pair strings 337 | - output_dir: the directory where matches_raw.h5 is saved 338 | - result_dir: the output directory to save processed matches and keypoints 339 | - sc_thres: the score threshold to filter less confident matches 340 | - qt_dthres, qt_psize: args for quantization, set qt_psize as 0 can skip the quantization. 341 | - qt_unique: flag to maintain the uniqueness of matches after quantization 342 | 343 | Nothing is returned. The processed keypoints are saved as keypoints.h5 and matches (represented by the 344 | keypoints ids) are saved as matches.h5 in the result directory. 345 | """ 346 | # Select matches from the raw matches and quantize keypoints 347 | if (result_dir / "keypoints.h5").exists(): 348 | logging.info("Result matches and keypoints already existed, skip") 349 | return 350 | 351 | match_file = h5py.File(output_dir / "matches_raw.h5", "r") 352 | all_kp_data = {} 353 | logging.info("Start parse matches and quantize keypoints ...") 354 | with h5py.File(result_dir / "matches.h5", "w") as res_fmatch: 355 | for pair in tqdm(pair_list, smoothing=0.1): 356 | name0, name1 = pair.split() 357 | pair = names_to_pair(name0, name1) 358 | 359 | if pair not in match_file: 360 | continue 361 | 362 | matches = match_file[pair]["matches"].__array__() 363 | scores = match_file[pair]["scores"].__array__() 364 | valid = np.where(scores >= sc_thres)[0] 365 | matches = matches[valid] 366 | scores = scores[valid] 367 | 368 | # Compute match ids and quantize keypoints 369 | match_ids = matches_to_keypoint_ids( 370 | matches, 371 | scores, 372 | name0, 373 | name1, 374 | all_kp_data, 375 | qt_dthres, 376 | qt_psize, 377 | qt_unique, 378 | ) 379 | 380 | # Save matches 381 | grp = res_fmatch.create_group(pair) 382 | grp.create_dataset("matches0", data=match_ids) 383 | num_pairs = len(res_fmatch.keys()) 384 | 385 | # Save keypoints 386 | with h5py.File(result_dir / "keypoints.h5", "w") as res_fkp: 387 | logging.info(f"Save keypoints from {len(all_kp_data)} images...") 388 | for name in tqdm(all_kp_data, smoothing=0.1): 389 | kps = np.array(all_kp_data[name]["kps"], dtype=np.float32) 390 | kgrp = res_fkp.create_group(name) 391 | kgrp.create_dataset("keypoints", data=kps) 392 | logging.info(f"Finished quantization, match pairs:{num_pairs}") 393 | match_file.close() 394 | 395 | 396 | def matches_to_keypoint_ids( 397 | matches, 398 | scores, 399 | name0, 400 | name1, 401 | all_kp_data, 402 | qt_dthres=-1, 403 | qt_psize=-1, 404 | qt_unique=True, 405 | ): 406 | if len(matches) == 0: 407 | return np.empty([0, 2], dtype=np.int32) 408 | 409 | if qt_psize > 0 and qt_dthres > 0: 410 | # Compute keypoints jointly with quantization 411 | if name0 not in all_kp_data: 412 | all_kp_data[name0] = {"kps": [], "kp_means": {}} 413 | if name1 not in all_kp_data: 414 | all_kp_data[name1] = {"kps": [], "kp_means": {}} 415 | 416 | id1s = quantize_keypoints( 417 | matches[:, 0:2], all_kp_data[name0], psize=qt_psize, dthres=qt_dthres 418 | ) 419 | id2s = quantize_keypoints( 420 | matches[:, 2:4], all_kp_data[name1], psize=qt_psize, dthres=qt_dthres 421 | ) 422 | match_ids = np.dstack([id1s, id2s]).reshape(-1, 2) # N, 2 423 | 424 | # Remove n-to-1 matches after quantization 425 | if qt_unique and len(match_ids) > 1: 426 | uids = get_unique_matches_ids(match_ids, scores) 427 | match_ids = match_ids[uids] 428 | uids = get_unique_matches_ids(match_ids, scores) 429 | match_ids = match_ids[uids] 430 | else: 431 | # Compute keypoints without quantization 432 | if name0 not in all_kp_data: 433 | all_kp_data[name0] = {"kps": [], "kpids": {}} 434 | if name1 not in all_kp_data: 435 | all_kp_data[name1] = {"kps": [], "kpids": {}} 436 | 437 | id1s = compute_keypoints(matches[:, 0:2], all_kp_data[name0]) 438 | id2s = compute_keypoints(matches[:, 2:4], all_kp_data[name1]) 439 | match_ids = np.dstack([id1s, id2s]).reshape(-1, 2) 440 | 441 | return match_ids 442 | -------------------------------------------------------------------------------- /immatch/utils/metrics.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | 4 | 5 | def check_data_hist(data_list, bins, tag="", return_hist=False): 6 | if not data_list: 7 | if return_hist: 8 | return "", None 9 | return "" 10 | hists = [] 11 | means = [] 12 | Ns = [] 13 | for data in data_list: 14 | N = len(data) 15 | Ns.append(N) 16 | if N == 0: 17 | continue 18 | counts = np.histogram(data, bins)[0] 19 | hists.append(counts / N) 20 | means.append(np.mean(data)) 21 | 22 | hist_print = f"{tag} Sample/N(mean/max/min)={len(data_list)}/{np.mean(Ns):.0f}/{np.max(Ns):.0f}/{np.min(Ns):.0f}\n" 23 | hist_print += f"Ratios(%): mean={np.mean(means):.2f}" 24 | mean_hists = np.mean(hists, axis=0) 25 | for val, low, high in zip(mean_hists, bins[0:-1], bins[1::]): 26 | hist_print += " [{},{})={:.2f}".format(low, high, 100 * val) 27 | if return_hist: 28 | return mean_hists, hist_print 29 | return hist_print 30 | 31 | 32 | def cal_relapose_auc(statis, thresholds=[5, 10, 20]): 33 | min_pose_err = np.maximum(np.array(statis["R_errs"]), np.array(statis["t_errs"])) 34 | auc = cal_error_auc(min_pose_err, thresholds) 35 | print(f"RelaPose AUC@{'/'.join(map(str, thresholds))}deg: {auc}%") 36 | return auc 37 | 38 | 39 | def cal_error_auc(errors, thresholds): 40 | if len(errors) == 0: 41 | return np.zeros(len(thresholds)) 42 | N = len(errors) 43 | errors = np.append([0.0], np.sort(errors)) 44 | recalls = np.arange(N + 1) / N 45 | aucs = [] 46 | for thres in thresholds: 47 | last_index = np.searchsorted(errors, thres) 48 | rcs_ = np.append(recalls[:last_index], recalls[last_index - 1]) 49 | errs_ = np.append(errors[:last_index], thres) 50 | aucs.append(np.trapz(rcs_, x=errs_) / thres) 51 | return 100 * np.array(aucs) 52 | 53 | 54 | def cal_abspose_error(R, R_gt, t, t_gt): 55 | if isinstance(R_gt, np.ndarray): 56 | R_gt = torch.from_numpy(R_gt).to(torch.float32) 57 | if isinstance(t_gt, np.ndarray): 58 | t_gt = torch.from_numpy(t_gt).to(torch.float32) 59 | R_err = ( 60 | torch.clip(0.5 * (torch.sum(R_gt * R_gt.new_tensor(R)) - 1), -1, 1).acos() 61 | * 180.0 62 | / pi 63 | ) 64 | t_err = torch.norm(t_gt.new_tensor(t) - t_gt) 65 | return R_err, t_err 66 | 67 | 68 | def cal_rot_error(R, R_gt): 69 | d = np.clip((np.trace(R.T.dot(R_gt)) - 1) / 2, -1.0, 1.0) 70 | err = np.rad2deg(np.arccos(d)) 71 | return err 72 | 73 | 74 | def cal_vec_angular_error(t, t_gt): 75 | norm = np.linalg.norm(t_gt) * np.linalg.norm(t) 76 | d = np.clip(np.dot(t, t_gt) / norm, -1.0, 1.0) 77 | err = np.rad2deg(np.arccos(d)) 78 | 79 | # This is what Aspanformer is using ... 80 | # t_err = np.minimum(t_err, 180 - t_err) # handle E ambiguity 81 | return err 82 | 83 | 84 | def cal_relapose_error(R, R_gt, t, t_gt): 85 | t_err = cal_vec_angular_error(t, t_gt) 86 | R_err = cal_rot_error(R, R_gt) 87 | return R_err, t_err 88 | 89 | 90 | def cal_reproj_dists_H(p1s, p2s, homography): 91 | """Compute the reprojection errors using the GT homography""" 92 | 93 | p1s_h = np.concatenate([p1s, np.ones([p1s.shape[0], 1])], axis=1) # Homogenous 94 | p2s_proj_h = np.transpose(np.dot(homography, np.transpose(p1s_h))) 95 | p2s_proj = p2s_proj_h[:, :2] / p2s_proj_h[:, 2:] 96 | dist = np.sqrt(np.sum((p2s - p2s_proj) ** 2, axis=1)) 97 | return dist 98 | -------------------------------------------------------------------------------- /immatch/utils/model_helper.py: -------------------------------------------------------------------------------- 1 | import os 2 | import yaml 3 | import immatch 4 | 5 | 6 | def parse_model_config(config, benchmark_name, root_dir="."): 7 | config_file = f"{root_dir}/configs/{config}.yml" 8 | with open(config_file, "r") as f: 9 | model_conf = yaml.load(f, Loader=yaml.FullLoader)[benchmark_name] 10 | 11 | # Update pretrained model path 12 | if "ckpt" in model_conf and root_dir != ".": 13 | model_conf["ckpt"] = os.path.join(root_dir, model_conf["ckpt"]) 14 | if "coarse" in model_conf and "ckpt" in model_conf["coarse"]: 15 | model_conf["coarse"]["ckpt"] = os.path.join( 16 | root_dir, model_conf["coarse"]["ckpt"] 17 | ) 18 | return model_conf 19 | 20 | 21 | def init_model(config, benchmark_name, root_dir="."): 22 | # Load model config 23 | model_conf = parse_model_config(config, benchmark_name, root_dir) 24 | 25 | # Initialize model 26 | class_name = model_conf["class"] 27 | model = immatch.__dict__[class_name](model_conf) 28 | print(f"Method:{class_name} Conf: {model_conf}") 29 | return model, model_conf 30 | -------------------------------------------------------------------------------- /outputs/hpatches/cache/CAPS_SIFT.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GrumpyZhou/image-matching-toolbox/470c9e1ba07bab28438afdff9f934c8589a92ed6/outputs/hpatches/cache/CAPS_SIFT.npy -------------------------------------------------------------------------------- /outputs/hpatches/cache/CAPS_SuperPoint_r4.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GrumpyZhou/image-matching-toolbox/470c9e1ba07bab28438afdff9f934c8589a92ed6/outputs/hpatches/cache/CAPS_SuperPoint_r4.npy -------------------------------------------------------------------------------- /outputs/hpatches/cache/D2Net.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GrumpyZhou/image-matching-toolbox/470c9e1ba07bab28438afdff9f934c8589a92ed6/outputs/hpatches/cache/D2Net.npy -------------------------------------------------------------------------------- /outputs/hpatches/cache/DoG1024-AffNet-HardNet.m0.95.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GrumpyZhou/image-matching-toolbox/470c9e1ba07bab28438afdff9f934c8589a92ed6/outputs/hpatches/cache/DoG1024-AffNet-HardNet.m0.95.npy -------------------------------------------------------------------------------- /outputs/hpatches/cache/NCNet.im1024.m0.9.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GrumpyZhou/image-matching-toolbox/470c9e1ba07bab28438afdff9f934c8589a92ed6/outputs/hpatches/cache/NCNet.im1024.m0.9.npy -------------------------------------------------------------------------------- /outputs/hpatches/cache/Patch2Pix.im1024.m0.5.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GrumpyZhou/image-matching-toolbox/470c9e1ba07bab28438afdff9f934c8589a92ed6/outputs/hpatches/cache/Patch2Pix.im1024.m0.5.npy -------------------------------------------------------------------------------- /outputs/hpatches/cache/Patch2Pix.im1024.m0.9.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GrumpyZhou/image-matching-toolbox/470c9e1ba07bab28438afdff9f934c8589a92ed6/outputs/hpatches/cache/Patch2Pix.im1024.m0.9.npy -------------------------------------------------------------------------------- /outputs/hpatches/cache/R2D2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GrumpyZhou/image-matching-toolbox/470c9e1ba07bab28438afdff9f934c8589a92ed6/outputs/hpatches/cache/R2D2.npy -------------------------------------------------------------------------------- /outputs/hpatches/cache/SparseNCNet_N2000.im3200.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GrumpyZhou/image-matching-toolbox/470c9e1ba07bab28438afdff9f934c8589a92ed6/outputs/hpatches/cache/SparseNCNet_N2000.im3200.npy -------------------------------------------------------------------------------- /outputs/hpatches/cache/SuperGlue_r4.m0.2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GrumpyZhou/image-matching-toolbox/470c9e1ba07bab28438afdff9f934c8589a92ed6/outputs/hpatches/cache/SuperGlue_r4.m0.2.npy -------------------------------------------------------------------------------- /outputs/hpatches/cache/SuperGlue_r4.m0.5.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GrumpyZhou/image-matching-toolbox/470c9e1ba07bab28438afdff9f934c8589a92ed6/outputs/hpatches/cache/SuperGlue_r4.m0.5.npy -------------------------------------------------------------------------------- /outputs/hpatches/cache/SuperGlue_r4.m0.9.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GrumpyZhou/image-matching-toolbox/470c9e1ba07bab28438afdff9f934c8589a92ed6/outputs/hpatches/cache/SuperGlue_r4.m0.9.npy -------------------------------------------------------------------------------- /outputs/hpatches/cache/SuperPoint_r4.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GrumpyZhou/image-matching-toolbox/470c9e1ba07bab28438afdff9f934c8589a92ed6/outputs/hpatches/cache/SuperPoint_r4.npy -------------------------------------------------------------------------------- /outputs/hpatches/cache/aslfeat.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GrumpyZhou/image-matching-toolbox/470c9e1ba07bab28438afdff9f934c8589a92ed6/outputs/hpatches/cache/aslfeat.npy -------------------------------------------------------------------------------- /outputs/hpatches/cache/delf.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GrumpyZhou/image-matching-toolbox/470c9e1ba07bab28438afdff9f934c8589a92ed6/outputs/hpatches/cache/delf.npy -------------------------------------------------------------------------------- /outputs/hpatches/cache/hesaff.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GrumpyZhou/image-matching-toolbox/470c9e1ba07bab28438afdff9f934c8589a92ed6/outputs/hpatches/cache/hesaff.npy -------------------------------------------------------------------------------- /outputs/hpatches/cache/hesaffnet.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GrumpyZhou/image-matching-toolbox/470c9e1ba07bab28438afdff9f934c8589a92ed6/outputs/hpatches/cache/hesaffnet.npy -------------------------------------------------------------------------------- /pretrained/download.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # CAPS 4 | mkdir -p caps 5 | gdown --id 1UVjtuhTDmlvvVuUlEq_M5oJVImQl6z1f -O caps/caps-pretrained.pth 6 | 7 | # D2Net 8 | mkdir -p d2net 9 | wget https://dsmn.ml/files/d2-net/d2_tf.pth -O d2net/d2_tf.pth 10 | wget https://dsmn.ml/files/d2-net/d2_tf_no_phototourism.pth -O d2net/d2_tf_no_phototourism.pth 11 | 12 | # R2D2 Symbolic links 13 | ln -s ../third_party/r2d2/models r2d2 14 | 15 | # SparseNCNet 16 | mkdir -p sparsencnet 17 | wget https://www.di.ens.fr/willow/research/sparse-ncnet/models/sparsencnet_k10.pth.tar -O sparsencnet/sparsencnet_k10.pth.tar 18 | 19 | # Patch2Pix Symbolic links 20 | ln -s ../third_party/patch2pix/pretrained patch2pix 21 | cd patch2pix 22 | bash download.sh 23 | cd .. 24 | 25 | # LoFTR 26 | mkdir -p loftr 27 | gdown --id 1M-VD35-qdB5Iw-AtbDBCKC7hPolFW9UY -O loftr/outdoor_ds.ckpt 28 | gdown --id 19s3QvcCWQ6g-N1PrYlDCg-2mOJZ3kkgS -O loftr/indoor_ds_new.ckpt 29 | 30 | # COTR 31 | mkdir -p cotr 32 | cd cotr/ 33 | wget https://www.cs.ubc.ca/research/kmyi_data/files/2021/cotr/default.zip 34 | unzip -j default.zip 35 | mv checkpoint.pth.tar cotr_default.pth.tar 36 | rm default.zip 37 | cd .. 38 | 39 | # ASpanFormer Symbolic links 40 | ln -s ../third_party/aspanformer/weights aspanformer -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | 3 | setup( 4 | name="immatch", 5 | version="0.1.0", 6 | install_requires=[ 7 | "transforms3d", 8 | "tqdm", 9 | "pyyaml", 10 | "einops", 11 | "kornia", 12 | "kornia_moons", 13 | "yacs", 14 | "pillow", 15 | "h5py", 16 | "gdown", 17 | "tables", 18 | "vispy", 19 | "glfw", 20 | "albumentations", 21 | ], 22 | packages=find_packages(), 23 | author="Qunjie Zhou", 24 | python_requires=">=3.7", 25 | classifiers=[ 26 | "License :: OSI Approved :: MIT License", 27 | "Operating System :: OS Independent", 28 | "Programming Language :: Python :: 3", 29 | "Programming Language :: Python :: 3.7", 30 | "Programming Language :: Python :: 3.8", 31 | "Programming Language :: Python :: 3 :: Only", 32 | "Topic :: Scientific/Engineering", 33 | ], 34 | license="MIT", 35 | keywords="image feature matching", 36 | ) 37 | --------------------------------------------------------------------------------