3dv24 cb1d86f982d1 cached
26 files
11.8 MB
53.1k tokens
119 symbols
1 requests
Download .txt
Showing preview only (219K chars total). Download the full file or copy to clipboard to get everything.
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

![teaser](media/teaser_wide.png)
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**![new](media/New.png)  
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.

![indoor6_sfm](media/indoor6_sfm.png)
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)
* <strike>scene2 (3021/283/284 images)</strike> 
* [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)
* <strike>scene4 (1942/272/272 images)</strike>
* [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.
![indoor6_sfm](media/run_inference_screenshot.png)

**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.
<br>  
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.  
<br>
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
================================================
<!-- BEGIN MICROSOFT SECURITY.MD V0.0.5 BLOCK -->

## 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).

<!-- END MICROSOFT SECURITY.MD BLOCK -->

================================================
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,
                #                                         b
Download .txt
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
Download .txt
SYMBOL INDEX (119 symbols across 15 files)

FILE: src/dataloader/indoor6.py
  class Indoor6 (line 21) | class Indoor6(Dataset):
    method __init__ (line 22) | def __init__(self, root_folder="",
    method original_image_name (line 98) | def original_image_name(self, index):
    method _modify_intrinsic (line 107) | def _modify_intrinsic(self, index, use_precomputed_focal_length=False):
    method _load_and_resize_image (line 145) | def _load_and_resize_image(self, index, W, H):
    method _load_pose (line 159) | def _load_pose(self, index):
    method __getitem__ (line 172) | def __getitem__(self, index):
    method __len__ (line 209) | def __len__(self):
  class Indoor6Patches (line 213) | class Indoor6Patches(Indoor6):
    method __init__ (line 214) | def __init__(self, root_folder="",
    method _extract_patch (line 237) | def _extract_patch(self, C_T_G, lm_idx, K, W_modified, H_modified, cen...
    method _project_landmarks_into_patch (line 276) | def _project_landmarks_into_patch(self, K, C_T_G, img_idx, _top, _bott...
    method __getitem__ (line 302) | def __getitem__(self, index):

FILE: src/inference.py
  function compute_error (line 15) | def compute_error(C_R_G, C_t_G, C_R_G_hat, C_t_G_hat):
  function compute_2d3d (line 23) | def compute_2d3d(opt, pred_heatmap, peak_threshold, landmark2d, landmark...
  function compute_pose (line 74) | def compute_pose(G_p_f, C_b_f_hm, weights, minimal_tight_thr, opt_tight_...
  function inference (line 94) | def inference(opt, minimal_tight_thr=1e-2, opt_tight_thr=5e-3, mode='tes...
  function inference_landmark_stats (line 266) | def inference_landmark_stats(opt, mode='test'):

FILE: src/local_inference.py
  function local_inference (line 9) | def local_inference():

FILE: src/local_training.py
  function launch_training (line 8) | def launch_training():

FILE: src/models/blocks.py
  function _make_encoder (line 6) | def _make_encoder(use_pretrained, exportable=True, output_downsample=4):
  function _make_pretrained_efficientnet_lite0 (line 21) | def _make_pretrained_efficientnet_lite0(use_pretrained, exportable=False):
  function _make_efficientnet_backbone (line 31) | def _make_efficientnet_backbone(effnet):
  function _make_resnet_backbone (line 44) | def _make_resnet_backbone(resnet):
  function _make_pretrained_resnext101_wsl (line 57) | def _make_pretrained_resnext101_wsl(use_pretrained):
  class Interpolate (line 62) | class Interpolate(nn.Module):
    method __init__ (line 66) | def __init__(self, scale_factor, mode, align_corners=False):
    method forward (line 80) | def forward(self, x):
  class ResidualConvUnit (line 97) | class ResidualConvUnit(nn.Module):
    method __init__ (line 101) | def __init__(self, features):
    method forward (line 119) | def forward(self, x):
  class FeatureFusionBlock (line 136) | class FeatureFusionBlock(nn.Module):
    method __init__ (line 140) | def __init__(self, features):
    method forward (line 151) | def forward(self, *xs):
  class ResidualConvUnit_custom (line 171) | class ResidualConvUnit_custom(nn.Module):
    method __init__ (line 175) | def __init__(self, features, activation, bn):
    method forward (line 203) | def forward(self, x):
  class FeatureFusionBlock_custom (line 231) | class FeatureFusionBlock_custom(nn.Module):
    method __init__ (line 235) | def __init__(self, features, activation, deconv=False, bn=False, expan...
    method forward (line 260) | def forward(self, *xs):

FILE: src/models/conv2d_layers.py
  function _ntuple (line 20) | def _ntuple(n):
  function _is_static_pad (line 34) | def _is_static_pad(kernel_size, stride=1, dilation=1, **_):
  function _get_padding (line 38) | def _get_padding(kernel_size, stride=1, dilation=1, **_):
  function _calc_same_pad (line 43) | def _calc_same_pad(i: int, k: int, s: int, d: int):
  function _same_pad_arg (line 47) | def _same_pad_arg(input_size, kernel_size, stride, dilation):
  function _split_channels (line 55) | def _split_channels(num_chan, num_groups):
  function conv2d_same (line 61) | def conv2d_same(
  class Conv2dSame (line 72) | class Conv2dSame(nn.Conv2d):
    method __init__ (line 77) | def __init__(self, in_channels, out_channels, kernel_size, stride=1,
    method forward (line 82) | def forward(self, x):
  class Conv2dSameExport (line 86) | class Conv2dSameExport(nn.Conv2d):
    method __init__ (line 93) | def __init__(self, in_channels, out_channels, kernel_size, stride=1, d...
    method forward (line 99) | def forward(self, x):
  function get_padding_value (line 112) | def get_padding_value(padding, kernel_size, **kwargs):
  function create_conv2d_pad (line 135) | def create_conv2d_pad(in_chs, out_chs, kernel_size, **kwargs):
  class MixedConv2d (line 149) | class MixedConv2d(nn.ModuleDict):
    method __init__ (line 155) | def __init__(self, in_channels, out_channels, kernel_size=3,
    method forward (line 175) | def forward(self, x):
  function get_condconv_initializer (line 182) | def get_condconv_initializer(initializer, num_experts, expert_shape):
  class CondConv2d (line 195) | class CondConv2d(nn.Module):
    method __init__ (line 204) | def __init__(self, in_channels, out_channels, kernel_size=3,
    method reset_parameters (line 234) | def reset_parameters(self):
    method forward (line 245) | def forward(self, x, routing_weights):
  function select_conv2d (line 286) | def select_conv2d(in_chs, out_chs, kernel_size, **kwargs):

FILE: src/models/efficientlitesld.py
  class ASPP (line 7) | class ASPP(nn.Module):
    method __init__ (line 8) | def __init__(self, in_ch, d1, d2, d3, d4, reduction=4):
    method forward (line 32) | def forward(self, x):
  class EfficientNetSLD (line 40) | class EfficientNetSLD(torch.nn.Module):
    method __init__ (line 44) | def __init__(self, path=None, num_landmarks=200, output_downsample=4, ...
    method forward (line 81) | def forward(self, x):

FILE: src/train.py
  function plotting (line 17) | def plotting(ROOT_FOLDER):
  function train (line 62) | def train(opt):
  function train_patches (line 193) | def train_patches(opt):

FILE: src/utils/generate_visibility_depth_normal.py
  function extract (line 16) | def extract(opt):

FILE: src/utils/heatmap.py
  function generate_heat_maps (line 5) | def generate_heat_maps(landmarks, visibility_mask, heatmap_size, K, sigm...
  function generate_heat_maps_gpu (line 59) | def generate_heat_maps_gpu(landmarks_2d, visibility_mask, heatmap_size, ...

FILE: src/utils/landmark_selection.py
  function ComputePerPointTimeSpan (line 9) | def ComputePerPointTimeSpan(image_ids, images):
  function ComputePerPointDepth (line 22) | def ComputePerPointDepth(pointInGlobal, image_ids, images):
  function ComputePerPointAngularSpan (line 35) | def ComputePerPointAngularSpan(pointInGlobal, image_ids, images):
  function SaveLandmarksAndVisibilityMask (line 51) | def SaveLandmarksAndVisibilityMask(selected_landmarks, points3D, images,...

FILE: src/utils/merge_landmark_files.py
  function save_landmark_visibility_mask (line 11) | def save_landmark_visibility_mask(landmarks, visibility_mask,

FILE: src/utils/pnp.py
  function Rotation2Quaternion (line 5) | def Rotation2Quaternion(R):
  function Quaternion2Rotation (line 43) | def Quaternion2Rotation(q):
  function skewsymm (line 78) | def skewsymm(x):
  function VectorizeInitialPose (line 91) | def VectorizeInitialPose(C_T_G):
  function MeasureReprojectionSinglePose (line 101) | def MeasureReprojectionSinglePose(z, p, b, w):
  function UpdatePose (line 118) | def UpdatePose(z):
  function P3PKe (line 131) | def P3PKe(m, X, inlier_thres=1e-5):
  function P3PKe_Ransac (line 259) | def P3PKe_Ransac(G_p_f, C_b_f_hm, w, thres=0.01):
  function RunPnPNL (line 290) | def RunPnPNL(C_T_G, G_p_f, C_b_f, w, cutoff=0.01):

FILE: src/utils/read_write_models.py
  class Image (line 52) | class Image(BaseImage):
    method qvec2rotmat (line 53) | def qvec2rotmat(self):
  function read_next_bytes (line 76) | def read_next_bytes(fid, num_bytes, format_char_sequence, endian_charact...
  function write_next_bytes (line 88) | def write_next_bytes(fid, data, format_char_sequence, endian_character="...
  function read_cameras_text (line 104) | def read_cameras_text(path):
  function read_cameras_binary (line 130) | def read_cameras_binary(path_to_model_file):
  function write_cameras_text (line 159) | def write_cameras_text(cameras, path):
  function write_cameras_binary (line 176) | def write_cameras_binary(cameras, path_to_model_file):
  function read_images_text (line 196) | def read_images_text(path):
  function read_images_binary (line 227) | def read_images_binary(path_to_model_file):
  function write_images_text (line 262) | def write_images_text(images, path):
  function write_images_binary (line 290) | def write_images_binary(images, path_to_model_file):
  function read_points3D_text (line 311) | def read_points3D_text(path):
  function read_points3D_binary (line 338) | def read_points3D_binary(path_to_model_file):
  function write_points3D_text (line 368) | def write_points3D_text(points3D, path):
  function write_points3D_binary (line 393) | def write_points3D_binary(points3D, path_to_model_file):
  function detect_model_format (line 412) | def detect_model_format(path, ext):
  function read_model (line 421) | def read_model(path, ext=""):
  function write_model (line 449) | def write_model(cameras, images, points3D, path, ext=".bin"):
  function qvec2rotmat (line 461) | def qvec2rotmat(qvec):
  function rotmat2qvec (line 474) | def rotmat2qvec(R):
  function main (line 488) | def main():

FILE: src/utils/select_additional_landmarks.py
  function load_landmark_files (line 21) | def load_landmark_files(landmark_path, visibility_path):
  function load_landmark_visibility_files (line 36) | def load_landmark_visibility_files(landmark_path, visibility_path):
  function visualize_keypoint_np (line 51) | def visualize_keypoint_np(image_, y, x, kp_color):
  function compute_error (line 63) | def compute_error(C_R_G, C_t_G, C_R_G_hat, C_t_G_hat):
  function compute_2d3d (line 71) | def compute_2d3d(opt, pred_heatmap, peak_threshold, landmark2d, landmark...
  function compute_pose (line 122) | def compute_pose(G_p_f, C_b_f_hm, weights, minimal_tight_thr, opt_tight_...
  function select_additional_landmarks (line 154) | def select_additional_landmarks(opt, minimal_tight_thr=1e-2, opt_tight_t...
Condensed preview — 26 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (222K chars).
[
  {
    "path": ".github/CODEOWNERS",
    "chars": 79,
    "preview": "# These code reviewers should be added by default.\n* @snsinha @omiksik @tien-d\n"
  },
  {
    "path": "CODE_OF_CONDUCT.md",
    "chars": 444,
    "preview": "# Microsoft Open Source Code of Conduct\n\nThis project has adopted the [Microsoft Open Source Code of Conduct](https://op"
  },
  {
    "path": "LICENSE",
    "chars": 18647,
    "preview": "Attribution 4.0 International\n\n=======================================================================\n\nCreative Common"
  },
  {
    "path": "LICENSE-CODE",
    "chars": 1141,
    "preview": "    MIT License\n\n    Copyright (c) Microsoft Corporation.\n\n    Permission is hereby granted, free of charge, to any pers"
  },
  {
    "path": "README.md",
    "chars": 14582,
    "preview": "# Scene Landmark Detection for Camera Localization\n\n## Introduction\n\n![teaser](media/teaser_wide.png)\nWe have devised a "
  },
  {
    "path": "SECURITY.md",
    "chars": 2780,
    "preview": "<!-- BEGIN MICROSOFT SECURITY.MD V0.0.5 BLOCK -->\n\n## Security\n\nMicrosoft takes the security of our software products an"
  },
  {
    "path": "src/dataloader/indoor6.py",
    "chars": 19520,
    "preview": "import argparse\nimport copy\nimport fnmatch\nimport numpy as np\nimport os\nimport pickle\nfrom PIL import Image\n\nimport sys\n"
  },
  {
    "path": "src/inference.py",
    "chars": 15870,
    "preview": "import copy\nimport numpy as np\nimport os\nimport torch\nfrom torch.utils.data import DataLoader\nfrom tqdm import tqdm\nimpo"
  },
  {
    "path": "src/local_inference.py",
    "chars": 2058,
    "preview": "# Copyright (c) Microsoft Corporation. All rights reserved.\n#from __future__ import print_function\nimport argparse\nimpor"
  },
  {
    "path": "src/local_training.py",
    "chars": 3295,
    "preview": "# Copyright (c) Microsoft Corporation. All rights reserved.\nimport argparse\nimport os\n#import re\n\nArgs = None\n\ndef launc"
  },
  {
    "path": "src/main.py",
    "chars": 8330,
    "preview": "import argparse\nfrom inference import *\nfrom train import *\n\nDEVICE = None\n# auto-detect default device\nif torch.backend"
  },
  {
    "path": "src/models/blocks.py",
    "chars": 6948,
    "preview": "import torch\nimport torch.nn as nn\nfrom .conv2d_layers import Conv2dSameExport\n\n\ndef _make_encoder(use_pretrained, expor"
  },
  {
    "path": "src/models/conv2d_layers.py",
    "chars": 12040,
    "preview": "\"\"\" Conv2D w/ SAME padding, CondConv, MixedConv\n\nA collection of conv layers and padding helpers needed by EfficientNet,"
  },
  {
    "path": "src/models/efficientlitesld.py",
    "chars": 3837,
    "preview": "import torch\nimport torch.nn as nn\n\nfrom .blocks import _make_encoder\n\n\nclass ASPP(nn.Module):\n    def __init__(self, in"
  },
  {
    "path": "src/requirements.txt",
    "chars": 252,
    "preview": "# Scene Landmarks Detector Requirements\n# Usage: pip install -r requirements.txt\n\nargparse\nmatplotlib>=3.2.2\nnumpy>=1.22"
  },
  {
    "path": "src/run_inference.py",
    "chars": 1571,
    "preview": "import os\nimport statistics as st\nimport sys\nimport torch\n\nif __name__ == '__main__':\n\n    home_dir = os.path.expanduser"
  },
  {
    "path": "src/run_training.py",
    "chars": 3428,
    "preview": "from math import exp\nimport os\nimport statistics as st\nfrom tabnanny import check\n\nif __name__ == '__main__':\n\n    home_"
  },
  {
    "path": "src/train.py",
    "chars": 15725,
    "preview": "from datetime import datetime\nimport logging\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport os\nimport pickle\n"
  },
  {
    "path": "src/utils/generate_visibility_depth_normal.py",
    "chars": 12428,
    "preview": "import argparse\nimport copy\nimport fnmatch\nimport numpy as np\nimport open3d as o3d\nimport os\nimport pickle\nfrom PIL impo"
  },
  {
    "path": "src/utils/heatmap.py",
    "chars": 3043,
    "preview": "import numpy as np\nimport torch\n\n\ndef generate_heat_maps(landmarks, visibility_mask, heatmap_size, K, sigma=3):\n    '''\n"
  },
  {
    "path": "src/utils/landmark_selection.py",
    "chars": 7376,
    "preview": "import argparse\r\nimport numpy as np\r\nimport os\r\nimport pickle\r\nfrom read_write_models import qvec2rotmat, read_model\r\nfr"
  },
  {
    "path": "src/utils/merge_landmark_files.py",
    "chars": 3042,
    "preview": "import argparse\nimport copy\nimport numpy as np\nimport os\n\nimport sys\nsys.path.append(os.path.join(sys.path[0], '..'))\n\nf"
  },
  {
    "path": "src/utils/pnp.py",
    "chars": 7573,
    "preview": "import numpy as np\nfrom scipy.optimize import least_squares\n\n\ndef Rotation2Quaternion(R):\n    \"\"\"\n    Convert a rotation"
  },
  {
    "path": "src/utils/read_write_models.py",
    "chars": 22339,
    "preview": "# Copyright (c) 2018, ETH Zurich and UNC Chapel Hill.\r\n# All rights reserved.\r\n#\r\n# Redistribution and use in source and"
  },
  {
    "path": "src/utils/select_additional_landmarks.py",
    "chars": 27717,
    "preview": "import argparse\nimport copy\nimport numpy as np\nimport os\nimport scipy.stats as stats\nimport torch\nfrom torch.utils.data "
  }
]

// ... and 1 more files (download for full content)

About this extraction

This page contains the full source code of the microsoft/SceneLandmarkLocalization GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 26 files (11.8 MB), approximately 53.1k tokens, and a symbol index with 119 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.

Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.

Copied to clipboard!