Repository: microsoft/SceneLandmarkLocalization
Branch: 3dv24
Commit: cb1d86f982d1
Files: 26
Total size: 11.8 MB
Directory structure:
gitextract_ttpd4pib/
├── .github/
│ └── CODEOWNERS
├── CODE_OF_CONDUCT.md
├── LICENSE
├── LICENSE-CODE
├── README.md
├── SECURITY.md
└── src/
├── dataloader/
│ └── indoor6.py
├── inference.py
├── local_inference.py
├── local_training.py
├── main.py
├── models/
│ ├── blocks.py
│ ├── conv2d_layers.py
│ └── efficientlitesld.py
├── pretrained_efficientnetlite0.net
├── requirements.txt
├── run_inference.py
├── run_training.py
├── train.py
└── utils/
├── generate_visibility_depth_normal.py
├── heatmap.py
├── landmark_selection.py
├── merge_landmark_files.py
├── pnp.py
├── read_write_models.py
└── select_additional_landmarks.py
================================================
FILE CONTENTS
================================================
================================================
FILE: .github/CODEOWNERS
================================================
# These code reviewers should be added by default.
* @snsinha @omiksik @tien-d
================================================
FILE: CODE_OF_CONDUCT.md
================================================
# Microsoft Open Source Code of Conduct
This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/).
Resources:
- [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/)
- [Microsoft Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/)
- Contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with questions or concerns
================================================
FILE: LICENSE
================================================
Attribution 4.0 International
=======================================================================
Creative Commons Corporation ("Creative Commons") is not a law firm and
does not provide legal services or legal advice. Distribution of
Creative Commons public licenses does not create a lawyer-client or
other relationship. Creative Commons makes its licenses and related
information available on an "as-is" basis. Creative Commons gives no
warranties regarding its licenses, any material licensed under their
terms and conditions, or any related information. Creative Commons
disclaims all liability for damages resulting from their use to the
fullest extent possible.
Using Creative Commons Public Licenses
Creative Commons public licenses provide a standard set of terms and
conditions that creators and other rights holders may use to share
original works of authorship and other material subject to copyright
and certain other rights specified in the public license below. The
following considerations are for informational purposes only, are not
exhaustive, and do not form part of our licenses.
Considerations for licensors: Our public licenses are
intended for use by those authorized to give the public
permission to use material in ways otherwise restricted by
copyright and certain other rights. Our licenses are
irrevocable. Licensors should read and understand the terms
and conditions of the license they choose before applying it.
Licensors should also secure all rights necessary before
applying our licenses so that the public can reuse the
material as expected. Licensors should clearly mark any
material not subject to the license. This includes other CC-
licensed material, or material used under an exception or
limitation to copyright. More considerations for licensors:
wiki.creativecommons.org/Considerations_for_licensors
Considerations for the public: By using one of our public
licenses, a licensor grants the public permission to use the
licensed material under specified terms and conditions. If
the licensor's permission is not necessary for any reason--for
example, because of any applicable exception or limitation to
copyright--then that use is not regulated by the license. Our
licenses grant only permissions under copyright and certain
other rights that a licensor has authority to grant. Use of
the licensed material may still be restricted for other
reasons, including because others have copyright or other
rights in the material. A licensor may make special requests,
such as asking that all changes be marked or described.
Although not required by our licenses, you are encouraged to
respect those requests where reasonable. More_considerations
for the public:
wiki.creativecommons.org/Considerations_for_licensees
=======================================================================
Creative Commons Attribution 4.0 International Public License
By exercising the Licensed Rights (defined below), You accept and agree
to be bound by the terms and conditions of this Creative Commons
Attribution 4.0 International Public License ("Public License"). To the
extent this Public License may be interpreted as a contract, You are
granted the Licensed Rights in consideration of Your acceptance of
these terms and conditions, and the Licensor grants You such rights in
consideration of benefits the Licensor receives from making the
Licensed Material available under these terms and conditions.
Section 1 -- Definitions.
a. Adapted Material means material subject to Copyright and Similar
Rights that is derived from or based upon the Licensed Material
and in which the Licensed Material is translated, altered,
arranged, transformed, or otherwise modified in a manner requiring
permission under the Copyright and Similar Rights held by the
Licensor. For purposes of this Public License, where the Licensed
Material is a musical work, performance, or sound recording,
Adapted Material is always produced where the Licensed Material is
synched in timed relation with a moving image.
b. Adapter's License means the license You apply to Your Copyright
and Similar Rights in Your contributions to Adapted Material in
accordance with the terms and conditions of this Public License.
c. Copyright and Similar Rights means copyright and/or similar rights
closely related to copyright including, without limitation,
performance, broadcast, sound recording, and Sui Generis Database
Rights, without regard to how the rights are labeled or
categorized. For purposes of this Public License, the rights
specified in Section 2(b)(1)-(2) are not Copyright and Similar
Rights.
d. Effective Technological Measures means those measures that, in the
absence of proper authority, may not be circumvented under laws
fulfilling obligations under Article 11 of the WIPO Copyright
Treaty adopted on December 20, 1996, and/or similar international
agreements.
e. Exceptions and Limitations means fair use, fair dealing, and/or
any other exception or limitation to Copyright and Similar Rights
that applies to Your use of the Licensed Material.
f. Licensed Material means the artistic or literary work, database,
or other material to which the Licensor applied this Public
License.
g. Licensed Rights means the rights granted to You subject to the
terms and conditions of this Public License, which are limited to
all Copyright and Similar Rights that apply to Your use of the
Licensed Material and that the Licensor has authority to license.
h. Licensor means the individual(s) or entity(ies) granting rights
under this Public License.
i. Share means to provide material to the public by any means or
process that requires permission under the Licensed Rights, such
as reproduction, public display, public performance, distribution,
dissemination, communication, or importation, and to make material
available to the public including in ways that members of the
public may access the material from a place and at a time
individually chosen by them.
j. Sui Generis Database Rights means rights other than copyright
resulting from Directive 96/9/EC of the European Parliament and of
the Council of 11 March 1996 on the legal protection of databases,
as amended and/or succeeded, as well as other essentially
equivalent rights anywhere in the world.
k. You means the individual or entity exercising the Licensed Rights
under this Public License. Your has a corresponding meaning.
Section 2 -- Scope.
a. License grant.
1. Subject to the terms and conditions of this Public License,
the Licensor hereby grants You a worldwide, royalty-free,
non-sublicensable, non-exclusive, irrevocable license to
exercise the Licensed Rights in the Licensed Material to:
a. reproduce and Share the Licensed Material, in whole or
in part; and
b. produce, reproduce, and Share Adapted Material.
2. Exceptions and Limitations. For the avoidance of doubt, where
Exceptions and Limitations apply to Your use, this Public
License does not apply, and You do not need to comply with
its terms and conditions.
3. Term. The term of this Public License is specified in Section
6(a).
4. Media and formats; technical modifications allowed. The
Licensor authorizes You to exercise the Licensed Rights in
all media and formats whether now known or hereafter created,
and to make technical modifications necessary to do so. The
Licensor waives and/or agrees not to assert any right or
authority to forbid You from making technical modifications
necessary to exercise the Licensed Rights, including
technical modifications necessary to circumvent Effective
Technological Measures. For purposes of this Public License,
simply making modifications authorized by this Section 2(a)
(4) never produces Adapted Material.
5. Downstream recipients.
a. Offer from the Licensor -- Licensed Material. Every
recipient of the Licensed Material automatically
receives an offer from the Licensor to exercise the
Licensed Rights under the terms and conditions of this
Public License.
b. No downstream restrictions. You may not offer or impose
any additional or different terms or conditions on, or
apply any Effective Technological Measures to, the
Licensed Material if doing so restricts exercise of the
Licensed Rights by any recipient of the Licensed
Material.
6. No endorsement. Nothing in this Public License constitutes or
may be construed as permission to assert or imply that You
are, or that Your use of the Licensed Material is, connected
with, or sponsored, endorsed, or granted official status by,
the Licensor or others designated to receive attribution as
provided in Section 3(a)(1)(A)(i).
b. Other rights.
1. Moral rights, such as the right of integrity, are not
licensed under this Public License, nor are publicity,
privacy, and/or other similar personality rights; however, to
the extent possible, the Licensor waives and/or agrees not to
assert any such rights held by the Licensor to the limited
extent necessary to allow You to exercise the Licensed
Rights, but not otherwise.
2. Patent and trademark rights are not licensed under this
Public License.
3. To the extent possible, the Licensor waives any right to
collect royalties from You for the exercise of the Licensed
Rights, whether directly or through a collecting society
under any voluntary or waivable statutory or compulsory
licensing scheme. In all other cases the Licensor expressly
reserves any right to collect such royalties.
Section 3 -- License Conditions.
Your exercise of the Licensed Rights is expressly made subject to the
following conditions.
a. Attribution.
1. If You Share the Licensed Material (including in modified
form), You must:
a. retain the following if it is supplied by the Licensor
with the Licensed Material:
i. identification of the creator(s) of the Licensed
Material and any others designated to receive
attribution, in any reasonable manner requested by
the Licensor (including by pseudonym if
designated);
ii. a copyright notice;
iii. a notice that refers to this Public License;
iv. a notice that refers to the disclaimer of
warranties;
v. a URI or hyperlink to the Licensed Material to the
extent reasonably practicable;
b. indicate if You modified the Licensed Material and
retain an indication of any previous modifications; and
c. indicate the Licensed Material is licensed under this
Public License, and include the text of, or the URI or
hyperlink to, this Public License.
2. You may satisfy the conditions in Section 3(a)(1) in any
reasonable manner based on the medium, means, and context in
which You Share the Licensed Material. For example, it may be
reasonable to satisfy the conditions by providing a URI or
hyperlink to a resource that includes the required
information.
3. If requested by the Licensor, You must remove any of the
information required by Section 3(a)(1)(A) to the extent
reasonably practicable.
4. If You Share Adapted Material You produce, the Adapter's
License You apply must not prevent recipients of the Adapted
Material from complying with this Public License.
Section 4 -- Sui Generis Database Rights.
Where the Licensed Rights include Sui Generis Database Rights that
apply to Your use of the Licensed Material:
a. for the avoidance of doubt, Section 2(a)(1) grants You the right
to extract, reuse, reproduce, and Share all or a substantial
portion of the contents of the database;
b. if You include all or a substantial portion of the database
contents in a database in which You have Sui Generis Database
Rights, then the database in which You have Sui Generis Database
Rights (but not its individual contents) is Adapted Material; and
c. You must comply with the conditions in Section 3(a) if You Share
all or a substantial portion of the contents of the database.
For the avoidance of doubt, this Section 4 supplements and does not
replace Your obligations under this Public License where the Licensed
Rights include other Copyright and Similar Rights.
Section 5 -- Disclaimer of Warranties and Limitation of Liability.
a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE
EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS
AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF
ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS,
IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION,
WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR
PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS,
ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT
KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT
ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU.
b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE
TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION,
NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT,
INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES,
COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR
USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN
ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR
DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR
IN PART, THIS LIMITATION MAY NOT APPLY TO YOU.
c. The disclaimer of warranties and limitation of liability provided
above shall be interpreted in a manner that, to the extent
possible, most closely approximates an absolute disclaimer and
waiver of all liability.
Section 6 -- Term and Termination.
a. This Public License applies for the term of the Copyright and
Similar Rights licensed here. However, if You fail to comply with
this Public License, then Your rights under this Public License
terminate automatically.
b. Where Your right to use the Licensed Material has terminated under
Section 6(a), it reinstates:
1. automatically as of the date the violation is cured, provided
it is cured within 30 days of Your discovery of the
violation; or
2. upon express reinstatement by the Licensor.
For the avoidance of doubt, this Section 6(b) does not affect any
right the Licensor may have to seek remedies for Your violations
of this Public License.
c. For the avoidance of doubt, the Licensor may also offer the
Licensed Material under separate terms or conditions or stop
distributing the Licensed Material at any time; however, doing so
will not terminate this Public License.
d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
License.
Section 7 -- Other Terms and Conditions.
a. The Licensor shall not be bound by any additional or different
terms or conditions communicated by You unless expressly agreed.
b. Any arrangements, understandings, or agreements regarding the
Licensed Material not stated herein are separate from and
independent of the terms and conditions of this Public License.
Section 8 -- Interpretation.
a. For the avoidance of doubt, this Public License does not, and
shall not be interpreted to, reduce, limit, restrict, or impose
conditions on any use of the Licensed Material that could lawfully
be made without permission under this Public License.
b. To the extent possible, if any provision of this Public License is
deemed unenforceable, it shall be automatically reformed to the
minimum extent necessary to make it enforceable. If the provision
cannot be reformed, it shall be severed from this Public License
without affecting the enforceability of the remaining terms and
conditions.
c. No term or condition of this Public License will be waived and no
failure to comply consented to unless expressly agreed to by the
Licensor.
d. Nothing in this Public License constitutes or may be interpreted
as a limitation upon, or waiver of, any privileges and immunities
that apply to the Licensor or You, including from the legal
processes of any jurisdiction or authority.
=======================================================================
Creative Commons is not a party to its public
licenses. Notwithstanding, Creative Commons may elect to apply one of
its public licenses to material it publishes and in those instances
will be considered the “Licensor.” The text of the Creative Commons
public licenses is dedicated to the public domain under the CC0 Public
Domain Dedication. Except for the limited purpose of indicating that
material is shared under a Creative Commons public license or as
otherwise permitted by the Creative Commons policies published at
creativecommons.org/policies, Creative Commons does not authorize the
use of the trademark "Creative Commons" or any other trademark or logo
of Creative Commons without its prior written consent including,
without limitation, in connection with any unauthorized modifications
to any of its public licenses or any other arrangements,
understandings, or agreements concerning use of licensed material. For
the avoidance of doubt, this paragraph does not form part of the
public licenses.
Creative Commons may be contacted at creativecommons.org.
================================================
FILE: LICENSE-CODE
================================================
MIT License
Copyright (c) Microsoft Corporation.
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE
================================================
FILE: README.md
================================================
# Scene Landmark Detection for Camera Localization
## Introduction
We have devised a new method to detect scene-specific _scene landmarks_ for localizing a camera within a pre-mapped scene. Our method is privacy-preserving, has low storage requirements and achieves high accuracy. **[Left]** Scene landmarks detected in a query image. **[Middle]** A CNN-based heatmap prediction architecture is trained. **[Right]** The 3D scene
landmarks (_in red_) and the estimated camera pose (_in blue_) are shown overlaid over the 3D point cloud (_in gray_). The 3D point
cloud is shown only for visualization. It is not actually used for camera localization.
---
## Papers
**Improved Scene Landmark Detection for Camera Localization**
Tien Do and Sudipta N. Sinha
International Conference on 3D Vision (**3DV**), 2024
[pdf](paper/DoSinha3DV2024.pdf)
**Learning to Detect Scene Landmarks for Camera Localization**
Tien Do, Ondrej Miksik, Joseph DeGol, Hyun Soo Park, and Sudipta N. Sinha
IEEE/CVF Conference on Computer Vision and Pattern Recognition (**CVPR**), 2022
[pdf](paper/DoEtalCVPR2022.pdf) [video](https://www.youtube.com/watch?v=HM2yLCLz5nY)
**Indoor6 Dataset**
[download](https://drive.google.com/drive/folders/1w7Adnd6MXmNOacT072JnQ6emHUeLrD71?usp=drive_link)
## Bibtex
If you find our work to be useful in your research, please consider citing our paper:
```
@InProceedings{Do_Sinha_2024_ImprovedSceneLandmarkLoc,
author = {Do, Tien and Sinha, Sudipta N.},
title = {Improved Scene Landmark Detection for Camera Localization},
booktitle = {Proceedings of the International Conference on 3D Vision (3DV)},
month = {March},
year = {2024}
}
@InProceedings{Do_2022_SceneLandmarkLoc,
author = {Do, Tien and Miksik, Ondrej and DeGol, Joseph and Park, Hyun Soo and Sinha, Sudipta N.},
title = {Learning to Detect Scene Landmarks for Camera Localization},
booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
month = {June},
year = {2022}
}
```
# Indoor-6 Dataset
The Indoor-6 dataset was created from multiple sessions captured in six indoor scenes over multiple days. The pseudo
ground truth (pGT) 3D point clouds and camera poses for each scene are computed using [COLMAP](https://colmap.github.io/). All training data uses only colmap reconstruction from training images. The figure below
shows the camera poses (in red) and point clouds (in gray) and for each scene, the number of video and images in the
training and test split respectively. Compared to [7-scenes](https://www.microsoft.com/en-us/research/project/rgb-d-dataset-7-scenes/), the scenes in Indoor-6 are larger, have multiple rooms,
contains illumination variations as the images span multiple days and different times of day.
Indoor-6 dataset SfM reconstructions. Train/val/test splits and download urls per scene are listed below:
* [scene1](https://drive.google.com/file/d/1AJhPh9nnZO0HJyxuXXZdtKtA7kFRi3LQ/view?usp=drive_link) (6289/798/799 images)
* scene2 (3021/283/284 images)
* [scene2a](https://drive.google.com/file/d/1DgTQ7fflZJ7DdbHDRZF-6gXdB_vJF7fY/view?usp=drive_link) (4890/256/257 images)
* [scene3](https://drive.google.com/file/d/12aER7rQkvGS_DPeugTHo_Ma_Fi7JuflS/view?usp=drive_link) (4181/313/315 images)
* scene4 (1942/272/272 images)
* [scene4a](https://drive.google.com/file/d/1gibneq5ixZ0lmeNAYTmY4Mh8a244T2nl/view?usp=drive_link) (2285/158/158 images)
* [scene5](https://drive.google.com/file/d/18wHn_69-eV22N4I8R0rWQkcSQ3EtCYMX/view?usp=drive_link) (4946/512/424 images)
* [scene6](https://drive.google.com/file/d/1mZYnoKo37KXRjREK5CKs5IzDox2G3Prt/view?usp=drive_link) (1761/322/323 images)
* [colmap](https://drive.google.com/file/d/1oMo552DYo2U5Fvjm5MrTYPMqpMjXEf7m/view?usp=drive_link) (colmap reconstructions for all scenes.)
**Note**: We added two new scenes (`scene2a` and `scene4a`) to the Indoor-6 dataset after our CVPR 2022 paper was published. This was because we were unable to release `scene2` and `scene4` from the original dataset due to privacy reasons.
The two new scenes have been included as replacements. Please refer to our 3DV 2024 paper for a quantitative evaluation of our method and several baselines on the latest version of the dataset.
# Source code
The repository contains all the source code for our project. The most recent version can be found in the `3dv24` git branch (which is now the default branch of the repository). The best performing pretrained models for `SLD-star` as proposed in our 3DV 2024 paper are also available (see below). It significantly outperforms the `SLD+NBE` approach proposed in our CVPR 2022 paper. The source code for the `SLD+NBE` method is not maintained anymore. The older version of the code (pre 3DV 2024) can be found in the `main` branch.
## Environment Setup
```
pip install -r requirements.txt
```
* Python 3.9.13 on Windows 11.
* CUDA version: release 11.8 (V11.8.89)
* PyTorch version: 2.1.0+cu118
For development purposes, training was tested to run on both CUDA and CPU on both Linux and Windows platforms, as well as using the latest experimental version of pyTorch with Metal Performance Shaders on Mac OS X (see below).
By default the code will select hardware acceleration for your device, if available.
### Experimental Mac OS Metal Performance Shaders (MPS)
To enable the MPS backend, make sure you are running the latest Apple Silicon compatible hardware and follow [these instructions](https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/) to get the latest Nightly build of pyTorch instead.
_NOTE_: MPS has max supported precision of FP32.
## Layout
The source code expects the following directory structure (currently in your home directory).
```
└── data
| └── outputs
| └── checkpoints
| └── indoor6
| └── scene1
| └── scene2a
| └── scene3
| └── scene4a
| └── scene5
| └── scene6
└── SceneLandmarkLocalization
└── src
└── README.md (this)
```
* Download the indoor6 dataset and place the contents in the `/data/indoor6/` folder, as indicated above.
* Download the pretrained models for `SLD-star` (see below) from our 3DV 2024 paper and place them in the `/data/checkpoints` folder, as indicated above.
[pretrained models](https://drive.google.com/file/d/1s8bUgAuy2LX4QMcKE8yKz6JRyhL3JgxZ/view?usp=drive_link)
* Clone this repo into `/SceneLandmarkLocalization`.
* Finally, create the folder `/data/outputs` for storing trained models and other files that will be created when training your own models using the training routine.
## Running Inference using Pre-trained Models
Instructions to test the `SLD-star` models from our 3DV 2024 paper are listed below.
**Step 1.** First, verify the contents of the checkpoints folder. You should see the following files and directories.
```
└── data
└── checkpoints
└── scene1_1000-125_v10
└── scene1_1000-125_v10.txt
└── scene2a_1000-125_v10
└── scene2a_1000-125_v10.txt
└── scene3_1000-125_v10
└── scene3_1000-125_v10.txt
└── scene4a_1000-125_v10
└── scene4a_1000-125_v10.txt
└── scene5_1000-125_v10
└── scene5_1000-125_v10.txt
└── scene6_1000-125_v10
└── scene6_1000-125_v10.txt
```
**Step 2.** For `1000-125_v10`, each scene has eight model checkpoints. For example, `scene6` has these files.
```
└── scene6_1000-125_v10
└── scene6-000-125
└── model-best_median.ckpt
└── scene6-125-250
└── model-best_median.ckpt
└── scene6-250-375
└── model-best_median.ckpt
└── scene6-375-500
└── model-best_median.ckpt
└── scene6-500-625
└── model-best_median.ckpt
└── scene6-625-750
└── model-best_median.ckpt
└── scene6-750-875
└── model-best_median.ckpt
└── scene6-875-1000
└── model-best_median.ckpt
```
**Step 3.** Each experiment file for the `1000-125_v10` experiment, for e.g. `scene6_1000-125_v10.txt` contains eight lines, one for each model checkpoint (or landmark subset). Each line contains various attributes for the associated model.
**Step 4.** Check the Python script `/SceneLandmarkLocalization/src/run_inference.py`. The relative paths hardcoded in the variables `checkpoint_dir` and `dataset_dir` both assume the directory layout that was described earlier. The variable `experiment` is set to `1000-125_v10` which corresponds to the `SLD-star` model trained for 1000 landmarks partitioned into eight subsets each with 125 landmarks. The suffix `v10` is a tag to keep track of the experiment and generated model checkpoints.
**Step 5.** Now, run the following script.
```
cd SceneLandmarkLocalization/src
python run_inference.py
```
**Step 6.** When the script finishes running, the following text will be displayed on the console. The final accuracy (5cm/5deg recall) in percent is printed alongwith the mean inference speed.
**Step 7.** The metrics are also written to the file `/data/checkpoints/RESULTS-1000-125_v10.txt`. Note that, `1000-125_v10` is the experiment name specified in the `run_inference.py` script.
## Training Models
We now discuss how to train an `SLD-star` model ensemble.
As proposed in our 3DV 2024 paper, the model ensemble is a set of models that share the same architecture (derived from an EfficientNet backbone), but have independent sets of model parameters.
Each model (or network) in the ensemble is trained on a different subset of scene landmarks.
In our implementation, we define the subsets by considering the ordered list of all the scene landmarks and partitioning that list into blocks of fixed size. For convenience, we choose block sizes that exactly divide the total number of landmarks to ensure that all the subsets have the same size.
For example, given 1000 scene landmarks and choosing a block size of 125, we will obtain eight subsets. The first subset will consist of landmarks with indices in the range `[0,125]` in the ordered list.
The second subset will have landmarks with indices in the range `[125,250]` and so on.
We will now discuss how to run the training code.
**Step 1.** Now, run the following script.
To train a single model in the ensemble (for a specific scene), you might need to edit certain variables and modify the default values hardcoded in the `SceneLandmarkLocalization/src/run_training.py` script.
Then, just run it as follows.
```
cd SceneLandmarkLocalization/src
python run_training.py
```
**Step 2.** Editing the script and modifying the parameter values.
The important hyperparameters and settings that might need to be modified are the follows.
1. ***Paths:*** The default values for the dataset path and output paths are as follows (based on the assumed directory structure). However, these can be modified as needed.
```
dataset_dir = '../../data/indoor6'
output_dir = '../../data/outputs'
```
2. ***Scene ID and landmarks:*** The names of the landmark and visibility files.
```
scene_name = 'scene6'
landmark_config = 'landmarks/landmarks-1000v10'
visibility_config = 'landmarks/visibility-1000v10_depth_normal'
```
3. ***Ensemble configuration:*** The number of landmarks and the block size of the ensemble. `subset_index` indicates which network within the ensemble will be trained. So in the following example, the value `0` indicates that the model will be trained for the landmarks in the index range of `[0,125]`. So for this `1000-125` ensemble, you will need to change `subset_index` to `1, 2, ..., 7` to train all eight networks.
```
num_landmarks = 1000
block_size = 125
subset_index = 0
```
4. ***Version No.:*** A string tag which is appended to the generated model names and experiment files. This helps us avoid nameclashes when training and testing multiple sets of models.
**Step 3.** When training completes, check the output directory, you should see a directory that contains the model checkpoint for the specified scene. There will also be an experiment text file with the same name.
Inside the scene directory are sub-directories, one for each network in the ensemble. For example, the subdirectories for the `1000-125` ensemble for `scene6` will be named as `scene6-000-125`, `scene6-125-250` and so on.
Look inside these subdirectories for the model checkpoint file `model-best_median.ckpt`.
# Contributing
This project welcomes contributions and suggestions. Most contributions require you to agree to a
Contributor License Agreement (CLA) declaring that you have the right to, and actually do, grant us
the rights to use your contribution. For details, visit https://cla.opensource.microsoft.com.
When you submit a pull request, a CLA bot will automatically determine whether you need to provide
a CLA and decorate the PR appropriately (e.g., status check, comment). Simply follow the instructions
provided by the bot. You will only need to do this once across all repos using our CLA.
This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/).
For more information see the [Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/) or
contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with any additional questions or comments.
# Legal Notices
Microsoft and any contributors grant you a license to the Microsoft documentation and other content
in this repository under the [Creative Commons Attribution 4.0 International Public License](https://creativecommons.org/licenses/by/4.0/legalcode),
see the [LICENSE](LICENSE) file, and grant you a license to any code in the repository under the [MIT License](https://opensource.org/licenses/MIT), see the
[LICENSE-CODE](LICENSE-CODE) file.
Microsoft, Windows, Microsoft Azure and/or other Microsoft products and services referenced in the documentation
may be either trademarks or registered trademarks of Microsoft in the United States and/or other countries.
The licenses for this project do not grant you rights to use any Microsoft names, logos, or trademarks.
Microsoft's general trademark guidelines can be found at http://go.microsoft.com/fwlink/?LinkID=254653.
Privacy information can be found at https://privacy.microsoft.com/en-us/
Microsoft and any contributors reserve all other rights, whether under their respective copyrights, patents,
or trademarks, whether by implication, estoppel or otherwise.
================================================
FILE: SECURITY.md
================================================
## Security
Microsoft takes the security of our software products and services seriously, which includes all source code repositories managed through our GitHub organizations, which include [Microsoft](https://github.com/Microsoft), [Azure](https://github.com/Azure), [DotNet](https://github.com/dotnet), [AspNet](https://github.com/aspnet), [Xamarin](https://github.com/xamarin), and [our GitHub organizations](https://opensource.microsoft.com/).
If you believe you have found a security vulnerability in any Microsoft-owned repository that meets [Microsoft's definition of a security vulnerability](https://docs.microsoft.com/en-us/previous-versions/tn-archive/cc751383(v=technet.10)), please report it to us as described below.
## Reporting Security Issues
**Please do not report security vulnerabilities through public GitHub issues.**
Instead, please report them to the Microsoft Security Response Center (MSRC) at [https://msrc.microsoft.com/create-report](https://msrc.microsoft.com/create-report).
If you prefer to submit without logging in, send email to [secure@microsoft.com](mailto:secure@microsoft.com). If possible, encrypt your message with our PGP key; please download it from the [Microsoft Security Response Center PGP Key page](https://www.microsoft.com/en-us/msrc/pgp-key-msrc).
You should receive a response within 24 hours. If for some reason you do not, please follow up via email to ensure we received your original message. Additional information can be found at [microsoft.com/msrc](https://www.microsoft.com/msrc).
Please include the requested information listed below (as much as you can provide) to help us better understand the nature and scope of the possible issue:
* Type of issue (e.g. buffer overflow, SQL injection, cross-site scripting, etc.)
* Full paths of source file(s) related to the manifestation of the issue
* The location of the affected source code (tag/branch/commit or direct URL)
* Any special configuration required to reproduce the issue
* Step-by-step instructions to reproduce the issue
* Proof-of-concept or exploit code (if possible)
* Impact of the issue, including how an attacker might exploit the issue
This information will help us triage your report more quickly.
If you are reporting for a bug bounty, more complete reports can contribute to a higher bounty award. Please visit our [Microsoft Bug Bounty Program](https://microsoft.com/msrc/bounty) page for more details about our active programs.
## Preferred Languages
We prefer all communications to be in English.
## Policy
Microsoft follows the principle of [Coordinated Vulnerability Disclosure](https://www.microsoft.com/en-us/msrc/cvd).
================================================
FILE: src/dataloader/indoor6.py
================================================
import argparse
import copy
import fnmatch
import numpy as np
import os
import pickle
from PIL import Image
import sys
sys.path.append('../utils')
import torch
from torch.utils.data.dataset import Dataset
from torch.utils.data import DataLoader
from torchvision import transforms
from utils.pnp import Quaternion2Rotation
np.random.seed(0)
class Indoor6(Dataset):
def __init__(self, root_folder="",
scene_id='', mode='all',
landmark_idx=[None], skip_image_index=1,
input_image_downsample=1, gray_image_output=False,
landmark_config='landmarks/landmarks-50',
visibility_config='landmarks/visibility-50',
use_precomputed_focal_length=False):
super(Indoor6, self).__init__()
self.to_tensor = transforms.ToTensor()
self.image_folder = os.path.join(root_folder,
scene_id,
'images')
image_files_all = fnmatch.filter(os.listdir(self.image_folder), '*.color.jpg')
image_files_all = sorted(image_files_all)[::skip_image_index]
self.image_files = []
if mode == 'train':
self.image_files = \
pickle.load(open('%s/%s/train_test_val.pkl' % (root_folder, scene_id), 'rb'))[
'train'][::skip_image_index]
self.image_indices = \
pickle.load(open('%s/%s/train_test_val.pkl' % (root_folder, scene_id), 'rb'))[
'train_idx'][::skip_image_index]
elif mode == 'test':
self.image_files = \
pickle.load(open('%s/%s/train_test_val.pkl' % (root_folder, scene_id), 'rb'))[
'test'][::skip_image_index]
self.image_indices = \
pickle.load(open('%s/%s/train_test_val.pkl' % (root_folder, scene_id), 'rb'))[
'test_idx'][::skip_image_index]
elif mode == 'val':
self.image_files = \
pickle.load(open('%s/%s/train_test_val.pkl' % (root_folder, scene_id), 'rb'))[
'val'][::skip_image_index]
self.image_indices = \
pickle.load(open('%s/%s/train_test_val.pkl' % (root_folder, scene_id), 'rb'))[
'val_idx'][::skip_image_index]
else:
self.image_files = image_files_all
self.image_indices = np.arange(0, len(image_files_all))
self.image_indices = np.asarray(self.image_indices)
self.num_images = len(self.image_files)
self.gray_image_output = gray_image_output
self.mode = mode
landmark_file = open(root_folder + '/' + scene_id
+ '/%s.txt' % landmark_config, 'r')
num_landmark = int(landmark_file.readline())
self.landmark = []
for l in range(num_landmark):
pl = landmark_file.readline().split()
pl = np.array([float(pl[i]) for i in range(len(pl))])
self.landmark.append(pl)
self.landmark = np.asarray(self.landmark)[:, 1:]
self.image_downsampled = input_image_downsample
visibility_file = root_folder + '/' + scene_id + '/%s.txt' % visibility_config
self.visibility = np.loadtxt(visibility_file).astype(bool)
if landmark_idx[0] != None:
self.landmark = self.landmark[landmark_idx]
self.visibility = self.visibility[landmark_idx]
self.landmark = self.landmark.transpose()
## Precomputed fixed focal length
self.precomputed_focal_length = None
if use_precomputed_focal_length:
PRECOMPUTED_FOCAL_LENGTH = {'scene1': 900, 'scene2a': 1100, 'scene3': 900, 'scene4a': 900, 'scene5': 900, 'scene6': 900}
self.precomputed_focal_length = PRECOMPUTED_FOCAL_LENGTH[scene_id]
def original_image_name(self, index):
intrinsics = open(os.path.join(self.image_folder,
self.image_files[index].replace('color.jpg', 'intrinsics.txt')))
intrinsics = intrinsics.readline().split()
return intrinsics[6]
def _modify_intrinsic(self, index, use_precomputed_focal_length=False):
W = None
H = None
K = None
K_inv = None
while K_inv is None:
try:
intrinsics = open(os.path.join(self.image_folder,
self.image_files[index].replace('color.jpg', 'intrinsics.txt')))
intrinsics = intrinsics.readline().split()
W = int(intrinsics[0]) // (self.image_downsampled * 32) * 32
H = int(intrinsics[1]) // (self.image_downsampled * 32) * 32
scale_factor_x = W / float(intrinsics[0])
scale_factor_y = H / float(intrinsics[1])
if use_precomputed_focal_length:
fx = self.precomputed_focal_length * scale_factor_x
fy = self.precomputed_focal_length * scale_factor_y
else:
fx = float(intrinsics[2]) * scale_factor_x
fy = float(intrinsics[2]) * scale_factor_y
cx = float(intrinsics[3]) * scale_factor_x
cy = float(intrinsics[4]) * scale_factor_y
K = np.array([[fx, 0., cx],
[0., fy, cy],
[0., 0., 1.]], dtype=float)
K_inv = np.linalg.inv(K)
except(RuntimeError, TypeError, NameError):
pass
return K, K_inv, W, H
def _load_and_resize_image(self, index, W, H):
color_img_rs = None
while color_img_rs is None:
try:
# Load color image
color_img = Image.open(os.path.join(self.image_folder, self.image_files[index]))
color_img_rs = color_img.resize((W, H), resample=Image.BILINEAR)
except(RuntimeError, TypeError, NameError):
pass
color_tensor = self.to_tensor(color_img_rs)
return color_tensor
def _load_pose(self, index):
pose = None
while pose is None:
try:
# Load 3x4 pose matrix and make it 4x4 by appending vector [0., 0., 0., 1.]
pose = np.loadtxt(os.path.join(self.image_folder, self.image_files[index].replace('color.jpg', 'pose.txt')))
except (RuntimeError, TypeError, NameError):
pass
pose_s = np.vstack((pose, np.array([0., 0., 0., 1.])))
return pose_s
def __getitem__(self, index):
K, K_inv, W_modified, H_modified = self._modify_intrinsic(index, use_precomputed_focal_length=False if self.precomputed_focal_length is None else True)
color_tensor = self._load_and_resize_image(index, W_modified, H_modified)
C_T_G = self._load_pose(index)
landmark3d = C_T_G @ np.vstack((self.landmark, np.ones((1, self.landmark.shape[1]))))
output = {'pose_gt': torch.tensor(C_T_G),
'image': color_tensor,
'intrinsics': torch.tensor(K, dtype=torch.float32, requires_grad=False),
'inv_intrinsics': torch.tensor(K_inv, dtype=torch.float32, requires_grad=False),
'landmark3d': torch.tensor(landmark3d[:3], dtype=torch.float32, requires_grad=False),
}
proj = K @ (C_T_G[:3, :3] @ self.landmark + C_T_G[:3, 3:])
landmark2d = proj / proj[2:]
output['landmark2d'] = landmark2d[:2]
inside_patch = (landmark2d[0] < W_modified) * \
(landmark2d[0] >= 0) * \
(landmark2d[1] < H_modified) * \
(landmark2d[1] >= 0) # L vector
# visible by propagated colmap visibility and inside image
_mask1 = self.visibility[:, self.image_indices[index]] * inside_patch
# outside patch
# _mask2 = ~inside_patch
# inside image but not visible by colmap
_mask3 = (self.visibility[:, self.image_indices[index]] == 0) * inside_patch
visibility_mask = 1.0 * _mask1 + 0.5 * _mask3
output['visibility'] = visibility_mask
return output
def __len__(self):
return self.num_images
class Indoor6Patches(Indoor6):
def __init__(self, root_folder="",
scene_id='', mode='all',
landmark_idx=[None], skip_image_index=1,
input_image_downsample=1, gray_image_output=False,
patch_size=96,
positive_samples=4, random_samples=4,
landmark_config='landmarks/landmarks-50',
visibility_config='landmarks/visibility-50',
augmentation=True):
super().__init__(root_folder=root_folder,
scene_id=scene_id, mode=mode,
landmark_idx=landmark_idx, skip_image_index=skip_image_index,
input_image_downsample=input_image_downsample, gray_image_output=gray_image_output,
landmark_config=landmark_config,
visibility_config=visibility_config)
self.patch_size = patch_size
self.positive_samples = positive_samples
self.random_samples = random_samples
self.landmark_idx = landmark_idx
self.augmentation = augmentation
self.num_landmarks = self.landmark.shape[1]
def _extract_patch(self, C_T_G, lm_idx, K, W_modified, H_modified, center=False, adjust_boundary=True):
proj = K @ (C_T_G[:3, :3] @ self.landmark[:, lm_idx:(lm_idx + 1)] + C_T_G[:3, 3:])
proj /= copy.copy(proj[2:])
# Extract patch
y = int(proj[1, 0])
x = int(proj[0, 0])
if center:
dy = -self.patch_size // 2
dx = -self.patch_size // 2
else:
dy = -np.random.rand(1) * self.patch_size
dx = -np.random.rand(1) * self.patch_size
_top = int(y + dy)
_bottom = _top + int(self.patch_size)
_left = int(x + dx)
_right = _left + int(self.patch_size)
if adjust_boundary:
# Adjust the boundary
if _top < 0:
_top = 0
_bottom = int(self.patch_size)
elif _bottom >= H_modified:
_top = H_modified - int(self.patch_size)
_bottom = H_modified
if _left < 0:
_left = 0
_right = int(self.patch_size)
elif _right >= W_modified:
_left = W_modified - int(self.patch_size)
_right = W_modified
return _left, _right, _top, _bottom
def _project_landmarks_into_patch(self, K, C_T_G, img_idx, _top, _bottom, _left, _right):
proj = K @ (C_T_G[:3, :3] @ self.landmark + C_T_G[:3, 3:])
in_front_of_camera = proj[2] > 0.0
proj /= copy.copy(proj[2:])
proj_patch = np.zeros_like(proj[:2])
proj_patch[0] = proj[0] - _left
proj_patch[1] = proj[1] - _top
# L vector
inside_patch = (proj[0] < _right) * (proj[0] >= _left) * (proj[1] < _bottom) * (
proj[1] >= _top) * in_front_of_camera
# visible by propagated colmap visibility and inside patch
_mask1 = self.visibility[:, self.image_indices[img_idx]] * inside_patch
# outside patch
# _mask2 = ~inside_patch
# inside patch but not visible by colmap
_mask3 = (self.visibility[:, self.image_indices[img_idx]] == 0) * inside_patch
visibility_mask = 1.0 * _mask1 + 0.5 * _mask3
return proj_patch, visibility_mask
def __getitem__(self, index):
patches = []
keypoint_locations = []
landmark_visibility_on_patch = []
L = self.landmark.shape[1] # number of keypoints
list_landmarks = np.random.permutation(L)[:self.positive_samples]
## Create positive examples
for lm_idx in list_landmarks:
## Randomly draw image index from visibility mask
training_img_ids_observe_lm_idx = self.visibility[lm_idx, self.image_indices].reshape(-1)
total_images_observed_this_lm = np.sum(training_img_ids_observe_lm_idx)
if total_images_observed_this_lm == 0:
print('no positive example')
img_idx_positive_sample_for_lm_idx = np.random.randint(self.num_images)
else:
# img_idx_observe_lm_idx = (index % int(np.sum(training_img_ids_observe_lm_idx)))
random_indices_observe_this_lm = np.random.randint(0, total_images_observed_this_lm)
img_idx_positive_sample_for_lm_idx = np.where(training_img_ids_observe_lm_idx==1)[0][random_indices_observe_this_lm]
K, K_inv, W_modified, H_modified = self._modify_intrinsic(img_idx_positive_sample_for_lm_idx)
C_T_G = self._load_pose(img_idx_positive_sample_for_lm_idx)
color_tensor = self._load_and_resize_image(img_idx_positive_sample_for_lm_idx, W_modified, H_modified)
if not self.augmentation:
_left, _right, _top, _bottom = self._extract_patch(C_T_G, lm_idx, K, W_modified, H_modified,
center=False, adjust_boundary=True)
color_patch = color_tensor.reshape(1, 3, H_modified, W_modified)[:, :, _top:_bottom, _left:_right]
Cg_T_G = C_T_G
K_scale = K
else:
## Random rotation, change K, T
q = np.random.rand(4) - 0.5
q[1] *= 0.1 # pitch
q[2] *= 0.1 # yaw
q[3] *= 0.1 # roll
q[0] = 1.0
q /= np.linalg.norm(q)
Cg_R_C = Quaternion2Rotation(q)
Cg_T_C = np.eye(4)
Cg_T_C[:3, :3] = Cg_R_C
Cg_T_G = Cg_T_C @ C_T_G
K_scale = K.copy()
K_scale[:2, :2] *= (0.9 + 0.2*np.random.rand())
K_scale_inv = np.linalg.inv(K_scale)
_left, _right, _top, _bottom = self._extract_patch(Cg_T_G, lm_idx, K_scale, W_modified, H_modified,
center=False, adjust_boundary=False)
## Extract patch
YY_patch, XX_patch = torch.meshgrid(torch.arange(_top, _bottom, 1),
torch.arange(_left, _right, 1))
XX_patch = XX_patch.reshape(1, self.patch_size, self.patch_size).float()
YY_patch = YY_patch.reshape(1, self.patch_size, self.patch_size).float()
in_H_out = K @ Cg_R_C.T @ K_scale_inv
in_H_out = torch.tensor(in_H_out, dtype=torch.float)
in_p_out = in_H_out @ torch.cat((XX_patch,
YY_patch,
torch.ones_like(XX_patch)), dim=1).reshape((3, self.patch_size**2))
in_p_out = in_p_out / in_p_out[2:].clone()
scale = torch.tensor([[2. / W_modified, 0.],
[0., 2. / H_modified]], dtype=torch.float).reshape(2, 2)
center = torch.tensor([0.5 * (W_modified - 1),
0.5 * (H_modified - 1)], dtype=torch.float).reshape(2, 1)
in_p_out_normalized = scale @ (in_p_out[:2] - center)
invalid_pixel_mask = (in_p_out_normalized[0] < -1) + \
(in_p_out_normalized[0] > 1) + \
(in_p_out_normalized[1] < -1) + \
(in_p_out_normalized[1] > 1)
if torch.sum(invalid_pixel_mask>0) > 0.25 * self.patch_size ** 2:
_left, _right, _top, _bottom = self._extract_patch(C_T_G, lm_idx, K, W_modified, H_modified,
center=False, adjust_boundary=True)
color_patch = color_tensor.reshape(1, 3, H_modified, W_modified)[:, :, _top:_bottom, _left:_right]
# Not using augmented transformation
K_scale = K.copy()
Cg_T_G = C_T_G
else:
grid_sampler = in_p_out_normalized.reshape(1, 2, self.patch_size, self.patch_size).permute(0, 2, 3, 1)
color_tensor = color_tensor.reshape(1, 3, H_modified, W_modified)
color_patch = torch.nn.functional.grid_sample(color_tensor, grid_sampler,
padding_mode='zeros', mode='bilinear', align_corners=False)
color_patch = torch.nn.functional.interpolate(color_patch, size=(self.patch_size, self.patch_size))
keypoints_2d, visibility_mask = self._project_landmarks_into_patch(K_scale, Cg_T_G, img_idx_positive_sample_for_lm_idx, _top, _bottom, _left, _right)
patches.append(color_patch)
keypoint_locations.append(keypoints_2d.reshape((1, 2, L)))
landmark_visibility_on_patch.append(visibility_mask.reshape((1, L)))
## Create random examples
patches_random = []
keypoint_locations_random = []
landmark_visibility_on_patch_random = []
C_T_G = self._load_pose(index)
K, K_inv, W_modified, H_modified = self._modify_intrinsic(index)
color_tensor = self._load_and_resize_image(index, W_modified, H_modified)
for _ in range(self.random_samples):
_top = int(np.random.rand(1) * (H_modified - self.patch_size))
_bottom = _top + self.patch_size
_left = int(np.random.rand(1) * (W_modified - self.patch_size))
_right = _left + self.patch_size
keypoints_2d, visibility_mask = self._project_landmarks_into_patch(K, C_T_G, index, _top, _bottom, _left, _right)
patches_random.append(color_tensor[:, _top:_bottom, _left:_right].clone().reshape(1, 3, self.patch_size, self.patch_size))
keypoint_locations_random.append(keypoints_2d.reshape((1, 2, L)))
landmark_visibility_on_patch_random.append(visibility_mask.reshape((1, L)))
patches = torch.cat(patches+patches_random, dim=0)
keypoint_locations = np.concatenate(keypoint_locations+keypoint_locations_random, axis=0)
landmark_visibility_on_patch = np.concatenate(landmark_visibility_on_patch+landmark_visibility_on_patch_random, axis=0)
## COLOR AUGMENTATION
if self.augmentation:
if torch.rand(1) > 0.5:
patches += 0.02 * (
torch.rand((patches.shape[0], patches.shape[1], 1, 1)) - 0.5) * torch.ones_like(patches)
else:
patches += 0.2 * (
torch.rand((patches.shape[0], 1, 1, 1)) - 0.5) * torch.ones_like(patches)
clipped_patches = torch.clip(patches, 0, 1)
output = {'patches': clipped_patches,
'landmark2d': torch.tensor(keypoint_locations, dtype=torch.float, requires_grad=False),
'visibility': torch.tensor(landmark_visibility_on_patch, requires_grad=False),
}
return output
================================================
FILE: src/inference.py
================================================
import copy
import numpy as np
import os
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm
import random
from datetime import datetime
from dataloader.indoor6 import Indoor6
from models.efficientlitesld import EfficientNetSLD
from utils.pnp import *
def compute_error(C_R_G, C_t_G, C_R_G_hat, C_t_G_hat):
rot_err = 180 / np.pi * np.arccos(np.clip(0.5 * (np.trace(C_R_G.T @ C_R_G_hat) - 1.0), a_min=-1., a_max=1.))
trans_err = np.linalg.norm(C_R_G_hat.T @ C_t_G_hat - C_R_G.T @ C_t_G)
return rot_err, trans_err
def compute_2d3d(opt, pred_heatmap, peak_threshold, landmark2d, landmark3d, C_b_f_gt, H_hm, W_hm, K_inv,
METRICS_LOGGING=None):
N = pred_heatmap.shape[0]
G_p_f = np.zeros((3, N))
C_b_f_hm = np.zeros((3, N))
weights = np.zeros(N)
validIdx = 0
pixel_error = []
angular_error = []
for l in range(N):
pred_heatmap_l = pred_heatmap[l]
max_pred_heatmap_l = np.max(pred_heatmap_l)
if max_pred_heatmap_l > peak_threshold:
peak_yx = np.unravel_index(np.argmax(pred_heatmap_l), np.array(pred_heatmap_l).shape)
peak_yx = np.array(peak_yx)
# Patch size extraction
P = int(min(1+2*np.min(np.array([peak_yx[0], H_hm-1.0-peak_yx[0], peak_yx[1], W_hm-1.0-peak_yx[1]])),
1+64//opt.output_downsample))
patch_peak_yx = pred_heatmap_l[peak_yx[0] - P // 2:peak_yx[0] + P // 2 + 1,
peak_yx[1] - P // 2:peak_yx[1] + P // 2 + 1]
xx_patch, yy_patch = np.meshgrid(np.arange(peak_yx[1] - P // 2, peak_yx[1] + P // 2 + 1, 1),
np.arange(peak_yx[0] - P // 2, peak_yx[0] + P // 2 + 1, 1))
refine_y = np.sum(patch_peak_yx * yy_patch) / np.sum(patch_peak_yx)
refine_x = np.sum(patch_peak_yx * xx_patch) / np.sum(patch_peak_yx)
pixel_error.append(np.linalg.norm(landmark2d[:2, l] -
opt.output_downsample * np.array([refine_x, refine_y])))
pred_bearing = K_inv @ np.array([refine_x, refine_y, 1])
pred_bearing = pred_bearing / np.linalg.norm(pred_bearing)
gt_bearing = C_b_f_gt[:, l]
gt_bearing = gt_bearing / np.linalg.norm(gt_bearing)
angular_error_batch = np.arccos(
np.clip(pred_bearing @ gt_bearing, a_min=-1, a_max=1)) * 180 / np.pi
angular_error.append(angular_error_batch)
weights[validIdx] = max_pred_heatmap_l
C_b_f_hm[:, validIdx] = pred_bearing
G_p_f[:, validIdx] = landmark3d[:, l]
validIdx += 1
return G_p_f[:, :validIdx], C_b_f_hm[:, :validIdx], weights[:validIdx], np.asarray(pixel_error), np.asarray(angular_error)
def compute_pose(G_p_f, C_b_f_hm, weights, minimal_tight_thr, opt_tight_thr):
Ndetected_landmarks = C_b_f_hm.shape[1]
if Ndetected_landmarks >= 4:
## P3P ransac
C_T_G_hat, PnP_inlier = P3PKe_Ransac(G_p_f, C_b_f_hm, weights,
thres=minimal_tight_thr)
if np.sum(PnP_inlier) >= 4:
C_T_G_opt = RunPnPNL(C_T_G_hat,
G_p_f[:, PnP_inlier],
C_b_f_hm[:, PnP_inlier],
weights[PnP_inlier],
cutoff=opt_tight_thr)
return np.sum(PnP_inlier), C_T_G_opt
return 0, None
def inference(opt, minimal_tight_thr=1e-2, opt_tight_thr=5e-3, mode='test'):
# random.seed(datetime.now().timestamp())
PRETRAINED_MODEL = opt.pretrained_model
device = opt.gpu_device
test_dataset = Indoor6(landmark_idx=np.arange(opt.landmark_indices[0], opt.landmark_indices[-1]),
scene_id=opt.scene_id,
mode=mode,
root_folder=opt.dataset_folder,
input_image_downsample=2,
landmark_config=opt.landmark_config,
visibility_config=opt.visibility_config,
skip_image_index=1,
use_precomputed_focal_length=opt.use_precomputed_focal_length)
test_dataloader = DataLoader(dataset=test_dataset, num_workers=1, batch_size=1, shuffle=False, pin_memory=True)
num_landmarks = test_dataset.landmark.shape[1]
landmark_data = test_dataset.landmark
cnns = []
nLandmarks = opt.landmark_indices
num_landmarks = opt.landmark_indices[-1] - opt.landmark_indices[0]
for idx, pretrained_model in enumerate(PRETRAINED_MODEL):
if opt.model == 'efficientnet':
cnn = EfficientNetSLD(num_landmarks=nLandmarks[idx+1]-nLandmarks[idx], output_downsample=opt.output_downsample).to(device=device)
cnn.load_state_dict(torch.load(pretrained_model))
cnn = cnn.to(device=device)
cnn.eval()
# Adding pretrained model
cnns.append(cnn)
peak_threshold = 3e-1
img_id = 0
METRICS_LOGGING = {'image_name': '',
'angular_error': 180.,
'pixel_error': 1800.,
'rot_err_all': 180.,
'trans_err_all': 180.,
'heatmap_peak': 0.0,
'ndetected': 0,
}
test_image_logging = []
with torch.no_grad():
## Only works for indoor-6
indoor6W = 640 // opt.output_downsample
indoor6H = 352 // opt.output_downsample
HH, WW = torch.meshgrid(torch.arange(indoor6H), torch.arange(indoor6W))
WW = WW.reshape(1, 1, indoor6H, indoor6W).to('cuda')
HH = HH.reshape(1, 1, indoor6H, indoor6W).to('cuda')
with tqdm(test_dataloader) as tq:
for idx, batch in enumerate(tq):
#for idx, batch in enumerate(tqdm(test_dataloader)):
image = batch['image'].to(device=device)
B, _, H, W = image.shape
K_inv = batch['inv_intrinsics'].to(device=device)
C_T_G_gt = batch['pose_gt'].cpu().numpy()
landmark2d = batch['intrinsics'] @ batch['landmark3d'].reshape(B, 3, num_landmarks)
landmark2d /= landmark2d[:, 2:].clone()
landmark2d = landmark2d.numpy()
pred_heatmap = []
for cnn in cnns:
pred = cnn(image)
pred_heatmap.append(pred['1'])
pred_heatmap = torch.cat(pred_heatmap, axis=1)
pred_heatmap *= (pred_heatmap > peak_threshold).float()
# tmp = torch.sqrt(pred_heatmap)
#
# w^{1.5}
# pred_heatmap *= tmp
#
# w^{2.5}
# pred_heatmap *= tmp
# pred_heatmap *= pred_heatmap
# w^2
pred_heatmap *= pred_heatmap
K_inv[:, :, :2] *= opt.output_downsample
## Compute 2D location of landmarks
P = torch.max(torch.max(pred_heatmap, dim=3)[0], dim=2)[0]
pred_normalized_heatmap = pred_heatmap / (torch.sum(pred_heatmap, axis=(2, 3), keepdim=True) + 1e-4)
projx = torch.sum(WW * pred_normalized_heatmap, axis=(2, 3)).reshape(B, 1, num_landmarks)
projy = torch.sum(HH * pred_normalized_heatmap, axis=(2, 3)).reshape(B, 1, num_landmarks)
xy1 = torch.cat((projx, projy, torch.ones_like(projx)), axis=1)
uv1 = K_inv @ xy1
C_B_f = uv1 / torch.sqrt(torch.sum(uv1 ** 2, axis=1, keepdim=True))
C_B_f = C_B_f.cpu().numpy()
P = P.cpu().numpy()
xy1 = xy1.cpu().numpy()
## Compute error
for b in range(B):
Pb = P[b]>peak_threshold
G_p_f = landmark_data[:, Pb]
C_b_f = C_B_f[b][:, Pb]
## MAKING THIS CHANGE FOR ABLATION STUDY IN PAPER: PLEASE REMOVE LATER!
## weights = np.ones_like(P[b][Pb])
weights = P[b][Pb]
xy1b = xy1[b][:2, Pb]
pnp_inlier, C_T_G_hat = compute_pose(G_p_f, C_b_f, weights,
minimal_tight_thr, opt_tight_thr)
rot_err, trans_err = 180., 1800.
if pnp_inlier >= 4:
rot_err, trans_err = compute_error(C_T_G_gt[b][:3, :3], C_T_G_gt[b][:3, 3],
C_T_G_hat[:3, :3], C_T_G_hat[:3, 3])
## Logging information
pixel_error = np.linalg.norm(landmark2d[b][:2, Pb] - opt.output_downsample * xy1b, axis=0)
C_b_f_gt = batch['landmark3d'][b]
C_b_f_gt = torch.nn.functional.normalize(C_b_f_gt, dim=0).cpu().numpy()
angular_error = np.arccos(np.clip(np.sum(C_b_f * C_b_f_gt[:, Pb], axis=0), -1, 1)) * 180. / np.pi
m = copy.deepcopy(METRICS_LOGGING)
m['image_name'] = test_dataset.image_files[img_id]
m['pixel_error'] = pixel_error
m['angular_error'] = angular_error
m['heatmap_peak'] = weights
m['rot_err_all'] = np.array([rot_err])
m['trans_err_all'] = np.array([trans_err])
test_image_logging.append(m)
img_id += 1
elapsedtime = tq.format_dict["elapsed"]
processing_speed = len(test_dataset)/elapsedtime
metrics_output = {'angular_error': [],
'pixel_error': [],
'heatmap_peak': [],
'rot_err_all': [],
'trans_err_all': []}
for k in metrics_output:
for imgdata in test_image_logging:
metrics_output[k].append(imgdata[k])
metrics_output[k] = np.concatenate(metrics_output[k])
metrics_output['r5'] = np.sum(metrics_output['rot_err_all'] < 5) / len(test_dataset)
metrics_output['r10'] = np.sum(metrics_output['rot_err_all'] < 10) / len(test_dataset)
metrics_output['p5'] = np.sum(metrics_output['trans_err_all'] < 0.05) / len(test_dataset)
metrics_output['p10'] = np.sum(metrics_output['trans_err_all'] < 0.1) / len(test_dataset)
metrics_output['r1p1'] = np.sum((metrics_output['rot_err_all'] < 1) * (metrics_output['trans_err_all'] < 0.01))/len(test_dataset)
metrics_output['r2p2'] = np.sum((metrics_output['rot_err_all'] < 2) * (metrics_output['trans_err_all'] < 0.02))/len(test_dataset)
metrics_output['r5p5'] = np.sum((metrics_output['rot_err_all'] < 5) * (metrics_output['trans_err_all'] < 0.05))/len(test_dataset)
metrics_output['r10p10'] = np.sum((metrics_output['rot_err_all'] < 10) * (metrics_output['trans_err_all'] < 0.1)) / len(test_dataset)
metrics_output['median_rot_error'] = np.median(metrics_output['rot_err_all'])
metrics_output['median_trans_error'] = np.median(metrics_output['trans_err_all'])
metrics_output['speed'] = processing_speed
return metrics_output
def inference_landmark_stats(opt, mode='test'):
import pickle
PRETRAINED_MODEL = opt.pretrained_model
device = opt.gpu_device
test_dataset = Indoor6(landmark_idx=np.arange(opt.landmark_indices[0], opt.landmark_indices[-1]),
scene_id=opt.scene_id,
mode=mode,
root_folder=opt.dataset_folder,
input_image_downsample=2,
landmark_config=opt.landmark_config,
visibility_config=opt.visibility_config,
skip_image_index=1)
test_dataloader = DataLoader(dataset=test_dataset, num_workers=1, batch_size=1, shuffle=False, pin_memory=True)
num_landmarks = test_dataset.landmark.shape[1]
cnns = []
nLandmarks = opt.landmark_indices
num_landmarks = opt.landmark_indices[-1] - opt.landmark_indices[0]
for idx, pretrained_model in enumerate(PRETRAINED_MODEL):
if opt.model == 'efficientnet':
cnn = EfficientNetSLD(num_landmarks=nLandmarks[idx+1]-nLandmarks[idx], output_downsample=opt.output_downsample).to(device=device)
cnn.load_state_dict(torch.load(pretrained_model))
cnn = cnn.to(device=device)
cnn.eval()
# Adding pretrained model
cnns.append(cnn)
peak_threshold = 2e-1
SINGLE_LANDMARK_STATS = {'image_idx': [],
'pixel_error': [],
}
landmark_stats = [copy.deepcopy(SINGLE_LANDMARK_STATS) for _ in range(num_landmarks)]
img_idx = 0
with torch.no_grad():
## Only works for indoor-6
indoor6W = 640 // opt.output_downsample
indoor6H = 352 // opt.output_downsample
HH, WW = torch.meshgrid(torch.arange(indoor6H), torch.arange(indoor6W))
WW = WW.reshape(1, 1, indoor6H, indoor6W).to('cuda')
HH = HH.reshape(1, 1, indoor6H, indoor6W).to('cuda')
for idx, batch in enumerate(tqdm(test_dataloader)):
image = batch['image'].to(device=device)
B, _, H, W = image.shape
landmark2d = batch['intrinsics'] @ batch['landmark3d'].reshape(B, 3, num_landmarks)
landmark2d /= landmark2d[:, 2:].clone()
landmark2d = landmark2d.numpy()
pred_heatmap = []
for cnn in cnns:
pred = cnn(image)
pred_heatmap.append(pred['1'])
pred_heatmap = torch.cat(pred_heatmap, axis=1)
pred_heatmap *= (pred_heatmap > peak_threshold).float()
## Compute 2D location of landmarks
P = torch.max(torch.max(pred_heatmap, dim=3)[0], dim=2)[0]
pred_normalized_heatmap = pred_heatmap / (torch.sum(pred_heatmap, axis=(2, 3), keepdim=True) + 1e-4)
projx = torch.sum(WW * pred_normalized_heatmap, axis=(2, 3)).reshape(B, 1, num_landmarks)
projy = torch.sum(HH * pred_normalized_heatmap, axis=(2, 3)).reshape(B, 1, num_landmarks)
xy1 = torch.cat((projx, projy, torch.ones_like(projx)), axis=1)
P = P.cpu().numpy()
xy1 = xy1.cpu().numpy()
## Compute error
for b in range(B):
for l in range(num_landmarks):
if P[b,l] > peak_threshold:
pixel_error = np.linalg.norm(landmark2d[b][:2, l] -
opt.output_downsample * xy1[b][:2, l])
landmark_stats[l]['pixel_error'].append(pixel_error)
landmark_stats[l]['image_idx'].append(test_dataset.image_indices[img_idx])
img_idx += 1
landmark_stats_np = np.zeros((num_landmarks, 5))
for l in range(num_landmarks):
landmark_stats_np[l, 0] = l
landmark_stats_np[l, 1] = len(landmark_stats[l]['image_idx'])
if landmark_stats_np[l, 1] > 0:
pixel_error = np.array(landmark_stats[l]['pixel_error'])
landmark_stats_np[l, 2] = np.mean(pixel_error)
landmark_stats_np[l, 3] = np.median(pixel_error)
landmark_stats_np[l, 4] = np.max(pixel_error)
np.savetxt(os.path.join(opt.output_folder, 'landmark_stats.txt'), landmark_stats_np)
pickle.dump(landmark_stats, open(os.path.join(opt.output_folder, 'landmark_stats.pkl'), 'wb'))
return
================================================
FILE: src/local_inference.py
================================================
# Copyright (c) Microsoft Corporation. All rights reserved.
#from __future__ import print_function
import argparse
import os
import time
Args = None
def local_inference():
cmd = 'python main.py --action test --dataset_folder %s --scene_id %s --landmark_config %s --visibility_config %s' % (Args.dataset_dir, Args.scene_id, Args.landmark_config, Args.visibility_config)
cmd += ' --output_downsample 8'
cmd += ' --landmark_indices 0'
for i in range(0, len(Args.landmark_indices)):
cmd += ' --landmark_indices %d' % (Args.landmark_indices[i])
for ckpt in Args.checkpoint_names:
cmd += ' --pretrained_model %s/%s/%s/model-best_median.ckpt' % (Args.checkpoint_dir, Args.experimentGroupName, ckpt)
cmd += ' --output_folder %s/%s' % (Args.checkpoint_dir, Args.experimentGroupName)
print("Running [" + cmd + "]")
os.system(cmd)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument(
'--experiment_file', default="", type=str, required=True,
help="Experiment file path.")
parser.add_argument(
'--dataset_dir', default="", type=str, required=True,
help="Dataset path.")
parser.add_argument(
'--checkpoint_dir', default="", type=str, required=True,
help="Checkpoints folder path.")
Args = parser.parse_args()
tmp = os.path.basename(Args.experiment_file)
Args.experimentGroupName = tmp[:tmp.rindex('.')]
Args.landmark_indices = []
Args.checkpoint_names = []
exp_file = os.path.join(Args.checkpoint_dir, Args.experiment_file)
fd = open(exp_file, 'r')
while True:
line = fd.readline()
if line == '':
break
split_line = line.split()
Args.scene_id = split_line[0]
expName = split_line[1]
Args.landmark_config = split_line[2]
Args.visibility_config = split_line[3]
Args.checkpoint_names.append(expName)
fields = expName.split('-')
Args.landmark_indices.append(int(fields[2]))
local_inference()
================================================
FILE: src/local_training.py
================================================
# Copyright (c) Microsoft Corporation. All rights reserved.
import argparse
import os
#import re
Args = None
def launch_training():
print("Experiment File: %s" % Args.experiment_file)
print("Model Dir: %s" % Args.model_dir)
cmd = 'python main.py --action train_patches'
cmd += ' --training_batch_size %d' % (Args.training_batch_size)
cmd += ' --output_downsample %d' % (Args.output_downsample)
cmd += ' --num_epochs %d' % (Args.num_epochs)
cmd += ' --dataset_folder %s' % (Args.dataset_dir)
cmd += ' --scene_id %s' % (Args.scene_id)
cmd += ' --landmark_config %s' % (Args.landmark_config)
cmd += ' --visibility_config %s' % (Args.visibility_config)
cmd += ' --output_folder %s' % (Args.model_dir)
cmd += ' --landmark_indices %d' % (Args.landmark_index_start)
cmd += ' --landmark_indices %d' % (Args.landmark_index_stop)
os.system(cmd)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument(
'--dataset_dir', type=str, required=True,
help="Dataset folder path.")
parser.add_argument(
'--experiment_file', type=str, required=True,
help="Experiment file path.")
parser.add_argument(
'--scene_id', type=str, required=True,
help="name of scene.")
parser.add_argument(
'--landmark_config', type=str, required=True,
help='Landmark configuration.')
parser.add_argument(
'--visibility_config', type=str, required=True,
help='Visibility configuration.')
parser.add_argument(
'--num_landmarks', type=int, required=True,
help='number of landmarks.')
parser.add_argument(
'--block_size', type=int, required=True,
help='number of landmarks in each block.')
parser.add_argument(
'--subset_index', type=int, required=True,
help='index of landmark subset (starts from 0).')
parser.add_argument(
'--output_dir', type=str, required=True,
help='folder to save experiment file in.')
parser.add_argument(
'--model_dir', type=str, required=True,
help='folder to save model ckpt file in.')
parser.add_argument(
'--training_batch_size', type=int, required=True,
help='batch size.')
parser.add_argument(
'--output_downsample', type=int, required=True,
help='Downsample factor for heat map resolution.')
parser.add_argument(
'--num_epochs', type=int, required=True,
help='the number of epochs used for training.')
Args = parser.parse_args()
# Write the experiment file
exp_fn = os.path.join(Args.output_dir, Args.experiment_file)
fd = open(exp_fn, "w")
for lid in range(0, Args.num_landmarks, Args.block_size):
Args.landmark_index_start = lid
Args.landmark_index_stop = lid + Args.block_size
str = '%s %s-%03d-%03d %s %s local' % (Args.scene_id, Args.scene_id, Args.landmark_index_start, Args.landmark_index_stop, Args.landmark_config, Args.visibility_config)
print(str, file=fd)
fd.close()
# Launch the training job for the specified subset only.
Args.landmark_index_start = Args.block_size * Args.subset_index
Args.landmark_index_stop = Args.block_size * (Args.subset_index + 1)
launch_training()
================================================
FILE: src/main.py
================================================
import argparse
from inference import *
from train import *
DEVICE = None
# auto-detect default device
if torch.backends.mps.is_available():
# Code to run on macOS
torch.backends.mps.enabled = True
DEVICE = "mps"
print ("MPS enabled")
elif torch.cuda.is_available():
# Windows or Linux GPU acceleration
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
DEVICE = "cuda"
print ("CUDA enabled")
else:
# CPU
torch.backends.cudnn.enabled = False
DEVICE = "cpu"
print ("CPU enabled")
if __name__ == '__main__':
parser = argparse.ArgumentParser(
description='Scene Landmark Detection',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument(
'--dataset_folder', type=str, required=True,
help='Root directory, where all data is stored')
parser.add_argument(
'--output_folder', type=str, required=True,
help='Output folder')
parser.add_argument(
'--landmark_config', type=str, default='landmarks/landmarks-300',
help='File containing scene-specific 3D landmarks.')
parser.add_argument(
'--landmark_indices', type=int, action='append',
help = 'Landmark indices, specify twice',
required=True)
parser.add_argument(
'--visibility_config', type=str, default='landmarks/visibility_aug-300',
help='File containing information about visibility of landmarks in cameras associated with training set.')
parser.add_argument(
'--scene_id', type=str, default='scene6',
help='Scene id')
parser.add_argument(
'--model', type=str, default='efficientnet',
help='Network architecture backbone.')
parser.add_argument(
'--output_downsample', type=int, default=4,
help='Down sampling factor for output resolution')
parser.add_argument(
'--gpu_device', type=str, default=DEVICE,
help='GPU device')
parser.add_argument(
'--pretrained_model', type=str, action='append', default=[],
help='Pretrained detector model')
parser.add_argument(
'--num_epochs', type=int, default=200,
help='Number of training epochs.')
parser.add_argument(
'--action', type=str, default='test',
help='train/train_patches/test')
parser.add_argument(
'--use_precomputed_focal_length', type=int, default=0)
parser.add_argument(
'--training_batch_size', type=int, default=8,
help='Batch size used during training.')
opt = parser.parse_args()
#print('scene_id: ', opt.scene_id)
#print('action: ', opt.action)
#print('training_batch_size: ', opt.training_batch_size)
#print('output downsample: ', opt.output_downsample)
if opt.action == 'train':
train(opt)
opt.pretrained_model = [opt.output_folder + '/model-best_median.ckpt']
eval_stats = inference(opt, minimal_tight_thr=1e-3, opt_tight_thr=1e-3)
print("{:>10} {:>30} {:>30} {:>20}".format('Scene ID',
'Median trans error (cm)',
'Median rotation error (deg)',
'Recall 5cm5deg (%)'))
print("{:>10} {:>30.4} {:>30.4} {:>20.2%}".format(opt.scene_id,
100. * eval_stats['median_trans_error'],
eval_stats['median_rot_error'],
eval_stats['r5p5']))
elif opt.action == 'train_patches':
train_patches(opt)
opt.pretrained_model = [opt.output_folder + '/model-best_median.ckpt']
eval_stats = inference(opt, minimal_tight_thr=1e-3, opt_tight_thr=1e-3)
print("{:>10} {:>30} {:>30} {:>20}".format('Scene ID',
'Median trans error (cm)',
'Median rotation error (deg)',
'Recall 5cm5deg (%)'))
print("{:>10} {:>30.4} {:>30.4} {:>20.2%}".format(opt.scene_id,
100. * eval_stats['median_trans_error'],
eval_stats['median_rot_error'],
eval_stats['r5p5']))
elif opt.action == 'landmark_stats':
inference_landmark_stats(opt, mode='train')
elif opt.action == 'test':
if opt.scene_id == 'all':
eval_stats = {}
pretrained_folder = opt.pretrained_model
output_folder = opt.output_folder
for scene_id in ['1', '2a', '3', '4a', '5', '6']:
opt.scene_id = 'scene' + scene_id
opt.pretrained_model = [pretrained_folder + 'scene%s.ckpt' % scene_id]
opt.output_folder = os.path.join(output_folder, 'scene' + scene_id)
eval_stats[opt.scene_id] = inference(opt, minimal_tight_thr=1e-3, opt_tight_thr=1e-3)
print("{:>10} {:>30} {:>30} {:>20}".format('Scene ID',
'Median trans error (cm)',
'Median rotation error (deg)',
'Recall 5cm5deg (%)'))
for x in eval_stats:
print("{:>10} {:>30.4} {:>30.4} {:>20.2%}".format(x,
100. * eval_stats[x]['median_trans_error'],
eval_stats[x]['median_rot_error'],
eval_stats[x]['r5p5']))
else:
eval_stats = inference(opt, minimal_tight_thr=1e-3, opt_tight_thr=1e-3)
metricsFilename = opt.output_folder + '/metrics.txt'
print(metricsFilename)
fd = open(metricsFilename, "w")
fd.write("%f\n" % (eval_stats['r5p5']))
fd.write("%f\n" % (eval_stats['speed']))
fd.close()
print("{:>10} {:>30} {:>30} {:>20} {:>15} {:>15} {:>15} {:>15} {:>20} {:>20}".format('Scene ID',
'Median trans error (cm)',
'Median rotation error (deg)',
'Recall 1cm1deg (%)',
'2cm2deg (%)',
'5cm5deg (%)',
'10cm10deg (%)',
'5deg (%)',
'Median Pixel Error',
'Median Angular Error'))
print("{:>10} {:>30.4} {:>30.4} {:>20.2%} {:>15.2%} {:>15.2%} {:>15.2%} {:>15.2%} {:>20.4} {:>20.4}".format(opt.scene_id,
100. * eval_stats['median_trans_error'],
eval_stats['median_rot_error'],
eval_stats['r1p1'],
eval_stats['r2p2'],
eval_stats['r5p5'],
eval_stats['r10p10'],
eval_stats['r5'],
np.median(eval_stats['pixel_error']),
np.median(eval_stats['angular_error'])))
================================================
FILE: src/models/blocks.py
================================================
import torch
import torch.nn as nn
from .conv2d_layers import Conv2dSameExport
def _make_encoder(use_pretrained, exportable=True, output_downsample=4):
# pretrained = _make_pretrained_efficientnet_lite0(use_pretrained, exportable=exportable)
pretrained = torch.load('pretrained_efficientnetlite0.net')
if output_downsample <= 16:
pretrained.layer2[0][0].conv_dw.stride = (1, 1)
if output_downsample <= 8:
pretrained.layer3[0][0].conv_dw.stride = (1, 1)
if output_downsample <= 4:
pretrained.layer4[0][0].conv_dw.stride = (1, 1)
return pretrained, None
def _make_pretrained_efficientnet_lite0(use_pretrained, exportable=False):
efficientnet = torch.hub.load(
"rwightman/gen-efficientnet-pytorch",
"tf_efficientnet_lite0",
pretrained=use_pretrained,
exportable=exportable
)
return _make_efficientnet_backbone(efficientnet)
def _make_efficientnet_backbone(effnet):
pretrained = nn.Module()
pretrained.layer1 = nn.Sequential(
effnet.conv_stem, effnet.bn1, effnet.act1, *effnet.blocks[0:2]
)
pretrained.layer2 = nn.Sequential(*effnet.blocks[2:3])
pretrained.layer3 = nn.Sequential(*effnet.blocks[3:5])
pretrained.layer4 = nn.Sequential(*effnet.blocks[5:9])
return pretrained
def _make_resnet_backbone(resnet):
pretrained = nn.Module()
pretrained.layer1 = nn.Sequential(
resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool, resnet.layer1
)
pretrained.layer2 = resnet.layer2
pretrained.layer3 = resnet.layer3
pretrained.layer4 = resnet.layer4
return pretrained
def _make_pretrained_resnext101_wsl(use_pretrained):
resnet = torch.hub.load("facebookresearch/WSL-Images", "resnext101_32x8d_wsl")
return _make_resnet_backbone(resnet)
class Interpolate(nn.Module):
"""Interpolation module.
"""
def __init__(self, scale_factor, mode, align_corners=False):
"""Init.
Args:
scale_factor (float): scaling
mode (str): interpolation mode
"""
super(Interpolate, self).__init__()
self.interp = nn.functional.interpolate
self.scale_factor = scale_factor
self.mode = mode
self.align_corners = align_corners
def forward(self, x):
"""Forward pass.
Args:
x (tensor): input
Returns:
tensor: interpolated data
"""
x = self.interp(
x, scale_factor=self.scale_factor, mode=self.mode, align_corners=self.align_corners
)
return x
class ResidualConvUnit(nn.Module):
"""Residual convolution module.
"""
def __init__(self, features):
"""Init.
Args:
features (int): number of features
"""
super().__init__()
self.conv1 = nn.Conv2d(
features, features, kernel_size=3, stride=1, padding=1, bias=True
)
self.conv2 = nn.Conv2d(
features, features, kernel_size=3, stride=1, padding=1, bias=True
)
self.relu = nn.ReLU(inplace=True)
def forward(self, x):
"""Forward pass.
Args:
x (tensor): input
Returns:
tensor: output
"""
out = self.relu(x)
out = self.conv1(out)
out = self.relu(out)
out = self.conv2(out)
return out + x
class FeatureFusionBlock(nn.Module):
"""Feature fusion block.
"""
def __init__(self, features):
"""Init.
Args:
features (int): number of features
"""
super(FeatureFusionBlock, self).__init__()
self.resConfUnit1 = ResidualConvUnit(features)
self.resConfUnit2 = ResidualConvUnit(features)
def forward(self, *xs):
"""Forward pass.
Returns:
tensor: output
"""
output = xs[0]
if len(xs) == 2:
output += self.resConfUnit1(xs[1])
output = self.resConfUnit2(output)
output = nn.functional.interpolate(
output, scale_factor=2, mode="bilinear", align_corners=True
)
return output
class ResidualConvUnit_custom(nn.Module):
"""Residual convolution module.
"""
def __init__(self, features, activation, bn):
"""Init.
Args:
features (int): number of features
"""
super().__init__()
self.bn = bn
self.groups = 1
self.conv1 = nn.Conv2d(
features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups
)
self.conv2 = nn.Conv2d(
features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups
)
if self.bn == True:
self.bn1 = nn.BatchNorm2d(features)
self.bn2 = nn.BatchNorm2d(features)
self.activation = activation
self.skip_add = nn.quantized.FloatFunctional()
def forward(self, x):
"""Forward pass.
Args:
x (tensor): input
Returns:
tensor: output
"""
out = self.activation(x)
out = self.conv1(out)
if self.bn == True:
out = self.bn1(out)
out = self.activation(out)
out = self.conv2(out)
if self.bn == True:
out = self.bn2(out)
if self.groups > 1:
out = self.conv_merge(out)
return self.skip_add.add(out, x)
# return out + x
class FeatureFusionBlock_custom(nn.Module):
"""Feature fusion block.
"""
def __init__(self, features, activation, deconv=False, bn=False, expand=False, align_corners=True):
"""Init.
Args:
features (int): number of features
"""
super(FeatureFusionBlock_custom, self).__init__()
self.deconv = deconv
self.align_corners = align_corners
self.groups = 1
self.expand = expand
out_features = features
if self.expand == True:
out_features = features // 2
self.out_conv = nn.Conv2d(features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1)
self.resConfUnit1 = ResidualConvUnit_custom(features, activation, bn)
self.resConfUnit2 = ResidualConvUnit_custom(features, activation, bn)
self.skip_add = nn.quantized.FloatFunctional()
def forward(self, *xs):
"""Forward pass.
Returns:
tensor: output
"""
output = xs[0]
if len(xs) == 2:
res = self.resConfUnit1(xs[1])
output = self.skip_add.add(output, res)
# output += res
output = self.resConfUnit2(output)
output = nn.functional.interpolate(
output, scale_factor=2, mode="bilinear", align_corners=self.align_corners
)
output = self.out_conv(output)
return output
================================================
FILE: src/models/conv2d_layers.py
================================================
""" Conv2D w/ SAME padding, CondConv, MixedConv
A collection of conv layers and padding helpers needed by EfficientNet, MixNet, and
MobileNetV3 models that maintain weight compatibility with original Tensorflow models.
Copyright 2020 Ross Wightman
"""
import collections.abc
import math
from functools import partial
from itertools import repeat
from typing import Tuple, Optional
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
# From PyTorch internals
def _ntuple(n):
def parse(x):
if isinstance(x, collections.abc.Iterable):
return x
return tuple(repeat(x, n))
return parse
_single = _ntuple(1)
_pair = _ntuple(2)
_triple = _ntuple(3)
_quadruple = _ntuple(4)
def _is_static_pad(kernel_size, stride=1, dilation=1, **_):
return stride == 1 and (dilation * (kernel_size - 1)) % 2 == 0
def _get_padding(kernel_size, stride=1, dilation=1, **_):
padding = ((stride - 1) + dilation * (kernel_size - 1)) // 2
return padding
def _calc_same_pad(i: int, k: int, s: int, d: int):
return max((-(i // -s) - 1) * s + (k - 1) * d + 1 - i, 0)
def _same_pad_arg(input_size, kernel_size, stride, dilation):
ih, iw = input_size
kh, kw = kernel_size
pad_h = _calc_same_pad(ih, kh, stride[0], dilation[0])
pad_w = _calc_same_pad(iw, kw, stride[1], dilation[1])
return [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2]
def _split_channels(num_chan, num_groups):
split = [num_chan // num_groups for _ in range(num_groups)]
split[0] += num_chan - sum(split)
return split
def conv2d_same(
x, weight: torch.Tensor, bias: Optional[torch.Tensor] = None, stride: Tuple[int, int] = (1, 1),
padding: Tuple[int, int] = (0, 0), dilation: Tuple[int, int] = (1, 1), groups: int = 1):
ih, iw = x.size()[-2:]
kh, kw = weight.size()[-2:]
pad_h = _calc_same_pad(ih, kh, stride[0], dilation[0])
pad_w = _calc_same_pad(iw, kw, stride[1], dilation[1])
x = F.pad(x, [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2])
return F.conv2d(x, weight, bias, stride, (0, 0), dilation, groups)
class Conv2dSame(nn.Conv2d):
""" Tensorflow like 'SAME' convolution wrapper for 2D convolutions
"""
# pylint: disable=unused-argument
def __init__(self, in_channels, out_channels, kernel_size, stride=1,
padding=0, dilation=1, groups=1, bias=True):
super(Conv2dSame, self).__init__(
in_channels, out_channels, kernel_size, stride, 0, dilation, groups, bias)
def forward(self, x):
return conv2d_same(x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
class Conv2dSameExport(nn.Conv2d):
""" ONNX export friendly Tensorflow like 'SAME' convolution wrapper for 2D convolutions
NOTE: This does not currently work with torch.jit.script
"""
# pylint: disable=unused-argument
def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, groups=1, bias=True):
super(Conv2dSameExport, self).__init__(
in_channels, out_channels, kernel_size, stride, 0, dilation, groups, bias)
self.pad = None
self.pad_input_size = (0, 0)
def forward(self, x):
input_size = x.size()[-2:]
if self.pad is None:
pad_arg = _same_pad_arg(input_size, self.weight.size()[-2:], self.stride, self.dilation)
self.pad = nn.ZeroPad2d(pad_arg)
self.pad_input_size = input_size
if self.pad is not None:
x = self.pad(x)
return F.conv2d(
x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
def get_padding_value(padding, kernel_size, **kwargs):
dynamic = False
if isinstance(padding, str):
# for any string padding, the padding will be calculated for you, one of three ways
padding = padding.lower()
if padding == 'same':
# TF compatible 'SAME' padding, has a performance and GPU memory allocation impact
if _is_static_pad(kernel_size, **kwargs):
# static case, no extra overhead
padding = _get_padding(kernel_size, **kwargs)
else:
# dynamic padding
padding = 0
dynamic = True
elif padding == 'valid':
# 'VALID' padding, same as padding=0
padding = 0
else:
# Default to PyTorch style 'same'-ish symmetric padding
padding = _get_padding(kernel_size, **kwargs)
return padding, dynamic
def create_conv2d_pad(in_chs, out_chs, kernel_size, **kwargs):
padding = kwargs.pop('padding', '')
kwargs.setdefault('bias', False)
padding, is_dynamic = get_padding_value(padding, kernel_size, **kwargs)
if is_dynamic:
if is_exportable():
assert not is_scriptable()
return Conv2dSameExport(in_chs, out_chs, kernel_size, **kwargs)
else:
return Conv2dSame(in_chs, out_chs, kernel_size, **kwargs)
else:
return nn.Conv2d(in_chs, out_chs, kernel_size, padding=padding, **kwargs)
class MixedConv2d(nn.ModuleDict):
""" Mixed Grouped Convolution
Based on MDConv and GroupedConv in MixNet impl:
https://github.com/tensorflow/tpu/blob/master/models/official/mnasnet/mixnet/custom_layers.py
"""
def __init__(self, in_channels, out_channels, kernel_size=3,
stride=1, padding='', dilation=1, depthwise=False, **kwargs):
super(MixedConv2d, self).__init__()
kernel_size = kernel_size if isinstance(kernel_size, list) else [kernel_size]
num_groups = len(kernel_size)
in_splits = _split_channels(in_channels, num_groups)
out_splits = _split_channels(out_channels, num_groups)
self.in_channels = sum(in_splits)
self.out_channels = sum(out_splits)
for idx, (k, in_ch, out_ch) in enumerate(zip(kernel_size, in_splits, out_splits)):
conv_groups = out_ch if depthwise else 1
self.add_module(
str(idx),
create_conv2d_pad(
in_ch, out_ch, k, stride=stride,
padding=padding, dilation=dilation, groups=conv_groups, **kwargs)
)
self.splits = in_splits
def forward(self, x):
x_split = torch.split(x, self.splits, 1)
x_out = [conv(x_split[i]) for i, conv in enumerate(self.values())]
x = torch.cat(x_out, 1)
return x
def get_condconv_initializer(initializer, num_experts, expert_shape):
def condconv_initializer(weight):
"""CondConv initializer function."""
num_params = np.prod(expert_shape)
if (len(weight.shape) != 2 or weight.shape[0] != num_experts or
weight.shape[1] != num_params):
raise (ValueError(
'CondConv variables must have shape [num_experts, num_params]'))
for i in range(num_experts):
initializer(weight[i].view(expert_shape))
return condconv_initializer
class CondConv2d(nn.Module):
""" Conditional Convolution
Inspired by: https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/condconv/condconv_layers.py
Grouped convolution hackery for parallel execution of the per-sample kernel filters inspired by this discussion:
https://github.com/pytorch/pytorch/issues/17983
"""
__constants__ = ['bias', 'in_channels', 'out_channels', 'dynamic_padding']
def __init__(self, in_channels, out_channels, kernel_size=3,
stride=1, padding='', dilation=1, groups=1, bias=False, num_experts=4):
super(CondConv2d, self).__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.kernel_size = _pair(kernel_size)
self.stride = _pair(stride)
padding_val, is_padding_dynamic = get_padding_value(
padding, kernel_size, stride=stride, dilation=dilation)
self.dynamic_padding = is_padding_dynamic # if in forward to work with torchscript
self.padding = _pair(padding_val)
self.dilation = _pair(dilation)
self.groups = groups
self.num_experts = num_experts
self.weight_shape = (self.out_channels, self.in_channels // self.groups) + self.kernel_size
weight_num_param = 1
for wd in self.weight_shape:
weight_num_param *= wd
self.weight = torch.nn.Parameter(torch.Tensor(self.num_experts, weight_num_param))
if bias:
self.bias_shape = (self.out_channels,)
self.bias = torch.nn.Parameter(torch.Tensor(self.num_experts, self.out_channels))
else:
self.register_parameter('bias', None)
self.reset_parameters()
def reset_parameters(self):
init_weight = get_condconv_initializer(
partial(nn.init.kaiming_uniform_, a=math.sqrt(5)), self.num_experts, self.weight_shape)
init_weight(self.weight)
if self.bias is not None:
fan_in = np.prod(self.weight_shape[1:])
bound = 1 / math.sqrt(fan_in)
init_bias = get_condconv_initializer(
partial(nn.init.uniform_, a=-bound, b=bound), self.num_experts, self.bias_shape)
init_bias(self.bias)
def forward(self, x, routing_weights):
B, C, H, W = x.shape
weight = torch.matmul(routing_weights, self.weight)
new_weight_shape = (B * self.out_channels, self.in_channels // self.groups) + self.kernel_size
weight = weight.view(new_weight_shape)
bias = None
if self.bias is not None:
bias = torch.matmul(routing_weights, self.bias)
bias = bias.view(B * self.out_channels)
# move batch elements with channels so each batch element can be efficiently convolved with separate kernel
x = x.view(1, B * C, H, W)
if self.dynamic_padding:
out = conv2d_same(
x, weight, bias, stride=self.stride, padding=self.padding,
dilation=self.dilation, groups=self.groups * B)
else:
out = F.conv2d(
x, weight, bias, stride=self.stride, padding=self.padding,
dilation=self.dilation, groups=self.groups * B)
out = out.permute([1, 0, 2, 3]).view(B, self.out_channels, out.shape[-2], out.shape[-1])
# Literal port (from TF definition)
# x = torch.split(x, 1, 0)
# weight = torch.split(weight, 1, 0)
# if self.bias is not None:
# bias = torch.matmul(routing_weights, self.bias)
# bias = torch.split(bias, 1, 0)
# else:
# bias = [None] * B
# out = []
# for xi, wi, bi in zip(x, weight, bias):
# wi = wi.view(*self.weight_shape)
# if bi is not None:
# bi = bi.view(*self.bias_shape)
# out.append(self.conv_fn(
# xi, wi, bi, stride=self.stride, padding=self.padding,
# dilation=self.dilation, groups=self.groups))
# out = torch.cat(out, 0)
return out
def select_conv2d(in_chs, out_chs, kernel_size, **kwargs):
assert 'groups' not in kwargs # only use 'depthwise' bool arg
if isinstance(kernel_size, list):
assert 'num_experts' not in kwargs # MixNet + CondConv combo not supported currently
# We're going to use only lists for defining the MixedConv2d kernel groups,
# ints, tuples, other iterables will continue to pass to normal conv and specify h, w.
m = MixedConv2d(in_chs, out_chs, kernel_size, **kwargs)
else:
depthwise = kwargs.pop('depthwise', False)
groups = out_chs if depthwise else 1
if 'num_experts' in kwargs and kwargs['num_experts'] > 0:
m = CondConv2d(in_chs, out_chs, kernel_size, groups=groups, **kwargs)
else:
m = create_conv2d_pad(in_chs, out_chs, kernel_size, groups=groups, **kwargs)
return m
================================================
FILE: src/models/efficientlitesld.py
================================================
import torch
import torch.nn as nn
from .blocks import _make_encoder
class ASPP(nn.Module):
def __init__(self, in_ch, d1, d2, d3, d4, reduction=4):
super(ASPP, self).__init__()
self.aspp_d1 = nn.Sequential(
nn.Conv2d(in_ch, in_ch // reduction, 3, padding=d1, dilation=d1),
nn.BatchNorm2d(in_ch // reduction),
nn.ReLU(inplace=True)
)
self.aspp_d2 = nn.Sequential(
nn.Conv2d(in_ch, in_ch // reduction, 3, padding=d2, dilation=d2),
nn.BatchNorm2d(in_ch // reduction),
nn.ReLU(inplace=True)
)
self.aspp_d3 = nn.Sequential(
nn.Conv2d(in_ch, in_ch // reduction, 3, padding=d3, dilation=d3),
nn.BatchNorm2d(in_ch // reduction),
nn.ReLU(inplace=True)
)
self.aspp_d4 = nn.Sequential(
nn.Conv2d(in_ch, in_ch // reduction, 3, padding=d4, dilation=d4),
nn.BatchNorm2d(in_ch // reduction),
nn.ReLU(inplace=True)
)
def forward(self, x):
d1 = self.aspp_d1(x)
d2 = self.aspp_d2(x)
d3 = self.aspp_d3(x)
d4 = self.aspp_d4(x)
return torch.cat((d1, d2, d3, d4), dim=1)
class EfficientNetSLD(torch.nn.Module):
"""Network for monocular depth estimation.
"""
def __init__(self, path=None, num_landmarks=200, output_downsample=4, features=320):
"""Init.
Args:
path (str, optional): Path to saved model. Defaults to None.
features (int, optional): Number of features. Defaults to 256.
backbone (str, optional): Backbone network for encoder. Defaults to efficientnetlite0
"""
super(EfficientNetSLD, self).__init__()
self.pretrained, _ = _make_encoder(use_pretrained=True, output_downsample=output_downsample)
self.aspp = nn.Sequential(
ASPP(in_ch=features, d1=1, d2=2, d3=3, d4=4, reduction=4),
)
self.heatmap_outputs_res1 = nn.Sequential(
nn.Conv2d(features, num_landmarks, kernel_size=1, stride=1, padding=0)
)
self.heatmap_outputs_res2 = None
if output_downsample == 2:
input_channels = features + num_landmarks
output_channels = features
self.heatmap_features_res2 = nn.Sequential(nn.ConvTranspose2d(in_channels=input_channels,
out_channels=output_channels,
kernel_size=4, stride=2, padding=1,
bias=False),
nn.BatchNorm2d(output_channels),
nn.ReLU(inplace=True)
)
self.heatmap_outputs_res2 = nn.Conv2d(output_channels, num_landmarks, kernel_size=1, stride=1, bias=False)
if path:
self.load(path)
def forward(self, x):
"""Forward pass.
Args:
x (tensor): input data (image)
Returns:
Heatmap prediction
['1']: quarter of input spatial dimension
['2']: half of input spatial dimension
"""
layer_1 = self.pretrained.layer1(x)
layer_2 = self.pretrained.layer2(layer_1)
layer_3 = self.pretrained.layer3(layer_2)
layer_4 = self.pretrained.layer4(layer_3)
y1 = self.aspp(layer_4)
z1 = self.heatmap_outputs_res1(y1)
z2 = None
if self.heatmap_outputs_res2 is not None:
y2 = self.heatmap_features_res2(torch.cat((y1, z1), dim=1))
z2 = self.heatmap_outputs_res2(y2)
return {'1': z1, '2': z2}
================================================
FILE: src/pretrained_efficientnetlite0.net
================================================
[File too large to display: 11.6 MB]
================================================
FILE: src/requirements.txt
================================================
# Scene Landmarks Detector Requirements
# Usage: pip install -r requirements.txt
argparse
matplotlib>=3.2.2
numpy>=1.22.3
Pillow>=8.2.0
scipy>=1.6.2
open3d
#torch==1.10.0+cu113
#torchvision==0.11.1+cu113
#torchaudio==0.10.0+cu113
tqdm>=4.59.0
geffnet
================================================
FILE: src/run_inference.py
================================================
import os
import statistics as st
import sys
import torch
if __name__ == '__main__':
home_dir = os.path.expanduser("~")
# specify dataset path, location of checkpoints and the experiment name.
checkpoint_dir = os.path.join(home_dir, 'data/checkpoints')
dataset_dir = os.path.join(home_dir, 'data/indoor6')
experiment = '1000-125_v10'
# run inference for all six scenes of the indoor6 dataset
for scene_name in ['scene1', 'scene2a', 'scene3', 'scene4a', 'scene5', 'scene6']:
command = 'python ./local_inference.py --experiment_file %s_%s.txt --dataset_dir %s --checkpoint_dir %s' % (scene_name, experiment, dataset_dir, checkpoint_dir)
os.system(command)
# calculate metrics
t1 = []
t2 = []
for scene_name in ['scene1', 'scene2a', 'scene3', 'scene4a', 'scene5', 'scene6']:
subfolder = '%s_%s' % (scene_name, experiment)
mfn = os.path.join(checkpoint_dir, subfolder, "metrics.txt")
mfd = open(mfn, 'r')
idx = 0
for line in mfd.readlines():
if (idx % 2 == 0):
t1.append(float(line))
else:
t2.append(float(line))
idx+=1
mfd.close();
print(t1)
print(t2)
metricPcnt = 100.0 * st.fmean(t1)
print(' mean = %s pcnt' % str(metricPcnt))
print(' rate = %s imgs./sec.' % str(st.fmean(t2)))
fname = 'RESULTS-%s.txt' % experiment
ffn = os.path.join(checkpoint_dir, fname)
ffd = open(ffn, 'w')
ffd.write(f"{metricPcnt}\n{st.fmean(t2)}\n")
ffd.close();
================================================
FILE: src/run_training.py
================================================
from math import exp
import os
import statistics as st
from tabnanny import check
if __name__ == '__main__':
home_dir = os.path.expanduser("~")
# Specify the paths to the dataset and the output folders.
dataset_dir = os.path.join(home_dir, "data/indoor6")
output_dir = os.path.join(home_dir, "data/outputs")
# Specify a version number which can be incremented when training multiple variants on
# the same scene.
version_no = 10
# Specify the scene name
scene_name = 'scene6'
# Specify the landmark file
landmark_config = 'landmarks/landmarks-1000v10'
# Specify the visibility file
visibility_config = 'landmarks/visibility-1000v10_depth_normal'
# Specify the batch size for the minibatches used for training.
training_batch_size = 8
# Specify the downsample factor for the output heatmap.
output_downsample = 8
# Specify the number of epochs to use during training.
num_epochs = 200
# Specify the number of landmarks and the block size. The number of landmarks should be
# identical to the number of landmarks in the landmark file specified for the
# landmark_config parameter.
num_landmarks = 1000
# Specify the number of landmarks that will be present in each subset when the set of
# landmarks is partitioned into mutually exclusive subsets. The value specified here
# should exactly divide the landmark count. For example, when num_landmarks = 1000 and
# block_size = 125, we get 1000/125 = 8 subsets of landmarks.
block_size = 125
# Specify which subset you want to train the model for. For example, when
# num_landmarks = 1000 and block_size = 125, then subset_index = 0 indicates that the
# range of indices of landmarks in the subset would be [0, 125]. If subset_index = 1,
# then the range of indices would be [125, 250].
subset_index = 0
# Format the experiment name.
experiment_name = '%s_%d-%d_v%d' % (scene_name, num_landmarks, block_size, version_no)
# Format the model_dir string
landmark_start_index = subset_index * block_size
landmark_stop_index = (subset_index + 1) * block_size
if landmark_start_index < 0 | landmark_stop_index > num_landmarks:
raise Exception('landmark indices are outside valid range!')
else:
tmp = '%s-%03d-%03d' % (scene_name, landmark_start_index, landmark_stop_index)
model_dir = os.path.join(output_dir, experiment_name, tmp)
# Create the model_dir folder.
os.makedirs(model_dir, exist_ok=True)
# Create the command line string for the training job.
cmd = 'python ./local_training.py'
cmd += ' --dataset_dir %s' % dataset_dir
cmd += ' --scene_id %s' % scene_name
cmd += ' --experiment_file %s.txt' % experiment_name
cmd += ' --num_landmarks %d' % num_landmarks
cmd += ' --block_size %d' % block_size
cmd += ' --landmark_config %s' % landmark_config
cmd += ' --visibility_config %s' % visibility_config
cmd += ' --subset_index %d' % subset_index
cmd += ' --output_dir %s' % output_dir
cmd += ' --model_dir %s' % model_dir
cmd += ' --training_batch_size %d' % training_batch_size
cmd += ' --output_downsample %d' % output_downsample
cmd += ' --num_epochs %d' % num_epochs
# Launch training
os.system(cmd)
================================================
FILE: src/train.py
================================================
from datetime import datetime
import logging
import matplotlib.pyplot as plt
import numpy as np
import os
import pickle
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm
from inference import *
from dataloader.indoor6 import *
from models.efficientlitesld import EfficientNetSLD
from utils.heatmap import generate_heat_maps_gpu
def plotting(ROOT_FOLDER):
data = pickle.load(open('%s/stats.pkl' % ROOT_FOLDER, 'rb'))
fig, axs = plt.subplots(4, 1)
t = 0
s = []
epoch = 0
for i in range(len(data['train'])-1):
if data['train'][i+1]['ep'] == epoch + 1:
epoch += 1
else:
t += 1
s.append(data['train'][i]['loss'])
t = np.arange(0, t)
s = np.array(s)
s = np.convolve(s, np.ones(10)/10., mode='same')
axs[0].plot(t, np.log(s))
axs[0].set(xlabel='iterations', ylabel='loss', title='')
axs[0].grid()
max_grad = np.array([data['train'][i]['max_grad'] for i in range(len(data['train']))])
axs[1].plot(np.arange(0, len(max_grad)), np.log10(max_grad))
axs[1].set(xlabel='iterations', ylabel='max gradient', title='')
axs[1].grid()
t = np.array([data['eval'][i]['ep'] for i in range(len(data['eval']))])
s = np.array([np.median(data['eval'][i]['pixel_error']) for i in range(len(data['eval']))])
axs[2].plot(t, s)
axs[2].set(xlabel='epoch', ylabel='Pixel error', title='')
axs[2].grid()
axs[2].set_yticks(np.arange(0, 20, 5), minor=False)
axs[2].set_ylim(0, 20)
r = np.array([data['eval'][i]['recall'] for i in range(len(data['eval']))])
axs[3].plot(t, r)
axs[3].set(xlabel='epoch', ylabel='recall', title='')
axs[3].grid()
plt.subplots_adjust(left=None, bottom=None, right=None, top=None, wspace=0.8, hspace=1.0)
plt.close()
fig.savefig('%s/curve_train_test.png' % ROOT_FOLDER, format='png', dpi=120)
def train(opt):
if not os.path.exists(opt.output_folder):
os.makedirs(opt.output_folder)
logging.basicConfig(filename='%s/training.log' % opt.output_folder, filemode='a', level=logging.DEBUG, format='')
logging.info("Scene Landmark Detector Training")
print('Start training ...')
stats_pkl_logging = {'train': [], 'eval': []}
device = opt.gpu_device
assert len(opt.landmark_indices) == 0 or len(opt.landmark_indices) == 2, "landmark indices must be empty or length 2"
train_dataset = Indoor6(landmark_idx=np.arange(opt.landmark_indices[0],
opt.landmark_indices[1]) if len(opt.landmark_indices) == 2 else [None],
scene_id=opt.scene_id,
mode='train',
root_folder=opt.dataset_folder,
input_image_downsample=2,
landmark_config=opt.landmark_config,
visibility_config=opt.visibility_config,
skip_image_index=1)
train_dataloader = DataLoader(dataset=train_dataset, num_workers=4, batch_size=opt.training_batch_size, shuffle=True,
pin_memory=True)
## Save the trained landmark configurations
np.savetxt(os.path.join(opt.output_folder, 'landmarks.txt'), train_dataset.landmark)
np.savetxt(os.path.join(opt.output_folder, 'visibility.txt'), train_dataset.visibility, fmt='%d')
num_landmarks = train_dataset.landmark.shape[1]
if opt.model == 'efficientnet':
cnn = EfficientNetSLD(num_landmarks=num_landmarks, output_downsample=opt.output_downsample).to(device=device)
optimizer = torch.optim.AdamW(cnn.parameters(), lr=1e-3, betas=(0.9, 0.999), eps=1e-4, weight_decay=0.01)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5)
lowest_median_angular_error = 1e6
for epoch in range(opt.num_epochs):
# Training
training_loss = 0
for idx, batch in enumerate(tqdm(train_dataloader)):
cnn.train()
images = batch['image'].to(device=device)
B, _, H, W = images.shape
visibility = batch['visibility'].reshape(B, num_landmarks).to(device=device)
landmark2d = batch['landmark2d'].reshape(B, 2, num_landmarks).to(device=device)
# Resolution configure
landmark2d /= opt.output_downsample
heat_map_size = [H // opt.output_downsample, W // opt.output_downsample]
gt = generate_heat_maps_gpu(landmark2d,
visibility,
heat_map_size,
sigma=torch.tensor([5.], dtype=torch.float, device=device, requires_grad=False))
gt.requires_grad = False
# Clear gradient
optimizer.zero_grad()
# CNN forward pass
pred = cnn(images)['1']
# Compute loss and do backward pass
losses = torch.sum((pred[visibility != 0.5] - gt[visibility != 0.5]) ** 2)
training_loss += losses.detach().clone().item()
losses.backward()
optimizer.step()
logging.info('epoch %d, iter %d, loss %4.4f' % (epoch, idx, losses.item()))
stats_pkl_logging['train'].append({'ep': epoch, 'iter': idx, 'loss': losses.item()})
# Saving the ckpt
path = '%s/model-latest.ckpt' % (opt.output_folder)
torch.save(cnn.state_dict(), path)
if scheduler.get_last_lr()[-1] > 5e-5:
scheduler.step()
opt.pretrained_model = path
eval_stats = inference(opt, opt_tight_thr=1e-3, minimal_tight_thr=1e-3, mode='val')
median_angular_error = np.median(eval_stats['angular_error'])
if (median_angular_error < lowest_median_angular_error):
lowest_median_angular_error = median_angular_error
path = '%s/model-best_median.ckpt' % (opt.output_folder)
torch.save(cnn.state_dict(), path)
# date time
ts = datetime.datetime.now().timestamp()
dt = datetime.datetime.fromtimestamp(ts)
datestring = dt.strftime("%Y-%m-%d_%H-%M-%S")
# Print, log and update plot
stats_pkl_logging['eval'].append(
{'ep': epoch,
'angular_error': eval_stats['angular_error'],
'pixel_error': eval_stats['pixel_error'],
'recall': eval_stats['r5p5']
})
str_log = 'epoch %3d: [%s] ' \
'tr_loss= %10.2f, ' \
'lowest_median= %8.4f deg. ' \
'recall= %2.4f ' \
'angular-err(deg.)= [%7.4f %7.4f %7.4f] ' \
'pixel-err= [%4.3f %4.3f %4.3f] [mean/med./min] ' % (epoch, datestring, training_loss,
lowest_median_angular_error,
eval_stats['r5p5'],
np.mean(eval_stats['angular_error']),
np.median(eval_stats['angular_error']),
np.min(eval_stats['angular_error']),
np.mean(eval_stats['pixel_error']),
np.median(eval_stats['pixel_error']),
np.min(eval_stats['pixel_error']))
print(str_log)
logging.info(str_log)
with open('%s/stats.pkl' % opt.output_folder, 'wb') as f:
pickle.dump(stats_pkl_logging, f)
plotting(opt.output_folder)
def train_patches(opt):
if not os.path.exists(opt.output_folder):
os.makedirs(opt.output_folder)
logging.basicConfig(filename='%s/training.log' % opt.output_folder, filemode='a', level=logging.DEBUG, format='')
logging.info("Scene Landmark Detector Training Patches")
stats_pkl_logging = {'train': [], 'eval': []}
device = opt.gpu_device
assert len(opt.landmark_indices) == 0 or len(opt.landmark_indices) == 2, "landmark indices must be empty or length 2"
train_dataset = Indoor6Patches(landmark_idx=np.arange(opt.landmark_indices[0],
opt.landmark_indices[1]) if len(opt.landmark_indices) == 2 else [None],
scene_id=opt.scene_id,
mode='train',
root_folder=opt.dataset_folder,
input_image_downsample=2,
landmark_config=opt.landmark_config,
visibility_config=opt.visibility_config,
skip_image_index=1)
train_dataloader = DataLoader(dataset=train_dataset, num_workers=4, batch_size=opt.training_batch_size, shuffle=True,
pin_memory=True)
## Save the trained landmark configurations
np.savetxt(os.path.join(opt.output_folder, 'landmarks.txt'), train_dataset.landmark)
np.savetxt(os.path.join(opt.output_folder, 'visibility.txt'), train_dataset.visibility, fmt='%d')
num_landmarks = train_dataset.landmark.shape[1]
if opt.model == 'efficientnet':
cnn = EfficientNetSLD(num_landmarks=num_landmarks, output_downsample=opt.output_downsample).to(device=device)
optimizer = torch.optim.AdamW(cnn.parameters(), lr=1e-3, betas=(0.9, 0.999), eps=1e-4, weight_decay=0.01)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=40, gamma=0.5)
lowest_median_angular_error = 1e6
for epoch in range(opt.num_epochs):
# Training
training_loss = 0
for idx, batch in enumerate(tqdm(train_dataloader)):
cnn.train()
B1, B2, _, H, W = batch['patches'].shape
B = B1 * B2
patches = batch['patches']
visibility = batch['visibility']
landmark2d = batch['landmark2d']
# highest supported precision for MPS is FP32
if device.lower() == 'mps':
patches = patches.float()
visibility = visibility.float()
landmark2d = landmark2d.float()
patches = patches.reshape(B, 3, H, W).to(device=device)
visibility = visibility.reshape(B, num_landmarks).to(device=device)
landmark2d = landmark2d.reshape(B, 2, num_landmarks).to(device=device)
# Batch randomization
input_batch_random = np.random.permutation(B)
landmark2d_rand = [landmark2d[input_batch_random[b:b + 1]] for b in range(B)]
patches_rand = [patches[input_batch_random[b:b + 1]] for b in range(B)]
visibility_rand = [visibility[input_batch_random[b:b + 1]] for b in range(B)]
landmark2d_rand = torch.cat(landmark2d_rand, dim=0)
patches_rand = torch.cat(patches_rand, dim=0)
visibility_rand = torch.cat(visibility_rand, axis=0)
# Resolution configure
landmark2d_rand /= opt.output_downsample
heat_map_size = [H // opt.output_downsample, W // opt.output_downsample]
gt = generate_heat_maps_gpu(landmark2d_rand,
visibility_rand,
heat_map_size,
sigma=torch.tensor([20. / opt.output_downsample], dtype=torch.float, device=device, requires_grad=False))
gt.requires_grad = False
# Clear gradient
optimizer.zero_grad()
# CNN forward pass
pred = cnn(patches_rand)['1']
# Compute loss and do backward pass
losses = torch.sum((pred[visibility_rand != 0.5] - gt[visibility_rand != 0.5]) ** 2)
training_loss += losses.detach().clone().item()
losses.backward()
m = torch.tensor([0.0]).to(device)
for p in cnn.parameters():
m = torch.max(torch.max(torch.abs(p.grad.data)), m)
## Ignore batch with large gradient element
if epoch == 0 or (epoch > 0 and m < 1e4):
optimizer.step()
else:
cnn.load_state_dict(torch.load('%s/model-best_median.ckpt' % (opt.output_folder)))
cnn.to(device=device)
logging.info('epoch %d, iter %d, loss %4.4f' % (epoch, idx, losses.item()))
stats_pkl_logging['train'].append({'ep': epoch, 'iter': idx, 'loss': losses.item(), 'max_grad': m.cpu().numpy()})
# Saving the ckpt
path = '%s/model-latest.ckpt' % (opt.output_folder)
torch.save(cnn.state_dict(), path)
if scheduler.get_last_lr()[-1] > 5e-5:
scheduler.step()
opt.pretrained_model = [path]
eval_stats = inference(opt, opt_tight_thr=1e-3, minimal_tight_thr=1e-3, mode='val')
median_angular_error = np.median(eval_stats['angular_error'])
path = '%s/model-best_median.ckpt' % (opt.output_folder)
if (median_angular_error < lowest_median_angular_error):
lowest_median_angular_error = median_angular_error
torch.save(cnn.state_dict(), path)
if (~os.path.exists(path) and len(eval_stats['angular_error']) == 0):
torch.save(cnn.state_dict(), path)
# date time
ts = datetime.now().timestamp()
dt = datetime.fromtimestamp(ts)
datestring = dt.strftime("%Y-%m-%d_%H-%M-%S")
# Print, log and update plot
stats_pkl_logging['eval'].append(
{'ep': epoch,
'angular_error': eval_stats['angular_error'],
'pixel_error': eval_stats['pixel_error'],
'recall': eval_stats['r5p5']
})
try:
str_log = 'epoch %3d: [%s] ' \
'tr_loss= %10.2f, ' \
'lowest_median= %8.4f deg. ' \
'recall= %2.4f ' \
'angular-err(deg.)= [%7.4f %7.4f %7.4f] ' \
'pixel-err= [%4.3f %4.3f %4.3f] [mean/med./min] ' % (epoch, datestring, training_loss,
lowest_median_angular_error,
eval_stats['r5p5'],
np.mean(eval_stats['angular_error']),
np.median(eval_stats['angular_error']),
np.min(eval_stats['angular_error']),
np.mean(eval_stats['pixel_error']),
np.median(eval_stats['pixel_error']),
np.min(eval_stats['pixel_error']))
print(str_log)
logging.info(str_log)
except ValueError: #raised if array is empty.
str_log = 'epoch %3d: [%s] ' \
'tr_loss= %10.2f, ' \
'No correspondences found' % (epoch, datestring, training_loss)
print(str_log)
logging.info(str_log)
with open('%s/stats.pkl' % opt.output_folder, 'wb') as f:
pickle.dump(stats_pkl_logging, f)
plotting(opt.output_folder)
================================================
FILE: src/utils/generate_visibility_depth_normal.py
================================================
import argparse
import copy
import fnmatch
import numpy as np
import open3d as o3d
import os
import pickle
from PIL import Image
from torch.utils.data import DataLoader
from tqdm import tqdm
import sys
sys.path.append(os.path.join(sys.path[0], '..'))
from dataloader.indoor6 import Indoor6
def extract(opt):
DATASET_FOLDER = os.path.join(opt.dataset_folder)
test_dataset = Indoor6(scene_id=opt.scene_id,
mode='all',
root_folder=DATASET_FOLDER,
input_image_downsample=1,
landmark_config=opt.landmark_config,
visibility_config=opt.visibility_config,
skip_image_index=1)
test_dataloader = DataLoader(dataset=test_dataset, num_workers=1, batch_size=1, shuffle=False, pin_memory=True)
return test_dataloader, test_dataset
if __name__ == '__main__':
parser = argparse.ArgumentParser(
description='Scene Landmark Detection',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument(
'--dataset_folder', type=str, required=False,
help='Root directory, where all data is stored')
parser.add_argument(
'--output_folder', type=str, required=False,
help='Output folder')
parser.add_argument(
'--landmark_config', type=str, default='landmarks/landmarks-300',
help='Landmark configuration.')
parser.add_argument(
'--visibility_config', type=str, default='landmarks/visibility-300',
help='Visibility configuration.')
parser.add_argument(
'--scene_id', type=str, default='scene1',
help='Scene id')
opt = parser.parse_args()
monodepth_folder = os.path.join(opt.dataset_folder, opt.scene_id, 'depth')
from read_write_models import *
cameras, images, points = read_model(os.path.join(opt.dataset_folder, 'indoor6-colmap/%s/sparse/0' % opt.scene_id), ext='.bin')
indoor6_name_2to_colmap_index = {}
for k in images:
indoor6_name_2to_colmap_index[images[k].name] = k
# print(images[k])
dataloader, data = extract(opt)
augmented_visibility = copy.deepcopy(data.visibility)
monodepth_folder = os.path.join(opt.dataset_folder,
opt.scene_id,
'depth')
count_invalid_images = 0
##############################################################
### Creating depth images and augment visibility based on ####
### the consistency between depth and 3D points from colmap ##
##############################################################
for idx, batch in enumerate(tqdm(dataloader)):
_, _, H, W = batch['image'].shape
# batch['intrinsic']
original_image_name = data.original_image_name(idx)
colmap_index = indoor6_name_2to_colmap_index[original_image_name]
if images[colmap_index].name != original_image_name:
print('indoor6 name: ', data.image_files[idx], ', original name ', original_image_name)
point3D_ids = images[colmap_index].point3D_ids
K = batch['intrinsics'][0].cpu().numpy()
R = batch['pose_gt'][0, :3, :3].cpu().numpy()
t = batch['pose_gt'][0, :3, 3].cpu().numpy()
xys = images[colmap_index].xys
monoscaled_depth_path = os.path.join(monodepth_folder, data.image_files[idx].replace('.jpg', '.scaled_depth.npy'))
dmonodense_scaled = None
if os.path.exists(monoscaled_depth_path):
dmonodense_scaled = np.load(monoscaled_depth_path)
# else:
# dmonodense = np.load(os.path.join(monodepth_folder, data.image_files[idx].replace('jpg', 'npy')))
# ds = np.zeros(len(point3D_ids))
# dmono = np.zeros(len(point3D_ids))
# validIdx = 0
# for i, k in enumerate(point3D_ids):
# if k != -1:
# Cp = R @ points[k].xyz + t
# xyz = K @ Cp
# proj_x = xyz[0] / xyz[2]
# proj_y = xyz[1] / xyz[2]
# px = xys[i][0]
# py = xys[i][1]
# if Cp[2] < 15.0 and proj_x >= 0 and proj_x < W and proj_y >= 0 and proj_y < H and np.abs(proj_x-px) < 5.0 and np.abs(proj_y-py) < 5.0:
# ds[validIdx] = Cp[2]
# dmono[validIdx] = dmonodense[int(proj_y), int(proj_x)]
# ## Doing sth here to compute surface normal
# validIdx += 1
# if validIdx < 10:
# dmonodense_scaled = None
# count_invalid_images += 1
# else:
# ds = ds[:validIdx]
# dmono = dmono[:validIdx]
# A = np.array([[np.sum(dmono**2), np.sum(dmono)], [np.sum(dmono), validIdx]])
# b = np.array([np.sum(dmono*ds), np.sum(ds)])
# k = np.linalg.solve(A, b)
# dmonodense_scaled = k[0] * dmonodense + k[1]
# np.save(monoscaled_depth_path, dmonodense_scaled)
if dmonodense_scaled is not None:
Cplm = batch['landmark3d'][0].cpu().numpy()
pixlm = K @ Cplm
px = pixlm[0] / pixlm[2]
py = pixlm[1] / pixlm[2]
infront_infrustum = (Cplm[2] > 0.3) * (Cplm[2] < 15.0) * (px >= 0) * (px < W) * (py >=0) * (py < H)
vis = copy.deepcopy(augmented_visibility[:, data.image_indices[idx]])
count_colmap_vs_depth_incompatibility = 0
count_infront_infrustum = 0
for l in range(data.landmark.shape[1]):
if infront_infrustum[l]:
count_infront_infrustum += 1
depth_from_scaled_mono = dmonodense_scaled[int(py[l]), int(px[l])]
depth_from_lm_proj = Cplm[2, l]
rel_depth = np.abs(depth_from_lm_proj - depth_from_scaled_mono) / depth_from_lm_proj
if vis[l]==0:
if rel_depth < 0.3: ## 30% depth compatible
vis[l] = True
augmented_visibility[:, data.image_indices[idx]] = vis
np.savetxt(os.path.join(opt.dataset_folder, opt.scene_id, opt.visibility_config + '_depth.txt'), augmented_visibility, fmt='%d')
#########################################################
### Adding visibility refinement using surface normal ###
#########################################################
root_folder=opt.dataset_folder
scene_id=opt.scene_id
data = pickle.load(open('%s/%s/train_test_val.pkl' % (root_folder, scene_id), 'rb'))
imgs = data['train'] + data['val'] + data['test']
idx = data['train_idx'] + data['val_idx'] + data['test_idx']
landmark_config = opt.landmark_config
visibility_config = opt.visibility_config
visibility_depth_config = visibility_config + '_depth'
np.random.seed(100)
landmark_colors = np.random.rand(10000, 3)
landmark_file = open(root_folder + '/' + scene_id + '/%s.txt' % landmark_config, 'r')
num_landmark = int(landmark_file.readline())
lm = []
for l in range(num_landmark):
pl = landmark_file.readline().split()
pl = np.array([float(pl[i]) for i in range(len(pl))])
lm.append(pl)
lm = np.asarray(lm)[:, 1:].T
visibility_file = root_folder + '/' + scene_id + '/%s.txt' % visibility_config
visibility = np.loadtxt(visibility_file).astype(bool)
visibility_file = root_folder + '/' + scene_id + '/%s.txt' % visibility_depth_config
visibility_depth = np.loadtxt(visibility_file).astype(bool)
new_visibility = copy.deepcopy(visibility_depth)
lm_spheres = []
mesh_arrows = []
mesh_arrows_ref = []
H = 720
W = 1280
WW, HH = np.meshgrid(np.arange(W), np.arange(H))
WW = WW.reshape(1, H, W)
HH = HH.reshape(1, H, W)
wh1 = np.concatenate((WW, HH, np.ones_like(HH)), axis=0)
lm_sn = np.zeros((num_landmark, 6))
lm_sn[:, :3] = lm.T
for lm_idx in tqdm(range(visibility.shape[0])):
## Observe from colmap
visibility_matrix_ids = [i for i in np.where(visibility[lm_idx, idx])[0]]
images_observe_lm = [imgs[i] for i in visibility_matrix_ids]
pose_paths = [os.path.join(root_folder, scene_id, 'images', ifile.replace('color.jpg', 'pose.txt')) for ifile in images_observe_lm]
depth_paths = [os.path.join(root_folder, scene_id, 'depth', ifile.replace('.jpg', '.scaled_depth.npy')) for ifile in images_observe_lm]
intrinsic_paths = [os.path.join(root_folder, scene_id, 'images', ifile.replace('color.jpg', 'intrinsics.txt')) for ifile in images_observe_lm]
depths = np.zeros((len(pose_paths), H, W))
Ts = np.zeros((len(pose_paths), 4, 4))
Ks = np.zeros((len(pose_paths), 3, 3))
for i, pp in enumerate(pose_paths):
T = np.loadtxt(pp)
T = np.concatenate( (T, np.array([[0, 0, 0, 1]])), axis=0)
Ts[i] = T
intrinsics = open(intrinsic_paths[i])
intrinsics = intrinsics.readline().split()
fx = float(intrinsics[2])
fy = float(intrinsics[2])
cx = float(intrinsics[3])
cy = float(intrinsics[4])
K = np.array([[fx, 0., cx],
[0., fy, cy],
[0., 0., 1.]])
Ks[i] = K
## First estimate for surface normal using just visibility vector
bsum = np.zeros(3)
for i in range(Ts.shape[0]):
Gpt = lm[:, lm_idx] + Ts[i, :3, :3].T @ Ts[i, :3, 3]
bsum -= (Gpt / np.linalg.norm(Gpt))
bsum /= np.linalg.norm(bsum)
## Refine the surface normal based on depth image
bref = np.zeros(3)
patch_size = 50
for i in range(Ts.shape[0]):
if os.path.exists(depth_paths[i]):
cp = Ts[i, :3, :3] @ lm[:, lm_idx] + Ts[i, :3, 3]
cp = Ks[i] @ cp
cp = cp.reshape(-1)
proj_x = int(cp[0] / cp[2])
proj_y = int(cp[1] / cp[2])
if proj_x >= patch_size and proj_x < W-patch_size and proj_y >= patch_size and proj_y < H-patch_size:
patch_x0, patch_x1 = proj_x-patch_size, proj_x+patch_size
patch_y0, patch_y1 = proj_y-patch_size, proj_y+patch_size
d = np.load(depth_paths[i])[patch_y0:patch_y1, patch_x0:patch_x1].reshape((1, patch_size * 2, patch_size * 2))
pcd = np.linalg.inv(Ks[i]) @ (wh1[:, patch_y0:patch_y1, patch_x0:patch_x1] * d).reshape(3, 4 * patch_size ** 2)
A = np.concatenate((pcd, np.ones((1, 4 * patch_size ** 2))), axis=0)
D, U = np.linalg.eig(A @ A.T)
sn = Ts[i, :3, :3].T @ U[:3, np.argsort(D)[0]]
sn /= np.linalg.norm(sn)
if np.sum(bsum * sn) > 0.0:
bref += sn
elif np.sum(bsum * sn) < 0.0:
bref -= sn
if np.linalg.norm(bref) == 0:
lm_sn[lm_idx, 3:] = bsum
else:
bref /= np.linalg.norm(bref)
lm_sn[lm_idx, 3:] = bref
visibility_matrix_ids = [i for i in np.where(visibility_depth[lm_idx, idx])[0]]
images_observe_lm = [imgs[i] for i in np.where(visibility_depth[lm_idx, idx])[0]]
pose_paths = [os.path.join(root_folder, scene_id, 'images', ifile.replace('color.jpg', 'pose.txt')) for ifile in images_observe_lm]
for i, pp in enumerate(pose_paths):
T = np.loadtxt(pp)
if visibility_depth[lm_idx, idx[visibility_matrix_ids[i]]]:
Gpt = lm[:, lm_idx] + T[:3, :3].T @ T[:3, 3]
Gpt /= np.linalg.norm(Gpt)
if np.sum(bref * Gpt) > -0.2: ## violate visibility direction
new_visibility[lm_idx, idx[visibility_matrix_ids[i]]] = 0
np.savetxt(os.path.join(root_folder, scene_id, '%s_normal.txt' % (landmark_config)), lm_sn)
np.savetxt(os.path.join(root_folder, scene_id, '%s_depth_normal.txt' % (visibility_config)), new_visibility, fmt='%d')
================================================
FILE: src/utils/heatmap.py
================================================
import numpy as np
import torch
def generate_heat_maps(landmarks, visibility_mask, heatmap_size, K, sigma=3):
'''
:param landmarks: [3, L]
:param visibility_mask: [L]
:return: hms, hms_weight(1: visible, 0: invisible)
'''
hms = np.zeros((landmarks.shape[1],
heatmap_size[0],
heatmap_size[1]),
dtype=np.float32)
hms_weights = np.ones((landmarks.shape[1]), dtype=np.float32)
tmp_size = sigma * 3
for lm_id in range(landmarks.shape[1]):
landmark_2d = K @ landmarks[:, lm_id]
landmark_2d /= landmark_2d[2]
mu_x = int(landmark_2d[0] + 0.5)
mu_y = int(landmark_2d[1] + 0.5)
# Check that any part of the gaussian is in-bounds
ul = [int(mu_y - tmp_size), int(mu_x - tmp_size)]
br = [int(mu_y + tmp_size + 1), int(mu_x + tmp_size + 1)]
if ul[0] >= heatmap_size[0] or ul[1] >= heatmap_size[1] \
or br[0] < 0 or br[1] < 0 or landmarks[2, lm_id] < 0:
continue
if visibility_mask[lm_id]:
## Generate gaussian
size = 2 * tmp_size + 1
x = np.arange(0, size, 1, np.float32)
y = x[:, np.newaxis]
x0 = y0 = size // 2
# The gaussian is not normalized, we want the center value to equal 1
g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2))
# Usable gaussian range
g_y = max(0, -ul[0]), min(br[0], heatmap_size[0]) - ul[0]
g_x = max(0, -ul[1]), min(br[1], heatmap_size[1]) - ul[1]
# Image range
img_y = max(0, ul[0]), min(br[0], heatmap_size[0])
img_x = max(0, ul[1]), min(br[1], heatmap_size[1])
hms[lm_id][img_y[0]:img_y[1], img_x[0]:img_x[1]] = \
g[g_y[0]:g_y[1], g_x[0]:g_x[1]]
else:
hms_weights[lm_id] = 0.0
return hms, hms_weights
def generate_heat_maps_gpu(landmarks_2d, visibility_mask, heatmap_size, sigma=3):
'''
gpu version of heat map generation
:param landmarks: [3, L]
:return: hms
'''
B, _, L = landmarks_2d.shape
H, W = heatmap_size[0], heatmap_size[1]
yy_grid, xx_grid = torch.meshgrid(torch.arange(0, heatmap_size[0]),
torch.arange(0, heatmap_size[1]))
xx_grid, yy_grid = xx_grid.to(device=landmarks_2d.device), yy_grid.to(device=landmarks_2d.device)
hms = torch.exp(-((xx_grid.reshape(1, 1, H, W)-landmarks_2d[:, 0].reshape(B, L, 1, 1))**2 +
(yy_grid.reshape(1, 1, H, W)-landmarks_2d[:, 1].reshape(B, L, 1, 1))**2)/(2*sigma**2))
hms_vis = hms * visibility_mask.reshape(B, L, 1, 1).float()
hms_vis[hms_vis < 0.1] = 0.0
normalizing_factor, _ = torch.max(hms_vis.reshape(B, L, -1), dim=2)
hms_vis[normalizing_factor > 0.5] = hms_vis[normalizing_factor > 0.5] / \
normalizing_factor.reshape(B, L, 1, 1)[normalizing_factor > 0.5]
return hms_vis
================================================
FILE: src/utils/landmark_selection.py
================================================
import argparse
import numpy as np
import os
import pickle
from read_write_models import qvec2rotmat, read_model
from tqdm import tqdm
def ComputePerPointTimeSpan(image_ids, images):
timespan = {}
for imageID in image_ids:
session_id = int(images[imageID].name.split('-')[0])
if session_id in timespan:
timespan[session_id] += 1
else:
timespan[session_id] = 1
return len(timespan)
def ComputePerPointDepth(pointInGlobal, image_ids, images):
d = np.zeros(len(image_ids))
for i, imageID in enumerate(image_ids):
R = qvec2rotmat(images[imageID].qvec)
t = images[imageID].tvec
pointInCamerai = R @ pointInGlobal + t
d[i] = pointInCamerai[2]
pointDepthMean, pointDepthStd = np.mean(d), np.std(d)
return pointDepthMean, pointDepthStd
def ComputePerPointAngularSpan(pointInGlobal, image_ids, images):
N = len(image_ids)
H = np.zeros((3, 3))
for i, imageID in enumerate(image_ids):
Ri = qvec2rotmat(images[imageID].qvec)
ti = images[imageID].tvec
bi = Ri.T @ (pointInGlobal - ti)
bi = bi / np.linalg.norm(bi)
H += (np.eye(3) - np.outer(bi, bi))
H /= N
eigH = np.linalg.eigvals(0.5*(H + H.T))
return np.arccos(np.clip(1 - 2.0 * np.min(eigH)/np.max(eigH), 0, 1))
def SaveLandmarksAndVisibilityMask(selected_landmarks, points3D, images, indoor6_imagename_to_index, num_images, root_path,
landmark_config, visibility_config, outformat):
num_landmarks = len(selected_landmarks['id'])
visibility_mask = np.zeros((num_landmarks, num_images), dtype=np.uint8)
for i, pid in enumerate(selected_landmarks['id']):
for imgid in points3D[pid].image_ids:
if images[imgid].name in indoor6_imagename_to_index:
visibility_mask[i, indoor6_imagename_to_index[images[imgid].name]] = 1
np.savetxt(os.path.join(root_path, '%s%s.txt' % (visibility_config, outformat)), visibility_mask, fmt='%d')
f = open(os.path.join(root_path, '%s%s.txt' % (landmark_config, outformat)), 'w')
f.write('%d\n' % num_landmarks)
for i in range(selected_landmarks['xyz'].shape[1]):
f.write('%d %4.4f %4.4f %4.4f\n' % (i,
selected_landmarks['xyz'][0, i],
selected_landmarks['xyz'][1, i],
selected_landmarks['xyz'][2, i]))
f.close()
if __name__ == '__main__':
parser = argparse.ArgumentParser(
description='Scene Landmark Detection',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument(
'--dataset_folder', type=str, required=True,
help='Root directory, where all data is stored')
parser.add_argument(
'--scene_id', type=str, default='scene6',
help='Scene id')
parser.add_argument(
'--num_landmarks', type=int, default=300,
help='Number of selected landmarks.')
parser.add_argument(
'--output_format', type=str, default='v2',
help='Landmark file output.')
opt = parser.parse_args()
opt.landmark_config = "landmarks/landmarks-%d" % (opt.num_landmarks)
opt.visibility_config = "landmarks/visibility-%d" % (opt.num_landmarks)
scene = opt.scene_id
path = os.path.join(opt.dataset_folder, 'indoor6-colmap/%s-tr/sparse/0/' % scene)
cameras, images, points3D = read_model(path, ext='.bin')
## Max number of sessions
sessions = {}
for i in images:
print(images[i].name)
session_id = int(images[i].name.split('-')[0])
sessions[session_id] = 1
maxSession = len(sessions)
## Initialization
numPoints3D = len(points3D)
points3D_ids = np.zeros(numPoints3D)
points3D_scores = np.zeros(numPoints3D)
validIdx = 0
## Compute score for each landmark
for i, k in enumerate(tqdm(points3D)):
pointInGlobal = points3D[k].xyz
image_ids = points3D[k].image_ids
trackLength = len(image_ids)
if trackLength > 25:
depthMean, depthStd = ComputePerPointDepth(pointInGlobal, image_ids, images)
timespan = ComputePerPointTimeSpan(image_ids, images)
anglespan = ComputePerPointAngularSpan(pointInGlobal, image_ids, images)
depthScore = min(1.0, depthStd / depthMean)
trackLengthScore = 0.25 * np.log2(trackLength)
timeSpanScore = timespan / maxSession
if timespan >= 1 and depthMean < 10.0 and anglespan > 0.3:
points3D_ids[validIdx] = k
points3D_scores[validIdx] = depthScore + trackLengthScore + timeSpanScore + anglespan
validIdx += 1
## Sort scores
points3D_ids = points3D_ids[:validIdx]
points3D_scores = points3D_scores[:validIdx]
sorted_indices = np.argsort(points3D_scores)
## Greedy selection
selected_landmarks = {'id': np.zeros(opt.num_landmarks),
'xyz': np.zeros((3, opt.num_landmarks)),
'score': np.zeros(opt.num_landmarks)}
## Selecting first point
selected_landmarks['id'][0] = points3D_ids[sorted_indices[-1]]
selected_landmarks['xyz'][:, 0] = points3D[selected_landmarks['id'][0]].xyz
selected_landmarks['score'][0] = points3D_scores[sorted_indices[-1]]
nselected = 1
radius = 5.0
while nselected < opt.num_landmarks:
for i in reversed(sorted_indices):
id = points3D_ids[i]
xyz = points3D[id].xyz
if np.sum(np.linalg.norm(xyz.reshape(3, 1) - selected_landmarks['xyz'][:, :nselected], axis=0) < radius):
continue
else:
selected_landmarks['id'][nselected] = id
selected_landmarks['xyz'][:, nselected] = xyz
selected_landmarks['score'][nselected] = points3D_scores[i]
nselected += 1
if nselected == opt.num_landmarks:
break
radius *= 0.5
## Saving
indoor6_images = pickle.load(open(os.path.join(opt.dataset_folder, '%s/train_test_val.pkl' % opt.scene_id), 'rb'))
indoor6_imagename_to_index = {}
for i, f in enumerate(indoor6_images['train']):
image_name = open(os.path.join(opt.dataset_folder,
opt.scene_id, 'images',
f.replace('color.jpg',
'intrinsics.txt'))).readline().split(' ')[-1][:-1]
indoor6_imagename_to_index[image_name] = indoor6_images['train_idx'][i]
num_images = len(indoor6_images['train']) + len(indoor6_images['val']) + len(indoor6_images['test'])
SaveLandmarksAndVisibilityMask(selected_landmarks, points3D, images, indoor6_imagename_to_index, num_images,
os.path.join(opt.dataset_folder, opt.scene_id),
opt.landmark_config, opt.visibility_config, opt.output_format)
================================================
FILE: src/utils/merge_landmark_files.py
================================================
import argparse
import copy
import numpy as np
import os
import sys
sys.path.append(os.path.join(sys.path[0], '..'))
from utils.select_additional_landmarks import load_landmark_visibility_files
def save_landmark_visibility_mask(landmarks, visibility_mask,
landmark_path, visibility_path):
num_landmarks = landmarks.shape[1]
np.savetxt(visibility_path, visibility_mask, fmt='%d')
f = open(landmark_path, 'w')
f.write('%d\n' % num_landmarks)
for i in range(num_landmarks):
f.write('%d %4.4f %4.4f %4.4f\n' % (i,
landmarks[0, i],
landmarks[1, i],
landmarks[2, i]))
f.close()
if __name__ == '__main__':
parser = argparse.ArgumentParser(
description='Scene Landmark Detection',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument(
'--dataset_folder', type=str, required=True,
help='Root directory, where all data is stored')
parser.add_argument(
'--scene_id', type=str, default='scene6',
help='Scene id')
parser.add_argument(
'--landmark_config', type=str, action='append',
help='File containing scene-specific 3D landmarks.')
parser.add_argument(
'--visibility_config', type=str, action='append',
help='File containing information about visibility of landmarks in cameras associated with training set.')
parser.add_argument(
'--output_format', type=str, required=True,
help='Output file format.')
opt = parser.parse_args()
assert len(opt.landmark_config) > 1
assert len(opt.landmark_config) == len(opt.visibility_config)
num_landmarks = 0
num_files = len(opt.landmark_config)
ls = []
vs = []
for (lp, vp) in zip(opt.landmark_config, opt.visibility_config):
landmark_path = os.path.join(opt.dataset_folder, opt.scene_id, lp + '.txt')
vis_path = os.path.join(opt.dataset_folder, opt.scene_id, vp + '.txt')
l, v = load_landmark_visibility_files(landmark_path=landmark_path,
visibility_path=vis_path)
num_landmarks += l.shape[1]
ls.append(l)
vs.append(v)
ls = np.concatenate(ls, axis=1)
vs = np.concatenate(vs, axis=0)
output_landmark_path = os.path.join(opt.dataset_folder, opt.scene_id, 'landmarks/landmarks-%d%s.txt' % (num_landmarks, opt.output_format))
if 'depth_normal' in opt.visibility_config[0]:
output_visibility_path = os.path.join(opt.dataset_folder, opt.scene_id, 'landmarks/visibility-%d%s_depth_normal.txt' % (num_landmarks, opt.output_format))
else:
output_visibility_path = os.path.join(opt.dataset_folder, opt.scene_id, 'landmarks/visibility-%d%s.txt' % (num_landmarks, opt.output_format))
save_landmark_visibility_mask(ls, vs, output_landmark_path, output_visibility_path)
================================================
FILE: src/utils/pnp.py
================================================
import numpy as np
from scipy.optimize import least_squares
def Rotation2Quaternion(R):
"""
Convert a rotation matrix to quaternion
Parameters
----------
R : ndarray of shape (3, 3)
Rotation matrix
Returns
-------
q : ndarray of shape (4,)
The unit quaternion (w, x, y, z)
"""
q = np.empty([4, ])
tr = np.trace(R)
if tr < 0:
i = R.diagonal().argmax()
j = (i + 1) % 3
k = (j + 1) % 3
q[i] = np.sqrt(1 - tr + 2 * R[i, i]) / 2
q[j] = (R[j, i] + R[i, j]) / (4 * q[i])
q[k] = (R[k, i] + R[i, k]) / (4 * q[i])
q[3] = (R[k, j] - R[j, k]) / (4 * q[i])
else:
q[3] = np.sqrt(1 + tr) / 2
q[0] = (R[2, 1] - R[1, 2]) / (4 * q[3])
q[1] = (R[0, 2] - R[2, 0]) / (4 * q[3])
q[2] = (R[1, 0] - R[0, 1]) / (4 * q[3])
q /= np.linalg.norm(q)
# Rearrange (x, y, z, w) to (w, x, y, z)
q = q[[3, 0, 1, 2]]
return q
def Quaternion2Rotation(q):
"""
Convert a quaternion to rotation matrix
Parameters
----------
q : ndarray of shape (4,)
Unit quaternion (w, x, y, z)
Returns
-------
R : ndarray of shape (3, 3)
The rotation matrix
"""
q /= np.linalg.norm(q)
w = q[0]
x = q[1]
y = q[2]
z = q[3]
R = np.empty([3, 3])
R[0, 0] = 1 - 2 * y ** 2 - 2 * z ** 2
R[0, 1] = 2 * (x * y - z * w)
R[0, 2] = 2 * (x * z + y * w)
R[1, 0] = 2 * (x * y + z * w)
R[1, 1] = 1 - 2 * x ** 2 - 2 * z ** 2
R[1, 2] = 2 * (y * z - x * w)
R[2, 0] = 2 * (x * z - y * w)
R[2, 1] = 2 * (y * z + x * w)
R[2, 2] = 1 - 2 * x ** 2 - 2 * y ** 2
return R
def skewsymm(x):
Sx = np.zeros((3, 3))
Sx[0, 1] = -x[2]
Sx[0, 2] = x[1]
Sx[1, 0] = x[2]
Sx[2, 0] = -x[1]
Sx[1, 2] = -x[0]
Sx[2, 1] = x[0]
return Sx
def VectorizeInitialPose(C_T_G):
R = C_T_G[:3, :3]
t = C_T_G[:3, 3]
q = Rotation2Quaternion(R)
z = np.concatenate([t, q])
return z
def MeasureReprojectionSinglePose(z, p, b, w):
n_points = b.shape[1]
q = z[3:7]
q_norm = np.sqrt(np.sum(q ** 2))
q = q / q_norm
R = Quaternion2Rotation(q)
t = z[:3]
b_hat = R @ p + t.reshape(3, 1)
b_hat_normalized = b_hat / np.sqrt(np.sum(b_hat ** 2, axis=0))
err = np.repeat(w, 3).reshape(n_points, 3).T * (b_hat_normalized - b)
return err.reshape(-1)
def UpdatePose(z):
p = z[0:7]
q = p[3:]
q = q / np.linalg.norm(q)
R = Quaternion2Rotation(q)
t = p[:3]
P_new = np.hstack([R, t[:, np.newaxis]])
return P_new
def P3PKe(m, X, inlier_thres=1e-5):
"""
Perspective-3-point algorithm from
Ke, T., & Roumeliotis, S. I. (CVPR'17). An efficient algebraic solution to the perspective-three-point problem.
Parameters
----------
m : ndarray of shape (3, 4)
unit bearing vectors to each landmarks w.r.t camera
X : ndarray of shape (3, 4)
3D points position w.r.t global
Returns
-------
R : ndarray of shape (3, 3)
t : ndarray of shape (3, 1)
(R, t) represents transformation from global to camera frame of reference
"""
w1 = X[:, 0]
w2 = X[:, 1]
w3 = X[:, 2]
u0 = w1 - w2
nu0 = np.linalg.norm(u0)
if nu0 < 1e-4:
return None, None
k1 = u0 / nu0
b1 = m[:, 0]
b2 = m[:, 1]
b3 = m[:, 2]
k3 = np.cross(b1, b2)
nk3 = np.linalg.norm(k3)
if nk3 < 1e-4:
return None, None
k3 = k3 / nk3
tz = np.cross(b1, k3)
v1 = np.cross(b1, b3)
v2 = np.cross(b2, b3)
u1 = w1 - w3
u1k1 = np.sum(u1 * k1)
k3b3 = np.sum(k3 * b3)
if np.abs(k3b3) < 1e-4:
return None, None
f11 = k3.T @ b3
f13 = k3.T @ v1
f15 = -u1k1 * f11
nl = np.cross(u1, k1)
delta = np.linalg.norm(nl)
if delta < 1e-4:
return None, None
nl = nl / delta
f11 = delta * f11
f13 = delta * f13
u2k1 = u1k1 - nu0
f21 = np.sum(tz * v2)
f22 = nk3 * k3b3
f23 = np.sum(k3 * v2)
f24 = u2k1 * f22
f25 = -u2k1 * f21
f21 = delta * f21
f22 = delta * f22
f23 = delta * f23
g1 = f13 * f22
g2 = f13 * f25 - f15 * f23
g3 = f11 * f23 - f13 * f21
g4 = -f13 * f24
g5 = f11 * f22
g6 = f11 * f25 - f15 * f21
g7 = -f15 * f24
alpha = np.array([g5 * g5 + g1 * g1 + g3 * g3,
2 * (g5 * g6 + g1 * g2 + g3 * g4),
g6 * g6 + 2 * g5 * g7 + g2 * g2 + g4 * g4 - g1 * g1 - g3 * g3,
2 * (g6 * g7 - g1 * g2 - g3 * g4),
g7 * g7 - g2 * g2 - g4 * g4])
if any(np.isnan(alpha)):
return None, None
sols = np.roots(alpha)
Ck1nl = np.vstack((k1, nl, np.cross(k1, nl))).T
Cb1k3tzT = np.vstack((b1, k3, tz))
b3p = (delta / k3b3) * b3
R = np.zeros((3, 3, 4))
t = np.zeros((3, 4))
for i in range(sols.shape[0]):
if np.imag(sols[i]) != 0:
continue
ctheta1p = np.real(sols[i])
if abs(ctheta1p) > 1:
continue
stheta1p = np.sqrt(1 - ctheta1p * ctheta1p)
if k3b3 < 0:
stheta1p = -stheta1p
ctheta3 = g1 * ctheta1p + g2
stheta3 = g3 * ctheta1p + g4
ntheta3 = stheta1p / ((g5 * ctheta1p + g6) * ctheta1p + g7)
ctheta3 = ntheta3 * ctheta3
stheta3 = ntheta3 * stheta3
C13 = np.array([[ctheta3, 0, -stheta3],
[stheta1p * stheta3, ctheta1p, stheta1p * ctheta3],
[ctheta1p * stheta3, -stheta1p, ctheta1p * ctheta3]])
Ri = (Ck1nl @ C13 @ Cb1k3tzT).T
pxstheta1p = stheta1p * b3p
ti = pxstheta1p - Ri @ w3
ti = ti.reshape(3, 1)
m_hat = Ri @ X + ti
m_hat = m_hat / np.linalg.norm(m_hat, axis=0)
if np.sum(np.sum(m_hat * m, axis=0) > 1.0 - inlier_thres) == 4:
return Ri, ti
return None, None
def P3PKe_Ransac(G_p_f, C_b_f_hm, w, thres=0.01):
inlier_thres = thres
C_T_G_best = None
inlier_best = np.zeros(G_p_f.shape[1], dtype=bool)
Nsample=4
inlier_score_best=0
for iter in range(125): #old value was 10
## Weighted sampling based on weight factor
min_set = np.argpartition(np.exp(w) * np.random.rand(w.shape[0]), -Nsample)[-Nsample:]
C_R_G_hat, C_t_G_hat = P3PKe(C_b_f_hm[:, min_set], G_p_f[:, min_set], inlier_thres=thres)
if C_R_G_hat is None or C_t_G_hat is None:
continue
# Get inlier
C_b_f_hat = C_R_G_hat @ G_p_f + C_t_G_hat
C_b_f_hat = C_b_f_hat / np.linalg.norm(C_b_f_hat, axis=0)
inlier_mask = np.sum(C_b_f_hat * C_b_f_hm, axis=0) > (1.0 - inlier_thres)
inlier_score = np.sum(w[inlier_mask])
if inlier_score > inlier_score_best:
inlier_best = inlier_mask
C_T_G_best = np.eye(4)
C_T_G_best[:3, :3] = C_R_G_hat
C_T_G_best[:3, 3:] = C_t_G_hat
inlier_score_best = inlier_score
return C_T_G_best, inlier_best
def RunPnPNL(C_T_G, G_p_f, C_b_f, w, cutoff=0.01):
'''
Weighted PnP based using weight w and bearing angular loss.
Return optimized P_new = optimized C_T_G.
'''
z0 = VectorizeInitialPose(C_T_G)
res = least_squares(
lambda x: MeasureReprojectionSinglePose(x, G_p_f, C_b_f, w),
z0,
verbose=0,
ftol=1e-4,
max_nfev=50,
xtol=1e-8,
loss='huber',
f_scale=cutoff
)
z = res.x
P_new = UpdatePose(z)
return P_new
================================================
FILE: src/utils/read_write_models.py
================================================
# Copyright (c) 2018, ETH Zurich and UNC Chapel Hill.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright
# notice, this list of conditions and the following disclaimer in the
# documentation and/or other materials provided with the distribution.
#
# * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of
# its contributors may be used to endorse or promote products derived
# from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#
# Author: Johannes L. Schoenberger (jsch-at-demuc-dot-de)
import os
import collections
import numpy as np
import struct
import argparse
import logging
logger = logging.getLogger(__name__)
CameraModel = collections.namedtuple(
"CameraModel", ["model_id", "model_name", "num_params"])
Camera = collections.namedtuple(
"Camera", ["id", "model", "width", "height", "params"])
BaseImage = collections.namedtuple(
"Image", ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"])
Point3D = collections.namedtuple(
"Point3D", ["id", "xyz", "rgb", "error", "image_ids", "point2D_idxs"])
class Image(BaseImage):
def qvec2rotmat(self):
return qvec2rotmat(self.qvec)
CAMERA_MODELS = {
CameraModel(model_id=0, model_name="SIMPLE_PINHOLE", num_params=3),
CameraModel(model_id=1, model_name="PINHOLE", num_params=4),
CameraModel(model_id=2, model_name="SIMPLE_RADIAL", num_params=4),
CameraModel(model_id=3, model_name="RADIAL", num_params=5),
CameraModel(model_id=4, model_name="OPENCV", num_params=8),
CameraModel(model_id=5, model_name="OPENCV_FISHEYE", num_params=8),
CameraModel(model_id=6, model_name="FULL_OPENCV", num_params=12),
CameraModel(model_id=7, model_name="FOV", num_params=5),
CameraModel(model_id=8, model_name="SIMPLE_RADIAL_FISHEYE", num_params=4),
CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5),
CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12)
}
CAMERA_MODEL_IDS = dict([(camera_model.model_id, camera_model)
for camera_model in CAMERA_MODELS])
CAMERA_MODEL_NAMES = dict([(camera_model.model_name, camera_model)
for camera_model in CAMERA_MODELS])
def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"):
"""Read and unpack the next bytes from a binary file.
:param fid:
:param num_bytes: Sum of combination of {2, 4, 8}, e.g. 2, 6, 16, 30, etc.
:param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}.
:param endian_character: Any of {@, =, <, >, !}
:return: Tuple of read and unpacked values.
"""
data = fid.read(num_bytes)
return struct.unpack(endian_character + format_char_sequence, data)
def write_next_bytes(fid, data, format_char_sequence, endian_character="<"):
"""pack and write to a binary file.
:param fid:
:param data: data to send, if multiple elements are sent at the same time,
they should be encapsuled either in a list or a tuple
:param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}.
should be the same length as the data list or tuple
:param endian_character: Any of {@, =, <, >, !}
"""
if isinstance(data, (list, tuple)):
bytes = struct.pack(endian_character + format_char_sequence, *data)
else:
bytes = struct.pack(endian_character + format_char_sequence, data)
fid.write(bytes)
def read_cameras_text(path):
"""
see: src/base/reconstruction.cc
void Reconstruction::WriteCamerasText(const std::string& path)
void Reconstruction::ReadCamerasText(const std::string& path)
"""
cameras = {}
with open(path, "r") as fid:
while True:
line = fid.readline()
if not line:
break
line = line.strip()
if len(line) > 0 and line[0] != "#":
elems = line.split()
camera_id = int(elems[0])
model = elems[1]
width = int(elems[2])
height = int(elems[3])
params = np.array(tuple(map(float, elems[4:])))
cameras[camera_id] = Camera(id=camera_id, model=model,
width=width, height=height,
params=params)
return cameras
def read_cameras_binary(path_to_model_file):
"""
see: src/base/reconstruction.cc
void Reconstruction::WriteCamerasBinary(const std::string& path)
void Reconstruction::ReadCamerasBinary(const std::string& path)
"""
cameras = {}
with open(path_to_model_file, "rb") as fid:
num_cameras = read_next_bytes(fid, 8, "Q")[0]
for _ in range(num_cameras):
camera_properties = read_next_bytes(
fid, num_bytes=24, format_char_sequence="iiQQ")
camera_id = camera_properties[0]
model_id = camera_properties[1]
model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name
width = camera_properties[2]
height = camera_properties[3]
num_params = CAMERA_MODEL_IDS[model_id].num_params
params = read_next_bytes(fid, num_bytes=8*num_params,
format_char_sequence="d"*num_params)
cameras[camera_id] = Camera(id=camera_id,
model=model_name,
width=width,
height=height,
params=np.array(params))
assert len(cameras) == num_cameras
return cameras
def write_cameras_text(cameras, path):
"""
see: src/base/reconstruction.cc
void Reconstruction::WriteCamerasText(const std::string& path)
void Reconstruction::ReadCamerasText(const std::string& path)
"""
HEADER = "# Camera list with one line of data per camera:\n" + \
"# CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n" + \
"# Number of cameras: {}\n".format(len(cameras))
with open(path, "w") as fid:
fid.write(HEADER)
for _, cam in cameras.items():
to_write = [cam.id, cam.model, cam.width, cam.height, *cam.params]
line = " ".join([str(elem) for elem in to_write])
fid.write(line + "\n")
def write_cameras_binary(cameras, path_to_model_file):
"""
see: src/base/reconstruction.cc
void Reconstruction::WriteCamerasBinary(const std::string& path)
void Reconstruction::ReadCamerasBinary(const std::string& path)
"""
with open(path_to_model_file, "wb") as fid:
write_next_bytes(fid, len(cameras), "Q")
for _, cam in cameras.items():
model_id = CAMERA_MODEL_NAMES[cam.model].model_id
camera_properties = [cam.id,
model_id,
cam.width,
cam.height]
write_next_bytes(fid, camera_properties, "iiQQ")
for p in cam.params:
write_next_bytes(fid, float(p), "d")
return cameras
def read_images_text(path):
"""
see: src/base/reconstruction.cc
void Reconstruction::ReadImagesText(const std::string& path)
void Reconstruction::WriteImagesText(const std::string& path)
"""
images = {}
with open(path, "r") as fid:
while True:
line = fid.readline()
if not line:
break
line = line.strip()
if len(line) > 0 and line[0] != "#":
elems = line.split()
image_id = int(elems[0])
qvec = np.array(tuple(map(float, elems[1:5])))
tvec = np.array(tuple(map(float, elems[5:8])))
camera_id = int(elems[8])
image_name = elems[9]
elems = fid.readline().split()
xys = np.column_stack([tuple(map(float, elems[0::3])),
tuple(map(float, elems[1::3]))])
point3D_ids = np.array(tuple(map(int, elems[2::3])))
images[image_id] = Image(
id=image_id, qvec=qvec, tvec=tvec,
camera_id=camera_id, name=image_name,
xys=xys, point3D_ids=point3D_ids)
return images
def read_images_binary(path_to_model_file):
"""
see: src/base/reconstruction.cc
void Reconstruction::ReadImagesBinary(const std::string& path)
void Reconstruction::WriteImagesBinary(const std::string& path)
"""
images = {}
with open(path_to_model_file, "rb") as fid:
num_reg_images = read_next_bytes(fid, 8, "Q")[0]
for _ in range(num_reg_images):
binary_image_properties = read_next_bytes(
fid, num_bytes=64, format_char_sequence="idddddddi")
image_id = binary_image_properties[0]
qvec = np.array(binary_image_properties[1:5])
tvec = np.array(binary_image_properties[5:8])
camera_id = binary_image_properties[8]
image_name = ""
current_char = read_next_bytes(fid, 1, "c")[0]
while current_char != b"\x00": # look for the ASCII 0 entry
image_name += current_char.decode("utf-8")
current_char = read_next_bytes(fid, 1, "c")[0]
num_points2D = read_next_bytes(fid, num_bytes=8,
format_char_sequence="Q")[0]
x_y_id_s = read_next_bytes(fid, num_bytes=24*num_points2D,
format_char_sequence="ddq"*num_points2D)
xys = np.column_stack([tuple(map(float, x_y_id_s[0::3])),
tuple(map(float, x_y_id_s[1::3]))])
point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3])))
images[image_id] = Image(
id=image_id, qvec=qvec, tvec=tvec,
camera_id=camera_id, name=image_name,
xys=xys, point3D_ids=point3D_ids)
return images
def write_images_text(images, path):
"""
see: src/base/reconstruction.cc
void Reconstruction::ReadImagesText(const std::string& path)
void Reconstruction::WriteImagesText(const std::string& path)
"""
if len(images) == 0:
mean_observations = 0
else:
mean_observations = sum((len(img.point3D_ids) for _, img in images.items()))/len(images)
HEADER = "# Image list with two lines of data per image:\n" + \
"# IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n" + \
"# POINTS2D[] as (X, Y, POINT3D_ID)\n" + \
"# Number of images: {}, mean observations per image: {}\n".format(len(images), mean_observations)
with open(path, "w") as fid:
fid.write(HEADER)
for _, img in images.items():
image_header = [img.id, *img.qvec, *img.tvec, img.camera_id, img.name]
first_line = " ".join(map(str, image_header))
fid.write(first_line + "\n")
points_strings = []
for xy, point3D_id in zip(img.xys, img.point3D_ids):
points_strings.append(" ".join(map(str, [*xy, point3D_id])))
fid.write(" ".join(points_strings) + "\n")
def write_images_binary(images, path_to_model_file):
"""
see: src/base/reconstruction.cc
void Reconstruction::ReadImagesBinary(const std::string& path)
void Reconstruction::WriteImagesBinary(const std::string& path)
"""
with open(path_to_model_file, "wb") as fid:
write_next_bytes(fid, len(images), "Q")
for _, img in images.items():
write_next_bytes(fid, img.id, "i")
write_next_bytes(fid, img.qvec.tolist(), "dddd")
write_next_bytes(fid, img.tvec.tolist(), "ddd")
write_next_bytes(fid, img.camera_id, "i")
for char in img.name:
write_next_bytes(fid, char.encode("utf-8"), "c")
write_next_bytes(fid, b"\x00", "c")
write_next_bytes(fid, len(img.point3D_ids), "Q")
for xy, p3d_id in zip(img.xys, img.point3D_ids):
write_next_bytes(fid, [*xy, p3d_id], "ddq")
def read_points3D_text(path):
"""
see: src/base/reconstruction.cc
void Reconstruction::ReadPoints3DText(const std::string& path)
void Reconstruction::WritePoints3DText(const std::string& path)
"""
points3D = {}
with open(path, "r") as fid:
while True:
line = fid.readline()
if not line:
break
line = line.strip()
if len(line) > 0 and line[0] != "#":
elems = line.split()
point3D_id = int(elems[0])
xyz = np.array(tuple(map(float, elems[1:4])))
rgb = np.array(tuple(map(int, elems[4:7])))
error = float(elems[7])
image_ids = np.array(tuple(map(int, elems[8::2])))
point2D_idxs = np.array(tuple(map(int, elems[9::2])))
points3D[point3D_id] = Point3D(id=point3D_id, xyz=xyz, rgb=rgb,
error=error, image_ids=image_ids,
point2D_idxs=point2D_idxs)
return points3D
def read_points3D_binary(path_to_model_file):
"""
see: src/base/reconstruction.cc
void Reconstruction::ReadPoints3DBinary(const std::string& path)
void Reconstruction::WritePoints3DBinary(const std::string& path)
"""
points3D = {}
with open(path_to_model_file, "rb") as fid:
num_points = read_next_bytes(fid, 8, "Q")[0]
for _ in range(num_points):
binary_point_line_properties = read_next_bytes(
fid, num_bytes=43, format_char_sequence="QdddBBBd")
point3D_id = binary_point_line_properties[0]
xyz = np.array(binary_point_line_properties[1:4])
rgb = np.array(binary_point_line_properties[4:7])
error = np.array(binary_point_line_properties[7])
track_length = read_next_bytes(
fid, num_bytes=8, format_char_sequence="Q")[0]
track_elems = read_next_bytes(
fid, num_bytes=8*track_length,
format_char_sequence="ii"*track_length)
image_ids = np.array(tuple(map(int, track_elems[0::2])))
point2D_idxs = np.array(tuple(map(int, track_elems[1::2])))
points3D[point3D_id] = Point3D(
id=point3D_id, xyz=xyz, rgb=rgb,
error=error, image_ids=image_ids,
point2D_idxs=point2D_idxs)
return points3D
def write_points3D_text(points3D, path):
"""
see: src/base/reconstruction.cc
void Reconstruction::ReadPoints3DText(const std::string& path)
void Reconstruction::WritePoints3DText(const std::string& path)
"""
if len(points3D) == 0:
mean_track_length = 0
else:
mean_track_length = sum((len(pt.image_ids) for _, pt in points3D.items()))/len(points3D)
HEADER = "# 3D point list with one line of data per point:\n" + \
"# POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n" + \
"# Number of points: {}, mean track length: {}\n".format(len(points3D), mean_track_length)
with open(path, "w") as fid:
fid.write(HEADER)
for _, pt in points3D.items():
point_header = [pt.id, *pt.xyz, *pt.rgb, pt.error]
fid.write(" ".join(map(str, point_header)) + " ")
track_strings = []
for image_id, point2D in zip(pt.image_ids, pt.point2D_idxs):
track_strings.append(" ".join(map(str, [image_id, point2D])))
fid.write(" ".join(track_strings) + "\n")
def write_points3D_binary(points3D, path_to_model_file):
"""
see: src/base/reconstruction.cc
void Reconstruction::ReadPoints3DBinary(const std::string& path)
void Reconstruction::WritePoints3DBinary(const std::string& path)
"""
with open(path_to_model_file, "wb") as fid:
write_next_bytes(fid, len(points3D), "Q")
for _, pt in points3D.items():
write_next_bytes(fid, pt.id, "Q")
write_next_bytes(fid, pt.xyz.tolist(), "ddd")
write_next_bytes(fid, pt.rgb.tolist(), "BBB")
write_next_bytes(fid, pt.error, "d")
track_length = pt.image_ids.shape[0]
write_next_bytes(fid, track_length, "Q")
for image_id, point2D_id in zip(pt.image_ids, pt.point2D_idxs):
write_next_bytes(fid, [image_id, point2D_id], "ii")
def detect_model_format(path, ext):
if os.path.isfile(os.path.join(path, "cameras" + ext)) and \
os.path.isfile(os.path.join(path, "images" + ext)) and \
os.path.isfile(os.path.join(path, "points3D" + ext)):
return True
return False
def read_model(path, ext=""):
# try to detect the extension automatically
if ext == "":
if detect_model_format(path, ".bin"):
ext = ".bin"
elif detect_model_format(path, ".txt"):
ext = ".txt"
else:
try:
cameras, images, points3D = read_model(os.path.join(path, "model/"))
logger.warning(
"This SfM file structure was deprecated in hloc v1.1")
return cameras, images, points3D
except FileNotFoundError:
raise FileNotFoundError(
f"Could not find binary or text COLMAP model at {path}")
if ext == ".txt":
cameras = read_cameras_text(os.path.join(path, "cameras" + ext))
images = read_images_text(os.path.join(path, "images" + ext))
points3D = read_points3D_text(os.path.join(path, "points3D") + ext)
else:
cameras = read_cameras_binary(os.path.join(path, "cameras" + ext))
images = read_images_binary(os.path.join(path, "images" + ext))
points3D = read_points3D_binary(os.path.join(path, "points3D") + ext)
return cameras, images, points3D
def write_model(cameras, images, points3D, path, ext=".bin"):
if ext == ".txt":
write_cameras_text(cameras, os.path.join(path, "cameras" + ext))
write_images_text(images, os.path.join(path, "images" + ext))
write_points3D_text(points3D, os.path.join(path, "points3D") + ext)
else:
write_cameras_binary(cameras, os.path.join(path, "cameras" + ext))
write_images_binary(images, os.path.join(path, "images" + ext))
write_points3D_binary(points3D, os.path.join(path, "points3D") + ext)
return cameras, images, points3D
def qvec2rotmat(qvec):
return np.array([
[1 - 2 * qvec[2]**2 - 2 * qvec[3]**2,
2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3],
2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2]],
[2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3],
1 - 2 * qvec[1]**2 - 2 * qvec[3]**2,
2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1]],
[2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2],
2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1],
1 - 2 * qvec[1]**2 - 2 * qvec[2]**2]])
def rotmat2qvec(R):
Rxx, Ryx, Rzx, Rxy, Ryy, Rzy, Rxz, Ryz, Rzz = R.flat
K = np.array([
[Rxx - Ryy - Rzz, 0, 0, 0],
[Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0],
[Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0],
[Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz]]) / 3.0
eigvals, eigvecs = np.linalg.eigh(K)
qvec = eigvecs[[3, 0, 1, 2], np.argmax(eigvals)]
if qvec[0] < 0:
qvec *= -1
return qvec
def main():
parser = argparse.ArgumentParser(description="Read and write COLMAP binary and text models")
parser.add_argument("--input_model", help="path to input model folder")
parser.add_argument("--input_format", choices=[".bin", ".txt"],
help="input model format", default="")
parser.add_argument("--output_model",
help="path to output model folder")
parser.add_argument("--output_format", choices=[".bin", ".txt"],
help="outut model format", default=".txt")
args = parser.parse_args()
cameras, images, points3D = read_model(path=args.input_model, ext=args.input_format)
print("num_cameras:", len(cameras))
print("num_images:", len(images))
print("num_points3D:", len(points3D))
if args.output_model is not None:
write_model(cameras, images, points3D, path=args.output_model, ext=args.output_format)
if __name__ == "__main__":
main()
================================================
FILE: src/utils/select_additional_landmarks.py
================================================
import argparse
import copy
import numpy as np
import os
import scipy.stats as stats
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm
import sys
sys.path.append(os.path.join(sys.path[0], '..'))
from dataloader.indoor6 import Indoor6
from models.efficientlitesld import EfficientNetSLD
from utils.pnp import *
from PIL import Image
# import open3d as o3d
def load_landmark_files(landmark_path, visibility_path):
landmark_file = open(landmark_path, 'r')
num_landmark = int(landmark_file.readline())
landmark = []
for l in range(num_landmark):
pl = landmark_file.readline().split()
pl = np.array([float(pl[i]) for i in range(len(pl))])
landmark.append(pl)
landmark = np.asarray(landmark)[:, 1:].T
visibility = np.loadtxt(visibility_path)
return landmark, visibility
def load_landmark_visibility_files(landmark_path, visibility_path):
landmark_file = open(landmark_path, 'r')
num_landmark = int(landmark_file.readline())
landmark = []
for l in range(num_landmark):
pl = landmark_file.readline().split()
pl = np.array([float(pl[i]) for i in range(len(pl))])
landmark.append(pl)
landmark = np.asarray(landmark)[:, 1:].T
visibility = np.loadtxt(visibility_path)
return landmark, visibility
def visualize_keypoint_np(image_, y, x, kp_color):
image = image_.copy()
if np.sum(kp_color) == 255:
square_size = 5
else:
square_size = 3
for c in range(3):
image[y - square_size:y + square_size, x - square_size:x + square_size, c] = kp_color[c]
return image
def compute_error(C_R_G, C_t_G, C_R_G_hat, C_t_G_hat):
rot_err = 180 / np.pi * np.arccos(np.clip(0.5 * (np.trace(C_R_G.T @ C_R_G_hat) - 1.0), a_min=-1., a_max=1.))
trans_err = np.linalg.norm(C_R_G_hat.T @ C_t_G_hat - C_R_G.T @ C_t_G)
return rot_err, trans_err
def compute_2d3d(opt, pred_heatmap, peak_threshold, landmark2d, landmark3d, C_b_f_gt, H_hm, W_hm, K_inv,
METRICS_LOGGING=None):
N = pred_heatmap.shape[0]
G_p_f = np.zeros((3, N))
C_b_f_hm = np.zeros((3, N))
weights = np.zeros(N)
validIdx = 0
pixel_error = []
angular_error = []
for l in range(N):
pred_heatmap_l = pred_heatmap[l]
max_pred_heatmap_l = np.max(pred_heatmap_l)
if max_pred_heatmap_l > peak_threshold:
peak_yx = np.unravel_index(np.argmax(pred_heatmap_l), np.array(pred_heatmap_l).shape)
peak_yx = np.array(peak_yx)
# Patch size extraction
P = int(min(1+2*np.min(np.array([peak_yx[0], H_hm-1.0-peak_yx[0], peak_yx[1], W_hm-1.0-peak_yx[1]])),
1+64//opt.output_downsample))
patch_peak_yx = pred_heatmap_l[peak_yx[0] - P // 2:peak_yx[0] + P // 2 + 1,
peak_yx[1] - P // 2:peak_yx[1] + P // 2 + 1]
xx_patch, yy_patch = np.meshgrid(np.arange(peak_yx[1] - P // 2, peak_yx[1] + P // 2 + 1, 1),
np.arange(peak_yx[0] - P // 2, peak_yx[0] + P // 2 + 1, 1))
refine_y = np.sum(patch_peak_yx * yy_patch) / np.sum(patch_peak_yx)
refine_x = np.sum(patch_peak_yx * xx_patch) / np.sum(patch_peak_yx)
pixel_error.append(np.linalg.norm(landmark2d[:2, l] -
opt.output_downsample * np.array([refine_x, refine_y])))
pred_bearing = K_inv @ np.array([refine_x, refine_y, 1])
pred_bearing = pred_bearing / np.linalg.norm(pred_bearing)
gt_bearing = C_b_f_gt[:, l]
gt_bearing = gt_bearing / np.linalg.norm(gt_bearing)
angular_error_batch = np.arccos(
np.clip(pred_bearing @ gt_bearing, a_min=-1, a_max=1)) * 180 / np.pi
angular_error.append(angular_error_batch)
weights[validIdx] = max_pred_heatmap_l
C_b_f_hm[:, validIdx] = pred_bearing
G_p_f[:, validIdx] = landmark3d[:, l]
validIdx += 1
return G_p_f[:, :validIdx], C_b_f_hm[:, :validIdx], weights[:validIdx], np.asarray(pixel_error), np.asarray(angular_error)
def compute_pose(G_p_f, C_b_f_hm, weights, minimal_tight_thr, opt_tight_thr, img_id, OUTPUT_FOLDER):
Ndetected_landmarks = C_b_f_hm.shape[1]
# ### Saving 2D-3D correspondences
# if Ndetected_landmarks > 0:
# if not os.path.exists(os.path.join(OUTPUT_FOLDER, 'sld2d3d')):
# os.makedirs(os.path.join(OUTPUT_FOLDER, 'sld2d3d'))
# np.savetxt('%s/sld2d3d/%06d.txt' % (OUTPUT_FOLDER, img_id),
# np.concatenate((C_b_f_hm, G_p_f), axis=0))
# else:
# C_b_f_hm = None
# G_p_f = None
# weights = None
if Ndetected_landmarks >= 4:
## P3P ransac
C_T_G_hat, PnP_inlier = P3PKe_Ransac(G_p_f, C_b_f_hm, weights,
thres=minimal_tight_thr)
# print('inlier: ', np.sum(PnP_inlier))
if np.sum(PnP_inlier) >= 4:
# C_T_G_opt = PnP(C_T_G_hat, G_p_f[:, PnP_inlier], C_b_f_hm[:, PnP_inlier], weights[PnP_inlier])
C_T_G_opt = RunPnPNL(C_T_G_hat,
G_p_f[:, PnP_inlier], C_b_f_hm[:, PnP_inlier],
weights[PnP_inlier],
cutoff=opt_tight_thr)
return np.sum(PnP_inlier), C_T_G_opt, PnP_inlier
return 0, None, np.empty((0))
def select_additional_landmarks(opt, minimal_tight_thr=1e-2, opt_tight_thr=5e-3, mode='test', peak_threshold=0.6):
PRETRAINED_MODEL = opt.pretrained_model
device = opt.gpu_device
test_dataset = Indoor6(landmark_idx=np.arange(opt.landmark_indices[0], opt.landmark_indices[-1]),
scene_id=opt.scene_id,
mode=mode,
root_folder=opt.dataset_folder,
input_image_downsample=2,
landmark_config=opt.landmark_config,
visibility_config=opt.visibility_config,
skip_image_index=1)
test_dataloader = DataLoader(dataset=test_dataset, num_workers=1, batch_size=1, shuffle=False, pin_memory=True)
landmark_data = test_dataset.landmark
cnns = []
nLandmarks = opt.landmark_indices
num_landmarks = opt.landmark_indices[-1]
if len(PRETRAINED_MODEL) == 0:
use_gt_2d3d = True
else:
use_gt_2d3d = False
for idx, pretrained_model in enumerate(PRETRAINED_MODEL):
if opt.model == 'efficientnet':
cnn = EfficientNetSLD(num_landmarks=nLandmarks[idx+1]-nLandmarks[idx], output_downsample=opt.output_downsample).to(device=device)
cnn.load_state_dict(torch.load(pretrained_model))
cnn = cnn.to(device=device)
cnn.eval()
# Adding pretrained model
cnns.append(cnn)
img_id = 0
METRICS_LOGGING = {'image_name': '',
'angular_error': [],
'pixel_error': [],
'rot_err_all': 180.,
'trans_err_all': 180.,
'heatmap_peak': 0.0,
'ndetected': 0,
'pnp_inlier': np.zeros(num_landmarks),
'pixel_inlier_error': np.array([1800.]),
}
test_image_logging = []
LANDMARKS_METRICS_LOGGING = {'image_name': [],
'angular_error': [],
'pixel_error': [],
'heatmap_peak': 0.0,
'ndetected': 0,
}
test_landmarks_logging = [copy.deepcopy(LANDMARKS_METRICS_LOGGING) for _ in range(num_landmarks)]
print(len(test_landmarks_logging))
with torch.no_grad():
## Only works for indoor-6
indoor6W = 640 // opt.output_downsample
indoor6H = 352 // opt.output_downsample
HH, WW = torch.meshgrid(torch.arange(indoor6H), torch.arange(indoor6W))
WW = WW.reshape(1, 1, indoor6H, indoor6W).to('cuda')
HH = HH.reshape(1, 1, indoor6H, indoor6W).to('cuda')
for idx, batch in enumerate(tqdm(test_dataloader)):
image = batch['image'].to(device=device)
B, _, H, W = image.shape
K_inv = batch['inv_intrinsics'].to(device=device)
C_T_G_gt = batch['pose_gt'].cpu().numpy()
landmark2d = batch['intrinsics'] @ batch['landmark3d'].reshape(B, 3, num_landmarks)
landmark2d /= landmark2d[:, 2:].clone()
landmark2d = landmark2d.numpy()
pred_heatmap = []
for cnn in cnns:
pred = cnn(image)
pred_heatmap.append(pred['1'])
pred_heatmap = torch.cat(pred_heatmap, axis=1)
pred_heatmap *= (pred_heatmap > peak_threshold).float()
K_inv[:, :, :2] *= opt.output_downsample
## Compute 2D location of landmarks
P = torch.max(torch.max(pred_heatmap, dim=3)[0], dim=2)[0]
pred_normalized_heatmap = pred_heatmap / (torch.sum(pred_heatmap, axis=(2, 3), keepdim=True) + 1e-4)
projx = torch.sum(WW * pred_normalized_heatmap, axis=(2, 3)).reshape(B, 1, num_landmarks)
projy = torch.sum(HH * pred_normalized_heatmap, axis=(2, 3)).reshape(B, 1, num_landmarks)
xy1 = torch.cat((projx, projy, torch.ones_like(projx)), axis=1)
uv1 = K_inv @ xy1
C_B_f = uv1 / torch.sqrt(torch.sum(uv1 ** 2, axis=1, keepdim=True))
C_B_f = C_B_f.cpu().numpy()
P = P.cpu().numpy()
xy1 = xy1.cpu().numpy()
## Compute error
for b in range(B):
# G_p_f, C_b_f, weights, pixel_error, angular_error = compute_2d3d(
# opt, pred_heatmap[b].cpu().numpy(),
# peak_threshold, landmark2d[b], landmark_data,
# batch['landmark3d'][b].cpu().numpy(),
# H_hm, W_hm, K_inv[b].cpu().numpy())
Pb = P[b]>peak_threshold
G_p_f = landmark_data[:, Pb]
C_b_f = C_B_f[b][:, Pb]
weights = P[b][Pb]
# xy1b = xy1[b][:2, Pb]
pnp_inlier, C_T_G_hat, pnp_inlier_mask = compute_pose(G_p_f, C_b_f, weights,
minimal_tight_thr, opt_tight_thr,
img_id, opt.output_folder)
rot_err, trans_err = 180., 1800.
if pnp_inlier >= 4:
rot_err, trans_err = compute_error(C_T_G_gt[b][:3, :3], C_T_G_gt[b][:3, 3],
C_T_G_hat[:3, :3], C_T_G_hat[:3, 3])
## Logging information
pixel_error = np.linalg.norm(landmark2d[b][:2, Pb] - opt.output_downsample * xy1[b][:2, Pb], axis=0)
C_b_f_gt = batch['landmark3d'][b]
C_b_f_gt = torch.nn.functional.normalize(C_b_f_gt, dim=0).cpu().numpy()
angular_error = np.arccos(np.clip(np.sum(C_b_f * C_b_f_gt[:, Pb], axis=0), -1, 1)) * 180. / np.pi
m = copy.deepcopy(METRICS_LOGGING)
m['image_name'] = test_dataset.image_files[img_id]
m['rgb'] = batch['image'][b].cpu().numpy().transpose(1, 2, 0)
m['pixel_error'] = pixel_error
m['angular_error'] = angular_error
m['heatmap_peak'] = P[b]
m['pixel_detected'] = xy1[b] * opt.output_downsample
m['pixel_gt'] = landmark2d[b]
m['visibility_gt'] = batch['visibility'][b] > 0.5
m['rot_err_all'] = np.array([rot_err])
m['trans_err_all'] = np.array([trans_err])
m['K'] = batch['intrinsics'][b].cpu().numpy()
m['C_T_G_gt'] = C_T_G_gt[b]
if len(pnp_inlier_mask):
m['pnp_inlier'][Pb] = pnp_inlier_mask
pixel_inlier_error = np.linalg.norm(landmark2d[b][:2, m['pnp_inlier']==1] -
opt.output_downsample * xy1[b][:2, m['pnp_inlier']==1], axis=0)
m['pixel_inlier_error'] = pixel_inlier_error
test_image_logging.append(m)
for l in range(num_landmarks):
if batch['visibility'][b, l] > 0.5:
test_landmarks_logging[l]['image_name'].append(test_dataset.image_files[img_id])
if P[b, l]:
test_landmarks_logging[l]['pixel_error'].append(np.linalg.norm(landmark2d[b][:2, l] -
opt.output_downsample * xy1[b][:2, l], axis=0))
else:
test_landmarks_logging[l]['pixel_error'].append(1e3)
test_landmarks_logging.append(m)
img_id += 1
## 2D visualization of images
# test_image_logging.sort(key = lambda x: x['trans_err_all'][0])
# for m in test_image_logging:
# print(m['image_name'], ': ', m['trans_err_all'][0])
# img = np.array(m['rgb'] * 255, dtype=np.uint8)
# for l in range(len(m['visibility_gt'])):
# if m['pnp_inlier'][l]:
# img = visualize_keypoint_np(img,
# int(m['pixel_detected'][1, l]),
# int(m['pixel_detected'][0, l]),
# np.array([0., 255., 0.]))
# if m['visibility_gt'][l]:
# img = visualize_keypoint_np(img,
# int(m['pixel_gt'][1, l]),
# int(m['pixel_gt'][0, l]),
# np.array([200., 0., 0.]))
# Image.fromarray(img).save('%s/%2.2f_%2.2f_%s.jpg' % (opt.output_folder,
# m['trans_err_all'][0],
# np.mean(m['pixel_inlier_error']),
# m['image_name']))
###########################################################################################
############################ Extra landmark selection analysis ############################
###########################################################################################
## Some more additional points to improve wacky poses
# lm_file = os.path.join(opt.dataset_folder, opt.scene_id, 'landmarks/landmarks-2000v8.txt')
# vis_file = os.path.join(opt.dataset_folder, opt.scene_id, 'landmarks/visibility-2000v8_depth_normal.txt')
# full_landmarks, full_vis = load_landmark_files(lm_file, vis_file)
# full_landmarks, full_vis = full_landmarks[:, :200], full_vis[:200]
## Colmap file
from utils.read_write_models import read_model
cameras, images, points = read_model(os.path.join(opt.dataset_folder, 'indoor6-colmap/%s/sparse/0' % opt.scene_id), ext='.bin')
indoor6_name_2to_colmap_index = {}
for k in images:
indoor6_name_2to_colmap_index[images[k].name] = k
## Images with bad poses
## Adding more landmarks on top
## For each test image, pick 10 landmarks that have highest score,
## adding to the high accuracy of camera position triangulation
additional_landmarks = set()
for idx, m in enumerate(test_image_logging):
if m['trans_err_all'][0] > 1.0:
## We want to add unseen points that isn't near the 2D detected points
img_vis_id = test_dataset.image_files.index(m['image_name'])
# print('---------------------------')
# print(m['image_name'])
# print(test_dataset.original_image_name(img_vis_id))
# print(images[indoor6_name_2to_colmap_index[test_dataset.original_image_name(img_vis_id)]].name)
xys = images[indoor6_name_2to_colmap_index[test_dataset.original_image_name(img_vis_id)]].xys
point3dids = images[indoor6_name_2to_colmap_index[test_dataset.original_image_name(img_vis_id)]].point3D_ids
xys = xys[point3dids != -1]
point3dids = point3dids[point3dids != -1]
img = np.array(m['rgb'] * 255, dtype=np.uint8)
for l in range(xys.shape[0]):
img = visualize_keypoint_np(img,
int(xys[l, 1] * 352 / 720),
int(xys[l, 0] * 0.5),
np.array([200., 0., 0.]))
xy_scaled = np.array([xys[l, 0] * 0.5, xys[l, 1] * 352 / 720])
if np.sum(m['pnp_inlier']) > 0:
dist_other_2d_kpts = np.linalg.norm(xy_scaled.reshape(2, 1) - m['pixel_detected'][:2, m['pnp_inlier']==1], axis=0)
if np.min(dist_other_2d_kpts) > 20: # 20 pixels, 1/10 of the image size
additional_landmarks.add(point3dids[l])
else:
additional_landmarks.add(point3dids[l])
# for l in range(len(m['visibility_gt'])):
# if m['pnp_inlier'][l]:
# img = visualize_keypoint_np(img,
# int(m['pixel_detected'][1, l]),
# int(m['pixel_detected'][0, l]),
# np.array([0., 255., 0.]))
# visible_landmarks_in_the_next_1k = np.where(full_vis[:, test_dataset.image_indices[img_vis_id]] == 1)[0]
# visible_landmarks_in_the_next_1k += 1000
# print(visible_landmarks_in_the_next_1k)
# img = np.array(m['rgb'] * 255, dtype=np.uint8)
# for l in visible_landmarks_in_the_next_1k:
# pix = m['K'] @ (m['C_T_G_gt'][:3, :3] @ full_landmarks[:, l] + m['C_T_G_gt'][:3, 3])
# img = visualize_keypoint_np(img,
# int(pix[1] / pix[2]),
# int(pix[0] / pix[2]),
# np.array([200., 0., 0.]))
## Re-do pnp, save new image with new translation error
# Image.fromarray(img).save('%s/%2.2f_%2.2f_%s_after.jpg' % (opt.output_folder,
# m['trans_err_all'][0],
# np.mean(m['pixel_inlier_error']),
# m['image_name']))
### Given additional set of landmarks, re-run the landmark selection to get 200 points
from landmark_selection import ComputePerPointAngularSpan, ComputePerPointDepth, SaveLandmarksAndVisibilityMask
### Adding a bank of new points
numPoints3D = len(additional_landmarks)
points3D_ids = np.zeros(numPoints3D)
points3D_scores = np.zeros(numPoints3D)
points3D_depth = np.zeros(numPoints3D)
points3D_tracklength = np.zeros(numPoints3D)
points3D_anglespan = np.zeros(numPoints3D)
validIdx = 0
## Compute score for each landmark
for i, k in enumerate(tqdm(additional_landmarks)):
pointInGlobal = points[k].xyz
image_ids = points[k].image_ids
trackLength = len(image_ids)
depthMean, depthStd = ComputePerPointDepth(pointInGlobal, image_ids, images)
# timespan = ComputePerPointTimeSpan(image_ids, images)
anglespan = ComputePerPointAngularSpan(pointInGlobal, image_ids, images)
if depthMean < 15.0 and trackLength > 3:
depthScore = min(1.0, depthStd / depthMean)
trackLengthScore = 0.25 * np.log2(trackLength)
points3D_depth[validIdx] = depthMean
points3D_tracklength[validIdx] = trackLength
points3D_anglespan[validIdx] = anglespan
points3D_ids[validIdx] = k
points3D_scores[validIdx] = depthScore + trackLengthScore + anglespan
validIdx += 1
points3D_depth = points3D_depth[:validIdx]
points3D_tracklength = points3D_tracklength[:validIdx]
points3D_anglespan = points3D_anglespan[:validIdx]
point3dids = points3D_ids[:validIdx]
points3D_scores = points3D_scores[:validIdx]
print('Number of additional points: ', validIdx)
print('[Depth mean] Max: %2.2f/Median: %2.2f/Mean: %2.2f/Min: %2.2f'
% (np.max(points3D_depth), np.median(points3D_depth), np.mean(points3D_depth), np.min(points3D_depth)))
print('[Track length] Max: %2.2f/Median: %2.2f/Mean: %2.2f/Min: %2.2f'
% (np.max(points3D_tracklength), np.median(points3D_tracklength), np.mean(points3D_tracklength), np.min(points3D_tracklength)))
print('[Angle span] Max: %2.2f/Median: %2.2f/Mean: %2.2f/Min: %2.2f'
% (np.max(points3D_anglespan), np.median(points3D_anglespan), np.mean(points3D_anglespan), np.min(points3D_anglespan)))
num_selected_landmark = opt.num_landmarks
## Sort scores
sorted_indices = np.argsort(points3D_scores)
## Greedy selection
selected_landmarks = {'id': np.zeros(num_selected_landmark),
'xyz': np.zeros((3, num_selected_landmark)),
'score': np.zeros(num_selected_landmark)}
## Selecting first point
selected_landmarks['id'][0] = points3D_ids[sorted_indices[-1]]
selected_landmarks['xyz'][:, 0] = points[selected_landmarks['id'][0]].xyz
selected_landmarks['score'][0] = points3D_scores[sorted_indices[-1]]
nselected = 1
radius = 5.0
while nselected < num_selected_landmark:
for i in reversed(sorted_indices):
id = points3D_ids[i]
xyz = points[id].xyz
if np.sum(np.linalg.norm(xyz.reshape(3, 1) - selected_landmarks['xyz'][:, :nselected], axis=0) < radius):
continue
else:
selected_landmarks['id'][nselected] = id
selected_landmarks['xyz'][:, nselected] = xyz
selected_landmarks['score'][nselected] = points3D_scores[i]
nselected += 1
if nselected == num_selected_landmark:
break
radius *= 0.5
## Saving
import pickle
indoor6_images = pickle.load(open(os.path.join(opt.dataset_folder, '%s/train_test_val.pkl' % opt.scene_id), 'rb'))
indoor6_imagename_to_index = {}
for i, f in enumerate(indoor6_images['train']):
image_name = open(os.path.join(opt.dataset_folder,
opt.scene_id, 'images',
f.replace('color.jpg',
'intrinsics.txt'))).readline().split(' ')[-1][:-1]
indoor6_imagename_to_index[image_name] = indoor6_images['train_idx'][i]
for i, f in enumerate(indoor6_images['val']):
image_name = open(os.path.join(opt.dataset_folder,
opt.scene_id, 'images',
f.replace('color.jpg',
'intrinsics.txt'))).readline().split(' ')[-1][:-1]
indoor6_imagename_to_index[image_name] = indoor6_images['val_idx'][i]
for i, f in enumerate(indoor6_images['test']):
image_name = open(os.path.join(opt.dataset_folder,
opt.scene_id, 'images',
f.replace('color.jpg',
'intrinsics.txt'))).readline().split(' ')[-1][:-1]
indoor6_imagename_to_index[image_name] = indoor6_images['test_idx'][i]
num_images = len(indoor6_images['train']) + len(indoor6_images['val']) + len(indoor6_images['test'])
SaveLandmarksAndVisibilityMask(selected_landmarks, points, images, indoor6_imagename_to_index, num_images,
os.path.join(opt.dataset_folder, opt.scene_id),
opt.landmark_config, opt.visibility_config, opt.output_format)
if __name__ == '__main__':
parser = argparse.ArgumentParser(
description='Scene Landmark Detection',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument(
'--dataset_folder', type=str, required=True,
help='Root directory, where all data is stored')
parser.add_argument(
'--scene_id', type=str, default='scene6',
help='Scene id')
parser.add_argument(
'--num_landmarks', type=int, default=300,
help='Number of selected landmarks.')
parser.add_argument(
'--output_format', type=str, default='',
help='Landmark file output.')
parser.add_argument(
'--output_folder', type=str, required=True,
help='Output folder')
parser.add_argument(
'--landmark_config', type=str, default='landmarks/landmarks-300',
help='File containing scene-specific 3D landmarks.')
parser.add_argument(
'--landmark_indices', type=int, action='append',
help = 'Landmark indices, specify twice',
required=True)
parser.add_argument(
'--visibility_config', type=str, default='landmarks/visibility_aug-300',
help='File containing information about visibility of landmarks in cameras associated with training set.')
parser.add_argument(
'--model', type=str, default='efficientnet',
help='Network architecture backbone.')
parser.add_argument(
'--output_downsample', type=int, default=4,
help='Down sampling factor for output resolution')
parser.add_argument(
'--gpu_device', type=str, default='cuda:0',
help='GPU device')
parser.add_argument(
'--pretrained_model', type=str, action='append', default=[],
help='Pretrained detector model')
opt = parser.parse_args()
select_additional_landmarks(opt, minimal_tight_thr=1e-3, opt_tight_thr=1e-3)