├── .gitignore ├── .gitmodules ├── EVAL.md ├── LICENSE ├── README.md ├── TRAIN.md ├── datasets ├── adelaide.py ├── hope.py ├── nyuvp.py ├── smh.py └── su3.py ├── environment.yml ├── networks └── cn_net.py ├── parsac.py └── utils ├── backward.py ├── evaluation.py ├── initialisation.py ├── inlier_counting.py ├── metrics.py ├── options.py ├── postprocessing.py ├── residual_functions.py ├── sampling.py ├── solvers.py ├── tee.py └── visualisation.py /.gitignore: -------------------------------------------------------------------------------- 1 | tmp/ 2 | results/ 3 | .idea/ 4 | __pycache__ -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "datasets/yud_plus"] 2 | path = datasets/yud_plus 3 | url = git@github.com:fkluger/yud_plus.git 4 | [submodule "datasets/nyu_vp"] 5 | path = datasets/nyu_vp 6 | url = git@github.com:fkluger/nyu_vp.git 7 | -------------------------------------------------------------------------------- /EVAL.md: -------------------------------------------------------------------------------- 1 | The `parsac.py` script can perform both training and evaluation. 2 | The `--eval` flag sets some default options to run evaluation on the test set. 3 | 4 | After following the instructions for downloading the datasets and pre-trained network weights ([README](README.md)), 5 | you can execute the commands below in order to reproduce the results from our paper. 6 | 7 | # Main Results 8 | ## Vanishing Points 9 | 10 | ### SU3 11 | ``` 12 | python parsac.py --eval --dataset su3 --data_path datasets/su3 --problem vp --load weights/main_results/vp_su3 --inlier_threshold 0.0001 --instances 8 --hypotheses 32 13 | ``` 14 | 15 | ### YUD 16 | ``` 17 | python parsac.py --eval --dataset yud --data_path datasets/yud_plus/data --problem vp --load weights/main_results/vp_su3 --inlier_threshold 0.0001 --instances 8 --hypotheses 32 18 | ``` 19 | 20 | ### NYU-VP 21 | ``` 22 | python parsac.py --eval --dataset nyuvp --data_path datasets/nyu_vp/data --problem vp --load weights/main_results/vp_nyu --inlier_threshold 0.0001 --instances 8 --hypotheses 32 23 | ``` 24 | 25 | 26 | ### YUD+ 27 | ``` 28 | python parsac.py --eval --dataset yudplus --data_path datasets/yud_plus/data --problem vp --load weights/main_results/vp_su3 --inlier_threshold 0.0001 --instances 8 --hypotheses 32 29 | ``` 30 | 31 | ## Fundamental Matrices 32 | ### HOPE-F 33 | ``` 34 | python parsac.py --eval --load ./weights/main_results/fundamental --dataset hope --data_path ./datasets/hope --problem fundamental --inlier_threshold 0.01 --assignment_threshold 0.02 --instances 4 --hypotheses 128 35 | ``` 36 | 37 | ### Adelaide 38 | ``` 39 | python parsac.py --eval --load ./weights/main_results/fundamental --dataset adelaide --data_path ./datasets/adelaide --problem fundamental --inlier_threshold 0.01 --assignment_threshold 0.02 --instances 4 --hypotheses 128 40 | ``` 41 | 42 | ## Homographies 43 | ### SMH 44 | ``` 45 | python parsac.py --eval --load weights/main_results/homography --dataset smh --data_path datasets/smh --problem homography --inlier_threshold 1e-6 --assignment_threshold 4e-6 --instances 24 --hypotheses 512 46 | ``` 47 | 48 | ### Adelaide 49 | ``` 50 | python parsac.py --eval --load weights/main_results/homography --dataset adelaide --data_path datasets/adelaide --problem homography --inlier_threshold 1e-4 --assignment_threshold 4e-3 --instances 24 --hypotheses 512 51 | ``` 52 | 53 | # Self-Supervised Learning 54 | 55 | ## Weighted Loss 56 | 57 | ### SU3 58 | 59 | ``` 60 | python parsac.py --eval --dataset su3 --data_path datasets/su3 --problem vp --inlier_threshold 0.0001 --instances 8 --hypotheses 32 --load weights/self_supervised/weighted_su3 61 | ``` 62 | 63 | ### NYU-VP 64 | ``` 65 | python parsac.py --eval --dataset nyuvp --data_path datasets/nyu_vp/data --problem vp --inlier_threshold 0.0001 --instances 8 --hypotheses 32 --load weights/self_supervised/weighted_nyu 66 | ``` 67 | 68 | ### HOPE-F 69 | ``` 70 | python parsac.py --eval --load ./weights/self_supervised/weighted_fundamental --dataset hope --data_path ./datasets/hope --problem fundamental --inlier_threshold 0.01 --assignment_threshold 0.02 --instances 4 --hypotheses 128 71 | ``` 72 | 73 | ### Adelaide-F 74 | ``` 75 | python parsac.py --eval --load ./weights/self_supervised/weighted_fundamental --dataset adelaide --data_path ./datasets/adelaide --problem fundamental --inlier_threshold 0.01 --assignment_threshold 0.02 --instances 4 --hypotheses 128 76 | ``` 77 | 78 | ## Unweighted Loss 79 | 80 | ### SU3 81 | ``` 82 | python parsac.py --eval --dataset su3 --data_path datasets/su3 --problem vp --inlier_threshold 0.0001 --instances 8 --hypotheses 32 --load weights/self_supervised/unweighted_su3 83 | ``` 84 | 85 | ### NYU-VP 86 | ``` 87 | python parsac.py --eval --dataset nyuvp --data_path datasets/nyu_vp/data --problem vp --inlier_threshold 0.0001 --instances 8 --hypotheses 32 --load weights/self_supervised/unweighted_nyu 88 | ``` 89 | 90 | ### HOPE-F 91 | ``` 92 | python parsac.py --eval --load ./weights/self_supervised/unweighted_fundamental --dataset hope --data_path ./datasets/hope --problem fundamental --inlier_threshold 0.01 --assignment_threshold 0.02 --instances 4 --hypotheses 128 93 | ``` 94 | 95 | ### Adelaide-F 96 | ``` 97 | python parsac.py --eval --load ./weights/self_supervised/unweighted_fundamental --dataset adelaide --data_path ./datasets/adelaide --problem fundamental --inlier_threshold 0.01 --assignment_threshold 0.02 --instances 4 --hypotheses 128 98 | ``` 99 | 100 | # Ablation Study: Number of Model Instances 101 | ``` 102 | python parsac.py --eval --dataset su3 --data_path datasets/su3 --problem vp --inlier_threshold 0.0001 --hypotheses 32 --load weights/ablation_instances/M_HAT --instances M_HAT 103 | ``` 104 | Replace `M_HAT` with the number of putative model instances. Valid values are {2, 3, 4, 6, 10, 12, 16}. 105 | 106 | # Ablation Study: Weighted Inlier Counting 107 | ## with weighted inlier counting 108 | See [main results](#main-results) 109 | 110 | ## w/o weighted inlier counting 111 | 112 | ### SU3 113 | ``` 114 | python parsac.py --eval --dataset su3 --data_path datasets/su3 --problem vp --load weights/ablation_unweighted/su3 --inlier_threshold 0.0001 --instances 8 --hypotheses 32 --inlier_counting unweighted 115 | ``` 116 | 117 | ### HOPE-F 118 | ``` 119 | python parsac.py --eval --dataset hope --data_path ./datasets/hope --problem fundamental --load ./weights/main_results/fundamental --inlier_threshold 0.01 --assignment_threshold 0.02 --instances 4 --hypotheses 128 --inlier_counting unweighted 120 | ``` 121 | 122 | ### SMH 123 | ``` 124 | python parsac.py --eval --load weights/main_results/homography --dataset smh --data_path datasets/smh --problem homography --inlier_threshold 1e-6 --assignment_threshold 4e-6 --instances 24 --hypotheses 512 --inlier_counting unweighted 125 | ``` 126 | 127 | # Ablation Study: Robustness to Noise and Outliers 128 | Note: the following options only work for SU3 and Adelaide. 129 | 130 | In order to add Gaussian noise with standard deviation `sigma` to the input observations, use the following parameter: 131 | ``` 132 | --ablation_noise sigma 133 | ``` 134 | 135 | In order to remove all ground truth outliers from the observations and then add synthetic outliers with an outlier rate of `outlier_rate`, use the following parameter: 136 | ``` 137 | --ablation_outlier_ratio outlier_rate 138 | ``` 139 | 140 | 141 | # Ablation Study: Feature Generalisation 142 | 143 | We provide line segments extracted with [DeepLSD](https://github.com/cvg/DeepLSD) for SU3, NYU-VP and YUD(+). 144 | Download and extract the following archive before you can run the ablation study experiments below:\ 145 | https://cloud.tnt.uni-hannover.de/index.php/s/M7TTyqGzbnCfiJX 146 | 147 | ## Train: LSD / Test: LSD 148 | See [main results](#main-results) 149 | 150 | ## Train: LSD / Test: DeepLSD 151 | ### SU3 152 | ``` 153 | python parsac.py --eval --dataset su3 --data_path datasets/su3 --problem vp --load weights/main_results/vp_su3 --inlier_threshold 0.0001 --instances 8 --hypotheses 32 --ablation_deeplsd_folder deeplsd_features/su3 154 | ``` 155 | 156 | ### NYU-VP 157 | ``` 158 | python parsac.py --eval --dataset nyuvp --data_path datasets/nyu_vp/data --problem vp --load weights/main_results/vp_nyu --inlier_threshold 0.0001 --instances 8 --hypotheses 32 --ablation_deeplsd_folder deeplsd_features/nyu 159 | ``` 160 | ### YUD 161 | ``` 162 | python parsac.py --eval --dataset yud --data_path datasets/yud_plus/data --problem vp --load weights/main_results/vp_su3 --inlier_threshold 0.0001 --instances 8 --hypotheses 32 --ablation_deeplsd_folder deeplsd_features/yud 163 | ``` 164 | ### YUD+ 165 | ``` 166 | python parsac.py --eval --dataset yudplus --data_path datasets/yud_plus/data --problem vp --load weights/main_results/vp_su3 --inlier_threshold 0.0001 --instances 8 --hypotheses 32 --ablation_deeplsd_folder deeplsd_features/yud 167 | ``` 168 | 169 | ## Train: DeepLSD / Test: LSD 170 | ### SU3 171 | ``` 172 | python parsac.py --eval --dataset su3 --data_path datasets/su3 --problem vp --load weights/ablation_features/su3_deeplsd/ --inlier_threshold 0.0001 --instances 8 --hypotheses 32 173 | ``` 174 | ### NYU-VP 175 | ``` 176 | python parsac.py --eval --dataset nyuvp --data_path datasets/nyu_vp/data --problem vp --load weights/ablation_features/nyu_deeplsd/ --inlier_threshold 0.0001 --instances 8 --hypotheses 32 177 | ``` 178 | ### YUD 179 | ``` 180 | python parsac.py --eval --dataset yud --data_path datasets/yud_plus/data --problem vp --load weights/ablation_features/su3_deeplsd --inlier_threshold 0.0001 --instances 8 --hypotheses 32 181 | ``` 182 | ### YUD+ 183 | ``` 184 | python parsac.py --eval --dataset yudplus --data_path datasets/yud_plus/data --problem vp --load weights/ablation_features/su3_deeplsd --inlier_threshold 0.0001 --instances 8 --hypotheses 32 185 | ``` 186 | 187 | ## Train: DeepLSD / Test: DeepLSD 188 | ### SU3 189 | ``` 190 | python parsac.py --eval --dataset su3 --data_path datasets/su3 --problem vp --load weights/ablation_features/su3_deeplsd/ --inlier_threshold 0.0001 --instances 8 --hypotheses 32 --ablation_deeplsd_folder deeplsd_features/su3 191 | ``` 192 | ### NYU-VP 193 | ``` 194 | python parsac.py --eval --dataset nyuvp --data_path datasets/nyu_vp/data --problem vp --load weights/ablation_features/nyu_deeplsd/ --inlier_threshold 0.0001 --instances 8 --hypotheses 32 --ablation_deeplsd_folder deeplsd_features/nyu 195 | ``` 196 | ### YUD 197 | ``` 198 | python parsac.py --eval --dataset yud --data_path datasets/yud_plus/data --problem vp --load weights/ablation_features/su3_deeplsd --inlier_threshold 0.0001 --instances 8 --hypotheses 32 --ablation_deeplsd_folder deeplsd_features/yud 199 | ``` 200 | ### YUD+ 201 | ``` 202 | python parsac.py --eval --dataset yudplus --data_path datasets/yud_plus/data --problem vp --load weights/ablation_features/su3_deeplsd --inlier_threshold 0.0001 --instances 8 --hypotheses 32 --ablation_deeplsd_folder deeplsd_features/yud 203 | ``` 204 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The copyright in this software is being made available under the BSD License, 2 | included below. This software may be subject to other third party and 3 | contributor rights, including patent rights, and no such rights are granted 4 | under this license. 5 | 6 | 7 | Copyright (c) 2023-present 8 | 9 | Leibniz University Hannover (LUH) 10 | Institut fuer Informationsverarbeitung (TNT) 11 | 12 | 13 | Contacts 14 | 15 | Florian Kluger 16 | 17 | All rights reserved. 18 | 19 | 20 | Redistribution and use in source and binary forms, with or without 21 | modification, are permitted provided that the following conditions are met: 22 | 1. Redistributions of source code must retain the above copyright notice, 23 | this list of conditions and the following disclaimer. 24 | 2. Redistributions in binary form must reproduce the above copyright notice, 25 | this list of conditions and the following disclaimer in the documentation 26 | and/or other materials provided with the distribution. 27 | 3. Neither the name of the copyright holder nor the names of its contributors 28 | may be used to endorse or promote products derived from this software 29 | without specific prior written permission. 30 | 31 | 32 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 33 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 34 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 35 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 36 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 37 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 38 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 39 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 40 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 41 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # PARSAC: Accelerating Robust Multi-Model Fitting with Parallel Sample Consensus 2 | 3 | The paper with supplementary material is available on arXiv: 4 | https://arxiv.org/abs/2401.14919 5 | 6 | If you use this code, please cite our paper: 7 | ``` 8 | @inproceedings{kluger2024parsac, 9 | title={PARSAC: Accelerating Robust Multi-Model Fitting with Parallel Sample Consensus}, 10 | author={Kluger, Florian and Rosenhahn, Bodo}, 11 | booktitle={Proceedings of the AAAI Conference on Artificial Intelligence}, 12 | year={2024} 13 | } 14 | ``` 15 | 16 | 17 | Related repositories: 18 | * [HOPE-F dataset](https://github.com/fkluger/hope-f) 19 | * [SMH dataset](https://github.com/fkluger/smh) 20 | * [NYU-VP dataset](https://github.com/fkluger/nyu_vp) 21 | * [YUD+ dataset](https://github.com/fkluger/yud_plus) 22 | * [CONSAC](https://github.com/fkluger/consac) 23 | * [Our J-/T-Linkage implementation for VP detection](https://github.com/fkluger/vp-linkage) 24 | 25 | ## Installation 26 | Get the code: 27 | ``` 28 | git clone --recurse-submodules https://github.com/fkluger/parsac.git 29 | cd parsac 30 | git submodule update --init --recursive 31 | ``` 32 | 33 | Set up the Python environment using [Anaconda](https://www.anaconda.com/): 34 | ``` 35 | conda env create -f environment.yml 36 | source activate parsac 37 | ``` 38 | 39 | ## Datasets 40 | 41 | ### HOPE-F 42 | Download the [HOPE-F dataset](https://github.com/fkluger/hope-f) and extract it inside the `datasets/hope` directory. 43 | The small dataset w/o images is sufficient for training and evaluation. 44 | 45 | ### Synthetic Metropolis Homographies 46 | Download the [SMH dataset](https://github.com/fkluger/smh) and extract it inside the `datasets/smh` directory. 47 | The small dataset w/o images is sufficient for training and evaluation. 48 | 49 | ### NYU-VP 50 | The vanishing point labels and pre-extracted line segments for the 51 | [NYU dataset](https://cs.nyu.edu/~silberman/datasets/nyu_depth_v2.html) are fetched automatically via the *nyu_vp* 52 | submodule. 53 | 54 | ### YUD and YUD+ 55 | Pre-extracted line segments and VP labels are fetched automatically via the *yud_plus* submodule. RGB images and camera 56 | calibration parameters, however, are not included. Download the original York Urban Dataset from the 57 | [Elder Laboratory's website](http://www.elderlab.yorku.ca/resources/york-urban-line-segment-database-information/) and 58 | store it under the ```datasets/yud_plus/data``` subfolder. 59 | 60 | 61 | ### Adelaide-H/-F 62 | 63 | We provide a mirror of the Adelaide dataset here: https://cloud.tnt.uni-hannover.de/index.php/s/egE6y9KRMxcLg6T. 64 | Download it and place the `.mat` files inside the `datasets/adelaide` directory. 65 | 66 | ## Evaluation 67 | 68 | In order to reproduce the results from the paper using our pre-trained network, 69 | first [download the neural network weights](https://cloud.tnt.uni-hannover.de/index.php/s/KMwborYZsbYAsd2) 70 | and then follow the instructions on the [EVAL page](EVAL.md). 71 | 72 | ## Training 73 | 74 | If you want to train PARSAC from scratch, please follow the instructions on the [TRAIN page](TRAIN.md). 75 | 76 | ## License 77 | [BSD License](LICENSE) 78 | -------------------------------------------------------------------------------- /TRAIN.md: -------------------------------------------------------------------------------- 1 | Use the commands below to train PARSAC according to the experiments in our paper. 2 | 3 | We use [Weights & Biases](https://wandb.ai/) to log the training progress. You can enable it with `--wandb online` to use online syncing or `--wandb offline` to save the logs in a folder offline. 4 | Use the options `--wandb_entity`, `--wandb_group` and `--wandb_dir` to set the entity, group and local directory for your logs. 5 | 6 | # Main Results 7 | ## Vanishing Points 8 | 9 | ### SU3 10 | ``` 11 | python parsac.py --hypotheses 32 --batch 64 --samplecount 8 --inlier_threshold 0.0001 --dataset su3 --problem vp --instances 8 --hypsamples 64 --data_path datasets/su3 --checkpoint_dir ./tmp/checkpoints --no_refine 12 | ``` 13 | 14 | ### NYU-VP 15 | ``` 16 | python parsac.py --hypotheses 32 --batch 64 --samplecount 8 --inlier_threshold 0.0001 --dataset nyuvp --problem vp --instances 8 --hypsamples 64 --data_path datasets/nyu_vp/data --checkpoint_dir ./tmp/checkpoints --no_refine 17 | ``` 18 | 19 | 20 | ## Fundamental Matrices 21 | ### HOPE-F 22 | ``` 23 | python parsac.py --hypotheses 32 --batch 32 --samplecount 16 --inlier_threshold 0.004 --assignment_threshold 0.02 --dataset hope --problem fundamental --instances 4 --hypsamples 128 --epochs 3000 --lr_steps 2500 --data_path datasets/hope --checkpoint_dir ./tmp/checkpoints 24 | ``` 25 | 26 | 27 | ## Homographies 28 | ### SMH 29 | ``` 30 | python parsac.py --hypotheses 32 --batch 4 --samplecount 8 --inlier_threshold 1e-6 --assignment_threshold 4e-6 --dataset smh --problem homography --instances 24 --hypsamples 64 --epochs 500 --lr_steps 350 --data_path datasets/smh --checkpoint_dir ./tmp/checkpoints 31 | ``` 32 | 33 | # Self-Supervised Learning 34 | 35 | ## Weighted Loss 36 | 37 | ### SU3 38 | ``` 39 | python parsac.py --hypotheses 32 --batch 64 --samplecount 8 --inlier_threshold 0.0001 --dataset su3 --problem vp --instances 8 --hypsamples 64 --data_path datasets/su3 --checkpoint_dir ./tmp/checkpoints --no_refine --self_supervised 40 | ``` 41 | 42 | ### NYU-VP 43 | ``` 44 | python parsac.py --hypotheses 32 --batch 64 --samplecount 8 --inlier_threshold 0.0001 --dataset nyuvp --problem vp --instances 8 --hypsamples 64 --data_path datasets/nyu_vp/data --checkpoint_dir ./tmp/checkpoints --no_refine --self_supervised 45 | ``` 46 | 47 | ### HOPE-F 48 | ``` 49 | python parsac.py --hypotheses 32 --batch 32 --samplecount 16 --inlier_threshold 0.004 --assignment_threshold 0.02 --dataset hope --problem fundamental --instances 4 --hypsamples 128 --epochs 3000 --lr_steps 2500 --data_path datasets/hope --checkpoint_dir ./tmp/checkpoints --self_supervised 50 | ``` 51 | 52 | 53 | ## Unweighted Loss 54 | 55 | ### SU3 56 | ``` 57 | python parsac.py --hypotheses 32 --batch 64 --samplecount 8 --inlier_threshold 0.0001 --dataset su3 --problem vp --instances 8 --hypsamples 64 --data_path datasets/su3 --checkpoint_dir ./tmp/checkpoints --no_refine --self_supervised --cumulative_loss -1 58 | ``` 59 | 60 | ### NYU-VP 61 | ``` 62 | python parsac.py --hypotheses 32 --batch 64 --samplecount 8 --inlier_threshold 0.0001 --dataset nyuvp --problem vp --instances 8 --hypsamples 64 --data_path datasets/nyu_vp/data --checkpoint_dir ./tmp/checkpoints --no_refine --self_supervised --cumulative_loss -1 63 | ``` 64 | 65 | 66 | ### HOPE-F 67 | ``` 68 | python parsac.py --hypotheses 32 --batch 32 --samplecount 16 --inlier_threshold 0.004 --assignment_threshold 0.02 --dataset hope --problem fundamental --instances 4 --hypsamples 128 --epochs 3000 --lr_steps 2500 --data_path datasets/hope --checkpoint_dir ./tmp/checkpoints --self_supervised --cumulative_loss -1 69 | ``` 70 | 71 | 72 | # Ablation Study: Number of Model Instances 73 | Replace `M` with the desired number of putative model instances: 74 | ``` 75 | python parsac.py --instances M --hypotheses 32 --batch 64 --samplecount 8 --inlier_threshold 0.0001 --dataset su3 --problem vp --hypsamples 64 --data_path datasets/su3 --checkpoint_dir ./tmp/checkpoints --no_refine 76 | ``` 77 | 78 | 79 | # Ablation Study: (Un-)weighted Inlier Counting 80 | 81 | ## w/o weighted inlier counting 82 | 83 | ### SU3 84 | ``` 85 | python parsac.py --hypotheses 32 --batch 64 --samplecount 8 --inlier_threshold 0.0001 --dataset su3 --problem vp --instances 8 --hypsamples 64 --data_path datasets/su3 --checkpoint_dir ./tmp/checkpoints --no_refine --inlier_counting unweighted 86 | ``` 87 | 88 | ### HOPE-F 89 | ``` 90 | python parsac.py --hypotheses 32 --batch 32 --samplecount 16 --inlier_threshold 0.004 --assignment_threshold 0.02 --dataset hope --problem fundamental --instances 4 --hypsamples 128 --epochs 3000 --lr_steps 2500 --data_path datasets/hope --checkpoint_dir ./tmp/checkpoints --inlier_counting unweighted 91 | ``` 92 | 93 | ### SMH 94 | ``` 95 | python parsac.py --hypotheses 32 --batch 4 --samplecount 8 --inlier_threshold 1e-6 --assignment_threshold 4e-6 --dataset smh --problem homography --instances 24 --hypsamples 64 --epochs 500 --lr_steps 350 --data_path datasets/smh --checkpoint_dir ./tmp/checkpoints --inlier_counting unweighted 96 | ``` 97 | 98 | 99 | # Ablation Study: Feature Generalisation 100 | 101 | ## Train on DeepLSD 102 | 103 | ### SU3 104 | ``` 105 | python parsac.py --hypotheses 32 --batch 64 --samplecount 8 --inlier_threshold 0.0001 --dataset su3 --problem vp --instances 8 --hypsamples 64 --data_path datasets/su3 --checkpoint_dir ./tmp/checkpoints --no_refine --ablation_deeplsd_folder deeplsd_features/su3 106 | ``` 107 | 108 | ### NYU-VP 109 | ``` 110 | python parsac.py --hypotheses 32 --batch 64 --samplecount 8 --inlier_threshold 0.0001 --dataset nyuvp --problem vp --instances 8 --hypsamples 64 --data_path datasets/nyu_vp/data --checkpoint_dir ./tmp/checkpoints --no_refine --ablation_deeplsd_folder deeplsd_features/nyu 111 | ``` 112 | 113 | -------------------------------------------------------------------------------- /datasets/adelaide.py: -------------------------------------------------------------------------------- 1 | import os.path 2 | import scipy.io 3 | from torch.utils.data import Dataset 4 | import numpy as np 5 | 6 | class AdelaideRMF: 7 | 8 | homography_sequences = [ 9 | "barrsmith.mat", 10 | "bonhall.mat", 11 | "bonython.mat", 12 | "elderhalla.mat", 13 | "elderhallb.mat", 14 | "hartley.mat", 15 | "johnsona.mat", 16 | "johnsonb.mat", 17 | "ladysymon.mat", 18 | "library.mat", 19 | "napiera.mat", 20 | "napierb.mat", 21 | "neem.mat", 22 | "nese.mat", 23 | "oldclassicswing.mat", 24 | "physics.mat", 25 | "sene.mat", 26 | "unihouse.mat", 27 | "unionhouse.mat", 28 | ] 29 | 30 | fundamental_sequences = [ 31 | "cube.mat", 32 | "book.mat", 33 | "biscuit.mat", 34 | "game.mat", 35 | "biscuitbook.mat", 36 | "breadcube.mat", 37 | "breadtoy.mat", 38 | "cubechips.mat", 39 | "cubetoy.mat", 40 | "gamebiscuit.mat", 41 | "breadtoycar.mat", 42 | "carchipscube.mat", 43 | "toycubecar.mat", 44 | "breadcubechips.mat", 45 | "biscuitbookbox.mat", 46 | "cubebreadtoychips.mat", 47 | "breadcartoychips.mat", 48 | "dinobooks.mat", 49 | "boardgame.mat" 50 | ] 51 | 52 | def __init__(self, data_dir, keep_in_mem=True, normalize_coords=False, return_images=False, problem="homography", 53 | ablation_outlier_ratio=-1, ablation_noise=0): 54 | 55 | self.data_dir = data_dir 56 | self.keep_in_mem = keep_in_mem 57 | self.normalize_coords = normalize_coords 58 | self.return_images = return_images 59 | self.ablation_outlier_ratio = ablation_outlier_ratio 60 | self.ablation_noise = ablation_noise 61 | 62 | self.dataset_files = [] 63 | 64 | if problem == "homography": 65 | self.sequences = self.homography_sequences 66 | elif problem == "fundamental": 67 | self.sequences = self.fundamental_sequences 68 | elif problem == "all": 69 | self.sequences = self.fundamental_sequences + self.homography_sequences 70 | else: 71 | assert False 72 | 73 | for seq in self.sequences: 74 | seq_ = os.path.join(self.data_dir, seq) 75 | self.dataset_files += [seq_] 76 | 77 | self.dataset = [None for _ in self.dataset_files] 78 | 79 | def __len__(self): 80 | return len(self.dataset) 81 | 82 | def __getitem__(self, key): 83 | 84 | filename = self.dataset_files[key] 85 | datum = self.dataset[key] 86 | 87 | if datum is None: 88 | 89 | data_mat = scipy.io.loadmat(filename, variable_names=["data", "label", "score", "img1", "img2"]) 90 | pts1 = np.transpose(data_mat["data"][0:2,:]) 91 | pts2 = np.transpose(data_mat["data"][3:5,:]) 92 | 93 | if self.ablation_noise > 0: 94 | pts1 += np.random.normal(0, self.ablation_noise, pts1.shape) 95 | pts2 += np.random.normal(0, self.ablation_noise, pts2.shape) 96 | 97 | sideinfo = data_mat["score"].squeeze() 98 | gt_label = data_mat["label"].squeeze() 99 | img1size = data_mat["img1"].shape[0:2] 100 | img2size = data_mat["img2"].shape[0:2] 101 | 102 | if self.ablation_outlier_ratio >= 0: 103 | pts1 = pts1[np.where(gt_label > 0)] 104 | pts2 = pts2[np.where(gt_label > 0)] 105 | sideinfo = sideinfo[np.where(gt_label > 0)] 106 | gt_label = gt_label[np.where(gt_label > 0)] 107 | 108 | if self.ablation_outlier_ratio > 0: 109 | N = pts1.shape[0] 110 | No = int(N * self.ablation_outlier_ratio / (1 - self.ablation_outlier_ratio)) 111 | 112 | opts1 = np.zeros((No, 2)).astype(np.float32) 113 | opts2 = np.zeros((No, 2)).astype(np.float32) 114 | osideinfo = np.zeros(No).astype(np.float32) 115 | ogt_label = np.zeros(No).astype(np.float32) 116 | opts1[:, 0] = np.random.uniform(0, img1size[1]-1, No) 117 | opts1[:, 1] = np.random.uniform(0, img1size[0]-1, No) 118 | opts2[:, 0] = np.random.uniform(0, img2size[1]-1, No) 119 | opts2[:, 1] = np.random.uniform(0, img2size[0]-1, No) 120 | 121 | pts1 = np.concatenate([pts1, opts1], axis=0) 122 | pts2 = np.concatenate([pts2, opts2], axis=0) 123 | sideinfo = np.concatenate([sideinfo, osideinfo], axis=0) 124 | gt_label = np.concatenate([gt_label, ogt_label], axis=0) 125 | 126 | if self.normalize_coords: 127 | scale1 = np.max(img1size) 128 | scale2 = np.max(img2size) 129 | 130 | pts1[:,0] -= img1size[1]/2. 131 | pts2[:,0] -= img2size[1]/2. 132 | pts1[:,1] -= img1size[0]/2. 133 | pts2[:,1] -= img2size[0]/2. 134 | pts1 /= (scale1/2.) 135 | pts2 /= (scale2/2.) 136 | 137 | datum = {'points1': pts1, 'points2': pts2, 'sideinfo': sideinfo, 'img1size': img1size, 'img2size': img2size, 138 | 'labels': gt_label, 'name': self.sequences[key][:-4]} 139 | 140 | if self.return_images: 141 | datum["img1"] = data_mat["img1"] 142 | datum["img2"] = data_mat["img2"] 143 | 144 | if self.keep_in_mem: 145 | self.dataset[key] = datum 146 | 147 | return datum 148 | 149 | 150 | class AdelaideRMFDataset(Dataset): 151 | 152 | def __init__(self, data_dir_path, max_num_points, keep_in_mem=True, ablation_outlier_ratio=-1, ablation_noise=0, 153 | permute_points=True, return_images=False, return_labels=True, problem="homography"): 154 | self.homdata = AdelaideRMF(data_dir_path, keep_in_mem, normalize_coords=True, return_images=return_images, 155 | problem=problem, ablation_outlier_ratio=ablation_outlier_ratio, 156 | ablation_noise=ablation_noise) 157 | self.max_num_points = max_num_points 158 | self.permute_points = permute_points 159 | self.return_images = return_images 160 | self.return_labels = return_labels 161 | 162 | def __len__(self): 163 | return len(self.homdata) 164 | 165 | def __getitem__(self, key): 166 | datum = self.homdata[key] 167 | 168 | if self.max_num_points <= 0: 169 | max_num_points = datum['points1'].shape[0] 170 | else: 171 | max_num_points = self.max_num_points 172 | 173 | if self.permute_points: 174 | perm = np.random.permutation(datum['points1'].shape[0]) 175 | datum['points1'] = datum['points1'][perm] 176 | datum['points2'] = datum['points2'][perm] 177 | datum['sideinfo'] = datum['sideinfo'][perm] 178 | datum['labels'] = datum['labels'][perm] 179 | 180 | points = np.zeros((max_num_points, 5)).astype(np.float32) 181 | mask = np.zeros((max_num_points, )).astype(int) 182 | labels = np.zeros((max_num_points, )).astype(int) 183 | 184 | num_actual_points = np.minimum(datum['points1'].shape[0], max_num_points) 185 | points[0:num_actual_points, 0:2] = datum['points1'][0:num_actual_points, :].copy() 186 | points[0:num_actual_points, 2:4] = datum['points2'][0:num_actual_points, :].copy() 187 | points[0:num_actual_points, 4] = datum['sideinfo'][0:num_actual_points].copy() 188 | labels[0:num_actual_points] = datum['labels'][0:num_actual_points].copy() 189 | 190 | mask[0:num_actual_points] = 1 191 | 192 | imgsize = np.array([datum['img1size'], datum['img2size']]) 193 | 194 | if num_actual_points < max_num_points: 195 | for i in range(num_actual_points, max_num_points, num_actual_points): 196 | rest = max_num_points-i 197 | num = min(rest, num_actual_points) 198 | points[i:i+num, :] = points[0:num, :].copy() 199 | labels[i:i+num] = labels[0:num].copy() 200 | 201 | return points, points, labels, 0, 0, imgsize 202 | 203 | 204 | def make_vis(): 205 | import matplotlib 206 | import matplotlib.pyplot as plt 207 | def rgb2gray(rgb): 208 | return np.dot(rgb[..., :3], [0.2989, 0.5870, 0.1140]) 209 | 210 | dataset = AdelaideRMF("../datasets/adelaide", keep_in_mem=False, normalize_coords=False, return_images=True, 211 | problem="fundamental", ablation_outlier_ratio=-1, ablation_noise=0) 212 | 213 | target_folder = "./tmp/fig/adelaide_f" 214 | 215 | os.makedirs(target_folder, exist_ok=True) 216 | 217 | for idx in range(len(dataset)): 218 | print("%d / %d" % (idx+1, len(dataset)), end="\r") 219 | sample = dataset[idx] 220 | img1 = sample["img1"].astype(np.uint8) 221 | img2 = sample["img2"].astype(np.uint8) 222 | pts1 = sample["points1"] 223 | pts2 = sample["points2"] 224 | y = sample["labels"] 225 | 226 | num_models = np.max(y) 227 | 228 | cb_hex = ["#000000", "#E69F00", "#56B4E9", "#009E73", "#F0E442", "#0072B2", "#D55E00", "#CC79A7", "#8e10b3", 229 | "#374009", "#aec8ea", "#56611b", "#64a8c6", "#99d8d4", "#745a50", "#46fa50", "#e09eea", "#5b2b1f", 230 | "#723f91", "#634418", "#7db0d0", "#1ae37c", "#aa462c", "#719bb7", "#463aa2", "#98f42e", "#32185b", 231 | "#364fcd", "#7e54c8", "#bb5f7f", "#d466d5", "#5a0382", "#443067", "#a76232", "#78dbc1", "#35a4b2", 232 | "#52d387", "#af5a7e", "#3ef57d", "#d6d993"] 233 | cb = np.array([matplotlib.colors.to_rgb(x) for x in cb_hex]) 234 | 235 | fig = plt.figure(figsize=(4 * 4, 4 * 2), dpi=256) 236 | axs = fig.subplots(nrows=1, ncols=2) 237 | for ax in axs: 238 | ax.set_aspect('equal', 'box') 239 | ax.axis('off') 240 | 241 | if len(img1.shape) == 3: 242 | img1g = rgb2gray(img1) * 0.5 + 128 243 | else: 244 | img1g = img1 * 0.5 + 128 245 | 246 | if len(img2.shape) == 3: 247 | img2g = rgb2gray(img2) * 0.5 + 128 248 | else: 249 | img2g = img2 * 0.5 + 128 250 | 251 | axs[0].imshow(img1g, cmap='Greys_r', vmin=0, vmax=255) 252 | axs[1].imshow(img2g, cmap='Greys_r', vmin=0, vmax=255) 253 | 254 | for j, pts in enumerate([pts1, pts2]): 255 | ax = axs[j] 256 | 257 | ms = np.where(y>0, 8, 4) 258 | 259 | c = cb[y] 260 | 261 | ax.scatter(pts[:, 0], pts[:, 1], c=c, s=ms**2) 262 | 263 | fig.tight_layout() 264 | plt.gca().set_axis_off() 265 | plt.subplots_adjust(top=1, bottom=0, right=1, left=0, 266 | hspace=0, wspace=0) 267 | plt.margins(0, 0) 268 | plt.savefig(os.path.join(target_folder, "%02d_%03d_vis.png" % (num_models, idx)), bbox_inches='tight', 269 | pad_inches=0) 270 | plt.close() 271 | 272 | fig = plt.figure(figsize=(4 * 4, 4 * 2), dpi=150) 273 | axs = fig.subplots(nrows=1, ncols=2) 274 | for ax in axs: 275 | ax.set_aspect('equal', 'box') 276 | ax.axis('off') 277 | 278 | if len(img1.shape) == 3: 279 | axs[0].imshow(img1) 280 | else: 281 | axs[0].imshow(img1, cmap="Greys_r") 282 | 283 | if len(img2.shape) == 3: 284 | axs[1].imshow(img2) 285 | else: 286 | axs[1].imshow(img2, cmap="Greys_r") 287 | 288 | fig.tight_layout() 289 | plt.gca().set_axis_off() 290 | plt.subplots_adjust(top=1, bottom=0, right=1, left=0, 291 | hspace=0, wspace=0) 292 | plt.margins(0, 0) 293 | plt.savefig(os.path.join(target_folder, "%02d_%03d_orig.png" % (num_models, idx)), bbox_inches='tight', 294 | pad_inches=0) 295 | plt.close() 296 | -------------------------------------------------------------------------------- /datasets/hope.py: -------------------------------------------------------------------------------- 1 | import os.path 2 | from torch.utils.data import Dataset 3 | import numpy as np 4 | import pickle 5 | import skimage 6 | import random 7 | 8 | def rgb2gray(rgb): 9 | return np.dot(rgb[...,:3], [0.2989, 0.5870, 0.1140]) 10 | 11 | class HopeF: 12 | 13 | def __init__(self, data_dir, split, cache_data=True, normalize_coords=True, return_images=False, return_gt_models=False): 14 | 15 | self.data_dir = data_dir 16 | self.cache_data = cache_data 17 | self.normalize_coords = normalize_coords 18 | self.return_images = return_images 19 | self.return_gt_models = return_gt_models 20 | 21 | if split == "train": 22 | data_ids = range(800) 23 | elif split == "val": 24 | data_ids = range(800, 900) 25 | elif split == "test": 26 | data_ids = range(900, 1000) 27 | elif split == "all": 28 | data_ids = range(0, 1000) 29 | else: 30 | assert False, "invalid split: %s" % split 31 | 32 | self.dataset_folders = [] 33 | 34 | for num_obj in [1, 2, 3, 4]: 35 | self.dataset_folders += [os.path.join(self.data_dir, "%d/%04d" % (num_obj, i)) for i in data_ids] 36 | 37 | self.img_size = (1024, 1024) 38 | 39 | cache_folders = ["/phys/ssd/tmp/hope", "/phys/ssd/slurmstorage/tmp/hope", "/tmp/hope", "/phys/intern/tmp/hope"] 40 | 41 | self.cache_dir = None 42 | if cache_data: 43 | for cache_folder in cache_folders: 44 | try: 45 | cache_folder = os.path.join(cache_folder, split) 46 | os.makedirs(cache_folder, exist_ok=True) 47 | self.cache_dir = cache_folder 48 | print("%s is cache folder" % cache_folder) 49 | break 50 | except: 51 | print("%s unavailable" % cache_folder) 52 | 53 | def __len__(self): 54 | return len(self.dataset_folders) 55 | 56 | def __getitem__(self, key): 57 | 58 | folder = self.dataset_folders[key] 59 | datum = None 60 | 61 | if self.cache_dir is not None: 62 | cache_path = os.path.join(self.cache_dir, "%09d.pkl" % key) 63 | if os.path.exists(cache_path): 64 | with open(cache_path, 'rb') as f: 65 | datum = pickle.load(f) 66 | 67 | if datum is None: 68 | 69 | features_and_gt = np.load(os.path.join(folder, "features_and_ground_truth.npz"), allow_pickle=True) 70 | 71 | gt_label = features_and_gt["labels"] 72 | pts1 = features_and_gt["points1"][:, :2] 73 | pts2 = features_and_gt["points2"][:, :2] 74 | sideinfo = features_and_gt["ratios"] 75 | 76 | if self.normalize_coords: 77 | scale = np.max(self.img_size) 78 | 79 | pts1[:,0] -= self.img_size[1]/2. 80 | pts2[:,0] -= self.img_size[1]/2. 81 | pts1[:,1] -= self.img_size[0]/2. 82 | pts2[:,1] -= self.img_size[0]/2. 83 | pts1 /= (scale/2.) 84 | pts2 /= (scale/2.) 85 | 86 | datum = {'points1': pts1, 'points2': pts2, 'sideinfo': sideinfo, 'img1size': self.img_size, 87 | 'img2size': self.img_size, 'labels': gt_label} 88 | 89 | if self.cache_dir is not None: 90 | cache_path = os.path.join(self.cache_dir, "%09d.pkl" % key) 91 | if not os.path.exists(cache_path): 92 | with open(cache_path, 'wb') as f: 93 | pickle.dump(datum, f, pickle.HIGHEST_PROTOCOL) 94 | 95 | if self.return_images: 96 | img1_path = os.path.join(folder, "render0.png") 97 | img2_path = os.path.join(folder, "render1.png") 98 | image1_rgb = skimage.io.imread(img1_path).astype(float)[:, :, :3] 99 | image2_rgb = skimage.io.imread(img2_path).astype(float)[:, :, :3] 100 | image1 = rgb2gray(image1_rgb) 101 | image2 = rgb2gray(image2_rgb) 102 | datum['image1'] = image1 103 | datum['image2'] = image2 104 | datum['img1'] = image1_rgb 105 | datum['img2'] = image2_rgb 106 | 107 | if self.return_gt_models: 108 | features_and_gt = np.load(os.path.join(folder, "features_and_ground_truth.npz"), allow_pickle=True) 109 | gt = features_and_gt["F"] 110 | datum["gt"] = gt 111 | 112 | return datum 113 | 114 | 115 | class HopeFDataset(Dataset): 116 | 117 | def __init__(self, data_dir_path, split, max_num_points, cache_data=True, normalize_coords=True, 118 | permute_points=True, return_images=False, return_labels=True, return_gt_models=False): 119 | self.dataset = HopeF(data_dir_path, split, cache_data=cache_data, normalize_coords=normalize_coords, 120 | return_images=return_images, return_gt_models=return_gt_models) 121 | self.max_num_points = max_num_points 122 | self.permute_points = permute_points 123 | self.return_images = return_images 124 | self.return_labels = return_labels 125 | 126 | def denormalise(self, X): 127 | scale = np.max(self.dataset.img_size) / 2.0 128 | off = (self.dataset.img_size[1] / 2.0, self.dataset.img_size[0] / 2.0) 129 | p1 = X[..., 0:2] * scale 130 | p2 = X[..., 0:2] * scale 131 | p1[..., 0] += off[0] 132 | p1[..., 1] += off[1] 133 | p2[..., 0] += off[0] 134 | p2[..., 1] += off[1] 135 | 136 | return p1, p2 137 | 138 | def __len__(self): 139 | return len(self.dataset) 140 | 141 | def __getitem__(self, key): 142 | datum = self.dataset[key] 143 | 144 | if self.max_num_points <= 0: 145 | max_num_points = datum['points1'].shape[0] 146 | else: 147 | max_num_points = self.max_num_points 148 | 149 | if self.permute_points: 150 | perm = np.random.permutation(datum['points1'].shape[0]) 151 | datum['points1'] = datum['points1'][perm] 152 | datum['points2'] = datum['points2'][perm] 153 | datum['sideinfo'] = datum['sideinfo'][perm] 154 | datum['labels'] = datum['labels'][perm] 155 | 156 | points = np.zeros((max_num_points, 5)).astype(np.float32) 157 | mask = np.zeros((max_num_points, )).astype(int) 158 | labels = np.zeros((max_num_points, )).astype(int) 159 | 160 | num_actual_points = np.minimum(datum['points1'].shape[0], max_num_points) 161 | points[0:num_actual_points, 0:2] = datum['points1'][0:num_actual_points, :].copy() 162 | points[0:num_actual_points, 2:4] = datum['points2'][0:num_actual_points, :].copy() 163 | points[0:num_actual_points, 4] = datum['sideinfo'][0:num_actual_points].copy() 164 | labels[0:num_actual_points] = datum['labels'][0:num_actual_points].copy() 165 | 166 | mask[0:num_actual_points] = 1 167 | 168 | if num_actual_points < max_num_points: 169 | for i in range(num_actual_points, max_num_points, num_actual_points): 170 | rest = max_num_points-i 171 | num = min(rest, num_actual_points) 172 | points[i:i+num, :] = points[0:num, :].copy() 173 | labels[i:i+num] = labels[0:num].copy() 174 | 175 | if 'img1' in datum.keys(): 176 | image = datum['img1'] 177 | else: 178 | image = 0 179 | 180 | imgsize = np.array([(1024, 1024), (1024, 1024)]) 181 | 182 | if self.dataset.return_gt_models: 183 | return points, points, labels, datum["gt"], image, imgsize 184 | else: 185 | return points, points, labels, 0, image, imgsize 186 | 187 | 188 | 189 | def make_vis(): 190 | import matplotlib 191 | import matplotlib.pyplot as plt 192 | 193 | random.seed(0) 194 | 195 | dataset = HopeF("./datasets/hope", "all", cache_data=False, normalize_coords=False, return_images=True) 196 | 197 | target_folder = "./tmp/fig/hopef" 198 | 199 | os.makedirs(target_folder, exist_ok=True) 200 | 201 | for idx in range(len(dataset)): 202 | print("%d / %d" % (idx+1, len(dataset)), end="\r") 203 | sample = dataset[idx] 204 | img1 = sample["img1"].astype(np.uint8) 205 | img2 = sample["img2"].astype(np.uint8) 206 | pts1 = sample["points1"] 207 | pts2 = sample["points2"] 208 | y = sample["labels"] 209 | 210 | num_models = np.max(y) 211 | cb_hex = ["#000000", "#E69F00", "#56B4E9", "#009E73", "#F0E442", "#0072B2", "#D55E00", "#CC79A7", "#8e10b3", 212 | "#374009", "#aec8ea", "#56611b", "#64a8c6", "#99d8d4", "#745a50", "#46fa50", "#e09eea", "#5b2b1f", 213 | "#723f91", "#634418", "#7db0d0", "#1ae37c", "#aa462c", "#719bb7", "#463aa2", "#98f42e", "#32185b", 214 | "#364fcd", "#7e54c8", "#bb5f7f", "#d466d5", "#5a0382", "#443067", "#a76232", "#78dbc1", "#35a4b2", 215 | "#52d387", "#af5a7e", "#3ef57d", "#d6d993"] 216 | cb = np.array([matplotlib.colors.to_rgb(x) for x in cb_hex]) 217 | 218 | fig = plt.figure(figsize=(4 * 4, 4 * 2), dpi=256) 219 | axs = fig.subplots(nrows=1, ncols=2) 220 | for ax in axs: 221 | ax.set_aspect('equal', 'box') 222 | ax.axis('off') 223 | 224 | img1g = rgb2gray(img1) * 0.5 + 128 225 | img2g = rgb2gray(img2) * 0.5 + 128 226 | 227 | axs[0].imshow(img1g, cmap='Greys_r', vmin=0, vmax=255) 228 | axs[1].imshow(img2g, cmap='Greys_r', vmin=0, vmax=255) 229 | 230 | for j, pts in enumerate([pts1, pts2]): 231 | ax = axs[j] 232 | 233 | ms = np.where(y>0, 8, 4) 234 | 235 | c = cb[y] 236 | 237 | ax.scatter(pts[:, 0], pts[:, 1], c=c, s=ms**2) 238 | 239 | fig.tight_layout() 240 | plt.gca().set_axis_off() 241 | plt.subplots_adjust(top=1, bottom=0, right=1, left=0, 242 | hspace=0, wspace=0) 243 | plt.margins(0, 0) 244 | plt.savefig(os.path.join(target_folder, "%02d_%03d_vis.png" % (num_models, idx)), bbox_inches='tight', 245 | pad_inches=0) 246 | plt.close() 247 | 248 | fig = plt.figure(figsize=(4 * 4, 4 * 2), dpi=150) 249 | axs = fig.subplots(nrows=1, ncols=2) 250 | for ax in axs: 251 | ax.set_aspect('equal', 'box') 252 | ax.axis('off') 253 | 254 | axs[0].imshow(img1) 255 | axs[1].imshow(img2) 256 | fig.tight_layout() 257 | plt.gca().set_axis_off() 258 | plt.subplots_adjust(top=1, bottom=0, right=1, left=0, 259 | hspace=0, wspace=0) 260 | plt.margins(0, 0) 261 | plt.savefig(os.path.join(target_folder, "%02d_%03d_orig.png" % (num_models, idx)), bbox_inches='tight', 262 | pad_inches=0) 263 | plt.close() 264 | -------------------------------------------------------------------------------- /datasets/nyuvp.py: -------------------------------------------------------------------------------- 1 | import torch.utils.data 2 | import numpy as np 3 | from datasets.nyu_vp import nyu 4 | from datasets.yud_plus import yud 5 | import utils.residual_functions 6 | 7 | def label_lines(vps, line_segments, threshold=1-np.cos(2.0*np.pi/180.0)): 8 | 9 | residuals = utils.residual_functions.vanishing_point(torch.from_numpy(line_segments)[None, ...].cuda(), torch.from_numpy(vps).cuda()).cpu().numpy() 10 | 11 | min_residuals = np.min(residuals, axis=0) 12 | 13 | inliers = min_residuals < threshold 14 | 15 | labels = np.argmin(residuals, axis=0) + 1 16 | labels *= inliers 17 | 18 | return labels 19 | 20 | def prepare_sample(sample, max_num_lines, max_num_vps, generate_labels=False): 21 | if max_num_lines < 0: 22 | max_num_lines = sample['line_segments'].shape[0] 23 | else: 24 | max_num_lines = max_num_lines 25 | 26 | lines = np.zeros((max_num_lines, 12)).astype(np.float32) 27 | vps = np.zeros((max_num_vps, 3)).astype(np.float32) 28 | 29 | np.random.shuffle(sample['line_segments']) 30 | 31 | num_actual_line_segments = np.minimum(sample['line_segments'].shape[0], max_num_lines) 32 | lines[0:num_actual_line_segments, :] = sample['line_segments'][0:num_actual_line_segments, :12].copy() 33 | if num_actual_line_segments < max_num_lines: 34 | rest = max_num_lines - num_actual_line_segments 35 | lines[num_actual_line_segments:num_actual_line_segments + rest, :] = lines[0:rest, :].copy() 36 | 37 | num_actual_vps = np.minimum(sample['VPs'].shape[0], max_num_vps) 38 | vps[0:num_actual_vps, :] = sample['VPs'][0:num_actual_vps] 39 | 40 | centroids = lines[:, 9:11] 41 | lengths = np.linalg.norm(lines[:, 0:3] - lines[:, 3:6], axis=-1)[:, None] 42 | vectors = lines[:, 0:3] - lines[:, 3:6] 43 | angles = np.abs(np.arctan2(vectors[:, 0], vectors[:, 1]))[:, None] 44 | 45 | features = np.concatenate([centroids, lengths, angles], axis=-1) 46 | 47 | if generate_labels: 48 | labels = label_lines(vps, lines) 49 | else: 50 | labels = 0 51 | 52 | return features, lines, labels, vps, 0, 0 53 | 54 | 55 | class NYUVP(torch.utils.data.Dataset): 56 | 57 | def __init__(self, split, max_num_lines=512, max_num_vps=8, use_yud=False, use_yud_plus=False, deeplsd_folder=None, 58 | cache=True, generate_labels=False): 59 | if use_yud: 60 | self.dataset = yud.YUDVP(split=split, normalize_coords=True, data_dir_path="./datasets/yud_plus/data", 61 | yudplus=use_yud_plus, keep_in_memory=cache, external_lines_folder=deeplsd_folder) 62 | else: 63 | self.dataset = nyu.NYUVP(split=split, normalise_coordinates=True, data_dir_path="./datasets/nyu_vp/data", 64 | keep_data_in_memory=cache, external_lines_folder=deeplsd_folder) 65 | self.max_num_lines = max_num_lines 66 | self.max_num_vps = max_num_vps 67 | self.generate_labels = generate_labels 68 | 69 | def __len__(self): 70 | return len(self.dataset) 71 | 72 | def __getitem__(self, k): 73 | sample = self.dataset[k] 74 | 75 | return prepare_sample(sample, self.max_num_lines, self.max_num_vps, self.generate_labels) 76 | -------------------------------------------------------------------------------- /datasets/smh.py: -------------------------------------------------------------------------------- 1 | import os.path 2 | from torch.utils.data import Dataset 3 | import numpy as np 4 | import pickle 5 | import skimage 6 | import random 7 | 8 | def rgb2gray(rgb): 9 | return np.dot(rgb[...,:3], [0.2989, 0.5870, 0.1140]) 10 | 11 | class SMH: 12 | 13 | def __init__(self, data_dir, split, cache_data=False, normalize_coords=True, return_images=False, shuffle=False): 14 | 15 | self.data_dir = data_dir 16 | self.cache_data = cache_data 17 | self.normalize_coords = normalize_coords 18 | self.return_images = return_images 19 | 20 | self.img_size = (1024, 1024) 21 | 22 | self.train_sequences = [0, 1, 2, 3, 4] 23 | self.val_sequences = [5] 24 | self.test_sequences = [6] 25 | 26 | self.pairs = [] 27 | 28 | if split == "train": 29 | self.coarse_paths = self.train_sequences 30 | elif split == "val": 31 | self.coarse_paths = self.val_sequences 32 | elif split == "test": 33 | self.coarse_paths = self.test_sequences 34 | elif split == "all": 35 | self.coarse_paths = self.train_sequences + self.val_sequences + self.test_sequences 36 | else: 37 | assert False, "invalid split: %s" % split 38 | 39 | os.makedirs("./tmp/smh_pairs", exist_ok=True) 40 | pairs_cache_file = os.path.join("./tmp/smh_pairs", split+".pkl") 41 | if os.path.exists(pairs_cache_file): 42 | with open(pairs_cache_file, 'rb') as f: 43 | self.pairs = pickle.load(f) 44 | else: 45 | print("loading SMH dataset for the first time, might take a few minutes.. ") 46 | for coarse_path in self.coarse_paths: 47 | for fine_path_dir in os.scandir(os.path.join(self.data_dir, "%d" % coarse_path)): 48 | if fine_path_dir.is_dir(): 49 | for pair_path_dir in os.scandir(fine_path_dir.path): 50 | if pair_path_dir.is_dir(): 51 | if os.path.exists(os.path.join(pair_path_dir.path, "features_and_ground_truth.npz")): 52 | self.pairs += [pair_path_dir.path] 53 | self.pairs.sort() 54 | with open(pairs_cache_file, 'wb') as f: 55 | pickle.dump(self.pairs, f, pickle.HIGHEST_PROTOCOL) 56 | 57 | if shuffle: 58 | random.shuffle(self.pairs) 59 | 60 | print("%s dataset: %d samples" % (split, len(self.pairs))) 61 | 62 | self.cache_dir = None 63 | if cache_data: 64 | cache_folders = ["/phys/ssd/tmp/smh", "/phys/ssd/slurmstorage/tmp/smh", "/tmp/smh", 65 | "/phys/intern/tmp/smh"] 66 | for cache_folder in cache_folders: 67 | try: 68 | cache_folder = os.path.join(cache_folder, split) 69 | os.makedirs(cache_folder, exist_ok=True) 70 | self.cache_dir = cache_folder 71 | print("%s is cache folder" % cache_folder) 72 | break 73 | except: 74 | print("%s unavailable" % cache_folder) 75 | 76 | def __len__(self): 77 | return len(self.pairs) 78 | 79 | def __getitem__(self, key): 80 | 81 | folder = self.pairs[key] 82 | datum = None 83 | 84 | if self.cache_dir is not None: 85 | cache_path = os.path.join(self.cache_dir, "%09d.pkl" % key) 86 | if os.path.exists(cache_path): 87 | with open(cache_path, 'rb') as f: 88 | datum = pickle.load(f) 89 | 90 | if datum is None: 91 | features_and_gt = np.load(os.path.join(folder, "features_and_ground_truth.npz"), allow_pickle=True) 92 | 93 | gt_label = features_and_gt["labels"] 94 | pts1 = features_and_gt["points1"][:, :2] 95 | pts2 = features_and_gt["points2"][:, :2] 96 | sideinfo = features_and_gt["ratios"] 97 | 98 | if self.normalize_coords: 99 | scale = np.max(self.img_size) 100 | 101 | pts1[:,0] -= self.img_size[1]/2. 102 | pts2[:,0] -= self.img_size[1]/2. 103 | pts1[:,1] -= self.img_size[0]/2. 104 | pts2[:,1] -= self.img_size[0]/2. 105 | pts1 /= (scale/2.) 106 | pts2 /= (scale/2.) 107 | 108 | datum = {'points1': pts1, 'points2': pts2, 'sideinfo': sideinfo, 'img1size': self.img_size, 'img2size': self.img_size, 109 | 'labels': gt_label} 110 | 111 | if self.cache_dir is not None: 112 | cache_path = os.path.join(self.cache_dir, "%09d.pkl" % key) 113 | if not os.path.exists(cache_path): 114 | with open(cache_path, 'wb') as f: 115 | pickle.dump(datum, f, pickle.HIGHEST_PROTOCOL) 116 | 117 | 118 | if self.return_images: 119 | img1_path = os.path.join(folder, "render0.png") 120 | img2_path = os.path.join(folder, "render1.png") 121 | image1_rgb = skimage.io.imread(img1_path).astype(float)[:, :, :3] 122 | image2_rgb = skimage.io.imread(img2_path).astype(float)[:, :, :3] 123 | image1 = rgb2gray(image1_rgb) 124 | image2 = rgb2gray(image2_rgb) 125 | 126 | datum['image1'] = image1 127 | datum['image2'] = image2 128 | datum['img1'] = image1_rgb 129 | datum['img2'] = image2_rgb 130 | 131 | 132 | return datum 133 | 134 | 135 | class SMHDataset(Dataset): 136 | 137 | def __init__(self, data_dir_path, split, max_num_points, keep_in_mem=True, 138 | permute_points=True, return_images=False, return_labels=True, max_num_models=28): 139 | self.homdata = SMH(data_dir_path, split, cache_data=keep_in_mem, normalize_coords=True, return_images=return_images) 140 | self.max_num_points = max_num_points 141 | self.permute_points = permute_points 142 | self.return_images = return_images 143 | self.return_labels = return_labels 144 | self.max_num_models = max_num_models 145 | self.split = split 146 | 147 | def denormalise(self, X): 148 | scale = np.max(self.homdata.img_size) / 2.0 149 | off = (self.homdata.img_size[1] / 2.0, self.homdata.img_size[0] / 2.0) 150 | p1 = X[..., 0:2] * scale 151 | p2 = X[..., 0:2] * scale 152 | p1[..., 0] += off[0] 153 | p1[..., 1] += off[1] 154 | p2[..., 0] += off[0] 155 | p2[..., 1] += off[1] 156 | 157 | return p1, p2 158 | 159 | def __len__(self): 160 | return len(self.homdata) 161 | 162 | def __getitem__(self, key): 163 | datum = self.homdata[key] 164 | 165 | if self.max_num_points <= 0: 166 | max_num_points = datum['points1'].shape[0] 167 | else: 168 | max_num_points = self.max_num_points 169 | 170 | if self.permute_points: 171 | 172 | perm = np.random.permutation(datum['points1'].shape[0]) 173 | datum['points1'] = datum['points1'][perm] 174 | datum['points2'] = datum['points2'][perm] 175 | datum['sideinfo'] = datum['sideinfo'][perm] 176 | datum['labels'] = datum['labels'][perm] 177 | 178 | points = np.zeros((max_num_points, 5)).astype(np.float32) 179 | mask = np.zeros((max_num_points, )).astype(int) 180 | labels = np.zeros((max_num_points, )).astype(int) 181 | 182 | num_actual_points = np.minimum(datum['points1'].shape[0], max_num_points) 183 | points[0:num_actual_points, 0:2] = datum['points1'][0:num_actual_points, :].copy() 184 | points[0:num_actual_points, 2:4] = datum['points2'][0:num_actual_points, :].copy() 185 | points[0:num_actual_points, 4] = datum['sideinfo'][0:num_actual_points].copy() 186 | labels[0:num_actual_points] = datum['labels'][0:num_actual_points].copy() 187 | 188 | mask[0:num_actual_points] = 1 189 | 190 | if num_actual_points < max_num_points: 191 | for i in range(num_actual_points, max_num_points, num_actual_points): 192 | rest = max_num_points-i 193 | num = min(rest, num_actual_points) 194 | points[i:i+num, :] = points[0:num, :].copy() 195 | labels[i:i+num] = labels[0:num].copy() 196 | 197 | if self.max_num_models: 198 | labels[np.nonzero(labels >= self.max_num_models)] = 0 199 | 200 | if 'img1' in datum.keys(): 201 | image = datum['img1'] 202 | else: 203 | image = 0 204 | 205 | imgsize = np.array([(1024, 1024), (1024, 1024)]) 206 | 207 | return points, points, labels, 0, image, imgsize 208 | 209 | 210 | 211 | def make_vis(): 212 | import matplotlib 213 | import matplotlib.pyplot as plt 214 | 215 | random.seed(0) 216 | 217 | dataset = SMH("../datasets/smh", "all", cache_data=False, normalize_coords=False, return_images=True, shuffle=True) 218 | 219 | target_folder = "./tmp/fig/smh" 220 | 221 | os.makedirs(target_folder, exist_ok=True) 222 | 223 | for idx in range(len(dataset)): 224 | print("%d / %d" % (idx+1, len(dataset)), end="\r") 225 | sample = dataset[idx] 226 | img1 = sample["img1"].astype(np.uint8) 227 | img2 = sample["img2"].astype(np.uint8) 228 | pts1 = sample["points1"] 229 | pts2 = sample["points2"] 230 | y = sample["labels"] 231 | 232 | num_models = np.max(y) 233 | 234 | cb_hex = ["#000000", "#E69F00", "#56B4E9", "#009E73", "#F0E442", "#0072B2", "#D55E00", "#CC79A7", "#8e10b3", "#374009", "#aec8ea", "#56611b", "#64a8c6", "#99d8d4", "#745a50", "#46fa50", "#e09eea", "#5b2b1f", "#723f91", "#634418", "#7db0d0", "#1ae37c", "#aa462c", "#719bb7", "#463aa2", "#98f42e", "#32185b", "#364fcd", "#7e54c8", "#bb5f7f", "#d466d5", "#5a0382", "#443067", "#a76232", "#78dbc1", "#35a4b2", "#52d387", "#af5a7e", "#3ef57d", "#d6d993"] 235 | cb = np.array([matplotlib.colors.to_rgb(x) for x in cb_hex]) 236 | 237 | fig = plt.figure(figsize=(4 * 4, 4 * 2), dpi=256) 238 | axs = fig.subplots(nrows=1, ncols=2) 239 | for ax in axs: 240 | ax.set_aspect('equal', 'box') 241 | ax.axis('off') 242 | 243 | img1g = rgb2gray(img1) * 0.5 + 128 244 | img2g = rgb2gray(img2) * 0.5 + 128 245 | 246 | axs[0].imshow(img1g, cmap='Greys_r', vmin=0, vmax=255) 247 | axs[1].imshow(img2g, cmap='Greys_r', vmin=0, vmax=255) 248 | 249 | for j, pts in enumerate([pts1, pts2]): 250 | ax = axs[j] 251 | 252 | ms = np.where(y>0, 8, 4) 253 | 254 | c = cb[y] 255 | 256 | ax.scatter(pts[:, 0], pts[:, 1], c=c, s=ms**2) 257 | 258 | fig.tight_layout() 259 | plt.gca().set_axis_off() 260 | plt.subplots_adjust(top=1, bottom=0, right=1, left=0, 261 | hspace=0, wspace=0) 262 | plt.margins(0, 0) 263 | plt.savefig(os.path.join(target_folder, "%02d_%03d_vis.png" % (num_models, idx)), bbox_inches='tight', 264 | pad_inches=0) 265 | plt.close() 266 | 267 | fig = plt.figure(figsize=(4 * 4, 4 * 2), dpi=150) 268 | axs = fig.subplots(nrows=1, ncols=2) 269 | for ax in axs: 270 | ax.set_aspect('equal', 'box') 271 | ax.axis('off') 272 | 273 | axs[0].imshow(img1) 274 | axs[1].imshow(img2) 275 | fig.tight_layout() 276 | plt.gca().set_axis_off() 277 | plt.subplots_adjust(top=1, bottom=0, right=1, left=0, 278 | hspace=0, wspace=0) 279 | plt.margins(0, 0) 280 | plt.savefig(os.path.join(target_folder, "%02d_%03d_orig.png" % (num_models, idx)), bbox_inches='tight', 281 | pad_inches=0) 282 | plt.close() 283 | -------------------------------------------------------------------------------- /datasets/su3.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch.utils.data 3 | import numpy as np 4 | import skimage 5 | from pylsd.lsd import lsd 6 | import pickle 7 | import csv 8 | from glob import glob 9 | import utils.residual_functions 10 | 11 | def rgb2gray(rgb): 12 | return np.dot(rgb[...,:3], [0.2989, 0.5870, 0.1140]) 13 | 14 | def augment_sample(datum): 15 | M = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) 16 | if np.random.uniform(0, 1) > 0.5: 17 | F = np.array([[-1, 0, 0], [0, 1, 0], [0, 0, 1]]) 18 | M = F @ M 19 | if np.random.uniform(0, 1) > 0.5: 20 | F = np.array([[1, 0, 0], [0, -1, 0], [0, 0, 1]]) 21 | M = F @ M 22 | Mi = np.linalg.inv(M).T 23 | 24 | lines = datum["line_segments"] 25 | p1 = (M @ lines[:, 0:3].T).T 26 | p2 = (M @ lines[:, 3:6].T).T 27 | l = (Mi @ lines[:, 6:9].T).T 28 | c = (M @ lines[:, 9:12].T).T 29 | lines[:, 0:3] = p1 30 | lines[:, 3:6] = p2 31 | lines[:, 6:9] = l 32 | lines[:, 9:12] = c 33 | 34 | vps = (M @ datum["VPs"].T).T 35 | 36 | datum["line_segments"] = lines 37 | datum["VPs"] = vps 38 | 39 | return datum 40 | 41 | 42 | def remove_outliers(vps, lsd_line_segments, threshold=1-np.cos(0.5*np.pi/180.0)): 43 | 44 | lsd_normed = lsd_line_segments.copy() 45 | lsd_normed[:, 0:4] -= 256 46 | lsd_normed[:, 0:4] /= 256. 47 | 48 | line_segments = np.zeros((lsd_line_segments.shape[0], 12)) 49 | for li in range(line_segments.shape[0]): 50 | p1 = np.array([lsd_normed[li, 0], lsd_normed[li, 1], 1]) 51 | p2 = np.array([lsd_normed[li, 2], lsd_normed[li, 3], 1]) 52 | centroid = 0.5 * (p1 + p2) 53 | line = np.cross(p1, p2) 54 | line /= np.linalg.norm(line[0:2]) 55 | line_segments[li, 0:3] = p1 56 | line_segments[li, 3:6] = p2 57 | line_segments[li, 6:9] = line 58 | line_segments[li, 9:12] = centroid 59 | 60 | residuals = utils.residual_functions.vanishing_point(torch.from_numpy(line_segments)[None, ...].cuda(), torch.from_numpy(vps).cuda()).cpu().numpy() 61 | 62 | min_residuals = np.min(residuals, axis=0) 63 | 64 | inliers = min_residuals < threshold 65 | 66 | lsd_line_segments = lsd_line_segments[np.where(inliers)] 67 | 68 | return lsd_line_segments 69 | 70 | 71 | def label_lines(vps, line_segments, threshold=1-np.cos(2.0*np.pi/180.0)): 72 | 73 | residuals = utils.residual_functions.vanishing_point(torch.from_numpy(line_segments)[None, ...].cuda(), torch.from_numpy(vps).cuda()).cpu().numpy() 74 | 75 | min_residuals = np.min(residuals, axis=0) 76 | 77 | inliers = min_residuals < threshold 78 | 79 | labels = np.argmin(residuals, axis=0) + 1 80 | labels *= inliers 81 | 82 | return labels 83 | 84 | 85 | def add_outliers(line_segments, outlier_percentage=0.0): 86 | 87 | N = line_segments.shape[0] 88 | 89 | No = int(N * outlier_percentage/(1-outlier_percentage)) 90 | 91 | if outlier_percentage == 0 or No < 1: 92 | return line_segments 93 | 94 | outlier = np.zeros((No, line_segments.shape[1])).astype(np.float32) 95 | outlier[:, 0:2] = np.random.uniform(0, 511, (No, 2)) 96 | outlier[:, 2:4] = outlier[:, 0:2] + np.random.uniform(-50, 50, (No, 2)) 97 | outlier = np.clip(outlier, a_min=0, a_max=512) 98 | 99 | lines = np.concatenate([line_segments, outlier], axis=0) 100 | 101 | return lines 102 | 103 | 104 | def prepare_sample(sample, max_num_lines, max_num_vps, augment=False): 105 | 106 | if augment: 107 | sample = augment_sample(sample) 108 | 109 | if max_num_lines < 0: 110 | max_num_lines = sample['line_segments'].shape[0] 111 | else: 112 | max_num_lines = max_num_lines 113 | 114 | lines = np.zeros((max_num_lines, 12)).astype(np.float32) 115 | vps = np.zeros((max_num_vps, 3)).astype(np.float32) 116 | 117 | np.random.shuffle(sample['line_segments']) 118 | 119 | num_actual_line_segments = np.minimum(sample['line_segments'].shape[0], max_num_lines) 120 | lines[0:num_actual_line_segments, :] = sample['line_segments'][0:num_actual_line_segments, :12].copy() 121 | if num_actual_line_segments < max_num_lines: 122 | rest = max_num_lines - num_actual_line_segments 123 | lines[num_actual_line_segments:num_actual_line_segments + rest, :] = lines[0:rest, :].copy() 124 | 125 | num_actual_vps = np.minimum(sample['VPs'].shape[0], max_num_vps) 126 | vps[0:num_actual_vps, :] = sample['VPs'][0:num_actual_vps] 127 | 128 | centroids = lines[:, 9:11] 129 | lengths = np.linalg.norm(lines[:, 0:3] - lines[:, 3:6], axis=-1)[:, None] 130 | vectors = lines[:, 0:3] - lines[:, 3:6] 131 | angles = np.abs(np.arctan2(vectors[:, 0], vectors[:, 1]))[:, None] 132 | 133 | features = np.concatenate([centroids, lengths, angles], axis=-1) 134 | 135 | return features, lines, sample['labels'], vps, sample['image'], 0 136 | 137 | 138 | class SU3(torch.utils.data.Dataset): 139 | 140 | def __init__(self, rootdir, split, max_num_lines=512, normalise_coords=False, augmentation=False, 141 | deeplsd_folder=None, cache=True, ablation_outlier_ratio=-1, ablation_noise=0, return_dict=False, 142 | generate_labels=False): 143 | 144 | self.rootdir = rootdir 145 | self.deeplsd_folder = deeplsd_folder 146 | self.ablation_noise = ablation_noise 147 | self.ablation_outlier_ratio = ablation_outlier_ratio 148 | filelist = sorted(glob(f"{rootdir}/*/*.png")) 149 | 150 | self.split = split 151 | division = int(len(filelist) * 0.1) 152 | print("num of valid/test", division) 153 | if split == "train": 154 | num_train = int(len(filelist) * 0.8 * 1) 155 | self.filelist = filelist[2 * division: 2 * division + num_train] 156 | self.size = len(self.filelist) 157 | print("subset for training: percentage ", 1, num_train) 158 | elif split == "valid": 159 | self.filelist = [f for f in filelist[division:division*2] if "a1" not in f] 160 | self.size = len(self.filelist) 161 | elif split == "test": 162 | self.filelist = [f for f in filelist[:division] if "a1" not in f] 163 | self.size = len(self.filelist) 164 | elif split == "all": 165 | self.filelist = [f for f in filelist if "a1" not in f] 166 | self.size = len(self.filelist) 167 | print(f"n{split}:", len(self.filelist)) 168 | 169 | self.augmentation = augmentation 170 | self.return_dict = return_dict 171 | self.max_num_lines = max_num_lines 172 | self.normalise_coords = normalise_coords 173 | self.generate_labels = generate_labels 174 | 175 | f = 2.1875 * 256 176 | c = 256 177 | self.K = np.array([[f, 0, c], [0, -f, c], [0, 0, 1]]) 178 | self.S = np.array([[1. / c, 0, -1.], [0, 1. / c, -1.], [0, 0, 1]]) 179 | if self.normalise_coords: 180 | self.SK = self.S @ self.K 181 | else: 182 | self.SK = self.K 183 | 184 | self.cache_dir = None 185 | if cache: 186 | cache_folders = ["/phys/ssd/tmp/su3_new", "/phys/ssd/slurmstorage/tmp/su3_new", "/tmp/su3_new", "/phys/intern/tmp/su3_new", ] 187 | for cache_folder in cache_folders: 188 | try: 189 | cache_folder = os.path.join(cache_folder, split) 190 | if deeplsd_folder is not None: 191 | cache_folder = os.path.join(cache_folder, "deeplsd") 192 | os.makedirs(cache_folder, exist_ok=True) 193 | self.cache_dir = cache_folder 194 | print("%s is cache folder" % cache_folder) 195 | break 196 | except: 197 | print("%s unavailable" % cache_folder) 198 | 199 | def denormalise(self, X): 200 | p1 = X[..., :2] * 256 + 256 201 | p2 = X[..., 3:5] * 256 + 256 202 | 203 | return p1, p2 204 | 205 | 206 | def __len__(self): 207 | return len(self.filelist) 208 | 209 | def __getitem__(self, k): 210 | 211 | if k >= len(self.filelist): 212 | raise IndexError 213 | 214 | sample = None 215 | 216 | if self.cache_dir is not None: 217 | cache_path = os.path.join(self.cache_dir, "%09d.pkl" % k) 218 | if os.path.exists(cache_path): 219 | with open(cache_path, 'rb') as f: 220 | sample = pickle.load(f) 221 | 222 | if sample is None or not ("image" in sample.keys()): 223 | iname = self.filelist[k % len(self.filelist)] 224 | image_rgb = skimage.io.imread(iname).astype(float)[:, :, :3] 225 | image = rgb2gray(image_rgb) 226 | 227 | with np.load(iname.replace(".png", "_label.npz")) as npz: 228 | vds = npz["vpts"] 229 | vps = (self.SK @ vds.T).T 230 | 231 | if self.deeplsd_folder: 232 | file_seq = iname.split("/")[-2] 233 | file_idx = iname.split("/")[-1].split(".")[-2] 234 | target_dir = os.path.join(self.deeplsd_folder, file_seq) 235 | deeplsd_file = os.path.join(target_dir, file_idx + ".csv") 236 | 237 | lsd_line_segments = [] 238 | with open(deeplsd_file, 'r') as csv_file: 239 | reader = csv.reader(csv_file, delimiter=' ') 240 | for line in reader: 241 | p1x = float(line[0]) 242 | p1y = float(line[1]) 243 | p2x = float(line[2]) 244 | p2y = float(line[3]) 245 | lsd_line_segments += [np.array([p1x, p1y, p2x, p2y])] 246 | lsd_line_segments = np.vstack(lsd_line_segments) 247 | else: 248 | lsd_line_segments = lsd(image) 249 | 250 | if self.ablation_outlier_ratio >= 0: 251 | lsd_line_segments = remove_outliers(vps, lsd_line_segments) 252 | lsd_line_segments = add_outliers(lsd_line_segments, self.ablation_outlier_ratio) 253 | 254 | if self.ablation_noise > 0: 255 | lsd_line_segments[:, :4] += np.random.normal(0, self.ablation_noise, lsd_line_segments[:, :4].shape) 256 | 257 | if self.normalise_coords: 258 | lsd_line_segments[:,0:4] -= 256 259 | lsd_line_segments[:,0:4] /= 256. 260 | 261 | line_segments = np.zeros((lsd_line_segments.shape[0], 12)) 262 | for li in range(line_segments.shape[0]): 263 | p1 = np.array([lsd_line_segments[li,0], lsd_line_segments[li,1], 1]) 264 | p2 = np.array([lsd_line_segments[li,2], lsd_line_segments[li,3], 1]) 265 | centroid = 0.5*(p1+p2) 266 | line = np.cross(p1, p2) 267 | line /= np.linalg.norm(line[0:2]) 268 | line_segments[li, 0:3] = p1 269 | line_segments[li, 3:6] = p2 270 | line_segments[li, 6:9] = line 271 | line_segments[li, 9:12] = centroid 272 | 273 | sample = {"line_segments": line_segments, "VPs": vps, "image": image_rgb, "labels": 0} 274 | 275 | if self.cache_dir is not None: 276 | cache_path = os.path.join(self.cache_dir, "%09d.pkl" % k) 277 | if not os.path.exists(cache_path): 278 | with open(cache_path, 'wb') as f: 279 | # Pickle the 'data' dictionary using the highest protocol available. 280 | pickle.dump(sample, f, pickle.HIGHEST_PROTOCOL) 281 | 282 | 283 | if self.generate_labels: 284 | vps = sample["VPs"] 285 | line_segments = sample["line_segments"] 286 | labels = label_lines(vps, line_segments) 287 | sample["labels"] = labels 288 | 289 | if self.return_dict: 290 | return sample 291 | 292 | data = prepare_sample(sample, self.max_num_lines, 3, augment=self.augmentation) 293 | 294 | return data 295 | 296 | def save_mat(split): 297 | import scipy.io 298 | dataset = SU3("/data/scene_understanding/SU3", split, return_dict=True) 299 | 300 | target_folder = "/data/kluger/datasets/SU3/matlab/%s" % split 301 | os.makedirs(target_folder, exist_ok=True) 302 | 303 | filename = os.path.join(target_folder, "../intrinsic_camera.mat") 304 | scipy.io.savemat(filename, {"K": dataset.K}) 305 | 306 | print("num images: ", len(dataset)) 307 | for idx, sample in enumerate(dataset): 308 | print("%d / %d" % (idx+1, len(dataset))) 309 | img1 = rgb2gray(sample["image"])[..., None] 310 | filename = os.path.join(target_folder, "%04d.mat" % idx) 311 | scipy.io.savemat(filename, {"image": img1, "lines": sample["line_segments"], "VPs": sample["VPs"]}) 312 | 313 | def make_vis(): 314 | import matplotlib 315 | import matplotlib.pyplot as plt 316 | 317 | def line_vp_distances(lines, vps): 318 | distances = np.zeros((lines.shape[0], vps.shape[0])) 319 | 320 | for li in range(lines.shape[0]): 321 | for vi in range(vps.shape[0]): 322 | vp = vps[vi, :] 323 | line = lines[li, 6:9] 324 | centroid = lines[li, 9:12] 325 | constrained_line = np.cross(vp, centroid) 326 | constrained_line /= np.linalg.norm(constrained_line[0:2]) 327 | 328 | distance = np.arccos(np.abs((line[0:2] * constrained_line[0:2]).sum(axis=0))) * 180.0 / np.pi 329 | 330 | distances[li, vi] = distance 331 | return distances 332 | 333 | dataset = SU3("datasets/su3", "test", cache=False, return_dict=True) 334 | 335 | target_folder = "./tmp/fig/su3_examples" 336 | 337 | os.makedirs(target_folder, exist_ok=True) 338 | 339 | for idx in range(len(dataset)): 340 | vps = dataset[idx]['VPs'] 341 | num_vps = vps.shape[0] 342 | 343 | ls = dataset[idx]['line_segments'] 344 | vp = dataset[idx]['VPs'] 345 | vp[:, 0] /= vp[:, 2] 346 | vp[:, 1] /= vp[:, 2] 347 | vp[:, 2] /= vp[:, 2] 348 | 349 | distances = line_vp_distances(ls, vp) 350 | closest_vps = np.argmin(distances, axis=1) + 1 351 | closest_vps[np.nonzero(np.min(distances, axis=1) > 5.0)[0]] = 0 352 | 353 | image = dataset[idx]['image'].astype(np.uint8) 354 | 355 | hues = [0, 338, 208, 45, 170, 99, 310, 255, 80, 190, 230, 120] 356 | 357 | fig2 = plt.figure(figsize=(4 * 2, 4 * 2)) 358 | ax2 = fig2.add_subplot() 359 | ax2.set_aspect('equal', 'box') 360 | ax2.axis('off') 361 | 362 | if image is not None: 363 | grey_image = rgb2gray(image) * 0.5 + 128 364 | ax2.imshow(grey_image, cmap='Greys_r', vmin=0, vmax=255) 365 | 366 | for li in range(ls.shape[0]): 367 | vpidx = closest_vps[li] 368 | if vpidx: 369 | hue = hues[vpidx % len(hues)] / 360.0 370 | sat = 1 371 | val = 1 if vpidx else 0 372 | c = matplotlib.colors.hsv_to_rgb([hue, sat, val]) 373 | 374 | ax2.plot([ls[li, 0], ls[li, 3]], [ls[li, 1], ls[li, 4]], c=c, lw=3) 375 | 376 | fig2.tight_layout() 377 | plt.gca().set_axis_off() 378 | plt.subplots_adjust(top=1, bottom=0, right=1, left=0, 379 | hspace=0, wspace=0) 380 | plt.margins(0, 0) 381 | plt.savefig(os.path.join(target_folder, "%03d_vp_%d.png" % (idx, num_vps)), bbox_inches='tight', 382 | pad_inches=0) 383 | plt.imsave(os.path.join(target_folder, "%03d_orig.png" % idx), image) 384 | plt.close() 385 | -------------------------------------------------------------------------------- /environment.yml: -------------------------------------------------------------------------------- 1 | name: parsac 2 | channels: 3 | - pytorch 4 | - nvidia 5 | - conda-forge 6 | - defaults 7 | dependencies: 8 | - matplotlib=3.8.0 9 | - numpy=1.26.3 10 | - pip 11 | - python=3.10 12 | - pytorch=1.13.1 13 | - pytorch-cuda=11.7 14 | - scikit-image=0.20.0 15 | - scikit-learn=1.3.0 16 | - scipy=1.11.4 17 | - torchvision=0.14.1 18 | - pip: 19 | - pylsd-nova 20 | - wandb 21 | 22 | -------------------------------------------------------------------------------- /networks/cn_net.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | import torch.nn.functional as F 4 | 5 | class CNNet(nn.Module): 6 | 7 | def __init__(self, input_dim, output_dim, blocks=5, batch_norm=True, separate_weights=True): 8 | 9 | super(CNNet, self).__init__() 10 | 11 | self.input_dim = input_dim 12 | self.output_dim = output_dim 13 | 14 | self.p_in = nn.Conv1d(self.input_dim, 128, 1, 1, 0) 15 | 16 | self.res_blocks = [] 17 | 18 | self.batch_norm = batch_norm 19 | self.separate_probs = separate_weights 20 | 21 | for i in range(0, blocks): 22 | if batch_norm: 23 | self.res_blocks.append(( 24 | nn.Conv1d(128, 128, 1, 1, 0), 25 | nn.BatchNorm1d(128), 26 | nn.Conv1d(128, 128, 1, 1, 0), 27 | nn.BatchNorm1d(128), 28 | )) 29 | else: 30 | self.res_blocks.append(( 31 | nn.Conv1d(128, 128, 1, 1, 0), 32 | nn.Conv1d(128, 128, 1, 1, 0), 33 | )) 34 | 35 | for i, r in enumerate(self.res_blocks): 36 | super(CNNet, self).add_module(str(i) + 's0', r[0]) 37 | super(CNNet, self).add_module(str(i) + 's1', r[1]) 38 | if batch_norm: 39 | super(CNNet, self).add_module(str(i) + 's2', r[2]) 40 | super(CNNet, self).add_module(str(i) + 's3', r[3]) 41 | 42 | self.p_out = nn.Conv1d(128, output_dim, 1, 1, 0) 43 | if self.separate_probs: 44 | self.p_out2 = nn.Conv1d(128, output_dim, 1, 1, 0) 45 | 46 | 47 | def forward(self, inputs): 48 | ''' 49 | Forward pass. 50 | 51 | inputs -- 3D data tensor (BxNxC) 52 | ''' 53 | inputs_ = torch.transpose(inputs, 1, 2) 54 | 55 | x = inputs_[:, 0:self.input_dim] 56 | 57 | x = F.relu(self.p_in(x)) 58 | 59 | for r in self.res_blocks: 60 | res = x 61 | if self.batch_norm: 62 | x = F.relu(r[1](F.instance_norm(r[0](x)))) 63 | x = F.relu(r[3](F.instance_norm(r[2](x)))) 64 | else: 65 | x = F.relu(F.instance_norm(r[0](x))) 66 | x = F.relu(F.instance_norm(r[1](x))) 67 | x = x + res 68 | 69 | log_ng = F.logsigmoid(self.p_out(x)) 70 | log_ng = torch.transpose(log_ng, 1, 2) 71 | normalizer = torch.logsumexp(log_ng, dim=-1, keepdim=True) 72 | log_ng = log_ng - normalizer 73 | 74 | if self.separate_probs: 75 | log_ng2 = F.logsigmoid(self.p_out2(x)) 76 | log_ng2 = torch.transpose(log_ng2, 1, 2) 77 | normalizer = torch.logsumexp(log_ng2, dim=-2, keepdim=True) 78 | log_ng2 = log_ng2 - normalizer 79 | else: 80 | log_ng2 = log_ng 81 | 82 | return log_ng, log_ng2 83 | -------------------------------------------------------------------------------- /parsac.py: -------------------------------------------------------------------------------- 1 | from utils import \ 2 | options, initialisation, sampling, backward, visualisation, evaluation, residual_functions, inlier_counting, metrics, postprocessing 3 | import torch 4 | import time 5 | 6 | opt = options.get_options() 7 | 8 | initialisation.seeds(opt) 9 | 10 | ckpt_dir, log = initialisation.setup_logging_and_checkpointing(opt) 11 | 12 | model, optimizer, scheduler, device = initialisation.get_model(opt) 13 | 14 | datasets = initialisation.get_dataset(opt) 15 | 16 | for epoch in range(opt.epochs): 17 | 18 | print("epoch %d / %d" % (epoch + 1, opt.epochs)) 19 | 20 | dataloaders = initialisation.get_dataloader(opt, datasets, shuffle_all=False) 21 | 22 | for mode in opt.modes: 23 | 24 | assert not (dataloaders[mode] is None), "no dataloader for %s available" % mode 25 | 26 | print("mode: %s" % mode) 27 | 28 | if mode == "train": 29 | model.train() 30 | else: 31 | model.eval() 32 | 33 | eval_metrics = {"loss": [], "time": [], "total_time": []} 34 | wandb_log_data = {} 35 | 36 | total_start = time.time() 37 | 38 | for batch_idx, (features, X, gt_labels, gt_models, image, image_size) in enumerate(dataloaders[mode]): 39 | 40 | for run_idx in range(opt.runcount): 41 | 42 | X = X.to(device) 43 | features = features.to(device) 44 | gt_labels = gt_labels.to(device) 45 | gt_models = gt_models.to(device) 46 | image_size = image_size.to(device) 47 | 48 | optimizer.zero_grad() 49 | 50 | start_time = time.time() 51 | 52 | log_inlier_weights, log_sample_weights = model(features) 53 | 54 | with torch.no_grad(): 55 | 56 | minimal_sets = sampling.sample_minimal_sets(opt, log_sample_weights) 57 | hypotheses = sampling.generate_hypotheses(opt, X, minimal_sets) 58 | residuals = residual_functions.compute_residuals(opt, X, hypotheses) 59 | 60 | weighted_inlier_ratios, inlier_scores = \ 61 | inlier_counting.count_inliers(opt, residuals, log_inlier_weights) 62 | 63 | log_p_M_S, sampled_inlier_scores, sampled_hypotheses, sampled_residuals = \ 64 | sampling.sample_hypotheses(opt, mode, hypotheses, weighted_inlier_ratios, inlier_scores, residuals) 65 | 66 | if opt.refine: 67 | if opt.problem == "vp": 68 | sampled_hypotheses, sampled_residuals, sampled_inlier_scores = \ 69 | postprocessing.refinement_with_inliers(opt, X, sampled_inlier_scores) 70 | 71 | ranked_choices, ranked_inlier_ratios, ranked_hypotheses, ranked_scores, labels, clusters = \ 72 | postprocessing.ranking_and_clustering(opt, sampled_inlier_scores, sampled_hypotheses, 73 | sampled_residuals) 74 | 75 | duration = (time.time() - start_time) * 1000 76 | 77 | eval_metrics["time"] += [duration] 78 | 79 | if not opt.self_supervised: 80 | with torch.no_grad(): 81 | if opt.problem == "vp": 82 | exp_losses, _ = metrics.vp_loss(gt_models, ranked_hypotheses, datasets["inverse_intrinsics"]) 83 | elif opt.problem == "fundamental" or opt.problem == "homography": 84 | exp_losses = metrics.classification_loss(opt, gt_labels, clusters) 85 | else: 86 | assert False 87 | else: 88 | cumulative_inlier_losses = inlier_counting.compute_cumulative_inliers(opt, ranked_scores) 89 | sample_inlier_counts = inlier_counting.combine_hypotheses_inliers(ranked_scores) 90 | exp_losses = backward.expected_self_losses(opt, sample_inlier_counts, cumulative_inlier_losses) 91 | 92 | if mode == "train": 93 | log_p_M = backward.log_probabilities(log_sample_weights, minimal_sets, log_p_M_S) 94 | _ = backward.backward_pass(opt, exp_losses, log_p_M, optimizer) 95 | else: 96 | eval_metrics = \ 97 | evaluation.compute_validation_metrics(opt, eval_metrics, ranked_hypotheses, 98 | ranked_inlier_ratios, gt_models, gt_labels, X, image_size, clusters, 99 | run_idx, datasets["inverse_intrinsics"], train=(mode == "train")) 100 | mean_loss = exp_losses.mean() 101 | eval_metrics["loss"] += [mean_loss.item()] 102 | 103 | total_duration = (time.time() - total_start) * 1000 104 | eval_metrics["total_time"] += [total_duration] 105 | total_start = time.time() 106 | 107 | if opt.visualise: 108 | visualisation.save_visualisation_plots(opt, X, ranked_choices, log_inlier_weights, 109 | log_sample_weights, ranked_scores, clusters, 110 | labels, gt_models, gt_labels, image, 111 | dataloaders[mode].dataset, metrics=eval_metrics) 112 | 113 | visualisation.log_wandb(wandb_log_data, eval_metrics, mode, epoch) 114 | if opt.eval: 115 | for key, val in wandb_log_data.items(): 116 | print(key, ":", val) 117 | 118 | if mode == "train": 119 | scheduler.step() 120 | 121 | if opt.ckpt_mode == "all": 122 | torch.save(model.state_dict(), '%s/model_weights_%06d.net' % (ckpt_dir, epoch)) 123 | torch.save(optimizer.state_dict(), '%s/optimizer_%06d.net' % (ckpt_dir, epoch)) 124 | elif opt.ckpt_mode == "last": 125 | torch.save(model.state_dict(), '%s/model_weights.net' % (ckpt_dir)) 126 | torch.save(optimizer.state_dict(), '%s/optimizer.net' % (ckpt_dir)) 127 | -------------------------------------------------------------------------------- /utils/backward.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def expected_self_losses(opt, inlier_counts, cumulative_losses): 5 | if opt.cumulative_loss > -1: 6 | return cumulative_losses 7 | else: 8 | return -inlier_counts 9 | 10 | 11 | def log_probabilities(log_sample_weights, choices, log_p_M_S): 12 | # implements Eq. 10 13 | B, K, S, M, mss = choices.size() 14 | B, N, Mo = log_sample_weights.size() 15 | 16 | log_p_x_j = log_sample_weights.transpose(1, 2).view(B, 1, 1, Mo, N).expand(B, K, S, Mo, N) 17 | 18 | log_p_h = torch.gather(log_p_x_j[..., :M, :], -1, choices) # B, K, S, M, mss 19 | log_p_h = log_p_h.sum(-1) # B, K, S, M 20 | 21 | log_p_S = log_p_h.sum(2) # B, K, M 22 | 23 | log_p_M = log_p_M_S + log_p_S.sum(2).view(B, K, 1) 24 | 25 | return log_p_M 26 | 27 | 28 | def backward_pass(opt, losses, log_p_M, optimizer): 29 | if opt.hypsamples > 0: 30 | baselines = losses.mean((-1, -2), keepdim=True) 31 | outputs = [log_p_M] 32 | gradients = [(losses - baselines).detach()] 33 | 34 | else: 35 | baselines = losses.mean(-1, keepdim=True) 36 | 37 | outputs = [log_p_M, losses] 38 | gradients = [(losses - baselines).detach(), torch.ones_like(losses, device=losses.device)] 39 | 40 | torch.autograd.backward(outputs, gradients) 41 | 42 | optimizer.step() 43 | 44 | return baselines.mean() 45 | 46 | 47 | -------------------------------------------------------------------------------- /utils/evaluation.py: -------------------------------------------------------------------------------- 1 | import utils.metrics 2 | import numpy as np 3 | import sklearn.metrics 4 | 5 | 6 | def calc_auc(error_array, cutoff=0.25): 7 | error_array = error_array.squeeze() 8 | error_array = np.sort(error_array) 9 | num_values = error_array.shape[0] 10 | 11 | plot_points = np.zeros((num_values, 2)) 12 | 13 | midfraction = 1. 14 | 15 | for i in range(num_values): 16 | fraction = (i + 1) * 1.0 / num_values 17 | value = error_array[i] 18 | plot_points[i, 1] = fraction 19 | plot_points[i, 0] = value 20 | if i > 0: 21 | lastvalue = error_array[i - 1] 22 | if lastvalue < cutoff < value: 23 | midfraction = (lastvalue * plot_points[i - 1, 1] + value * fraction) / (value + lastvalue) 24 | 25 | if plot_points[-1, 0] < cutoff: 26 | plot_points = np.vstack([plot_points, np.array([cutoff, 1])]) 27 | else: 28 | plot_points = np.vstack([plot_points, np.array([cutoff, midfraction])]) 29 | 30 | sorting = np.argsort(plot_points[:, 0]) 31 | plot_points = plot_points[sorting, :] 32 | 33 | auc = sklearn.metrics.auc(plot_points[plot_points[:, 0] <= cutoff, 0], 34 | plot_points[plot_points[:, 0] <= cutoff, 1]) 35 | auc = auc / cutoff 36 | 37 | return auc, plot_points 38 | 39 | 40 | def compute_validation_metrics(opt, metrics, models, counts_total, gt_models, gt_labels, X, image_size, clusters, run_id, inv_intrinsics, train=False): 41 | if not opt.eval: 42 | if not ("inlier_count" in metrics.keys()): 43 | metrics["inlier_count"] = [] 44 | 45 | metrics["inlier_count"] += [counts_total.mean().item()] 46 | 47 | if opt.problem == "vp": 48 | 49 | B, K, M, _, D = models.size() 50 | 51 | if not ("vp_error" in metrics.keys()): 52 | metrics["vp_error"] = [[] for _ in range(opt.runcount)] 53 | 54 | _, min_costs = utils.metrics.vp_loss(gt_models, models, inv_intrinsics, max_error=None) 55 | 56 | for bi in range(B): 57 | for ki in range(K): 58 | vps_true = gt_models[bi].cpu().detach().numpy() 59 | n = np.linalg.norm(vps_true, axis=-1) 60 | num_true = np.sum(n > 1e-8) 61 | 62 | errors = min_costs[bi, ki, 0, 0, :num_true].cpu().detach().numpy().tolist() 63 | 64 | metrics["vp_error"][run_id] += errors 65 | 66 | if (opt.dataset == "adelaide" or opt.dataset == "hope" or opt.dataset == "smh") and not train: 67 | 68 | if not ("misclassification_error" in metrics.keys()): 69 | metrics["misclassification_error"] = [] 70 | if not ("geometric_error" in metrics.keys()): 71 | metrics["geometric_error"] = [] 72 | 73 | ge = utils.metrics.geometric_errors(opt, X, image_size, gt_labels, models[..., 0, :]) 74 | metrics["geometric_error"] += ge.flatten().tolist() 75 | 76 | classification_losses = utils.metrics.classification_loss(opt, gt_labels, clusters) * 100.0 77 | 78 | misclassification_errors = classification_losses.view(-1).cpu().numpy().tolist() 79 | metrics["misclassification_error"] += misclassification_errors 80 | 81 | return metrics 82 | 83 | 84 | -------------------------------------------------------------------------------- /utils/initialisation.py: -------------------------------------------------------------------------------- 1 | from networks.cn_net import CNNet 2 | from utils.tee import Tee 3 | from datasets.nyuvp import NYUVP 4 | from datasets.su3 import SU3 5 | from datasets.adelaide import AdelaideRMFDataset 6 | from datasets.hope import HopeFDataset 7 | from datasets.smh import SMHDataset 8 | import torch.optim as optim 9 | import torch 10 | import os 11 | import json 12 | import wandb 13 | import numpy as np 14 | import random 15 | 16 | 17 | def seeds(opt): 18 | torch.manual_seed(opt.seed) 19 | np.random.seed(opt.seed) 20 | random.seed(opt.seed) 21 | 22 | 23 | def seed_worker(worker_id): 24 | worker_seed = torch.initial_seed() % 2 ** 32 25 | np.random.seed(worker_seed) 26 | random.seed(worker_seed) 27 | 28 | 29 | def setup_logging_and_checkpointing(opt): 30 | wandb.init(project="parsac", entity=opt.wandb_entity, group=opt.wandb_group, mode=opt.wandb, dir=opt.wandb_dir) 31 | 32 | parent_ckpt_dir = os.path.join(opt.checkpoint_dir, opt.wandb_group) 33 | os.makedirs(parent_ckpt_dir, exist_ok=True) 34 | 35 | dir_success = False 36 | 37 | while not dir_success: 38 | if opt.debug_session: 39 | ckpt_dir = os.path.join(opt.checkpoint_dir, opt.group, "debug_session") 40 | os.makedirs(ckpt_dir, exist_ok=True) 41 | dir_success = True 42 | else: 43 | run_name = wandb.run.name 44 | if run_name is None or run_name == "": 45 | if not (opt.jobid == ""): 46 | run_name = opt.jobid 47 | else: 48 | run_name = "%08d" % np.random.randint(0, 99999999) 49 | ckpt_dir = os.path.join(opt.checkpoint_dir, opt.wandb_group, run_name) 50 | try: 51 | os.makedirs(ckpt_dir, exist_ok=False) 52 | dir_success = True 53 | except FileExistsError as err: 54 | print(err) 55 | print("%s exists, try again.." % ckpt_dir) 56 | 57 | print("saving models to: ", ckpt_dir) 58 | 59 | log_file = os.path.join(ckpt_dir, "output.log") 60 | print("log file: ", log_file) 61 | log = Tee(log_file, "w", file_only=False) 62 | 63 | with open(os.path.join(ckpt_dir, 'commandline_args.txt'), 'w') as f: 64 | json.dump(opt.__dict__, f, indent=2) 65 | 66 | wandb.config.update(vars(opt)) 67 | wandb.config.update({"checkpoints": ckpt_dir}) 68 | 69 | return ckpt_dir, log 70 | 71 | 72 | def get_model(opt): 73 | device_id = int(opt.gpu) 74 | if device_id is None or device_id < 0 or not torch.cuda.is_available(): 75 | device = torch.device('cpu') 76 | else: 77 | device = torch.device('cuda', device_id) 78 | 79 | input_dim_map = { 80 | "nyuvp": 4, 81 | "yudplus": 4, 82 | "yud": 4, 83 | "su3": 4, 84 | "adelaide": 4, 85 | "hope": 4, 86 | "smh": 4, 87 | } 88 | 89 | input_dim = input_dim_map[opt.dataset] 90 | 91 | model = CNNet(input_dim, opt.instances+1, opt.network_layers, batch_norm=True, 92 | separate_weights=opt.separate_weights) 93 | 94 | if opt.load is not None and len(opt.load) > 0: 95 | print("load weights from ", opt.load) 96 | model.load_state_dict(torch.load(os.path.join(opt.load, "model_weights.net"), map_location=device), strict=True) 97 | 98 | model = model.to(device) 99 | 100 | optimizer = optim.Adam(model.parameters(), lr=opt.lr, eps=1e-4, weight_decay=1e-4) 101 | if opt.load is not None and len(opt.load) > 0: 102 | optimizer.load_state_dict(torch.load(os.path.join(opt.load, "optimizer.net"), map_location=device)) 103 | 104 | scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=opt.lr_steps, gamma=0.1) 105 | 106 | total_params = sum(p.numel() for p in model.parameters()) 107 | print(f"Number of network parameters: {total_params}") 108 | 109 | return model, optimizer, scheduler, device 110 | 111 | 112 | def get_dataset(opt): 113 | test_dataset = None 114 | val_dataset = None 115 | train_dataset = None 116 | 117 | return_images = opt.visualise and not ("train" in opt.modes) 118 | cache = True 119 | 120 | if opt.dataset == "nyuvp": 121 | train_dataset = NYUVP("train", opt.max_num_points) 122 | val_dataset = NYUVP("val", opt.max_num_points) 123 | test_dataset = NYUVP("test", opt.max_num_points, deeplsd_folder=opt.ablation_deeplsd_folder, cache=False) 124 | elif opt.dataset == "su3": 125 | train_dataset = SU3(opt.data_path, "train", opt.max_num_points, normalise_coords=True, augmentation=opt.augment, 126 | deeplsd_folder=opt.ablation_deeplsd_folder) 127 | val_dataset = SU3(opt.data_path, "valid", opt.max_num_points, normalise_coords=True, 128 | deeplsd_folder=opt.ablation_deeplsd_folder) 129 | test_dataset = SU3(opt.data_path, "test", opt.max_num_points, normalise_coords=True, cache=True, 130 | deeplsd_folder=opt.ablation_deeplsd_folder, ablation_outlier_ratio=opt.ablation_outlier_ratio, 131 | ablation_noise=opt.ablation_noise) 132 | elif opt.dataset == "yudplus" or opt.dataset == "yud": 133 | train_dataset = NYUVP("train", opt.max_num_points, use_yud=True, use_yud_plus=(opt.dataset == "yudplus")) 134 | test_dataset = NYUVP("test", opt.max_num_points, use_yud=True, use_yud_plus=(opt.dataset == "yudplus"), 135 | deeplsd_folder=opt.ablation_deeplsd_folder, cache=False) 136 | elif opt.dataset == "adelaide": 137 | test_dataset = AdelaideRMFDataset(opt.data_path, opt.max_num_points, problem=opt.problem, permute_points=False, 138 | ablation_outlier_ratio=opt.ablation_outlier_ratio, 139 | ablation_noise=opt.ablation_noise) 140 | elif opt.dataset == "hope": 141 | train_dataset = HopeFDataset(opt.data_path, "train", opt.max_num_points) 142 | val_dataset = HopeFDataset(opt.data_path, "val", opt.max_num_points, return_images=return_images) 143 | test_dataset = HopeFDataset(opt.data_path, "test", opt.max_num_points, return_images=return_images, 144 | cache_data=cache) 145 | elif opt.dataset == "smh": 146 | train_dataset = SMHDataset(opt.data_path, "train", opt.max_num_points, keep_in_mem=cache) 147 | val_dataset = SMHDataset(opt.data_path, "val", opt.max_num_points, keep_in_mem=cache, 148 | return_images=return_images) 149 | test_dataset = SMHDataset(opt.data_path, "test", opt.max_num_points, keep_in_mem=cache, 150 | return_images=return_images) 151 | else: 152 | assert False, "unknown dataset %s" % opt.dataset 153 | 154 | if opt.dataset == "nyuvp": 155 | fx_rgb = 5.1885790117450188e+02 156 | fy_rgb = 5.1946961112127485e+02 157 | cx_rgb = 3.2558244941119034e+02 158 | cy_rgb = 2.5373616633400465e+02 159 | intrinsics = torch.tensor([[fx_rgb, 0, cx_rgb], [0, fy_rgb, cy_rgb], [0, 0, 1]]) 160 | scale_matrix = torch.tensor([[1. / 320., 0, -1.], [0, 1. / 320., -.75], [0, 0, 1]]) 161 | SKmat = scale_matrix @ intrinsics 162 | inv_intrinsics = torch.inverse(SKmat) 163 | elif opt.dataset == "yudplus" or opt.dataset == "yud": 164 | f = 6.053170589753693 165 | ps = 0.00896875 166 | pp = [307.55130528, 251.45424496] 167 | intrinsics = torch.tensor([[f / ps, 0, pp[0]], [0, f / ps, pp[1]], [0, 0, 1]]) 168 | scale_matrix = torch.tensor([[1. / 320., 0, -1.], [0, 1. / 320., -.75], [0, 0, 1]]) 169 | SKmat = scale_matrix @ intrinsics 170 | inv_intrinsics = torch.inverse(SKmat) 171 | elif opt.dataset == "su3": 172 | f = 2.1875 * 256 173 | c = 256 174 | intrinsics = torch.tensor([[f, 0, c], [0, -f, c], [0, 0, 1]]) 175 | scale_matrix = torch.tensor([[1. / c, 0, -1.], [0, 1. / c, -1.], [0, 0, 1]]) 176 | SKmat = scale_matrix @ intrinsics 177 | inv_intrinsics = torch.inverse(SKmat) 178 | else: 179 | inv_intrinsics = None 180 | 181 | print("initialised %s dataset" % opt.dataset) 182 | 183 | return {"train": train_dataset, "val": val_dataset, "test": test_dataset, "inverse_intrinsics": inv_intrinsics} 184 | 185 | 186 | def get_dataloader(opt, datasets, shuffle_all=False): 187 | g = torch.Generator() 188 | g.manual_seed(opt.seed) 189 | 190 | dataloaders = {} 191 | for split, dataset in datasets.items(): 192 | if dataset is None: 193 | loader = None 194 | else: 195 | loader = torch.utils.data.DataLoader(dataset, shuffle=(split == "train" or shuffle_all), 196 | num_workers=opt.num_workers, generator=g, worker_init_fn=seed_worker, 197 | batch_size=opt.batch, drop_last=False) 198 | dataloaders[split] = loader 199 | return dataloaders 200 | 201 | 202 | def get_parameter_sweep_range(opt): 203 | 204 | if opt.parameter_sweep == "S": 205 | values = range(opt.hypotheses//8, opt.hypotheses*2+1, opt.hypotheses//8) 206 | elif opt.parameter_sweep == "tau": 207 | if opt.problem == "homography" and opt.dataset == "adelaide": 208 | values = np.logspace(np.log10(opt.inlier_threshold)-3, np.log10(opt.inlier_threshold)+1, num=100, endpoint=True) 209 | elif opt.problem == "homography" and opt.dataset == "smh": 210 | values = np.logspace(np.log10(opt.inlier_threshold)+1, np.log10(opt.inlier_threshold)+3, num=50, endpoint=True) 211 | else: 212 | values = np.logspace(np.log10(opt.inlier_threshold)-1, np.log10(opt.inlier_threshold)+1, num=50, endpoint=True) 213 | elif opt.parameter_sweep == "taua": 214 | if opt.problem == "homography" and opt.dataset == "adelaide": 215 | values = np.logspace(np.log10(opt.assignment_threshold)-4, np.log10(opt.assignment_threshold)+1, num=150, endpoint=True) 216 | elif opt.problem == "homography" and opt.dataset == "smh": 217 | values = np.logspace(np.log10(opt.assignment_threshold)+3, np.log10(opt.assignment_threshold)+4, num=25, endpoint=True) 218 | else: 219 | values = np.logspace(np.log10(opt.assignment_threshold)-1, np.log10(opt.assignment_threshold)+1, num=50, endpoint=True) 220 | elif opt.parameter_sweep == "beta": 221 | values = np.logspace(np.log10(opt.inlier_softness)-1, np.log10(opt.inlier_softness)+1, num=50, endpoint=True) 222 | else: 223 | assert False 224 | 225 | if opt.dataset in ["hope", "smh", "nyuvp", "su3"]: 226 | opt.modes = ["val"] 227 | else: 228 | opt.modes = ["test"] 229 | 230 | return values -------------------------------------------------------------------------------- /utils/inlier_counting.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def soft_inlier_fun_gen(beta, tau, inverse=False): 5 | def f(d): 6 | if inverse: 7 | return torch.sigmoid(beta * d / tau - beta) 8 | else: 9 | return 1 - torch.sigmoid(beta * d / tau - beta) 10 | return f 11 | 12 | 13 | def hard_inlier_fun_gen(beta, tau, inverse=False): 14 | def f(d): 15 | inlier = (d < tau).float() 16 | return 1-inlier if inverse else inlier 17 | return f 18 | 19 | 20 | inlier_functions = {"hard": hard_inlier_fun_gen, "soft": soft_inlier_fun_gen} 21 | 22 | 23 | def count_inliers(opt, residuals, log_weights): 24 | 25 | weights = torch.softmax(log_weights, dim=2) 26 | B, N, Mo = weights.size() 27 | _, K, S, M, N = residuals.size() 28 | weights_e = weights[..., :M].transpose(1, 2).view(B, 1, 1, M, N).expand(B, K, S, M, N) 29 | 30 | inlier_fun = inlier_functions[opt.inlier_function](opt.inlier_softness, opt.inlier_threshold) 31 | 32 | inlier_scores = inlier_fun(residuals) 33 | 34 | if opt.inlier_counting == "unweighted": 35 | inlier_scores_weighted = inlier_scores 36 | else: 37 | inlier_scores_weighted = inlier_scores * weights_e 38 | 39 | inlier_ratios = inlier_scores_weighted.sum(-1) * 1.0 / N 40 | 41 | return inlier_ratios, inlier_scores 42 | 43 | 44 | def compute_cumulative_inliers(opt, scores): 45 | 46 | B, K, M, H, N = scores.size() 47 | 48 | combined_scores = torch.zeros_like(scores, device=scores.device) 49 | for mi in range(M): 50 | if mi == 0: 51 | combined_scores[:, :, mi] = scores[:, :, mi] 52 | else: 53 | combined_scores[:, :, mi] = torch.max(scores[:, :, 0:(mi+1)], dim=2)[0] * (opt.cumulative_loss ** mi) 54 | 55 | inlier_loss = -combined_scores.sum((2, 4)) * 1.0 / N 56 | 57 | return inlier_loss 58 | 59 | 60 | def combine_hypotheses_inliers(inlier_scores): 61 | 62 | B, K, M, H, N = inlier_scores.size() 63 | 64 | combined_inliers = torch.max(inlier_scores, dim=2)[0] 65 | 66 | inlier_counts = combined_inliers.sum(-1) * 1.0 / N # B, K, H 67 | 68 | return inlier_counts 69 | -------------------------------------------------------------------------------- /utils/metrics.py: -------------------------------------------------------------------------------- 1 | import itertools 2 | import math 3 | 4 | import numpy as np 5 | import scipy.optimize 6 | import torch 7 | 8 | from utils import residual_functions 9 | 10 | 11 | def hungarian_loss(cost, num_models=None, max_error=None): 12 | B, K, H, M1, M2 = cost.size() 13 | losses = np.zeros((B, K, H), dtype=np.float32) 14 | 15 | cost_np = cost.cpu().detach().numpy() 16 | for bi in range(B): 17 | for ki in range(K): 18 | for hi in range(H): 19 | if num_models is not None: 20 | num_true = num_models[bi] 21 | cost_mat = cost_np[bi, ki, hi, :num_true, :num_true] 22 | else: 23 | cost_mat = cost_np[bi, ki, hi] 24 | row_ind, col_ind = scipy.optimize.linear_sum_assignment(cost_mat) 25 | errors = cost_mat[row_ind, col_ind] 26 | if max_error is None: 27 | loss = errors.sum() 28 | else: 29 | loss = np.clip(errors, a_max=max_error, a_min=None).sum() 30 | 31 | losses[bi, ki, hi] = loss 32 | return torch.from_numpy(losses).to(cost.device) 33 | 34 | 35 | def hungarian_loss_torch(cost, num_models=None, max_error=None, max_gpu_alloc=6): 36 | B, K, H, M1, M2 = cost.size() 37 | if M1 > M2: 38 | cost = cost.transpose(-1, -2) 39 | B, K, H, M1, M2 = cost.size() 40 | 41 | perm = itertools.permutations(range(M2), M1) 42 | perm_list = list(perm) 43 | P = len(perm_list) 44 | 45 | if B * K * H * P * M1 * 4 > max_gpu_alloc * (1024 ** 3): 46 | return hungarian_loss(cost, num_models, max_error) 47 | 48 | perm_list = np.array(perm_list) 49 | 50 | range1 = torch.arange(0, M1).view(1, M1).expand(B, M1).to(cost.device) 51 | range2 = torch.arange(0, M2).view(1, M2).expand(B, M2).to(cost.device) 52 | 53 | if num_models is not None: 54 | mask1 = range1 < num_models.view(B, 1) 55 | mask2 = range2 < num_models.view(B, 1) 56 | mask = torch.logical_and(mask1.view(B, M1, 1), mask2.view(B, 1, M2)).view(B, 1, 1, M1, M2).float() 57 | cost = cost * mask 58 | mask = torch.logical_xor(mask1.view(B, M1, 1), mask2.view(B, 1, M2)).view(B, 1, 1, M1, M2).float() 59 | cost = cost + mask * 10000.0 60 | 61 | permutations = torch.from_numpy(perm_list).to(cost.device) 62 | permutations = permutations.view(1, 1, 1, P, M1, 1).expand(B, K, H, P, M1, 1) 63 | 64 | cost_e = cost.view(B, K, H, 1, M1, M2).expand(B, K, H, P, M1, M2) 65 | 66 | selected_cost = torch.gather(cost_e, -1, permutations).squeeze(-1) 67 | 68 | cost_sums = selected_cost.sum(-1) 69 | min_idx = torch.argmin(cost_sums, dim=3, keepdim=True) 70 | 71 | min_costs = torch.gather(selected_cost, 3, min_idx.view(B, K, H, 1, 1).expand(B, K, H, 1, M1)) 72 | 73 | if max_error is not None: 74 | losses = torch.clip(min_costs, max=max_error).sum(-1).squeeze(-1) 75 | else: 76 | losses = min_costs.sum(-1).squeeze(-1) 77 | 78 | return losses, min_costs 79 | 80 | 81 | def vp_loss(true_vps, estm_vps, inv_intrinsics, max_error=10.): 82 | B, Mt, D = true_vps.size() 83 | B, K, M, H, D = estm_vps.size() 84 | 85 | inv_intrinsics = inv_intrinsics.to(true_vps.device) 86 | 87 | true_vds = (inv_intrinsics @ true_vps.unsqueeze(-1)).squeeze(-1) 88 | estm_vds = (inv_intrinsics @ estm_vps.unsqueeze(-1)).squeeze(-1) 89 | 90 | true_vds = true_vds / torch.clip(torch.norm(true_vds, dim=-1, keepdim=True), min=1e-8) 91 | estm_vds = estm_vds / torch.clip(torch.norm(estm_vds, dim=-1, keepdim=True), min=1e-8) 92 | 93 | estm_vds = estm_vds.transpose(2, 3).view(B, K, H, 1, M, D) 94 | true_vds = true_vds.view(B, 1, 1, Mt, 1, D) 95 | cosines = (true_vds * estm_vds).sum(-1) 96 | cosines = torch.clip(cosines, min=-1.0 + 1e-8, max=1.0 - 1e-8) 97 | 98 | cost_matrix = torch.arccos(torch.abs(cosines)) * 180. / math.pi 99 | 100 | true_vps_norms = torch.norm(true_vps, dim=-1) 101 | 102 | losses, min_costs = hungarian_loss_torch(cost_matrix, (true_vps_norms > 1e-8).sum(-1), max_error) 103 | 104 | return losses, min_costs 105 | 106 | 107 | def classification_loss(opt, true_labels, clusters): 108 | with torch.no_grad(): 109 | 110 | inliers_only = opt.ablation_outlier_ratio >= 0 111 | if inliers_only: 112 | num_true_inliers = (true_labels > 0).sum(-1)[0] 113 | true_labels = true_labels[..., :num_true_inliers] 114 | clusters = clusters[..., :num_true_inliers] 115 | 116 | B, K, H, Mo, N = clusters.size() 117 | M = opt.instances 118 | 119 | num_true_labels = torch.max(true_labels).item() + 1 120 | 121 | true_clusters = torch.zeros((B, num_true_labels, N), dtype=torch.bool, device=true_labels.device) 122 | 123 | for li in range(num_true_labels): 124 | true_clusters[:, li] = torch.where(true_labels == li, 125 | torch.ones((B, N), dtype=torch.bool, device=true_labels.device), 126 | true_clusters[:, li]) 127 | 128 | true_clusters = true_clusters.view(B, 1, 1, num_true_labels, N) 129 | 130 | correct_assignments = torch.logical_and(clusters[..., None, :], true_clusters[..., None, :, :]) 131 | 132 | num_correct_per_class = correct_assignments.int().sum(-1) 133 | if num_true_labels > 5 or M > 9: 134 | num_correct_total = -hungarian_loss(-num_correct_per_class) 135 | else: 136 | num_correct_total, _ = hungarian_loss_torch(-num_correct_per_class) 137 | num_correct_total = -num_correct_total 138 | 139 | losses = (true_clusters.sum((-1, -2)) - num_correct_total) / true_clusters.sum((-1, -2)) 140 | 141 | return losses 142 | 143 | 144 | def geometric_errors(opt, X, image_size, ground_truth, models): 145 | x = torch.clone(X[..., :4]) 146 | 147 | scale = torch.max(image_size, dim=-1)[0] 148 | 149 | x[..., 0:2] *= scale[:, 0][:, None, None] / 2.0 150 | x[..., 2:4] *= scale[:, 1][:, None, None] / 2.0 151 | 152 | x[..., 0] += image_size[..., 0, 1][:, None] / 2.0 153 | x[..., 1] += image_size[..., 0, 0][:, None] / 2.0 154 | x[..., 2] += image_size[..., 1, 1][:, None] / 2.0 155 | x[..., 3] += image_size[..., 1, 0][:, None] / 2.0 156 | 157 | B, N, C = x.size() 158 | _, K, M, D = models.size() 159 | 160 | xe = x.view(B, 1, 1, N, C).expand(B, K, M, N, C) 161 | 162 | if opt.problem == "fundamental": 163 | T1 = torch.zeros((B, 3, 3), device=X.device) 164 | T2 = torch.zeros((B, 3, 3), device=X.device) 165 | 166 | T1[:, 0, 0] = 2.0 / scale[:, 0] 167 | T1[:, 1, 1] = 2.0 / scale[:, 0] 168 | T1[:, 2, 2] = 1 169 | 170 | T2[:, 0, 0] = 2.0 / scale[:, 1] 171 | T2[:, 1, 1] = 2.0 / scale[:, 1] 172 | T2[:, 2, 2] = 1 173 | 174 | T1[:, 0, 2] = -image_size[..., 0, 1] / scale[:, 0] 175 | T1[:, 1, 2] = -image_size[..., 0, 0] / scale[:, 0] 176 | T2[:, 0, 2] = -image_size[..., 1, 1] / scale[:, 1] 177 | T2[:, 1, 2] = -image_size[..., 1, 0] / scale[:, 1] 178 | 179 | T1 = T1[:, None, None, ...] 180 | T2 = T2[:, None, None, ...].transpose(-1, -2) 181 | 182 | F = torch.clone(models.view(B, K, M, 3, 3)) 183 | F = T2 @ F @ T1 184 | 185 | m = F.view(B, K, M, 9) 186 | 187 | elif opt.problem == "homography": 188 | T1 = torch.zeros((B, 3, 3), device=X.device) 189 | T2 = torch.zeros((B, 3, 3), device=X.device) 190 | 191 | T1[:, 0, 0] = 2.0 / scale[:, 0] 192 | T1[:, 1, 1] = 2.0 / scale[:, 0] 193 | T1[:, 2, 2] = 1 194 | 195 | T2[:, 0, 0] = scale[:, 1] / 2.0 196 | T2[:, 1, 1] = scale[:, 1] / 2.0 197 | T2[:, 2, 2] = 1 198 | 199 | T1[:, 0, 2] = -image_size[..., 0, 1] / scale[:, 0] 200 | T1[:, 1, 2] = -image_size[..., 0, 0] / scale[:, 0] 201 | T2[:, 0, 2] = image_size[..., 1, 1] / 2.0 202 | T2[:, 1, 2] = image_size[..., 1, 0] / 2.0 203 | 204 | T1 = T1[:, None, None, ...] 205 | T2 = T2[:, None, None, ...] 206 | 207 | H = torch.clone(models.view(B, K, M, 3, 3)) 208 | H = T2 @ H @ T1 209 | 210 | m = H.view(B, K, M, 9) 211 | 212 | else: 213 | m = models 214 | 215 | r = residual_functions.mapping[opt.problem](xe, m) 216 | 217 | if opt.problem == "homography": 218 | r = torch.sqrt(r) 219 | 220 | r = torch.clamp(r, min=None, max=torch.max(image_size)) 221 | 222 | r = r.detach().cpu().numpy() 223 | 224 | r_avg = np.zeros((B, K)) 225 | 226 | gt = ground_truth.cpu().numpy() 227 | 228 | for bi in range(B): 229 | for ki in range(K): 230 | all_selected_ge = [] 231 | num_classes = gt[bi].max() 232 | for ci in range(num_classes): 233 | ge_sel = r[bi, ki, :, :][:num_classes, np.where(gt[bi] == ci + 1)[0]] 234 | 235 | selected_ge = np.min(ge_sel, axis=0) 236 | 237 | all_selected_ge += selected_ge.flatten().tolist() 238 | 239 | r_avg[bi, ki] = np.mean(all_selected_ge) 240 | 241 | return r_avg 242 | -------------------------------------------------------------------------------- /utils/options.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | 4 | def get_options(): 5 | parser = argparse.ArgumentParser( 6 | description='PARSAC: Accelerating Robust Multi-Model Fitting with Parallel Sample Consensus', 7 | formatter_class=argparse.ArgumentDefaultsHelpFormatter) 8 | parser.add_argument('--eval', dest='eval', action='store_true', help='set some default values to run evaluation only') 9 | # general: 10 | parser.add_argument('--problem', default="vp", choices=["vp", "homography", "fundamental"], help='type of problem') 11 | parser.add_argument('--dataset', default="nyuvp", 12 | choices=["nyuvp", "yudplus", "yud", "su3", "adelaide", "hope", "smh"], 13 | help='name of dataset to use') 14 | parser.add_argument('--data_path', default="", help='path to dataset') 15 | parser.add_argument('--modes', nargs="+", default=["val", "train"], choices=["train", "val", "test"]) 16 | parser.add_argument('--checkpoint_dir', default='./tmp/checkpoints', 17 | help='directory for storing neural network weight checkpoints') 18 | parser.add_argument('--load', default="", type=str, help='load pretrained neural network weights from folder') 19 | parser.add_argument('--gpu', default="0", help='GPU ID to use (-1 for CPU)') 20 | parser.add_argument('--seed', default=0, type=int, help='random seed') 21 | parser.add_argument('--jobid', default="", help='Custom job ID (for setting the checkpoint directory)') 22 | 23 | # hyper-parameters: 24 | parser.add_argument('--instances', type=int, default=4, help='M-hat - number of putative model instances') 25 | parser.add_argument('--inlier_threshold', type=float, default=1e-4, help='tau - inlier threshold') 26 | parser.add_argument('--inlier_softness', default=5, type=float, help='beta - inlier softness') 27 | parser.add_argument('--assignment_threshold', type=float, default=1e-2, help='tau_a - assignment threshold') 28 | parser.add_argument('--lr', type=float, default=1e-4, help='learning rate for training') 29 | parser.add_argument('--softmax_alpha', type=float, default=1000., help='alpha_s - softmax scale factor') 30 | parser.add_argument('--epochs', type=int, default=2000, help='N_e - number of training epochs') 31 | parser.add_argument('--lr_steps', nargs="+", default=[1500], 32 | help='N_lr - number of epochs before reducing LR by factor of 10') 33 | parser.add_argument('--batch', type=int, default=48, help='B - batch size') 34 | parser.add_argument('--samplecount', type=int, default=1, help='K - hypotheses set samples') 35 | parser.add_argument('--hypsamples', type=int, default=0, 36 | help='K_hat - model set samples (0 = select argmax of inlier count)') 37 | parser.add_argument('--hypotheses', type=int, default=4, help='S - number of model hypotheses') 38 | parser.add_argument('--max_num_points', default=512, type=int, 39 | help='|X| - max. number of observations to select for input (set to -1 in order to use all observations)') 40 | parser.add_argument('--network_layers', default=6, type=int, help='number of residual blocks') 41 | 42 | # training only: 43 | parser.add_argument('--num_workers', type=int, default=4, help='number of workers for data loader') 44 | parser.add_argument('--ckpt_mode', default="last", type=str, 45 | help='', choices=["all", "last", "disabled"]) 46 | parser.add_argument('--self_supervised', action='store_true', help='use self-supervised loss function') 47 | parser.add_argument('--cumulative_loss', type=float, default=0.3, help='gamma - weight for self-supervised loss') 48 | 49 | # evaluation only: 50 | parser.add_argument('--runcount', type=int, default=1, help='perform multiple runs and compute mean and standard deviation of all metrics') 51 | 52 | # ablation 53 | parser.add_argument('--ablation_outlier_ratio', type=float, default=-1, help='synthetic outlier ratio') 54 | parser.add_argument('--ablation_noise', type=float, default=0, help='gaussian noise sigma') 55 | parser.add_argument('--ablation_deeplsd_folder', 56 | default=None, type=str, help='folder with pre-computed DeepLSD features') 57 | parser.add_argument('--inlier_counting', default="weighted", type=str, 58 | choices=["weighted", "unweighted"], help="use weighted or unweighted inlier counting") 59 | parser.add_argument('--inlier_function', default="soft", choices=["hard", "soft"], 60 | help='use hard or soft inlier scoring function') 61 | parser.add_argument('--no_refine', dest='refine', action='store_false', 62 | help='disable refinement for vanishing points') 63 | parser.add_argument('--no_separate_weights', dest='separate_weights', action='store_false', 64 | help="only predict one set of log weights for sample and inlier weights") 65 | 66 | # other 67 | parser.add_argument('--mss', default=0, type=int, 68 | help='minimal set size (overrides value given by the problem type)') 69 | parser.add_argument('--debug_session', action='store_true', help='') 70 | parser.add_argument('--augment', action='store_true', help='') 71 | parser.add_argument('--visualise', action='store_true', help='') 72 | 73 | # logging 74 | parser.add_argument('--wandb_group', default="", type=str, help='Weights and Biases group') 75 | parser.add_argument('--wandb', default="disabled", choices=["online", "offline", "disabled"], 76 | help='Weights and Biases mode') 77 | parser.add_argument('--wandb_entity', default="tnt", help='Weights and Biases entity') 78 | parser.add_argument('--wandb_dir', default="./tmp", type=str, 79 | help='Weights and Biases offline storage folder') 80 | 81 | opt = parser.parse_args() 82 | 83 | if opt.mss < 2: 84 | mss_map = {"lines": 2, "vp": 2, "homography": 4, "fundamental": 7} 85 | opt.mss = mss_map[opt.problem] 86 | 87 | if opt.eval: 88 | opt.modes = ["test"] 89 | opt.batch = 1 90 | opt.max_num_points = -1 91 | opt.num_workers = 0 92 | opt.epochs = 1 93 | opt.hypsamples = 0 94 | opt.samplecount = 1 95 | opt.wandb = "disabled" 96 | opt.runcount = 5 97 | 98 | return opt 99 | -------------------------------------------------------------------------------- /utils/postprocessing.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from utils.residual_functions import compute_residuals 4 | from utils.inlier_counting import inlier_functions 5 | 6 | 7 | def assign_cluster_labels(opt, residuals, counts): 8 | if len(residuals.size()) == 5: 9 | mask = (counts < opt.mss).float()[..., None].transpose(2, 3) 10 | res = residuals.transpose(2, 3) 11 | B, K, M, H, N = residuals.size() 12 | batch_dims = [B, K, H] 13 | else: 14 | mask = (counts < opt.mss).float()[..., None] 15 | res = residuals 16 | B, K, M, N = residuals.size() 17 | batch_dims = [B, K] 18 | 19 | estm_labels = -1 * torch.ones(batch_dims + [N], device=residuals.device, dtype=torch.long) 20 | ones = torch.ones(batch_dims + [N], device=residuals.device, dtype=torch.long) 21 | min_dists = torch.ones(batch_dims + [N], device=residuals.device, dtype=torch.float32) * opt.inlier_threshold 22 | 23 | res = res + mask * 1000.0 24 | 25 | for mi in range(M): 26 | condition = torch.logical_and(res[..., mi, :] < opt.assignment_threshold, estm_labels == -1) | \ 27 | (res[..., mi, :] < min_dists) 28 | 29 | estm_labels = torch.where(condition, ones * mi, estm_labels) 30 | min_dists = torch.minimum(min_dists, res[..., mi, :]) 31 | 32 | estm_labels += 1 33 | 34 | estm_clusters = torch.zeros(batch_dims + [M+1, N], dtype=torch.bool, device=residuals.device) 35 | for mi in range(M+1): 36 | estm_clusters[..., mi, :] = torch.where(estm_labels == mi, ones.bool(), estm_clusters[..., mi, :]) 37 | 38 | return estm_labels, estm_clusters 39 | 40 | 41 | def ranking_and_clustering(opt, inlier_scores, hypotheses, residuals): 42 | scores = torch.clone(inlier_scores.squeeze(2)) 43 | 44 | all_best_models = [] 45 | all_best_counts = [] 46 | 47 | if len(scores.size()) == 4: 48 | B, K, M, N = scores.size() 49 | selected_scores = torch.zeros((B, K, 1, N), dtype=torch.float32, device=scores.device) 50 | elif len(scores.size()) == 5: 51 | B, K, M, H, N = scores.size() 52 | selected_scores = torch.zeros((B, K, 1, H, N), dtype=torch.float32, device=scores.device) 53 | else: 54 | assert False 55 | 56 | for mi in range(M): 57 | overlap = scores * selected_scores 58 | unique = scores * (1-selected_scores) 59 | 60 | ranked_counts_per_model = unique.sum(-1) - overlap.sum(-1) 61 | 62 | if len(all_best_models) > 0: 63 | bm = torch.concatenate(all_best_models, dim=2) 64 | ranked_counts_per_model.scatter_(2, bm, torch.ones_like(ranked_counts_per_model, device=ranked_counts_per_model.device) * (-1e8)) 65 | 66 | best_counts, best_models = torch.max(ranked_counts_per_model, dim=2, keepdim=True) 67 | 68 | all_best_models += [best_models] 69 | all_best_counts += [best_counts] 70 | 71 | if len(scores.size()) == 4: 72 | best_scores = torch.gather(scores, 2, best_models.unsqueeze(-1).expand(B, K, 1, N)) 73 | 74 | elif len(scores.size()) == 5: 75 | best_scores = torch.gather(scores, 2, best_models.unsqueeze(-1).expand(B, K, 1, H, N)) 76 | else: 77 | assert False 78 | 79 | selected_scores = torch.maximum(best_scores, selected_scores) 80 | 81 | models = torch.concatenate(all_best_models, dim=2) 82 | ranked_counts_per_model = torch.concatenate(all_best_counts, dim=2) 83 | 84 | D = hypotheses.size(-1) 85 | ranked_hypotheses = torch.gather(hypotheses, 2, models.view(B, K, M, H, 1).expand(B, K, M, H, D)) 86 | ranked_residuals = torch.gather(residuals, 2, models.view(B, K, M, H, 1).expand(B, K, M, H, N)) 87 | ranked_scores = torch.gather(inlier_scores, 2, models.view(B, K, M, H, 1).expand(B, K, M, H, N)) 88 | total_counts = selected_scores.sum(-1) 89 | 90 | if not opt.problem == "vp": 91 | labels, clusters = assign_cluster_labels(opt, ranked_residuals, ranked_counts_per_model) 92 | else: 93 | labels = None 94 | clusters = None 95 | 96 | return models, total_counts, ranked_hypotheses, ranked_scores, labels, clusters 97 | 98 | 99 | def refinement_with_inliers(opt, X, inliers): 100 | B, K, M, H, N = inliers.size() 101 | 102 | if opt.problem == "vp": 103 | lines = X[:, :, 6:9] 104 | lines_repeated = lines.view(B, 1, 1, 1, N, 3).expand(B, K, M, H, N, 3) 105 | weights_repeated = inliers.view(B, K, M, H, N, 1).expand(B, K, M, H, N, 3) 106 | Mat = (lines_repeated * weights_repeated) 107 | MatSq = Mat.transpose(-1, -2) @ Mat 108 | _, _, V = torch.svd(MatSq) 109 | models = V[..., 2] 110 | 111 | else: 112 | assert False, "refinement for %s is not implemented yet" % opt.problem 113 | 114 | residuals = compute_residuals(opt, X, models) 115 | 116 | inlier_fun = inlier_functions[opt.inlier_function](opt.inlier_softness, opt.inlier_threshold) 117 | 118 | inlier_scores = inlier_fun(residuals) 119 | 120 | return models, residuals, inlier_scores -------------------------------------------------------------------------------- /utils/residual_functions.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def vanishing_point(X, vp): 5 | lines = X[..., 6:9] 6 | centroids = X[..., 9:12] 7 | 8 | constrained_lines = torch.cross(centroids, vp.unsqueeze(-2), dim=-1) 9 | con_line_norms = torch.norm(constrained_lines[..., 0:2], dim=-1, keepdim=True) 10 | constrained_lines = constrained_lines / (con_line_norms + 1e-8) 11 | line_norms = torch.norm(lines[..., 0:2], dim=-1, keepdim=True) 12 | lines = lines / (line_norms + 1e-8) 13 | distances = 1 - torch.abs((lines[..., 0:2] * constrained_lines[..., 0:2]).sum(dim=-1)) 14 | 15 | return distances 16 | 17 | 18 | def transfer_error(X, h): 19 | h_dims = h.size() 20 | batch_dims = list(h_dims)[:-1] 21 | H_dims = tuple(list(h_dims)[:-1] + [3, 3]) 22 | 23 | N = X.size(-2) 24 | 25 | H = h.view(H_dims) 26 | I = torch.eye(3, device=X.device).view([1]*(len(H_dims)-2) + [3, 3]) 27 | det = torch.det(H).view(batch_dims + [1, 1]) 28 | H_ = torch.where(det < 1e-6, I, H) 29 | 30 | X1 = torch.ones(batch_dims + [N, 3, 1], device=X.device) 31 | X2 = torch.ones(batch_dims + [N, 3, 1], device=X.device) 32 | X1[..., 0:2, 0] = X[..., 0:2] 33 | X2[..., 0:2, 0] = X[..., 2:4] 34 | 35 | HX1 = H.unsqueeze(-3) @ X1 36 | HX1[..., 0, 0] /= torch.clamp(HX1[..., 2, 0], min=1e-8) 37 | HX1[..., 1, 0] /= torch.clamp(HX1[..., 2, 0], min=1e-8) 38 | HX1[..., 2, 0] /= torch.clamp(HX1[..., 2, 0], min=1e-8) 39 | 40 | HX2 = torch.linalg.solve(H_.unsqueeze(-3), X2) 41 | HX2[..., 0, 0] /= torch.clamp(HX2[..., 2, 0], min=1e-8) 42 | HX2[..., 1, 0] /= torch.clamp(HX2[..., 2, 0], min=1e-8) 43 | HX2[..., 2, 0] /= torch.clamp(HX2[..., 2, 0], min=1e-8) 44 | 45 | signed_distances_1 = HX1 - X2 46 | distances_1 = (signed_distances_1 * signed_distances_1).sum(dim=-2).squeeze(-1) 47 | signed_distances_2 = HX2 - X1 48 | distances_2 = (signed_distances_2 * signed_distances_2).sum(dim=-2).squeeze(-1) 49 | 50 | distances = distances_1 + distances_2 51 | 52 | return distances 53 | 54 | 55 | def sampson_distance(X, f, squared=False): 56 | f_dims = f.size() 57 | batch_dims = list(f_dims)[:-1] 58 | F_dims = tuple(list(f_dims)[:-1] + [3, 3]) 59 | 60 | degenerate = torch.norm(f, dim=-1) < 1e-8 61 | 62 | N = X.size(-2) 63 | 64 | F = f.view(F_dims)[..., None, :, :] 65 | 66 | x1 = X[..., 0:2] 67 | x2 = X[..., 2:4] 68 | X1 = torch.ones(batch_dims + [N, 3, 1], device=X.device) 69 | X2 = torch.ones(batch_dims + [N, 3, 1], device=X.device) 70 | X1[..., 0:2, 0] = x1 71 | X2[..., 0:2, 0] = x2 72 | 73 | Fx1 = F @ X1 74 | Fx2 = F.transpose(-1, -2) @ X2 75 | 76 | xFx = X2 * Fx1 77 | xFx = torch.sum(xFx[..., 0], dim=-1) ** 2 78 | 79 | denom = Fx1[..., 0, 0] ** 2 + Fx1[..., 1, 0] ** 2 + Fx2[..., 0, 0] ** 2 + Fx2[..., 1, 0] ** 2 80 | denom = torch.clamp(denom, min=1e-8) 81 | 82 | sq_distances = xFx / denom 83 | 84 | sq_distances = sq_distances + degenerate[..., None].float() * 1e6 85 | 86 | if squared: 87 | return sq_distances 88 | else: 89 | return torch.sqrt(sq_distances) 90 | 91 | 92 | mapping = {"vp": vanishing_point, "homography": transfer_error, "fundamental": sampson_distance} 93 | 94 | 95 | def compute_residuals(opt, X, hypotheses): 96 | 97 | B, N, C = X.size() 98 | _, K, S, M, D = hypotheses.size() 99 | 100 | X_e = X.view(B, 1, 1, 1, N, C).expand(B, K, S, M, N, C) 101 | 102 | residuals = mapping[opt.problem](X_e, hypotheses) 103 | 104 | return residuals -------------------------------------------------------------------------------- /utils/sampling.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import utils.solvers 3 | 4 | 5 | def sample_minimal_sets(opt, log_probs): 6 | 7 | # batch x N x instances 8 | probs = torch.softmax(log_probs, dim=1) 9 | B, N, Mo = probs.size() 10 | M = opt.instances 11 | 12 | choice_weights = probs[..., :M].view(B, 1, 1, N, M).expand(B, opt.samplecount, opt.hypotheses, N, M) 13 | choice_weights = choice_weights.transpose(-1, -2).contiguous().view(-1, N) 14 | choice_batched = torch.multinomial(choice_weights, opt.mss, replacement=True) 15 | choices = choice_batched.view(B, opt.samplecount, opt.hypotheses, M, opt.mss) 16 | 17 | return choices 18 | 19 | 20 | def generate_hypotheses(opt, X, choices): 21 | 22 | B, N, C = X.size() 23 | _, K, S, M, mss = choices.size() 24 | 25 | X_e = X.view(B, 1, 1, 1, N, C).expand(B, K, S, M, N, C) 26 | choices_e = choices.view(B, K, S, M, mss, 1).expand(B, K, S, M, mss, C) 27 | 28 | X_samples = torch.gather(X_e, -2, choices_e) 29 | 30 | hypotheses = utils.solvers.minimal_solver[opt.problem](X_samples) 31 | 32 | return hypotheses 33 | 34 | 35 | def sample_hypotheses(opt, mode, hypotheses, weighted_inlier_counts, inlier_scores, residuals): 36 | B, K, S, M = weighted_inlier_counts.size() 37 | 38 | softmax_input = opt.softmax_alpha * weighted_inlier_counts 39 | 40 | hyp_selection_weights = torch.softmax(softmax_input, dim = 2) 41 | log_p_h_S = torch.nn.functional.log_softmax(softmax_input, dim=2) # B, K, S, M 42 | 43 | choice_weights = hyp_selection_weights.transpose(-1, -2).contiguous().view(-1, S) 44 | 45 | if opt.hypsamples > 0 and mode == "train": 46 | choice_batched = torch.multinomial(choice_weights, opt.hypsamples, replacement=True) 47 | choices = choice_batched.view(B, K, M, opt.hypsamples) 48 | H = opt.hypsamples 49 | else: 50 | choice_batched = torch.argmax(choice_weights, dim=-1) 51 | choices = choice_batched.view(B, K, M, 1) 52 | H = 1 53 | 54 | hyp_choices_e = choices.view(B, K, 1, M, H) 55 | log_p_e = log_p_h_S.view(B, K, S, M, 1).expand(B, K, S, M, H) 56 | selected_log_p = torch.gather(log_p_e, 2, hyp_choices_e).squeeze(2) # B, K, M, H 57 | log_p_M_S = selected_log_p.sum(2) # B, K, H 58 | 59 | B, K, S, M, D = hypotheses.size() 60 | B, K, M, H = choices.size() 61 | B, K, S, M, N = inlier_scores.size() 62 | hypotheses_e = hypotheses.view(B, K, S, M, 1, D).expand(B, K, S, M, H, D) 63 | inlier_scores_e = inlier_scores.view(B, K, S, M, 1, N).expand(B, K, S, M, H, N) 64 | residuals_e = residuals.view(B, K, S, M, 1, N).expand(B, K, S, M, H, N) 65 | hyp_choices_e = choices.view(B, K, 1, M, H, 1).expand(B, K, 1, M, H, D) 66 | selected_hypotheses = torch.gather(hypotheses_e, 2, hyp_choices_e).squeeze(2) 67 | hyp_choices_e = choices.view(B, K, 1, M, H, 1).expand(B, K, 1, M, H, N) 68 | selected_inlier_scores = torch.gather(inlier_scores_e, 2, hyp_choices_e).squeeze(2) # B, K, M, H, N 69 | selected_residuals = torch.gather(residuals_e, 2, hyp_choices_e).squeeze(2) 70 | 71 | return log_p_M_S, selected_inlier_scores, selected_hypotheses, selected_residuals 72 | 73 | 74 | -------------------------------------------------------------------------------- /utils/solvers.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import utils.residual_functions 3 | 4 | 5 | def vanishing_point(X): 6 | l1 = X[..., 0, 6:9] 7 | l2 = X[..., 1, 6:9] 8 | 9 | vp = torch.cross(l1, l2, dim=-1) 10 | 11 | vp_norm = torch.norm(vp, keepdim=True, dim=-1) + 1e-8 12 | vp = vp / vp_norm 13 | 14 | return vp 15 | 16 | 17 | def homography(X): 18 | 19 | dims = X.size() 20 | A_dims = tuple(list(dims)[:-2] + [9, 9]) 21 | 22 | A = torch.zeros(A_dims, device=X.device) 23 | 24 | A[..., 0:8:2, 5] = -1 25 | A[..., 0:8:2, 3:5] = -X[..., :, 0:2] 26 | A[..., 0:8:2, 6:8] = X[..., :, 0:2] 27 | A[..., 0:8:2, 6] *= X[..., :, 3] 28 | A[..., 0:8:2, 7] *= X[..., :, 3] 29 | A[..., 0:8:2, 8] = X[..., :, 3] 30 | A[..., 1:8:2, 2] = 1 31 | A[..., 1:8:2, 0:2] = X[..., :, 0:2] 32 | A[..., 1:8:2, 6:8] = -X[..., :, 0:2] 33 | A[..., 1:8:2, 6] *= X[..., :, 2] 34 | A[..., 1:8:2, 7] *= X[..., :, 2] 35 | A[..., 1:8:2, 8] = -X[..., :, 2] 36 | 37 | try: 38 | u, s, v = torch.svd(A) 39 | h = v[..., -1].to(X.device) 40 | except: 41 | h = torch.zeros(list(dims)[:-2] + [9], device=X.device) 42 | 43 | return h 44 | 45 | 46 | def fundamental_7point(X): 47 | 48 | # adapted from https://imkaywu.github.io/blog/2017/06/fundamental-matrix/ 49 | 50 | dims = X.size() 51 | A_dims = tuple(list(dims)[:-2] + [7, 9]) 52 | 53 | A = torch.ones(A_dims, device=X.device) 54 | 55 | A[..., 0:2] = X[..., 0:2] 56 | A[..., 0:2] *= X[..., 2][..., None] 57 | A[..., 3:5] = X[..., 0:2] 58 | A[..., 3:5] *= X[..., 3][..., None] 59 | A[..., 6:8] = X[..., 0:2] 60 | 61 | B = A.cpu() 62 | u, s, v = torch.svd(B) 63 | 64 | degenerate = s[..., -3].to(X.device).abs() < 1e-5 65 | 66 | fs = v[..., -2:].transpose(-1, -2).to(X.device) 67 | fs = fs / torch.clamp(torch.norm(fs, dim=-1, keepdim=True), min=1e-9) 68 | Fs = fs.contiguous().view(list(dims)[:-2] + [2, 3, 3]) 69 | 70 | D = torch.zeros(list(dims)[:-2] + [2, 2, 2], device=Fs.device) 71 | D_tmp = torch.zeros(list(dims)[:-2] + [3, 3], device=D.device) 72 | for i1 in range(2): 73 | for i2 in range(2): 74 | for i3 in range(2): 75 | D_tmp[..., 0] = Fs[..., i1, :, 0] 76 | D_tmp[..., 1] = Fs[..., i2, :, 1] 77 | D_tmp[..., 2] = Fs[..., i3, :, 2] 78 | D[..., i1, i2, i3] = torch.det(D_tmp) 79 | 80 | poly_coeffs = torch.zeros(list(dims)[:-2] + [4], device=X.device) 81 | poly_coeffs[..., 3] = -D[..., 1, 0, 0]+D[..., 0, 1, 1]+D[..., 0, 0, 0]+D[..., 1, 1, 0]+D[..., 1, 0, 1]-D[..., 0, 1, 0]-D[..., 0, 0, 1]-D[..., 1, 1, 1] 82 | poly_coeffs[..., 2] = D[..., 0, 0, 1]-2*D[..., 0, 1, 1]-2*D[..., 1, 0, 1]+D[..., 1, 0, 0]-2*D[..., 1, 1, 0]+D[..., 0, 1, 0]+3*D[..., 1, 1, 1] 83 | poly_coeffs[..., 1] = D[..., 1, 1, 0]+D[..., 0, 1, 1]+D[..., 1, 0, 1]-3*D[..., 1, 1, 1] 84 | poly_coeffs[..., 0] = D[..., 1, 1, 1] 85 | poly_coeffs = poly_coeffs / torch.clamp(torch.abs(poly_coeffs[..., 3][..., None]), min=1e-9) 86 | poly_coeffs = poly_coeffs * torch.sign(poly_coeffs[..., 3][..., None]) 87 | 88 | # find roots of polynomial det(a*F1 + (1-a)*F2) 89 | # https://en.wikipedia.org/wiki/Companion_matrix 90 | companion = torch.diag_embed(torch.ones(2, dtype=torch.float32, device=X.device), offset=-1) 91 | companion = companion.view([1]*(len(dims)-2) + [3, 3]).expand(list(dims[:-2]) + [3, 3]) 92 | companion = torch.clone(companion) 93 | companion[..., 0, -1] -= poly_coeffs[..., 0] 94 | companion[..., 1, -1] -= poly_coeffs[..., 1] 95 | companion[..., 2, -1] -= poly_coeffs[..., 2] 96 | 97 | eigvals_complex, _ = torch.linalg.eig(companion) 98 | eigvals = eigvals_complex.real.to(X.device) 99 | 100 | isnotreal = torch.logical_not(torch.isreal(eigvals_complex.to(X.device))) 101 | 102 | # compute residuals for all three solutions and select the best 103 | f = fs[..., 0, None, :] * eigvals[..., None] + fs[..., 1, None, :] * (1 - eigvals[..., None]) 104 | X_ = X[..., None, :, :].expand(list(dims)[:-2] + [3] + list(dims)[-2:]) 105 | res = utils.residual_functions.sampson_distance(X_, f) 106 | mean_res = res.mean(-1) + isnotreal.float() * res.max() * 10 107 | best_fs = torch.argmin(mean_res, dim=-1, keepdim=True)[..., None].expand(list(dims)[:-2] + [1, 9]) 108 | h = torch.gather(f, -2, best_fs).squeeze(-2) 109 | 110 | h = h * torch.logical_not(degenerate).float()[..., None] 111 | 112 | return h 113 | 114 | 115 | minimal_solver = {"vp": vanishing_point, "homography": homography, "fundamental": fundamental_7point} -------------------------------------------------------------------------------- /utils/tee.py: -------------------------------------------------------------------------------- 1 | import sys 2 | class Tee(object): 3 | def __init__(self, name, mode, file_only=False): 4 | self.file = open(name, mode) 5 | self.stdout = sys.stdout 6 | self.stderr = sys.stderr 7 | sys.stdout = self 8 | sys.stderr = self 9 | self.file_only = file_only 10 | 11 | def __del__(self): 12 | sys.stdout = self.stdout 13 | sys.stderr = self.stderr 14 | self.file.close() 15 | 16 | def write(self, data): 17 | self.file.write(data) 18 | if not self.file_only: 19 | self.stdout.write(data) 20 | 21 | def flush(self): 22 | self.file.flush() 23 | -------------------------------------------------------------------------------- /utils/visualisation.py: -------------------------------------------------------------------------------- 1 | import matplotlib.collections 2 | import matplotlib.pyplot as plt 3 | from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas 4 | import torch 5 | import numpy as np 6 | import wandb 7 | import utils.evaluation 8 | 9 | 10 | def rgb2gray(rgb): 11 | return np.dot(rgb[..., :3], [0.2989, 0.5870, 0.1140]) 12 | 13 | 14 | def generate_plot_single_batch(opt, X, log_inlier_weights, log_sample_weights, inlier_scores, clusters, labels, 15 | gt_models, gt_labels, image, dataset, max_num_samples=6, metrics=None): 16 | B, N, Mo = log_inlier_weights.size() 17 | M = opt.instances 18 | inlier_weights = torch.softmax(log_inlier_weights, dim=2) 19 | sample_weights = torch.softmax(log_sample_weights, dim=1) 20 | inlier_weights = inlier_weights / torch.max(inlier_weights, dim=1, keepdim=True)[0] 21 | sample_weights = sample_weights / torch.max(sample_weights, dim=1, keepdim=True)[0] 22 | 23 | batch = min(B, max_num_samples) 24 | 25 | num_cols = 4 if opt.problem == "vp" else 5 26 | 27 | if True: 28 | fig, all_axes = plt.subplots(nrows=Mo, ncols=batch * num_cols, figsize=(batch * 4 * num_cols, Mo * 4)) 29 | canvas = FigureCanvas(fig) 30 | plot_funs = {"vp": plot_vps, "homography": plot_fh, "fundamental": plot_fh} 31 | 32 | for mi, axes in enumerate(all_axes): 33 | for ai, ax in enumerate(axes): 34 | bi = int(ai / num_cols) 35 | 36 | x = X[bi].cpu().detach().numpy() 37 | 38 | size = 3 39 | 40 | if ai % num_cols == 0: 41 | w = inlier_weights[bi, :, mi].cpu().detach().numpy() 42 | c = '#1E88E5' 43 | if opt.problem == "vp": 44 | ax.set_facecolor('k') 45 | elif ai % num_cols == 1: 46 | w = sample_weights[bi, :, mi].cpu().detach().numpy() 47 | c = '#FFC107' 48 | if opt.problem == "vp": 49 | ax.set_facecolor('k') 50 | elif ai % num_cols == 2: 51 | if mi < M: 52 | w = inlier_scores[bi, mi].cpu().detach().numpy() 53 | else: 54 | if inlier_scores.size(1) == Mo: 55 | w = inlier_scores[bi, 0].cpu().detach().numpy() 56 | else: 57 | w = 0 58 | c = '#D81B60' 59 | 60 | if opt.problem == "vp": 61 | ax.set_facecolor('k') 62 | elif ai % num_cols == 3: 63 | w = clusters[bi, (mi + 1) % Mo].cpu().detach().numpy() 64 | 65 | c = '#D81B60' 66 | 67 | if opt.problem == "vp": 68 | ax.set_facecolor('k') 69 | elif ai % num_cols == 4: 70 | if mi < M: 71 | w = torch.where(gt_labels == (mi + 1), 72 | torch.ones(N, dtype=torch.float32, device=inlier_weights.device), 73 | torch.zeros(N, dtype=torch.float32, 74 | device=inlier_weights.device)).cpu().detach().numpy() 75 | else: 76 | w = torch.where(gt_labels == 0, 77 | torch.ones(N, dtype=torch.float32, device=inlier_weights.device), 78 | torch.zeros(N, dtype=torch.float32, 79 | device=inlier_weights.device)).cpu().detach().numpy() 80 | c = 'g' 81 | else: 82 | pass 83 | 84 | error = None 85 | if mi == 0 and ai == 0: 86 | try: 87 | error = metrics["misclassification_error"][-1] 88 | except: 89 | pass 90 | 91 | plot_funs[opt.problem](ax, x, w, c, error, size) 92 | 93 | ax.set_xticklabels([]) 94 | ax.set_xticks([]) 95 | ax.set_yticklabels([]) 96 | ax.set_yticks([]) 97 | plt.tight_layout() 98 | canvas.draw() 99 | plot_image = np.frombuffer(canvas.tostring_rgb(), dtype='uint8') 100 | plot_image = plot_image.reshape(fig.canvas.get_width_height()[::-1] + (3,)) 101 | 102 | plt.close() 103 | 104 | if not opt.dataset == "adelaide": 105 | 106 | labels_ = torch.clamp(labels.transpose(-1, -2) - 1, min=0) 107 | probs = torch.softmax(log_sample_weights, dim=1) 108 | probs = probs / torch.max(probs, dim=1)[0][:, None, :] 109 | sample_weights = torch.gather(probs[0], 1, labels_).detach().cpu().numpy()[:, 0] 110 | probs = torch.softmax(log_inlier_weights, dim=2) 111 | probs = probs / torch.max(probs, dim=2)[0][..., None] 112 | inlier_weights = torch.gather(probs[0], 1, labels_).detach().cpu().numpy()[:, 0] 113 | 114 | labels_ = labels - 1 115 | 116 | p1, p2 = dataset.denormalise(X) 117 | p1 = p1.cpu().numpy()[0] 118 | p2 = p2.cpu().numpy()[0] 119 | 120 | fig, axs = plt.subplots(nrows=1, ncols=5, figsize=(5 * 4, 4), ) 121 | canvas = FigureCanvas(fig) 122 | 123 | for ax_id in [0, 2, 3, 4]: 124 | axs[ax_id].set_facecolor('k') 125 | 126 | image_rgb = image[0].cpu().numpy() 127 | image_gray = rgb2gray(image_rgb) 128 | 129 | axs[0].imshow(image_rgb / 255.0) 130 | 131 | for ax in [axs[4]]: 132 | img = (image_gray.astype(float) / 255.0) * 0.7 + 0.3 133 | ax.imshow(img, cmap="Greys_r", vmin=0) 134 | img = np.zeros_like(img) 135 | axs[2].imshow(img, cmap="Greys_r", vmin=0, vmax=1) 136 | axs[3].imshow(img, cmap="Greys_r", vmin=0, vmax=1) 137 | 138 | colors_sw = [] 139 | colors_iw = [] 140 | colors_lb = [] 141 | hues = [338, 208, 45, 170, 99, 310, 255, 80, 190, 230, 120] 142 | lws = [] 143 | for ni in range(N): 144 | label = labels_[0, ni].item() 145 | if label == -1: 146 | hue = 0 147 | sat = 1 148 | val = 0 149 | val2 = 0 150 | lw = 0.1 151 | else: 152 | hue = hues[label % len(hues)] / 360.0 153 | sat = 1 154 | val = inlier_weights[ni] 155 | val2 = sample_weights[ni] 156 | lw = 4 157 | 158 | lws += [lw] 159 | color = matplotlib.colors.hsv_to_rgb([hue, sat, val]) 160 | colors_iw += [color] 161 | color = matplotlib.colors.hsv_to_rgb([hue, sat, val2]) 162 | colors_sw += [color] 163 | color = matplotlib.colors.hsv_to_rgb([hue, sat, 0 if label == -1 else 1]) 164 | colors_lb += [color] 165 | 166 | if opt.problem == "vp": 167 | 168 | lines = [((a[0], a[1]), (b[0], b[1])) for a, b, in zip(p1.tolist(), p2.tolist())] 169 | lines2 = [((a[0], -a[1]), (b[0], -b[1])) for a, b, in zip(p1.tolist(), p2.tolist())] 170 | 171 | axs[1].set_xlim(0, 512) 172 | axs[1].set_ylim(-512, 0) 173 | 174 | lw = 3 175 | lc1 = matplotlib.collections.LineCollection(lines2, linewidths=lw, colors='k') 176 | lc2 = matplotlib.collections.LineCollection(lines, linewidths=lw, colors=colors_sw) 177 | lc3 = matplotlib.collections.LineCollection(lines, linewidths=lw, colors=colors_iw) 178 | lc4 = matplotlib.collections.LineCollection(lines, linewidths=lws, colors=colors_lb) 179 | 180 | axs[1].add_collection(lc1) 181 | axs[2].add_collection(lc2) 182 | axs[3].add_collection(lc3) 183 | axs[4].add_collection(lc4) 184 | else: 185 | lw = 4 186 | axs[1].scatter(p1[:, 0], p1[:, 1], s=lw, c='k') 187 | lw = 4 188 | axs[2].scatter(p1[:, 0], p1[:, 1], s=lw, c=colors_sw) 189 | axs[3].scatter(p1[:, 0], p1[:, 1], s=lw, c=colors_iw) 190 | axs[4].scatter(p1[:, 0], p1[:, 1], s=lws, c=colors_lb) 191 | 192 | for ax in axs: 193 | ax.set_xlim(np.min(p1) - 10, np.max(p1) + 10) 194 | ax.set_ylim(np.max(p1) + 10, np.min(p1) - 10) 195 | 196 | for ax in axs: 197 | ax.set_xticklabels([]) 198 | ax.set_xticks([]) 199 | ax.set_yticklabels([]) 200 | ax.set_yticks([]) 201 | plt.tight_layout() 202 | canvas.draw() 203 | plot_image2 = np.frombuffer(canvas.tostring_rgb(), dtype='uint8') 204 | plot_image2 = plot_image2.reshape(fig.canvas.get_width_height()[::-1] + (3,)) 205 | plt.close() 206 | 207 | return plot_image, plot_image2 208 | 209 | 210 | def plot_fh(ax, X, weights, color, error, size): 211 | ax.set_xlim([-1, 1]) 212 | ax.set_ylim([-1, 1]) 213 | ax.set_aspect('equal') 214 | 215 | ax.scatter(X[:, 0], X[:, 1], s=1, c='k') 216 | ax.scatter(X[:, 0], X[:, 1], s=weights * 20, c=color) 217 | 218 | if error is not None: 219 | ax.set_title("%.1f" % error) 220 | 221 | 222 | def plot_vps(ax, X, weights, color, metrics, size): 223 | 224 | ax.set_xlim([-1, 1]) 225 | ax.set_ylim([-1, 1]) 226 | ax.set_aspect('equal') 227 | 228 | p1 = X[:, 0:3] 229 | p2 = X[:, 3:6] 230 | p1 = p1 / p1[:, 2][:, None] 231 | p2 = p2 / p2[:, 2][:, None] 232 | 233 | lines = [((a[0], -a[1]), (b[0], -b[1])) for a, b, in zip(p1.tolist(), p2.tolist())] 234 | 235 | if weights is None: 236 | colors = 'k' 237 | else: 238 | cmap = plt.get_cmap('bone') 239 | colors = cmap(weights) 240 | 241 | lc1 = matplotlib.collections.LineCollection(lines, linewidths=size, colors=colors) 242 | ax.add_collection(lc1) 243 | 244 | 245 | def save_visualisation_plots(opt, X, ranked_choices, log_probs_m, log_probs_n, inlier_scores, clusters, labels, 246 | gt_models, gt_labels, image, dataset, **kwargs): 247 | N = inlier_scores.size(-1) 248 | B, K, M, H = ranked_choices.size() 249 | 250 | log_probs_e = log_probs_m.transpose(-1, -2).contiguous().view(B, 1, M + 1, 1, N).expand(B, K, M + 1, H, N) 251 | ranked_log_probs_m = torch.gather(log_probs_e, 2, ranked_choices.view(B, K, M, H, 1).expand(B, K, M, H, N)) 252 | ranked_log_probs_m = torch.concatenate([ranked_log_probs_m, log_probs_e[..., -1, :, :].view(B, K, 1, H, N)], dim=2) 253 | log_probs_e = log_probs_n.transpose(-1, -2).contiguous().view(B, 1, M + 1, 1, N).expand(B, K, M + 1, H, N) 254 | ranked_log_probs_n = torch.gather(log_probs_e, 2, ranked_choices.view(B, K, M, H, 1).expand(B, K, M, H, N)) 255 | ranked_log_probs_n = torch.concatenate([ranked_log_probs_n, log_probs_e[..., -1, :, :].view(B, K, 1, H, N)], dim=2) 256 | 257 | plot_image, plot_image2 = generate_plot_single_batch(opt, X, ranked_log_probs_m[:, 0, :, 0].transpose(-1, -2), 258 | ranked_log_probs_n[:, 0, :, 0].transpose(-1, -2), inlier_scores[:, 0, :, 0], 259 | clusters[:, 0, 0], labels[:, 0, 0], gt_models, gt_labels, image, dataset, 260 | max_num_samples=6, **kwargs) 261 | 262 | assert False, "TODO: save plots to somewhere" 263 | 264 | 265 | 266 | def log_wandb(log_data, metrics, mode, epoch, images=[]): 267 | print_metrics = ["loss", "time", "total_time"] 268 | 269 | print_string = " | ".join([("%s: %.3f" % (key, np.mean(metrics[key]))) for key in print_metrics]) 270 | print(print_string) 271 | 272 | for key, value in metrics.items(): 273 | if "vp_err" in key: 274 | for cutoff in (1, 3, 5, 10): 275 | aucs = [] 276 | for errors in value: 277 | try: 278 | auc, plot_points = utils.evaluation.calc_auc(np.array(errors), cutoff=cutoff) 279 | except: 280 | auc = 0 281 | aucs += [auc] 282 | log_data["%s/%s_auc_%d_avg" % (mode, key, cutoff)] = np.mean(aucs) 283 | log_data["%s/%s_auc_%d_std" % (mode, key, cutoff)] = np.std(aucs) 284 | 285 | log_data["%s/%s_avg" % (mode, key)] = np.mean(value) 286 | log_data["%s/%s_std" % (mode, key)] = np.std(value) 287 | log_data["%s/%s_med" % (mode, key)] = np.median(value) 288 | else: 289 | try: 290 | log_data["%s/%s_avg" % (mode, key)] = np.mean(value) 291 | log_data["%s/%s_std" % (mode, key)] = np.std(value) 292 | log_data["%s/%s_med" % (mode, key)] = np.median(value) 293 | except: 294 | pass 295 | 296 | wandb.log(log_data, step=epoch) 297 | if len(images) > 0: 298 | for idx, wandb_image in enumerate(images): 299 | wandb.log({"%s/examples" % mode: wandb_image}, step=idx + 1) 300 | --------------------------------------------------------------------------------