Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/fisheye-skyline.jpg filter=lfs diff=lfs merge=lfs -text
+assets/teaser.gif filter=lfs diff=lfs merge=lfs -text
+demo.ipynb filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,138 @@
+# folders
+data/**
+outputs/**
+weights/**
+**.DS_Store
+.vscode/**
+wandb/**
+third_party/**
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/

.pre-commit-config.yaml

@@ -0,0 +1,35 @@
+default_stages: [commit]
+default_language_version:
+  python: python3.10
+repos:
+  - repo: https://github.com/psf/black
+    rev: 23.9.1
+    hooks:
+      - id: black
+        args: [--line-length=100]
+        exclude: ^(venv/|docs/)
+        types: [python]
+  - repo: https://github.com/PyCQA/flake8
+    rev: 6.1.0
+    hooks:
+      - id: flake8
+        additional_dependencies: [flake8-docstrings]
+        args:
+          [
+            --max-line-length=100,
+            --docstring-convention=google,
+            --ignore=E203 W503 E402 E731,
+          ]
+        exclude: ^(venv/|docs/|.*__init__.py)
+        types: [python]
+
+  - repo: https://github.com/pycqa/isort
+    rev: 5.12.0
+    hooks:
+      - id: isort
+        args: [--line-length=100, --profile=black, --atomic]
+
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.1.1
+    hooks:
+      - id: mypy

LICENSE

@@ -0,0 +1,190 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   Copyright 2024 ETH Zurich
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,12 +1,587 @@
 ---
 title: GeoCalib
-emoji: 👀
-colorFrom: pink
-colorTo: yellow
+app_file: gradio_app.py
 sdk: gradio
-sdk_version: 4.43.0
-app_file: app.py
-pinned: false
+sdk_version: 4.38.1
 ---
+<p align="center">
+  <h1 align="center"><ins>GeoCalib</ins> 📸<br>Single-image Calibration with Geometric Optimization</h1>
+  <p align="center">
+    <a href="https://www.linkedin.com/in/alexander-veicht/">Alexander Veicht</a>
+    ·
+    <a href="https://psarlin.com/">Paul-Edouard&nbsp;Sarlin</a>
+    ·
+    <a href="https://www.linkedin.com/in/philipplindenberger/">Philipp Lindenberger</a>
+    ·
+    <a href="https://www.microsoft.com/en-us/research/people/mapoll/">Marc&nbsp;Pollefeys</a>
+  </p>
+  <h2 align="center">
+    <p>ECCV 2024</p>
+    <a href="" align="center">Paper</a> | <!--TODO: update link-->
+    <a href="https://colab.research.google.com/drive/1oMzgPGppAPAIQxe-s7SRd_q8r7dVfnqo#scrollTo=etdzQZQzoo-K" align="center">Colab</a> | 
+    <a href="https://huggingface.co/spaces/veichta/GeoCalib" align="center">Demo 🤗</a>
+  </h2>
+  
+</p>
+<p align="center">
+    <a href=""><img src="assets/teaser.gif" alt="example" width=80%></a> <!--TODO: update link-->
+    <br>
+    <em>
+      GeoCalib accurately estimates the camera intrinsics and gravity direction from a single image 
+      <br>
+      by combining geometric optimization with deep learning.
+    </em>
+</p>
+
+##
+
+GeoCalib is a an algoritm for single-image calibration: it estimates the camera intrinsics and gravity direction from a single image only. By combining geometric optimization with deep learning, GeoCalib provides a more flexible and accurate calibration compared to previous approaches. This repository hosts the [inference](#setup-and-demo), [evaluation](#evaluation), and [training](#training) code for GeoCalib and instructions to download our training set [OpenPano](#openpano-dataset).
+
+
+## Setup and demo 
+
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1oMzgPGppAPAIQxe-s7SRd_q8r7dVfnqo#scrollTo=etdzQZQzoo-K)
+[![Hugging Face](https://img.shields.io/badge/Gradio-Demo-blue)](https://huggingface.co/spaces/veichta/GeoCalib)
+
+We provide a small inference package [`geocalib`](geocalib) that requires only minimal dependencies and Python >= 3.9. First clone the repository and install the dependencies:
+
+```bash
+git clone https://github.com/cvg/GeoCalib.git && cd GeoCalib
+python -m pip install -e .
+# OR
+python -m pip install -e "git+https://github.com/cvg/GeoCalib#egg=geocalib"
+```
+
+Here is a minimal usage example:
+
+```python
+from geocalib import GeoCalib
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = GeoCalib().to(device)
+
+# load image as tensor in range [0, 1] with shape [C, H, W]
+img = model.load_image("path/to/image.jpg").to(device)
+result = model.calibrate(img)
+
+print("camera:", result["camera"])
+print("gravity:", result["gravity"])
+```
+
+When either the intrinsics or the gravity are already know, they can be provided:
+
+```python
+# known intrinsics:
+result = model.calibrate(img, priors={"focal": focal_length_tensor})
+
+# known gravity:
+result = model.calibrate(img, priors={"gravity": gravity_direction_tensor})
+```
+
+The default model is optimized for pinhole images. To handle lens distortion, use the following:
+
+```python
+model = GeoCalib(weights="distorted")  # default is "pinhole"
+result = model.calibrate(img, camera_model="simple_radial")  # or pinhole, simple_divisional
+```
+
+Check out our [demo notebook](demo.ipynb) for a full working example.
+
+<details>
+<summary><b>[Interactive demo for your webcam - click to expand]</b></summary>
+Run the following command:
+  
+```bash
+python -m geocalib.interactive_demo --camera_id 0
+```
+
+The demo will open a window showing the camera feed and the calibration results. If `--camera_id` is not provided, the demo will ask for the IP address of a [droidcam](https://droidcam.app) camera.
+
+Controls:
+
+>Toggle the different features using the following keys:
+>
+>- ```h```: Show the estimated horizon line
+>- ```u```: Show the estimated up-vectors
+>- ```l```: Show the estimated latitude heatmap
+>- ```c```: Show the confidence heatmap for the up-vectors and latitudes
+>- ```d```: Show undistorted image, will overwrite the other features
+>- ```g```: Shows a virtual grid of points
+>- ```b```: Shows a virtual box object
+>
+>Change the camera model using the following keys:
+>
+>- ```1```: Pinhole -> Simple and fast
+>- ```2```: Simple Radial -> For small distortions
+>- ```3```: Simple Divisional -> For large distortions
+>
+>Press ```q``` to quit the demo.
+
+</details>
+
+
+<details>
+<summary><b>[Load GeoCalib with torch hub - click to expand]</b></summary>
+
+```python
+model = torch.hub.load("cvg/GeoCalib", "GeoCalib", trust_repo=True)
+```
+
+</details>
+
+## Evaluation
+
+The full evaluation and training code is provided in the single-image calibration library [`siclib`](siclib), which can be installed as:
+```bash
+python -m pip install -e siclib
+```
+
+Running the evaluation commands will write the results to `outputs/results/`.
+
+### LaMAR
+
+Running the evaluation commands will download the dataset to ```data/lamar2k``` which will take around 400 MB of disk space.
+
+<details>
+<summary>[Evaluate GeoCalib]</summary>
+
+To evaluate GeoCalib trained on the OpenPano dataset, run:
+
+```bash
+python -m siclib.eval.lamar2k --conf geocalib-pinhole --tag geocalib --overwrite
+```
+
+</details>
+
+<details>
+<summary>[Evaluate DeepCalib]</summary>
+
+To evaluate DeepCalib trained on the OpenPano dataset, run:
+
+```bash
+python -m siclib.eval.lamar2k --conf deepcalib --tag deepcalib --overwrite
+```
+
+</details>
+
+<details>
+<summary>[Evaluate Perspective Fields]</summary>
+
+Coming soon!
+
+</details>
+
+<details>
+<summary>[Evaluate UVP]</summary>
+
+To evaluate UVP, install the [VP-Estimation-with-Prior-Gravity](https://github.com/cvg/VP-Estimation-with-Prior-Gravity) under ```third_party/VP-Estimation-with-Prior-Gravity```. Then run:
+
+```bash
+python -m siclib.eval.lamar2k --conf uvp --tag uvp --overwrite data.preprocessing.edge_divisible_by=null
+```
+
+</details>
+
+<details>
+<summary>[Evaluate your own model]</summary>
+
+If you have trained your own model, you can evaluate it by running:
+
+```bash
+python -m siclib.eval.lamar2k --checkpoint <experiment name> --tag <eval name> --overwrite
+```
+
+</details>
+
+
+<details>
+<summary>[Results]</summary>
+
+Here are the results for the Area Under the Curve (AUC) for the roll, pitch and field of view (FoV) errors at 1/5/10 degrees for the different methods:
+
+| Approach  | Roll               | Pitch              | FoV                |
+| --------- | ------------------ | ------------------ | ------------------ |
+| DeepCalib | 44.1 / 73.9 / 84.8 | 10.8 / 28.3 / 49.8 | 0.7 / 13.0 / 24.0  |
+| ParamNet  | 51.7 / 77.0 / 86.0 | 27.0 / 52.7 / 70.2 | 02.8 / 06.8 / 14.3 |
+| UVP       | 72.7 / 81.8 / 85.7 | 42.3 / 59.9 / 69.4 | 15.6 / 30.6 / 43.5 |
+| GeoCalib  | 86.4 / 92.5 / 95.0 | 55.0 / 76.9 / 86.2 | 19.1 / 41.5 / 60.0 |
+</details>
+
+### MegaDepth
+
+Running the evaluation commands will download the dataset to ```data/megadepth2k``` or ```data/memegadepth2k-radial``` which will take around 2.1 GB and 1.47 GB of disk space respectively.
+
+<details>
+<summary>[Evaluate GeoCalib]</summary>
+
+To evaluate GeoCalib trained on the OpenPano dataset, run:
+
+```bash
+python -m siclib.eval.megadepth2k --conf geocalib-pinhole --tag geocalib --overwrite
+```
+
+To run the eval on the radial distorted images, run:
+
+```bash
+python -m siclib.eval.megadepth2k_radial --conf geocalib-pinhole --tag geocalib --overwrite model.camera_model=simple_radial
+```
+
+</details>
+
+<details>
+<summary>[Evaluate DeepCalib]</summary>
+
+To evaluate DeepCalib trained on the OpenPano dataset, run:
+
+```bash
+python -m siclib.eval.megadepth2k --conf deepcalib --tag deepcalib --overwrite
+```
+
+</details>
+
+<details>
+<summary>[Evaluate Perspective Fields]</summary>
+
+Coming soon!
+
+</details>
+
+<details>
+<summary>[Evaluate UVP]</summary>
+
+To evaluate UVP, install the [VP-Estimation-with-Prior-Gravity](https://github.com/cvg/VP-Estimation-with-Prior-Gravity) under ```third_party/VP-Estimation-with-Prior-Gravity```. Then run:
+
+```bash
+python -m siclib.eval.megadepth2k --conf uvp --tag uvp --overwrite data.preprocessing.edge_divisible_by=null
+```
+
+</details>
+
+<details>
+<summary>[Evaluate your own model]</summary>
+
+If you have trained your own model, you can evaluate it by running:
+
+```bash
+python -m siclib.eval.megadepth2k --checkpoint <experiment name> --tag <eval name> --overwrite
+```
+
+</details>
+
+<details>
+<summary>[Results]</summary>
+
+Here are the results for the Area Under the Curve (AUC) for the roll, pitch and field of view (FoV) errors at 1/5/10 degrees for the different methods:
+
+| Approach  | Roll               | Pitch              | FoV                |
+| --------- | ------------------ | ------------------ | ------------------ |
+| DeepCalib | 34.6 / 65.4 / 79.4 | 11.9 / 27.8 / 44.8 | 5.6 / 12.1 / 22.9  |
+| ParamNet  | 43.4 / 70.7 / 82.2 | 15.4 / 34.5 / 53.3 | 3.2 / 10.1 / 21.3  |
+| UVP       | 69.2 / 81.6 / 86.9 | 21.6 / 36.2 / 47.4 | 8.2 / 18.7 / 29.8  |
+| GeoCalib  | 82.6 / 90.6 / 94.0 | 32.4 / 53.3 / 67.5 | 13.6 / 31.7 / 48.2 |
+</details>
+
+### TartanAir
+
+Running the evaluation commands will download the dataset to ```data/tartanair``` which will take around 1.85 GB of disk space.
+
+<details>
+<summary>[Evaluate GeoCalib]</summary>
+
+To evaluate GeoCalib trained on the OpenPano dataset, run:
+
+```bash
+python -m siclib.eval.tartanair --conf geocalib-pinhole --tag geocalib --overwrite
+```
+
+</details>
+
+<details>
+<summary>[Evaluate DeepCalib]</summary>
+
+To evaluate DeepCalib trained on the OpenPano dataset, run:
+
+```bash
+python -m siclib.eval.tartanair --conf deepcalib --tag deepcalib --overwrite
+```
+
+</details>
+
+<details>
+<summary>[Evaluate Perspective Fields]</summary>
+
+Coming soon!
+
+</details>
+
+<details>
+<summary>[Evaluate UVP]</summary>
+
+To evaluate UVP, install the [VP-Estimation-with-Prior-Gravity](https://github.com/cvg/VP-Estimation-with-Prior-Gravity) under ```third_party/VP-Estimation-with-Prior-Gravity```. Then run:
+
+```bash
+python -m siclib.eval.tartanair --conf uvp --tag uvp --overwrite data.preprocessing.edge_divisible_by=null
+```
+
+</details>
+
+<details>
+<summary>[Evaluate your own model]</summary>
+
+If you have trained your own model, you can evaluate it by running:
+
+```bash
+python -m siclib.eval.tartanair --checkpoint <experiment name> --tag <eval name> --overwrite
+```
+
+</details>
+
+<details>
+<summary>[Results]</summary>
+
+Here are the results for the Area Under the Curve (AUC) for the roll, pitch and field of view (FoV) errors at 1/5/10 degrees for the different methods:
+
+| Approach  | Roll               | Pitch              | FoV                |
+| --------- | ------------------ | ------------------ | ------------------ |
+| DeepCalib | 24.7 / 55.4 / 71.5 | 16.3 / 38.8 / 58.5 | 1.5 / 8.8 / 27.2   |
+| ParamNet  | 34.5 / 59.2 / 73.9 | 19.4 / 42.0 / 60.3 | 6.0 / 16.8 / 31.6  |
+| UVP       | 52.1 / 64.8 / 71.9 | 36.2 / 48.8 / 58.6 | 15.8 / 25.8 / 35.7 |
+| GeoCalib  | 71.3 / 83.8 / 89.8 | 38.2 / 62.9 / 76.6 | 14.1 / 30.4 / 47.6 |
+</details>
+
+### Stanford2D3D
+
+Before downloading and running the evaluation, you will need to agree to the [terms of use](https://docs.google.com/forms/d/e/1FAIpQLScFR0U8WEUtb7tgjOhhnl31OrkEs73-Y8bQwPeXgebqVKNMpQ/viewform?c=0&w=1) for the Stanford2D3D dataset.
+Running the evaluation commands will download the dataset to ```data/stanford2d3d``` which will take around 885 MB of disk space.
+
+<details>
+<summary>[Evaluate GeoCalib]</summary>
+
+To evaluate GeoCalib trained on the OpenPano dataset, run:
+
+```bash
+python -m siclib.eval.stanford2d3d --conf geocalib-pinhole --tag geocalib --overwrite
+```
+
+</details>
+
+<details>
+<summary>[Evaluate DeepCalib]</summary>
+
+To evaluate DeepCalib trained on the OpenPano dataset, run:
+
+```bash
+python -m siclib.eval.stanford2d3d --conf deepcalib --tag deepcalib --overwrite
+```
+
+</details>
+
+<details>
+<summary>[Evaluate Perspective Fields]</summary>
+
+Coming soon!
+
+</details>
+
+<details>
+<summary>[Evaluate UVP]</summary>
+
+To evaluate UVP, install the [VP-Estimation-with-Prior-Gravity](https://github.com/cvg/VP-Estimation-with-Prior-Gravity) under ```third_party/VP-Estimation-with-Prior-Gravity```. Then run:
+
+```bash
+python -m siclib.eval.stanford2d3d --conf uvp --tag uvp --overwrite data.preprocessing.edge_divisible_by=null
+```
+
+</details>
+
+<details>
+<summary>[Evaluate your own model]</summary>
+
+If you have trained your own model, you can evaluate it by running:
+
+```bash
+python -m siclib.eval.stanford2d3d --checkpoint <experiment name> --tag <eval name> --overwrite
+```
+
+</details>
+
+<details>
+<summary>[Results]</summary>
+
+Here are the results for the Area Under the Curve (AUC) for the roll, pitch and field of view (FoV) errors at 1/5/10 degrees for the different methods:
+
+| Approach  | Roll               | Pitch              | FoV                |
+| --------- | ------------------ | ------------------ | ------------------ |
+| DeepCalib | 33.8 / 63.9 / 79.2 | 21.6 / 46.9 / 65.7 | 8.1 / 20.6 / 37.6  |
+| ParamNet  | 44.6 / 73.9 / 84.8 | 29.2 / 56.7 / 73.1 | 5.8 / 14.3 / 27.8  |
+| UVP       | 65.3 / 74.6 / 79.1 | 51.2 / 63.0 / 69.2 | 22.2 / 39.5 / 51.3 |
+| GeoCalib  | 83.1 / 91.8 / 94.8 | 52.3 / 74.8 / 84.6 | 17.4 / 40.0 / 59.4 |
+
+</details>
+
+### Evaluation options
+
+If you want to provide priors during the evaluation, you can add one or multiple of the following flags:
+
+```bash
+python -m siclib.eval.<benchmark> --conf <config> \
+    --tag <tag> \
+    data.use_prior_focal=true \
+    data.use_prior_gravity=true \
+    data.use_prior_k1=true
+```
+
+<details>
+<summary>[Visual inspection]</summary>
+
+To visually inspect the results of the evaluation, you can run the following command:
+
+```bash
+python -m siclib.eval.inspect <benchmark> <one or multiple tags>
+
+```
+For example, to inspect the results of the evaluation of the GeoCalib model on the LaMAR dataset, you can run:
+```bash
+python -m siclib.eval.inspect lamar2k geocalib
+```
+</details>
+
+## OpenPano Dataset
+
+The OpenPano dataset is a new dataset for single-image calibration which contains about 2.8k panoramas from various sources, namely [HDRMAPS](https://hdrmaps.com/hdris/), [PolyHaven](https://polyhaven.com/hdris), and the [Laval Indoor HDR dataset](http://hdrdb.com/indoor/#presentation). While this dataset is smaller than previous ones, it is publicly available and it provides a better balance between indoor and outdoor scenes.
+
+<details>
+<summary>[Downloading and preparing the dataset]</summary>
+
+In order to assemble the training set, first download the Laval dataset following the instructions on [the corresponding project page](http://hdrdb.com/indoor/#presentation) and place the panoramas in ```data/indoorDatasetCalibrated```. Then, tonemap the HDR images using the following command:
+
+```bash
+python -m siclib.datasets.utils.tonemapping --hdr_dir data/indoorDatasetCalibrated --out_dir data/laval-tonemap
+```
+
+We provide a script to download the PolyHaven and HDRMAPS panos. The script will create folders ```data/openpano/panoramas/{split}``` containing the panoramas specified by the ```{split}_panos.txt``` files. To run the script, execute the following commands:
+
+```bash
+python -m siclib.datasets.utils.download_openpano --name openpano --laval_dir data/laval-tonemap
+```
+Alternatively, you can download the PolyHaven and HDRMAPS panos from [here](https://cvg-data.inf.ethz.ch/GeoCalib_ECCV2024/).
+
+
+After downloading the panoramas, you can create the training set by running the following command:
+
+```bash
+python -m siclib.datasets.create_dataset_from_pano --config-name openpano
+```
+
+The dataset creation can be sped up by using multiple workers and a GPU. To do so, add the following arguments to the command:
+
+```bash
+python -m siclib.datasets.create_dataset_from_pano --config-name openpano n_workers=10 device=cuda
+```
+
+This will create the training set in ```data/openpano/openpano``` with about 37k images for training, 2.1k for validation, and 2.1k for testing.
+
+<details>
+<summary>[Distorted OpenPano]</summary>
+
+To create the OpenPano dataset with radial distortion, run the following command:
+
+```bash
+python -m siclib.datasets.create_dataset_from_pano --config-name openpano_radial
+```
+
+</details>
+
+</details>
+
+## Training
+
+As for the evaluation, the training code is provided in the single-image calibration library [`siclib`](siclib), which can be installed by:
+
+```bash
+python -m pip install -e siclib
+```
+
+Once the [OpenPano Dataset](#openpano-dataset) has been downloaded and prepared, we can train GeoCalib with it:
+
+First download the pre-trained weights for the [MSCAN-B](https://cloud.tsinghua.edu.cn/d/c15b25a6745946618462/) backbone:
+
+```bash
+mkdir weights
+wget "https://cloud.tsinghua.edu.cn/d/c15b25a6745946618462/files/?p=%2Fmscan_b.pth&dl=1" -O weights/mscan_b.pth
+```
+
+Then, start the training with the following command:
+
+```bash
+python -m siclib.train geocalib-pinhole-openpano --conf geocalib --distributed
+```
+
+Feel free to use any other experiment name. By default, the checkpoints will be written to ```outputs/training/```. The default batch size is 24 which requires 2x 4090 GPUs with 24GB of VRAM each. Configurations are managed by [Hydra](https://hydra.cc/) and can be overwritten from the command line.
+For example, to train GeoCalib on a single GPU with a batch size of 5, run:
+
+```bash
+python -m siclib.train geocalib-pinhole-openpano \
+    --conf geocalib \
+    data.train_batch_size=5 # for 1x 2080 GPU
+```
+
+Be aware that this can impact the overall performance. You might need to adjust the learning rate and number of training steps accordingly.
+
+If you want to log the training progress to [tensorboard](https://www.tensorflow.org/tensorboard) or [wandb](https://wandb.ai/), you can set the ```train.writer``` option:
+
+```bash
+python -m siclib.train geocalib-pinhole-openpano \
+    --conf geocalib \
+    --distributed \
+    train.writer=tensorboard
+```
+
+The model can then be evaluated using its experiment name:
+
+```bash
+python -m siclib.eval.<benchmark> --checkpoint geocalib-pinhole-openpano \
+    --tag geocalib-retrained
+```
+
+<details>
+<summary>[Training DeepCalib]</summary>
+
+To train DeepCalib on the OpenPano dataset, run:
+
+```bash
+python -m siclib.train deepcalib-openpano --conf deepcalib --distributed
+```
+
+Make sure that you have generated the [OpenPano Dataset](#openpano-dataset) with radial distortion or add
+the flag ```data=openpano``` to the command to train on the pinhole images.
+
+</details>
+
+<details>
+<summary>[Training Perspective Fields]</summary>
+
+Coming soon!
+
+</details>
+
+## BibTeX citation
+
+If you use any ideas from the paper or code from this repo, please consider citing:
+
+```bibtex
+@inproceedings{veicht2024geocalib,
+  author    = {Alexander Veicht and
+               Paul-Edouard Sarlin and
+               Philipp Lindenberger and
+               Marc Pollefeys},
+  title     = {{GeoCalib: Single-image Calibration with Geometric Optimization}},
+  booktitle = {ECCV},
+  year      = {2024}
+}
+```
+
+## License
+
+The code is provided under the [Apache-2.0 License](LICENSE) while the weights of the trained model are provided under the [Creative Commons Attribution 4.0 International Public License](https://creativecommons.org/licenses/by/4.0/legalcode). Thanks to the authors of the [Laval Indoor HDR dataset](http://hdrdb.com/indoor/#presentation) for allowing this.
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

demo.ipynb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:115886d06182eca375251eab9e301180e86465fd0ed152e917d43d7eb4cbd722
+size 13275966

geocalib/__init__.py

@@ -0,0 +1,17 @@
+import logging
+
+from geocalib.extractor import GeoCalib  # noqa
+
+formatter = logging.Formatter(
+    fmt="[%(asctime)s %(name)s %(levelname)s] %(message)s", datefmt="%m/%d/%Y %H:%M:%S"
+)
+handler = logging.StreamHandler()
+handler.setFormatter(formatter)
+handler.setLevel(logging.INFO)
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+logger.addHandler(handler)
+logger.propagate = False
+
+__module_name__ = __name__

geocalib/camera.py

@@ -0,0 +1,774 @@
+"""Implementation of the pinhole, simple radial, and simple divisional camera models."""
+
+from abc import abstractmethod
+from typing import Dict, Optional, Tuple, Union
+
+import torch
+from torch.func import jacfwd, vmap
+from torch.nn import functional as F
+
+from geocalib.gravity import Gravity
+from geocalib.misc import TensorWrapper, autocast
+from geocalib.utils import deg2rad, focal2fov, fov2focal, rad2rotmat
+
+# flake8: noqa: E741
+# mypy: ignore-errors
+
+
+class BaseCamera(TensorWrapper):
+    """Camera tensor class."""
+
+    eps = 1e-3
+
+    @autocast
+    def __init__(self, data: torch.Tensor):
+        """Camera parameters with shape (..., {w, h, fx, fy, cx, cy, *dist}).
+
+        Tensor convention: (..., {w, h, fx, fy, cx, cy, pitch, roll, *dist}) where
+        - w, h: image size in pixels
+        - fx, fy: focal lengths in pixels
+        - cx, cy: principal points in normalized image coordinates
+        - dist: distortion parameters
+
+        Args:
+            data (torch.Tensor): Camera parameters with shape (..., {6, 7, 8}).
+        """
+        # w, h, fx, fy, cx, cy, dist
+        assert data.shape[-1] in {6, 7, 8}, data.shape
+
+        pad = data.new_zeros(data.shape[:-1] + (8 - data.shape[-1],))
+        data = torch.cat([data, pad], -1) if data.shape[-1] != 8 else data
+        super().__init__(data)
+
+    @classmethod
+    def from_dict(cls, param_dict: Dict[str, torch.Tensor]) -> "BaseCamera":
+        """Create a Camera object from a dictionary of parameters.
+
+        Args:
+            param_dict (Dict[str, torch.Tensor]): Dictionary of parameters.
+
+        Returns:
+            Camera: Camera object.
+        """
+        for key, value in param_dict.items():
+            if not isinstance(value, torch.Tensor):
+                param_dict[key] = torch.tensor(value)
+
+        h, w = param_dict["height"], param_dict["width"]
+        cx, cy = param_dict.get("cx", w / 2), param_dict.get("cy", h / 2)
+
+        if "f" in param_dict:
+            f = param_dict["f"]
+        elif "vfov" in param_dict:
+            vfov = param_dict["vfov"]
+            f = fov2focal(vfov, h)
+        else:
+            raise ValueError("Focal length or vertical field of view must be provided.")
+
+        if "dist" in param_dict:
+            k1, k2 = param_dict["dist"][..., 0], param_dict["dist"][..., 1]
+        elif "k1_hat" in param_dict:
+            k1 = param_dict["k1_hat"] * (f / h) ** 2
+
+            k2 = param_dict.get("k2", torch.zeros_like(k1))
+        else:
+            k1 = param_dict.get("k1", torch.zeros_like(f))
+            k2 = param_dict.get("k2", torch.zeros_like(f))
+
+        fx, fy = f, f
+        if "scales" in param_dict:
+            fx = fx * param_dict["scales"][..., 0] / param_dict["scales"][..., 1]
+
+        params = torch.stack([w, h, fx, fy, cx, cy, k1, k2], dim=-1)
+        return cls(params)
+
+    def pinhole(self):
+        """Return the pinhole camera model."""
+        return self.__class__(self._data[..., :6])
+
+    @property
+    def size(self) -> torch.Tensor:
+        """Size (width height) of the images, with shape (..., 2)."""
+        return self._data[..., :2]
+
+    @property
+    def f(self) -> torch.Tensor:
+        """Focal lengths (fx, fy) with shape (..., 2)."""
+        return self._data[..., 2:4]
+
+    @property
+    def vfov(self) -> torch.Tensor:
+        """Vertical field of view in radians."""
+        return focal2fov(self.f[..., 1], self.size[..., 1])
+
+    @property
+    def hfov(self) -> torch.Tensor:
+        """Horizontal field of view in radians."""
+        return focal2fov(self.f[..., 0], self.size[..., 0])
+
+    @property
+    def c(self) -> torch.Tensor:
+        """Principal points (cx, cy) with shape (..., 2)."""
+        return self._data[..., 4:6]
+
+    @property
+    def K(self) -> torch.Tensor:
+        """Returns the self intrinsic matrix with shape (..., 3, 3)."""
+        shape = self.shape + (3, 3)
+        K = self._data.new_zeros(shape)
+        K[..., 0, 0] = self.f[..., 0]
+        K[..., 1, 1] = self.f[..., 1]
+        K[..., 0, 2] = self.c[..., 0]
+        K[..., 1, 2] = self.c[..., 1]
+        K[..., 2, 2] = 1
+        return K
+
+    def update_focal(self, delta: torch.Tensor, as_log: bool = False):
+        """Update the self parameters after changing the focal length."""
+        f = torch.exp(torch.log(self.f) + delta) if as_log else self.f + delta
+
+        # clamp focal length to a reasonable range for stability during training
+        min_f = fov2focal(self.new_ones(self.shape[0]) * deg2rad(150), self.size[..., 1])
+        max_f = fov2focal(self.new_ones(self.shape[0]) * deg2rad(5), self.size[..., 1])
+        min_f = min_f.unsqueeze(-1).expand(-1, 2)
+        max_f = max_f.unsqueeze(-1).expand(-1, 2)
+        f = f.clamp(min=min_f, max=max_f)
+
+        # make sure focal ration stays the same (avoid inplace operations)
+        fx = f[..., 1] * self.f[..., 0] / self.f[..., 1]
+        f = torch.stack([fx, f[..., 1]], -1)
+
+        dist = self.dist if hasattr(self, "dist") else self.new_zeros(self.f.shape)
+        return self.__class__(torch.cat([self.size, f, self.c, dist], -1))
+
+    def scale(self, scales: Union[float, int, Tuple[Union[float, int]]]):
+        """Update the self parameters after resizing an image."""
+        scales = (scales, scales) if isinstance(scales, (int, float)) else scales
+        s = scales if isinstance(scales, torch.Tensor) else self.new_tensor(scales)
+
+        dist = self.dist if hasattr(self, "dist") else self.new_zeros(self.f.shape)
+        return self.__class__(torch.cat([self.size * s, self.f * s, self.c * s, dist], -1))
+
+    def crop(self, pad: Tuple[float]):
+        """Update the self parameters after cropping an image."""
+        pad = pad if isinstance(pad, torch.Tensor) else self.new_tensor(pad)
+        size = self.size + pad.to(self.size)
+        c = self.c + pad.to(self.c) / 2
+
+        dist = self.dist if hasattr(self, "dist") else self.new_zeros(self.f.shape)
+        return self.__class__(torch.cat([size, self.f, c, dist], -1))
+
+    @autocast
+    def in_image(self, p2d: torch.Tensor):
+        """Check if 2D points are within the image boundaries."""
+        assert p2d.shape[-1] == 2
+        size = self.size.unsqueeze(-2)
+        return torch.all((p2d >= 0) & (p2d <= (size - 1)), -1)
+
+    @autocast
+    def project(self, p3d: torch.Tensor) -> Tuple[torch.Tensor]:
+        """Project 3D points into the self plane and check for visibility."""
+        z = p3d[..., -1]
+        valid = z > self.eps
+        z = z.clamp(min=self.eps)
+        p2d = p3d[..., :-1] / z.unsqueeze(-1)
+        return p2d, valid
+
+    def J_project(self, p3d: torch.Tensor):
+        """Jacobian of the projection function."""
+        x, y, z = p3d[..., 0], p3d[..., 1], p3d[..., 2]
+        zero = torch.zeros_like(z)
+        z = z.clamp(min=self.eps)
+        J = torch.stack([1 / z, zero, -x / z**2, zero, 1 / z, -y / z**2], dim=-1)
+        J = J.reshape(p3d.shape[:-1] + (2, 3))
+        return J  # N x 2 x 3
+
+    def undo_scale_crop(self, data: Dict[str, torch.Tensor]):
+        """Undo transforms done during scaling and cropping."""
+        camera = self.crop(-data["crop_pad"]) if "crop_pad" in data else self
+        return camera.scale(1.0 / data["scales"])
+
+    @abstractmethod
+    def distort(self, pts: torch.Tensor, return_scale: bool = False) -> Tuple[torch.Tensor]:
+        """Distort normalized 2D coordinates and check for validity of the distortion model."""
+        raise NotImplementedError("distort() must be implemented.")
+
+    def J_distort(self, p2d: torch.Tensor, wrt: str = "pts") -> torch.Tensor:
+        """Jacobian of the distortion function."""
+        if wrt == "scale2pts":  # (..., 2)
+            J = [
+                vmap(jacfwd(lambda x: self[idx].distort(x, return_scale=True)[0]))(p2d[idx])[None]
+                for idx in range(p2d.shape[0])
+            ]
+
+            return torch.cat(J, dim=0).squeeze(-3, -2)
+
+        elif wrt == "scale2dist":  # (..., 1)
+            J = []
+            for idx in range(p2d.shape[0]):  # loop to batch pts dimension
+
+                def func(x):
+                    params = torch.cat([self._data[idx, :6], x[None]], -1)
+                    return self.__class__(params).distort(p2d[idx], return_scale=True)[0]
+
+                J.append(vmap(jacfwd(func))(self[idx].dist))
+
+            return torch.cat(J, dim=0)
+
+        else:
+            raise NotImplementedError(f"Jacobian not implemented for wrt={wrt}")
+
+    @abstractmethod
+    def undistort(self, pts: torch.Tensor) -> Tuple[torch.Tensor]:
+        """Undistort normalized 2D coordinates and check for validity of the distortion model."""
+        raise NotImplementedError("undistort() must be implemented.")
+
+    def J_undistort(self, p2d: torch.Tensor, wrt: str = "pts") -> torch.Tensor:
+        """Jacobian of the undistortion function."""
+        if wrt == "pts":  # (..., 2, 2)
+            J = [
+                vmap(jacfwd(lambda x: self[idx].undistort(x)[0]))(p2d[idx])[None]
+                for idx in range(p2d.shape[0])
+            ]
+
+            return torch.cat(J, dim=0).squeeze(-3)
+
+        elif wrt == "dist":  # (..., 1)
+            J = []
+            for batch_idx in range(p2d.shape[0]):  # loop to batch pts dimension
+
+                def func(x):
+                    params = torch.cat([self._data[batch_idx, :6], x[None]], -1)
+                    return self.__class__(params).undistort(p2d[batch_idx])[0]
+
+                J.append(vmap(jacfwd(func))(self[batch_idx].dist))
+
+            return torch.cat(J, dim=0)
+        else:
+            raise NotImplementedError(f"Jacobian not implemented for wrt={wrt}")
+
+    @autocast
+    def up_projection_offset(self, p2d: torch.Tensor) -> torch.Tensor:
+        """Compute the offset for the up-projection."""
+        return self.J_distort(p2d, wrt="scale2pts")  # (B, N, 2)
+
+    def J_up_projection_offset(self, p2d: torch.Tensor, wrt: str = "uv") -> torch.Tensor:
+        """Jacobian of the distortion offset for up-projection."""
+        if wrt == "uv":  # (B, N, 2, 2)
+            J = [
+                vmap(jacfwd(lambda x: self[idx].up_projection_offset(x)[0, 0]))(p2d[idx])[None]
+                for idx in range(p2d.shape[0])
+            ]
+
+            return torch.cat(J, dim=0)
+
+        elif wrt == "dist":  # (B, N, 2)
+            J = []
+            for batch_idx in range(p2d.shape[0]):  # loop to batch pts dimension
+
+                def func(x):
+                    params = torch.cat([self._data[batch_idx, :6], x[None]], -1)[None]
+                    return self.__class__(params).up_projection_offset(p2d[batch_idx][None])
+
+                J.append(vmap(jacfwd(func))(self[batch_idx].dist))
+
+            return torch.cat(J, dim=0).squeeze(1)
+        else:
+            raise NotImplementedError(f"Jacobian not implemented for wrt={wrt}")
+
+    @autocast
+    def denormalize(self, p2d: torch.Tensor) -> torch.Tensor:
+        """Convert normalized 2D coordinates into pixel coordinates."""
+        return p2d * self.f.unsqueeze(-2) + self.c.unsqueeze(-2)
+
+    def J_denormalize(self):
+        """Jacobian of the denormalization function."""
+        return torch.diag_embed(self.f)  # ..., 2 x 2
+
+    @autocast
+    def normalize(self, p2d: torch.Tensor) -> torch.Tensor:
+        """Convert pixel coordinates into normalized 2D coordinates."""
+        return (p2d - self.c.unsqueeze(-2)) / (self.f.unsqueeze(-2))
+
+    def J_normalize(self, p2d: torch.Tensor, wrt: str = "f"):
+        """Jacobian of the normalization function."""
+        # ... x N x 2 x 2
+        if wrt == "f":
+            J_f = -(p2d - self.c.unsqueeze(-2)) / ((self.f.unsqueeze(-2)) ** 2)
+            return torch.diag_embed(J_f)
+        elif wrt == "pts":
+            J_pts = 1 / self.f
+            return torch.diag_embed(J_pts)
+        else:
+            raise NotImplementedError(f"Jacobian not implemented for wrt={wrt}")
+
+    def pixel_coordinates(self) -> torch.Tensor:
+        """Pixel coordinates in self frame.
+
+        Returns:
+            torch.Tensor: Pixel coordinates as a tensor of shape (B, h * w, 2).
+        """
+        w, h = self.size[0].unbind(-1)
+        h, w = h.round().to(int), w.round().to(int)
+
+        # create grid
+        x = torch.arange(0, w, dtype=self.dtype, device=self.device)
+        y = torch.arange(0, h, dtype=self.dtype, device=self.device)
+        x, y = torch.meshgrid(x, y, indexing="xy")
+        xy = torch.stack((x, y), dim=-1).reshape(-1, 2)  # shape (h * w, 2)
+
+        # add batch dimension (normalize() would broadcast but we make it explicit)
+        B = self.shape[0]
+        xy = xy.unsqueeze(0).expand(B, -1, -1)  # if B > 0 else xy
+
+        return xy.to(self.device).to(self.dtype)
+
+    @autocast
+    def pixel_bearing_many(self, p3d: torch.Tensor) -> torch.Tensor:
+        """Get the bearing vectors of pixel coordinates by normalizing them."""
+        return F.normalize(p3d, dim=-1)
+
+    @autocast
+    def world2image(self, p3d: torch.Tensor) -> Tuple[torch.Tensor]:
+        """Transform 3D points into 2D pixel coordinates."""
+        p2d, visible = self.project(p3d)
+        p2d, mask = self.distort(p2d)
+        p2d = self.denormalize(p2d)
+        valid = visible & mask & self.in_image(p2d)
+        return p2d, valid
+
+    @autocast
+    def J_world2image(self, p3d: torch.Tensor):
+        """Jacobian of the world2image function."""
+        p2d_proj, valid = self.project(p3d)
+
+        J_dnorm = self.J_denormalize()
+        J_dist = self.J_distort(p2d_proj)
+        J_proj = self.J_project(p3d)
+
+        J = torch.einsum("...ij,...jk,...kl->...il", J_dnorm, J_dist, J_proj)
+        return J, valid
+
+    @autocast
+    def image2world(self, p2d: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Transform point in the image plane to 3D world coordinates."""
+        p2d = self.normalize(p2d)
+        p2d, valid = self.undistort(p2d)
+        ones = p2d.new_ones(p2d.shape[:-1] + (1,))
+        p3d = torch.cat([p2d, ones], -1)
+        return p3d, valid
+
+    @autocast
+    def J_image2world(self, p2d: torch.Tensor, wrt: str = "f") -> Tuple[torch.Tensor, torch.Tensor]:
+        """Jacobian of the image2world function."""
+        if wrt == "dist":
+            p2d_norm = self.normalize(p2d)
+            return self.J_undistort(p2d_norm, wrt)
+        elif wrt == "f":
+            J_norm2f = self.J_normalize(p2d, wrt)
+            p2d_norm = self.normalize(p2d)
+            J_dist2norm = self.J_undistort(p2d_norm, "pts")
+
+            return torch.einsum("...ij,...jk->...ik", J_dist2norm, J_norm2f)
+        else:
+            raise ValueError(f"Unknown wrt: {wrt}")
+
+    @autocast
+    def undistort_image(self, img: torch.Tensor) -> torch.Tensor:
+        """Undistort an image using the distortion model."""
+        assert self.shape[0] == 1, "Batch size must be 1."
+        W, H = self.size.unbind(-1)
+        H, W = H.int().item(), W.int().item()
+
+        x, y = torch.meshgrid(torch.arange(0, W), torch.arange(0, H), indexing="xy")
+        coords = torch.stack((x, y), dim=-1).reshape(-1, 2)
+
+        p3d, _ = self.pinhole().image2world(coords.to(self.device).to(self.dtype))
+        p2d, _ = self.world2image(p3d)
+
+        mapx, mapy = p2d[..., 0].reshape((1, H, W)), p2d[..., 1].reshape((1, H, W))
+        grid = torch.stack((mapx, mapy), dim=-1)
+        grid = 2.0 * grid / torch.tensor([W - 1, H - 1]).to(grid) - 1
+        return F.grid_sample(img, grid, align_corners=True)
+
+    def get_img_from_pano(
+        self,
+        pano_img: torch.Tensor,
+        gravity: Gravity,
+        yaws: torch.Tensor = 0.0,
+        resize_factor: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Render an image from a panorama.
+
+        Args:
+            pano_img (torch.Tensor): Panorama image of shape (3, H, W) in [0, 1].
+            gravity (Gravity): Gravity direction of the camera.
+            yaws (torch.Tensor | list, optional): Yaw angle in radians. Defaults to 0.0.
+            resize_factor (torch.Tensor, optional): Resize the panorama to be a multiple of the
+            field of view. Defaults to 1.
+
+        Returns:
+            torch.Tensor: Image rendered from the panorama.
+        """
+        B = self.shape[0]
+        if B > 0:
+            assert self.size[..., 0].unique().shape[0] == 1, "All images must have the same width."
+            assert self.size[..., 1].unique().shape[0] == 1, "All images must have the same height."
+
+        w, h = self.size[0].unbind(-1)
+        h, w = h.round().to(int), w.round().to(int)
+
+        if isinstance(yaws, (int, float)):
+            yaws = [yaws]
+        if isinstance(resize_factor, (int, float)):
+            resize_factor = [resize_factor]
+
+        yaws = (
+            yaws.to(self.dtype).to(self.device)
+            if isinstance(yaws, torch.Tensor)
+            else self.new_tensor(yaws)
+        )
+
+        if isinstance(resize_factor, torch.Tensor):
+            resize_factor = resize_factor.to(self.dtype).to(self.device)
+        elif resize_factor is not None:
+            resize_factor = self.new_tensor(resize_factor)
+
+        assert isinstance(pano_img, torch.Tensor), "Panorama image must be a torch.Tensor."
+        pano_img = pano_img if pano_img.dim() == 4 else pano_img.unsqueeze(0)  # B x H x W x 3
+
+        pano_imgs = []
+        for i, yaw in enumerate(yaws):
+            if resize_factor is not None:
+                # resize the panorama such that the fov of the panorama has the same height as the
+                # image
+                vfov = self.vfov[i] if B != 0 else self.vfov
+                scale = torch.pi / float(vfov) * float(h) / pano_img.shape[0] * resize_factor[i]
+                pano_shape = (int(pano_img.shape[0] * scale), int(pano_img.shape[1] * scale))
+
+                mode = "bicubic" if scale >= 1 else "area"
+                resized_pano = F.interpolate(pano_img, size=pano_shape, mode=mode)
+            else:
+                # make sure to copy: resized_pano = pano_img
+                resized_pano = pano_img
+                pano_shape = pano_img.shape[-2:][::-1]
+
+            pano_imgs.append((resized_pano, pano_shape))
+
+        xy = self.pixel_coordinates()
+        uv1, _ = self.image2world(xy)
+        bearings = self.pixel_bearing_many(uv1)
+
+        # rotate bearings
+        R_yaw = rad2rotmat(self.new_zeros(yaw.shape), self.new_zeros(yaw.shape), yaws)
+        rotated_bearings = bearings @ gravity.R @ R_yaw
+
+        # spherical coordinates
+        lon = torch.atan2(rotated_bearings[..., 0], rotated_bearings[..., 2])
+        lat = torch.atan2(
+            rotated_bearings[..., 1], torch.norm(rotated_bearings[..., [0, 2]], dim=-1)
+        )
+
+        images = []
+        for idx, (resized_pano, pano_shape) in enumerate(pano_imgs):
+            min_lon, max_lon = -torch.pi, torch.pi
+            min_lat, max_lat = -torch.pi / 2.0, torch.pi / 2.0
+            min_x, max_x = 0, pano_shape[0] - 1.0
+            min_y, max_y = 0, pano_shape[1] - 1.0
+
+            # map Spherical Coordinates to Panoramic Coordinates
+            nx = (lon[idx] - min_lon) / (max_lon - min_lon) * (max_x - min_x) + min_x
+            ny = (lat[idx] - min_lat) / (max_lat - min_lat) * (max_y - min_y) + min_y
+
+            # reshape and cast to numpy for remap
+            mapx, mapy = nx.reshape((1, h, w)), ny.reshape((1, h, w))
+
+            grid = torch.stack((mapx, mapy), dim=-1)  # Add batch dimension
+            # Normalize to [-1, 1]
+            grid = 2.0 * grid / torch.tensor([pano_shape[-2] - 1, pano_shape[-1] - 1]).to(grid) - 1
+            # Apply grid sample
+            image = F.grid_sample(resized_pano, grid, align_corners=True)
+            images.append(image)
+
+        return torch.concatenate(images, 0) if B > 0 else images[0]
+
+    def __repr__(self):
+        """Print the Camera object."""
+        return f"{self.__class__.__name__} {self.shape} {self.dtype} {self.device}"
+
+
+class Pinhole(BaseCamera):
+    """Implementation of the pinhole camera model.
+
+    Use this model for undistorted images.
+    """
+
+    def distort(self, p2d: torch.Tensor, return_scale: bool = False) -> Tuple[torch.Tensor]:
+        """Distort normalized 2D coordinates."""
+        if return_scale:
+            return p2d.new_ones(p2d.shape[:-1] + (1,))
+
+        return p2d, p2d.new_ones((p2d.shape[0], 1)).bool()
+
+    def J_distort(self, p2d: torch.Tensor, wrt: str = "pts") -> torch.Tensor:
+        """Jacobian of the distortion function."""
+        if wrt == "pts":
+            return torch.eye(2, device=p2d.device, dtype=p2d.dtype).expand(p2d.shape[:-1] + (2, 2))
+
+        raise ValueError(f"Unknown wrt: {wrt}")
+
+    def undistort(self, pts: torch.Tensor) -> Tuple[torch.Tensor]:
+        """Undistort normalized 2D coordinates."""
+        return pts, pts.new_ones((pts.shape[0], 1)).bool()
+
+    def J_undistort(self, p2d: torch.Tensor, wrt: str = "pts") -> torch.Tensor:
+        """Jacobian of the undistortion function."""
+        if wrt == "pts":
+            return torch.eye(2, device=p2d.device, dtype=p2d.dtype).expand(p2d.shape[:-1] + (2, 2))
+
+        raise ValueError(f"Unknown wrt: {wrt}")
+
+    def J_up_projection_offset(self, p2d: torch.Tensor, wrt: str = "uv") -> torch.Tensor:
+        """Jacobian of the up-projection offset."""
+        if wrt == "uv":
+            return torch.zeros(p2d.shape[:-1] + (2, 2), device=p2d.device, dtype=p2d.dtype)
+
+        raise ValueError(f"Unknown wrt: {wrt}")
+
+
+class SimpleRadial(BaseCamera):
+    """Implementation of the simple radial camera model.
+
+    Use this model for weakly distorted images.
+
+    The distortion model is 1 + k1 * r^2 where r^2 = x^2 + y^2.
+    The undistortion model is 1 - k1 * r^2 estimated as in
+    "An Exact Formula for Calculating Inverse Radial Lens Distortions" by Pierre Drap.
+    """
+
+    @property
+    def dist(self) -> torch.Tensor:
+        """Distortion parameters, with shape (..., 1)."""
+        return self._data[..., 6:]
+
+    @property
+    def k1(self) -> torch.Tensor:
+        """Distortion parameters, with shape (...)."""
+        return self._data[..., 6]
+
+    def update_dist(self, delta: torch.Tensor, dist_range: Tuple[float, float] = (-0.7, 0.7)):
+        """Update the self parameters after changing the k1 distortion parameter."""
+        delta_dist = self.new_ones(self.dist.shape) * delta
+        dist = (self.dist + delta_dist).clamp(*dist_range)
+        data = torch.cat([self.size, self.f, self.c, dist], -1)
+        return self.__class__(data)
+
+    @autocast
+    def check_valid(self, p2d: torch.Tensor) -> torch.Tensor:
+        """Check if the distorted points are valid."""
+        return p2d.new_ones(p2d.shape[:-1]).bool()
+
+    def distort(self, p2d: torch.Tensor, return_scale: bool = False) -> Tuple[torch.Tensor]:
+        """Distort normalized 2D coordinates and check for validity of the distortion model."""
+        r2 = torch.sum(p2d**2, -1, keepdim=True)
+        radial = 1 + self.k1[..., None, None] * r2
+
+        if return_scale:
+            return radial, None
+
+        return p2d * radial, self.check_valid(p2d)
+
+    def J_distort(self, p2d: torch.Tensor, wrt: str = "pts"):
+        """Jacobian of the distortion function."""
+        if wrt == "scale2dist":  # (..., 1)
+            return torch.sum(p2d**2, -1, keepdim=True)
+        elif wrt == "scale2pts":  # (..., 2)
+            return 2 * self.k1[..., None, None] * p2d
+        else:
+            return super().J_distort(p2d, wrt)
+
+    @autocast
+    def undistort(self, p2d: torch.Tensor) -> Tuple[torch.Tensor]:
+        """Undistort normalized 2D coordinates and check for validity of the distortion model."""
+        b1 = -self.k1[..., None, None]
+        r2 = torch.sum(p2d**2, -1, keepdim=True)
+        radial = 1 + b1 * r2
+        return p2d * radial, self.check_valid(p2d)
+
+    @autocast
+    def J_undistort(self, p2d: torch.Tensor, wrt: str = "pts") -> torch.Tensor:
+        """Jacobian of the undistortion function."""
+        b1 = -self.k1[..., None, None]
+        r2 = torch.sum(p2d**2, -1, keepdim=True)
+        if wrt == "dist":
+            return -r2 * p2d
+        elif wrt == "pts":
+            radial = 1 + b1 * r2
+            radial_diag = torch.diag_embed(radial.expand(radial.shape[:-1] + (2,)))
+            ppT = torch.einsum("...i,...j->...ij", p2d, p2d)  # (..., 2, 2)
+            return (2 * b1[..., None] * ppT) + radial_diag
+        else:
+            return super().J_undistort(p2d, wrt)
+
+    def J_up_projection_offset(self, p2d: torch.Tensor, wrt: str = "uv") -> torch.Tensor:
+        """Jacobian of the up-projection offset."""
+        if wrt == "uv":  # (..., 2, 2)
+            return torch.diag_embed((2 * self.k1[..., None, None]).expand(p2d.shape[:-1] + (2,)))
+        elif wrt == "dist":
+            return 2 * p2d  # (..., 2)
+        else:
+            return super().J_up_projection_offset(p2d, wrt)
+
+
+class SimpleDivisional(BaseCamera):
+    """Implementation of the simple divisional camera model.
+
+    Use this model for strongly distorted images.
+
+    The distortion model is (1 - sqrt(1 - 4 * k1 * r^2)) / (2 * k1 * r^2) where r^2 = x^2 + y^2.
+    The undistortion model is 1 / (1 + k1 * r^2).
+    """
+
+    @property
+    def dist(self) -> torch.Tensor:
+        """Distortion parameters, with shape (..., 1)."""
+        return self._data[..., 6:]
+
+    @property
+    def k1(self) -> torch.Tensor:
+        """Distortion parameters, with shape (...)."""
+        return self._data[..., 6]
+
+    def update_dist(self, delta: torch.Tensor, dist_range: Tuple[float, float] = (-3.0, 3.0)):
+        """Update the self parameters after changing the k1 distortion parameter."""
+        delta_dist = self.new_ones(self.dist.shape) * delta
+        dist = (self.dist + delta_dist).clamp(*dist_range)
+        data = torch.cat([self.size, self.f, self.c, dist], -1)
+        return self.__class__(data)
+
+    @autocast
+    def check_valid(self, p2d: torch.Tensor) -> torch.Tensor:
+        """Check if the distorted points are valid."""
+        return p2d.new_ones(p2d.shape[:-1]).bool()
+
+    def distort(self, p2d: torch.Tensor, return_scale: bool = False) -> Tuple[torch.Tensor]:
+        """Distort normalized 2D coordinates and check for validity of the distortion model."""
+        r2 = torch.sum(p2d**2, -1, keepdim=True)
+        radial = 1 - torch.sqrt((1 - 4 * self.k1[..., None, None] * r2).clamp(min=0))
+        denom = 2 * self.k1[..., None, None] * r2
+
+        ones = radial.new_ones(radial.shape)
+        radial = torch.where(denom == 0, ones, radial / denom.masked_fill(denom == 0, 1e6))
+
+        if return_scale:
+            return radial, None
+
+        return p2d * radial, self.check_valid(p2d)
+
+    def J_distort(self, p2d: torch.Tensor, wrt: str = "pts"):
+        """Jacobian of the distortion function."""
+        r2 = torch.sum(p2d**2, -1, keepdim=True)
+        t0 = torch.sqrt((1 - 4 * self.k1[..., None, None] * r2).clamp(min=1e-6))
+        if wrt == "scale2pts":  # (B, N, 2)
+            d1 = t0 * 2 * r2
+            d2 = self.k1[..., None, None] * r2**2
+            denom = d1 * d2
+            return p2d * (4 * d2 - (1 - t0) * d1) / denom.masked_fill(denom == 0, 1e6)
+
+        elif wrt == "scale2dist":
+            d1 = 2 * self.k1[..., None, None] * t0
+            d2 = 2 * r2 * self.k1[..., None, None] ** 2
+            denom = d1 * d2
+            return (2 * d2 - (1 - t0) * d1) / denom.masked_fill(denom == 0, 1e6)
+
+        else:
+            return super().J_distort(p2d, wrt)
+
+    @autocast
+    def undistort(self, p2d: torch.Tensor) -> Tuple[torch.Tensor]:
+        """Undistort normalized 2D coordinates and check for validity of the distortion model."""
+        r2 = torch.sum(p2d**2, -1, keepdim=True)
+        denom = 1 + self.k1[..., None, None] * r2
+        radial = 1 / denom.masked_fill(denom == 0, 1e6)
+        return p2d * radial, self.check_valid(p2d)
+
+    def J_undistort(self, p2d: torch.Tensor, wrt: str = "pts") -> torch.Tensor:
+        """Jacobian of the undistortion function."""
+        # return super().J_undistort(p2d, wrt)
+        r2 = torch.sum(p2d**2, -1, keepdim=True)
+        k1 = self.k1[..., None, None]
+        if wrt == "dist":
+            denom = (1 + k1 * r2) ** 2
+            return -r2 / denom.masked_fill(denom == 0, 1e6) * p2d
+        elif wrt == "pts":
+            t0 = 1 + k1 * r2
+            t0 = t0.masked_fill(t0 == 0, 1e6)
+            ppT = torch.einsum("...i,...j->...ij", p2d, p2d)  # (..., 2, 2)
+            J = torch.diag_embed((1 / t0).expand(p2d.shape[:-1] + (2,)))
+            return J - 2 * k1[..., None] * ppT / t0[..., None] ** 2  # (..., N, 2, 2)
+
+        else:
+            return super().J_undistort(p2d, wrt)
+
+    def J_up_projection_offset(self, p2d: torch.Tensor, wrt: str = "uv") -> torch.Tensor:
+        """Jacobian of the up-projection offset.
+
+        func(uv, dist) = 4 / (2 * norm2(uv)^2 * (1-4*k1*norm2(uv)^2)^0.5) * uv
+        - (1-(1-4*k1*norm2(uv)^2)^0.5) / (k1 * norm2(uv)^4) * uv
+        """
+        k1 = self.k1[..., None, None]
+        r2 = torch.sum(p2d**2, -1, keepdim=True)
+        t0 = (1 - 4 * k1 * r2).clamp(min=1e-6)
+        t1 = torch.sqrt(t0)
+        if wrt == "dist":
+            denom = 4 * t0 ** (3 / 2)
+            denom = denom.masked_fill(denom == 0, 1e6)
+            J = 16 / denom
+
+            denom = r2 * t1 * k1
+            denom = denom.masked_fill(denom == 0, 1e6)
+            J = J - 2 / denom
+
+            denom = (r2 * k1) ** 2
+            denom = denom.masked_fill(denom == 0, 1e6)
+            J = J + (1 - t1) / denom
+
+            return J * p2d
+        elif wrt == "uv":
+            # ! unstable (gradient checker might fail), rewrite to use single division (by denom)
+            ppT = torch.einsum("...i,...j->...ij", p2d, p2d)  # (..., 2, 2)
+
+            denom = 2 * r2 * t1
+            denom = denom.masked_fill(denom == 0, 1e6)
+            J = torch.diag_embed((4 / denom).expand(p2d.shape[:-1] + (2,)))
+
+            denom = 4 * t1 * r2**2
+            denom = denom.masked_fill(denom == 0, 1e6)
+            J = J - 16 / denom[..., None] * ppT
+
+            denom = 4 * r2 * t0 ** (3 / 2)
+            denom = denom.masked_fill(denom == 0, 1e6)
+            J = J + (32 * k1[..., None]) / denom[..., None] * ppT
+
+            denom = r2**2 * t1
+            denom = denom.masked_fill(denom == 0, 1e6)
+            J = J - 4 / denom[..., None] * ppT
+
+            denom = k1 * r2**3
+            denom = denom.masked_fill(denom == 0, 1e6)
+            J = J + (4 * (1 - t1) / denom)[..., None] * ppT
+
+            denom = k1 * r2**2
+            denom = denom.masked_fill(denom == 0, 1e6)
+            J = J - torch.diag_embed(((1 - t1) / denom).expand(p2d.shape[:-1] + (2,)))
+
+            return J
+        else:
+            return super().J_up_projection_offset(p2d, wrt)
+
+
+camera_models = {
+    "pinhole": Pinhole,
+    "simple_radial": SimpleRadial,
+    "simple_divisional": SimpleDivisional,
+}

geocalib/extractor.py

@@ -0,0 +1,126 @@
+"""Simple interface for GeoCalib model."""
+
+from pathlib import Path
+from typing import Dict, Optional
+
+import torch
+import torch.nn as nn
+from torch.nn.functional import interpolate
+
+from geocalib.camera import BaseCamera
+from geocalib.geocalib import GeoCalib as Model
+from geocalib.utils import ImagePreprocessor, load_image
+
+
+class GeoCalib(nn.Module):
+    """Simple interface for GeoCalib model."""
+
+    def __init__(self, weights: str = "pinhole"):
+        """Initialize the model with optional config overrides.
+
+        Args:
+            weights (str): trained variant, "pinhole" (default) or "distorted".
+        """
+        super().__init__()
+        if weights == "pinhole":
+            url = "https://github.com/cvg/GeoCalib/releases/download/v1.0/geocalib-pinhole.tar"
+        elif weights == "distorted":
+            url = (
+                "https://github.com/cvg/GeoCalib/releases/download/v1.0/geocalib-simple_radial.tar"
+            )
+        else:
+            raise ValueError(f"Unknown weights: {weights}")
+
+        # load checkpoint
+        model_dir = f"{torch.hub.get_dir()}/geocalib"
+        state_dict = torch.hub.load_state_dict_from_url(
+            url, model_dir, map_location="cpu", file_name=f"{weights}.tar"
+        )
+
+        self.model = Model()
+        self.model.flexible_load(state_dict["model"])
+        self.model.eval()
+
+        self.image_processor = ImagePreprocessor({"resize": 320, "edge_divisible_by": 32})
+
+    def load_image(self, path: Path) -> torch.Tensor:
+        """Load image from path."""
+        return load_image(path)
+
+    def _post_process(
+        self, camera: BaseCamera, img_data: dict[str, torch.Tensor], out: dict[str, torch.Tensor]
+    ) -> tuple[BaseCamera, dict[str, torch.Tensor]]:
+        """Post-process model output by undoing scaling and cropping."""
+        camera = camera.undo_scale_crop(img_data)
+
+        w, h = camera.size.unbind(-1)
+        h = h[0].round().int().item()
+        w = w[0].round().int().item()
+
+        for k in ["latitude_field", "up_field"]:
+            out[k] = interpolate(out[k], size=(h, w), mode="bilinear")
+        for k in ["up_confidence", "latitude_confidence"]:
+            out[k] = interpolate(out[k][:, None], size=(h, w), mode="bilinear")[:, 0]
+
+        inverse_scales = 1.0 / img_data["scales"]
+        zero = camera.new_zeros(camera.f.shape[0])
+        out["focal_uncertainty"] = out.get("focal_uncertainty", zero) * inverse_scales[1]
+        return camera, out
+
+    @torch.no_grad()
+    def calibrate(
+        self,
+        img: torch.Tensor,
+        camera_model: str = "pinhole",
+        priors: Optional[Dict[str, torch.Tensor]] = None,
+        shared_intrinsics: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        """Perform calibration with online resizing.
+
+        Assumes input image is in range [0, 1] and in RGB format.
+
+        Args:
+            img (torch.Tensor): Input image, shape (C, H, W) or (1, C, H, W)
+            camera_model (str, optional): Camera model. Defaults to "pinhole".
+            priors (Dict[str, torch.Tensor], optional): Prior parameters. Defaults to {}.
+            shared_intrinsics (bool, optional): Whether to share intrinsics. Defaults to False.
+
+        Returns:
+            Dict[str, torch.Tensor]: camera and gravity vectors and uncertainties.
+        """
+        if len(img.shape) == 3:
+            img = img[None]  # add batch dim
+        if not shared_intrinsics:
+            assert len(img.shape) == 4 and img.shape[0] == 1
+
+        img_data = self.image_processor(img)
+
+        if priors is None:
+            priors = {}
+
+        prior_values = {}
+        if prior_focal := priors.get("focal"):
+            prior_focal = prior_focal[None] if len(prior_focal.shape) == 0 else prior_focal
+            prior_values["prior_focal"] = prior_focal * img_data["scales"][1]
+
+        if "gravity" in priors:
+            prior_gravity = priors["gravity"]
+            prior_gravity = prior_gravity[None] if len(prior_gravity.shape) == 0 else prior_gravity
+            prior_values["prior_gravity"] = prior_gravity
+
+        self.model.optimizer.set_camera_model(camera_model)
+        self.model.optimizer.shared_intrinsics = shared_intrinsics
+
+        out = self.model(img_data | prior_values)
+
+        camera, gravity = out["camera"], out["gravity"]
+        camera, out = self._post_process(camera, img_data, out)
+
+        return {
+            "camera": camera,
+            "gravity": gravity,
+            "covariance": out["covariance"],
+            **{k: out[k] for k in out.keys() if "field" in k},
+            **{k: out[k] for k in out.keys() if "confidence" in k},
+            **{k: out[k] for k in out.keys() if "uncertainty" in k},
+        }

geocalib/geocalib.py

@@ -0,0 +1,150 @@
+"""GeoCalib model definition."""
+
+import logging
+from typing import Dict
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from geocalib.lm_optimizer import LMOptimizer
+from geocalib.modules import MSCAN, ConvModule, LightHamHead
+
+# mypy: ignore-errors
+
+logger = logging.getLogger(__name__)
+
+
+class LowLevelEncoder(nn.Module):
+    """Very simple low-level encoder."""
+
+    def __init__(self):
+        """Simple low-level encoder."""
+        super().__init__()
+        self.in_channel = 3
+        self.feat_dim = 64
+
+        self.conv1 = ConvModule(self.in_channel, self.feat_dim, kernel_size=3, padding=1)
+        self.conv2 = ConvModule(self.feat_dim, self.feat_dim, kernel_size=3, padding=1)
+
+    def forward(self, data: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        """Forward pass."""
+        x = data["image"]
+
+        assert (
+            x.shape[-1] % 32 == 0 and x.shape[-2] % 32 == 0
+        ), "Image size must be multiple of 32 if not using single image input."
+
+        c1 = self.conv1(x)
+        c2 = self.conv2(c1)
+
+        return {"features": c2}
+
+
+class UpDecoder(nn.Module):
+    """Minimal implementation of UpDecoder."""
+
+    def __init__(self):
+        """Up decoder."""
+        super().__init__()
+        self.decoder = LightHamHead()
+        self.linear_pred_up = nn.Conv2d(self.decoder.out_channels, 2, kernel_size=1)
+
+    def forward(self, data: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        """Forward pass."""
+        x, log_confidence = self.decoder(data["features"])
+        up = self.linear_pred_up(x)
+        return {"up_field": F.normalize(up, dim=1), "up_confidence": torch.sigmoid(log_confidence)}
+
+
+class LatitudeDecoder(nn.Module):
+    """Minimal implementation of LatitudeDecoder."""
+
+    def __init__(self):
+        """Latitude decoder."""
+        super().__init__()
+        self.decoder = LightHamHead()
+        self.linear_pred_latitude = nn.Conv2d(self.decoder.out_channels, 1, kernel_size=1)
+
+    def forward(self, data: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        """Forward pass."""
+        x, log_confidence = self.decoder(data["features"])
+        eps = 1e-5  # avoid nan in backward of asin
+        lat = torch.tanh(self.linear_pred_latitude(x))
+        lat = torch.asin(torch.clamp(lat, -1 + eps, 1 - eps))
+        return {"latitude_field": lat, "latitude_confidence": torch.sigmoid(log_confidence)}
+
+
+class PerspectiveDecoder(nn.Module):
+    """Minimal implementation of PerspectiveDecoder."""
+
+    def __init__(self):
+        """Perspective decoder wrapping up and latitude decoders."""
+        super().__init__()
+        self.up_head = UpDecoder()
+        self.latitude_head = LatitudeDecoder()
+
+    def forward(self, data: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        """Forward pass."""
+        return self.up_head(data) | self.latitude_head(data)
+
+
+class GeoCalib(nn.Module):
+    """GeoCalib inference model."""
+
+    def __init__(self, **optimizer_options):
+        """Initialize the GeoCalib inference model.
+
+        Args:
+            optimizer_options: Options for the lm optimizer.
+        """
+        super().__init__()
+        self.backbone = MSCAN()
+        self.ll_enc = LowLevelEncoder()
+        self.perspective_decoder = PerspectiveDecoder()
+
+        self.optimizer = LMOptimizer({**optimizer_options})
+
+    def forward(self, data: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        """Forward pass."""
+        features = {"hl": self.backbone(data)["features"], "ll": self.ll_enc(data)["features"]}
+        out = self.perspective_decoder({"features": features})
+
+        out |= {
+            k: data[k]
+            for k in ["image", "scales", "prior_gravity", "prior_focal", "prior_k1"]
+            if k in data
+        }
+
+        out |= self.optimizer(out)
+
+        return out
+
+    def flexible_load(self, state_dict: Dict[str, torch.Tensor]) -> None:
+        """Load a checkpoint with flexible key names."""
+        dict_params = set(state_dict.keys())
+        model_params = set(map(lambda n: n[0], self.named_parameters()))
+
+        if dict_params == model_params:  # perfect fit
+            logger.info("Loading all parameters of the checkpoint.")
+            self.load_state_dict(state_dict, strict=True)
+            return
+        elif len(dict_params & model_params) == 0:  # perfect mismatch
+            strip_prefix = lambda x: ".".join(x.split(".")[:1] + x.split(".")[2:])
+            state_dict = {strip_prefix(n): p for n, p in state_dict.items()}
+            dict_params = set(state_dict.keys())
+            if len(dict_params & model_params) == 0:
+                raise ValueError(
+                    "Could not manage to load the checkpoint with"
+                    "parameters:" + "\n\t".join(sorted(dict_params))
+                )
+        common_params = dict_params & model_params
+        left_params = dict_params - model_params
+        left_params = [
+            p for p in left_params if "running" not in p and "num_batches_tracked" not in p
+        ]
+        logger.debug("Loading parameters:\n\t" + "\n\t".join(sorted(common_params)))
+        if left_params:
+            # ignore running stats of batchnorm
+            logger.warning("Could not load parameters:\n\t" + "\n\t".join(sorted(left_params)))
+        self.load_state_dict(state_dict, strict=False)

geocalib/gravity.py

@@ -0,0 +1,131 @@
+"""Tensor class for gravity vector in camera frame."""
+
+import torch
+from torch.nn import functional as F
+
+from geocalib.misc import EuclideanManifold, SphericalManifold, TensorWrapper, autocast
+from geocalib.utils import rad2rotmat
+
+# mypy: ignore-errors
+
+
+class Gravity(TensorWrapper):
+    """Gravity vector in camera frame."""
+
+    eps = 1e-4
+
+    @autocast
+    def __init__(self, data: torch.Tensor) -> None:
+        """Create gravity vector from data.
+
+        Args:
+            data (torch.Tensor): gravity vector as 3D vector in camera frame.
+        """
+        assert data.shape[-1] == 3, data.shape
+
+        data = F.normalize(data, dim=-1)
+
+        super().__init__(data)
+
+    @classmethod
+    def from_rp(cls, roll: torch.Tensor, pitch: torch.Tensor) -> "Gravity":
+        """Create gravity vector from roll and pitch angles."""
+        if not isinstance(roll, torch.Tensor):
+            roll = torch.tensor(roll)
+        if not isinstance(pitch, torch.Tensor):
+            pitch = torch.tensor(pitch)
+
+        sr, cr = torch.sin(roll), torch.cos(roll)
+        sp, cp = torch.sin(pitch), torch.cos(pitch)
+        return cls(torch.stack([-sr * cp, -cr * cp, sp], dim=-1))
+
+    @property
+    def vec3d(self) -> torch.Tensor:
+        """Return the gravity vector in the representation."""
+        return self._data
+
+    @property
+    def x(self) -> torch.Tensor:
+        """Return first component of the gravity vector."""
+        return self._data[..., 0]
+
+    @property
+    def y(self) -> torch.Tensor:
+        """Return second component of the gravity vector."""
+        return self._data[..., 1]
+
+    @property
+    def z(self) -> torch.Tensor:
+        """Return third component of the gravity vector."""
+        return self._data[..., 2]
+
+    @property
+    def roll(self) -> torch.Tensor:
+        """Return the roll angle of the gravity vector."""
+        roll = torch.asin(-self.x / (torch.sqrt(1 - self.z**2) + self.eps))
+        offset = -torch.pi * torch.sign(self.x)
+        return torch.where(self.y < 0, roll, -roll + offset)
+
+    def J_roll(self) -> torch.Tensor:
+        """Return the Jacobian of the roll angle of the gravity vector."""
+        cp, _ = torch.cos(self.pitch), torch.sin(self.pitch)
+        cr, sr = torch.cos(self.roll), torch.sin(self.roll)
+        Jr = self.new_zeros(self.shape + (3,))
+        Jr[..., 0] = -cr * cp
+        Jr[..., 1] = sr * cp
+        return Jr
+
+    @property
+    def pitch(self) -> torch.Tensor:
+        """Return the pitch angle of the gravity vector."""
+        return torch.asin(self.z)
+
+    def J_pitch(self) -> torch.Tensor:
+        """Return the Jacobian of the pitch angle of the gravity vector."""
+        cp, sp = torch.cos(self.pitch), torch.sin(self.pitch)
+        cr, sr = torch.cos(self.roll), torch.sin(self.roll)
+
+        Jp = self.new_zeros(self.shape + (3,))
+        Jp[..., 0] = sr * sp
+        Jp[..., 1] = cr * sp
+        Jp[..., 2] = cp
+        return Jp
+
+    @property
+    def rp(self) -> torch.Tensor:
+        """Return the roll and pitch angles of the gravity vector."""
+        return torch.stack([self.roll, self.pitch], dim=-1)
+
+    def J_rp(self) -> torch.Tensor:
+        """Return the Jacobian of the roll and pitch angles of the gravity vector."""
+        return torch.stack([self.J_roll(), self.J_pitch()], dim=-1)
+
+    @property
+    def R(self) -> torch.Tensor:
+        """Return the rotation matrix from the gravity vector."""
+        return rad2rotmat(roll=self.roll, pitch=self.pitch)
+
+    def J_R(self) -> torch.Tensor:
+        """Return the Jacobian of the rotation matrix from the gravity vector."""
+        raise NotImplementedError
+
+    def update(self, delta: torch.Tensor, spherical: bool = False) -> "Gravity":
+        """Update the gravity vector by adding a delta."""
+        if spherical:
+            data = SphericalManifold.plus(self.vec3d, delta)
+            return self.__class__(data)
+
+        data = EuclideanManifold.plus(self.rp, delta)
+        return self.from_rp(data[..., 0], data[..., 1])
+
+    def J_update(self, spherical: bool = False) -> torch.Tensor:
+        """Return the Jacobian of the update."""
+        return (
+            SphericalManifold.J_plus(self.vec3d)
+            if spherical
+            else EuclideanManifold.J_plus(self.vec3d)
+        )
+
+    def __repr__(self):
+        """Print the Camera object."""
+        return f"{self.__class__.__name__} {self.shape} {self.dtype} {self.device}"

geocalib/interactive_demo.py

@@ -0,0 +1,450 @@
+import argparse
+import logging
+import queue
+import threading
+import time
+from time import time
+
+import cv2
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+from geocalib.extractor import GeoCalib
+from geocalib.perspective_fields import get_perspective_field
+from geocalib.utils import get_device, rad2deg
+
+# flake8: noqa
+# mypy: ignore-errors
+
+
+description = """
+-------------------------
+GeoCalib Interactive Demo
+-------------------------
+
+This script is an interactive demo for GeoCalib. It will open a window showing the camera feed and 
+the calibration results. 
+
+Arguments:
+- '--camera_id': Camera ID to use. If none, will ask for ip of droidcam (https://droidcam.app)
+
+You can toggle different features using the following keys:
+
+- 'h': Toggle horizon line
+- 'u': Toggle up vector field
+- 'l': Toggle latitude heatmap
+- 'c': Toggle confidence heatmap
+- 'd': Toggle undistorted image
+- 'g': Toggle grid of points
+- 'b': Toggle box object
+
+You can also change the camera model using the following keys:
+
+- '1': Pinhole
+- '2': Simple Radial
+- '3': Simple Divisional
+
+Press 'q' to quit the demo.
+"""
+
+
+# Custom VideoCapture class to get the most recent frame instead FIFO
+class VideoCapture:
+    def __init__(self, name):
+        self.cap = cv2.VideoCapture(name)
+        self.q = queue.Queue()
+        t = threading.Thread(target=self._reader)
+        t.daemon = True
+        t.start()
+
+    # read frames as soon as they are available, keeping only most recent one
+    def _reader(self):
+        while True:
+            ret, frame = self.cap.read()
+            if not ret:
+                break
+            if not self.q.empty():
+                try:
+                    self.q.get_nowait()  # discard previous (unprocessed) frame
+                except queue.Empty:
+                    pass
+            self.q.put(frame)
+
+    def read(self):
+        return 1, self.q.get()
+
+    def isOpened(self):
+        return self.cap.isOpened()
+
+
+def add_text(frame, text, align_left=True, align_top=True):
+    """Add text to a plot."""
+    h, w = frame.shape[:2]
+    sc = min(h / 640.0, 2.0)
+    Ht = int(40 * sc)  # text height
+
+    for i, l in enumerate(text.split("\n")):
+        max_line = len(max([l for l in text.split("\n")], key=len))
+        x = int(8 * sc if align_left else w - (max_line) * sc * 18)
+        y = Ht * (i + 1) if align_top else h - Ht * (len(text.split("\n")) - i - 1) - int(8 * sc)
+
+        c_back, c_front = (0, 0, 0), (255, 255, 255)
+        font, style = cv2.FONT_HERSHEY_DUPLEX, cv2.LINE_AA
+        cv2.putText(frame, l, (x, y), font, 1.0 * sc, c_back, int(6 * sc), style)
+        cv2.putText(frame, l, (x, y), font, 1.0 * sc, c_front, int(1 * sc), style)
+    return frame
+
+
+def is_corner(p, h, w):
+    """Check if a point is a corner."""
+    return p in [(0, 0), (0, h - 1), (w - 1, 0), (w - 1, h - 1)]
+
+
+def plot_latitude(frame, latitude):
+    """Plot latitude heatmap."""
+    if not isinstance(latitude, np.ndarray):
+        latitude = latitude.cpu().numpy()
+
+    cmap = plt.get_cmap("seismic")
+    h, w = frame.shape[0], frame.shape[1]
+    sc = min(h / 640.0, 2.0)
+
+    vmin, vmax = -90, 90
+    latitude = (latitude - vmin) / (vmax - vmin)
+
+    colors = (cmap(latitude)[..., :3] * 255).astype(np.uint8)[..., ::-1]
+    frame = cv2.addWeighted(frame, 1 - 0.4, colors, 0.4, 0)
+
+    for contour_line in np.linspace(vmin, vmax, 15):
+        contour_line = (contour_line - vmin) / (vmax - vmin)
+
+        mask = (latitude > contour_line).astype(np.uint8)
+        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+        for contour in contours:
+            color = (np.array(cmap(contour_line))[:3] * 255).astype(np.uint8)[::-1]
+
+            # remove corners
+            contour = [p for p in contour if not is_corner(tuple(p[0]), h, w)]
+            for index, item in enumerate(contour[:-1]):
+                cv2.line(frame, item[0], contour[index + 1][0], color.tolist(), int(5 * sc))
+
+    return frame
+
+
+def draw_horizon_line(frame, heatmap):
+    """Draw a horizon line."""
+    if not isinstance(heatmap, np.ndarray):
+        heatmap = heatmap.cpu().numpy()
+
+    h, w = frame.shape[0], frame.shape[1]
+    sc = min(h / 640.0, 2.0)
+
+    color = (0, 255, 255)
+    vmin, vmax = -90, 90
+    heatmap = (heatmap - vmin) / (vmax - vmin)
+
+    contours, _ = cv2.findContours(
+        (heatmap > 0.5).astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
+    )
+    if contours:
+        contour = [p for p in contours[0] if not is_corner(tuple(p[0]), h, w)]
+        for index, item in enumerate(contour[:-1]):
+            cv2.line(frame, item[0], contour[index + 1][0], color, int(5 * sc))
+    return frame
+
+
+def plot_confidence(frame, confidence):
+    """Plot confidence heatmap."""
+    if not isinstance(confidence, np.ndarray):
+        confidence = confidence.cpu().numpy()
+
+    confidence = np.log10(confidence.clip(1e-6)).clip(-4)
+    confidence = (confidence - confidence.min()) / (confidence.max() - confidence.min())
+
+    cmap = plt.get_cmap("turbo")
+    colors = (cmap(confidence)[..., :3] * 255).astype(np.uint8)[..., ::-1]
+    return cv2.addWeighted(frame, 1 - 0.4, colors, 0.4, 0)
+
+
+def plot_vector_field(frame, vector_field):
+    """Plot a vector field."""
+    if not isinstance(vector_field, np.ndarray):
+        vector_field = vector_field.cpu().numpy()
+
+    H, W = frame.shape[:2]
+    sc = min(H / 640.0, 2.0)
+
+    subsample = min(W, H) // 10
+    offset_x = ((W % subsample) + subsample) // 2
+    samples_x = np.arange(offset_x, W, subsample)
+    samples_y = np.arange(int(subsample * 0.9), H, subsample)
+
+    vec_len = 40 * sc
+    x_grid, y_grid = np.meshgrid(samples_x, samples_y)
+    x, y = vector_field[:, samples_y][:, :, samples_x]
+    for xi, yi, xi_dir, yi_dir in zip(x_grid.ravel(), y_grid.ravel(), x.ravel(), y.ravel()):
+        start = (xi, yi)
+        end = (int(xi + xi_dir * vec_len), int(yi + yi_dir * vec_len))
+        cv2.arrowedLine(
+            frame, start, end, (0, 255, 0), int(5 * sc), line_type=cv2.LINE_AA, tipLength=0.3
+        )
+
+    return frame
+
+
+def plot_box(frame, gravity, camera):
+    """Plot a box object."""
+    pts = np.array(
+        [[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1], [1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 1]]
+    )
+    pts = pts - np.array([0.5, 1, 0.5])
+    rotation_vec = cv2.Rodrigues(gravity.R.numpy()[0])[0]
+    t = np.array([0, 0, 1], dtype=float)
+    K = camera.K[0].cpu().numpy().astype(float)
+    dist = np.zeros(4, dtype=float)
+    axis_points, _ = cv2.projectPoints(
+        0.1 * pts.reshape(-1, 3).astype(float), rotation_vec, t, K, dist
+    )
+
+    h = frame.shape[0]
+    sc = min(h / 640.0, 2.0)
+
+    color = (85, 108, 228)
+    for p in axis_points:
+        center = tuple((int(p[0][0]), int(p[0][1])))
+        frame = cv2.circle(frame, center, 10, color, -1, cv2.LINE_AA)
+
+    for i in range(0, 4):
+        p1 = axis_points[i].astype(int)
+        p2 = axis_points[i + 4].astype(int)
+        frame = cv2.line(frame, tuple(p1[0]), tuple(p2[0]), color, int(5 * sc), cv2.LINE_AA)
+
+        p1 = axis_points[i].astype(int)
+        p2 = axis_points[(i + 1) % 4].astype(int)
+        frame = cv2.line(frame, tuple(p1[0]), tuple(p2[0]), color, int(5 * sc), cv2.LINE_AA)
+
+        p1 = axis_points[i + 4].astype(int)
+        p2 = axis_points[(i + 1) % 4 + 4].astype(int)
+        frame = cv2.line(frame, tuple(p1[0]), tuple(p2[0]), color, int(5 * sc), cv2.LINE_AA)
+
+    return frame
+
+
+def plot_grid(frame, gravity, camera, grid_size=0.2, num_points=5):
+    """Plot a grid of points."""
+    h = frame.shape[0]
+    sc = min(h / 640.0, 2.0)
+
+    samples = np.linspace(-grid_size, grid_size, num_points)
+    xz = np.meshgrid(samples, samples)
+    pts = np.stack((xz[0].ravel(), np.zeros_like(xz[0].ravel()), xz[1].ravel()), axis=-1)
+
+    # project points
+    rotation_vec = cv2.Rodrigues(gravity.R.numpy()[0])[0]
+    t = np.array([0, 0, 1], dtype=float)
+    K = camera.K[0].cpu().numpy().astype(float)
+    dist = np.zeros(4, dtype=float)
+    axis_points, _ = cv2.projectPoints(pts.reshape(-1, 3).astype(float), rotation_vec, t, K, dist)
+
+    color = (192, 77, 58)
+    # draw points
+    for p in axis_points:
+        center = tuple((int(p[0][0]), int(p[0][1])))
+        frame = cv2.circle(frame, center, 10, color, -1, cv2.LINE_AA)
+
+    # draw lines
+    for i in range(num_points):
+        for j in range(num_points - 1):
+            p1 = axis_points[i * num_points + j].astype(int)
+            p2 = axis_points[i * num_points + j + 1].astype(int)
+            frame = cv2.line(frame, tuple(p1[0]), tuple(p2[0]), color, int(5 * sc), cv2.LINE_AA)
+
+            p1 = axis_points[j * num_points + i].astype(int)
+            p2 = axis_points[(j + 1) * num_points + i].astype(int)
+            frame = cv2.line(frame, tuple(p1[0]), tuple(p2[0]), color, int(5 * sc), cv2.LINE_AA)
+
+    return frame
+
+
+def undistort_image(img, camera, padding=0.3):
+    """Undistort an image."""
+    W, H = camera.size.unbind(-1)
+    H, W = H.int().item(), W.int().item()
+
+    pad_h, pad_w = int(H * padding), int(W * padding)
+    x, y = torch.meshgrid(torch.arange(0, W + pad_w), torch.arange(0, H + pad_h), indexing="xy")
+    coords = torch.stack((x, y), dim=-1).reshape(-1, 2) - torch.tensor([pad_w / 2, pad_h / 2])
+
+    p3d, _ = camera.pinhole().image2world(coords.to(camera.device).to(camera.dtype))
+    p2d, _ = camera.world2image(p3d)
+
+    p2d = p2d.float().numpy().reshape(H + pad_h, W + pad_w, 2)
+    img = cv2.remap(img, p2d[..., 0], p2d[..., 1], cv2.INTER_LINEAR, borderValue=(254, 254, 254))
+    return cv2.resize(img, (W, H))
+
+
+class InteractiveDemo:
+    def __init__(self, capture: VideoCapture, device: str) -> None:
+        self.cap = capture
+
+        self.device = torch.device(device)
+        self.model = GeoCalib().to(device)
+
+        self.up_toggle = False
+        self.lat_toggle = False
+        self.conf_toggle = False
+
+        self.hl_toggle = False
+        self.grid_toggle = False
+        self.box_toggle = False
+
+        self.undist_toggle = False
+
+        self.camera_model = "pinhole"
+
+    def render_frame(self, frame, calibration):
+        """Render the frame with the calibration results."""
+        camera, gravity = calibration["camera"].cpu(), calibration["gravity"].cpu()
+
+        if self.undist_toggle:
+            return undistort_image(frame, camera)
+
+        up, lat = get_perspective_field(camera, gravity)
+
+        if gravity.pitch[0] > 0:
+            frame = plot_box(frame, gravity, camera) if self.box_toggle else frame
+            frame = plot_grid(frame, gravity, camera) if self.grid_toggle else frame
+        else:
+            frame = plot_grid(frame, gravity, camera) if self.grid_toggle else frame
+            frame = plot_box(frame, gravity, camera) if self.box_toggle else frame
+
+        frame = draw_horizon_line(frame, lat[0, 0]) if self.hl_toggle else frame
+
+        if self.conf_toggle and self.up_toggle:
+            frame = plot_confidence(frame, calibration["up_confidence"][0])
+        frame = plot_vector_field(frame, up[0]) if self.up_toggle else frame
+
+        if self.conf_toggle and self.lat_toggle:
+            frame = plot_confidence(frame, calibration["latitude_confidence"][0])
+        frame = plot_latitude(frame, rad2deg(lat)[0, 0]) if self.lat_toggle else frame
+
+        return frame
+
+    def format_results(self, calibration):
+        """Format the calibration results."""
+        camera, gravity = calibration["camera"].cpu(), calibration["gravity"].cpu()
+
+        vfov, focal = camera.vfov[0].item(), camera.f[0, 0].item()
+        fov_unc = rad2deg(calibration["vfov_uncertainty"].item())
+        f_unc = calibration["focal_uncertainty"].item()
+
+        roll, pitch = gravity.rp[0].unbind(-1)
+        roll, pitch, vfov = rad2deg(roll), rad2deg(pitch), rad2deg(vfov)
+        roll_unc = rad2deg(calibration["roll_uncertainty"].item())
+        pitch_unc = rad2deg(calibration["pitch_uncertainty"].item())
+
+        text = f"{self.camera_model.replace('_', ' ').title()}\n"
+        text += f"Roll:  {roll:.2f} (+- {roll_unc:.2f})\n"
+        text += f"Pitch: {pitch:.2f} (+- {pitch_unc:.2f})\n"
+        text += f"vFoV:  {vfov:.2f} (+- {fov_unc:.2f})\n"
+        text += f"Focal: {focal:.2f} (+- {f_unc:.2f})"
+
+        if hasattr(camera, "k1"):
+            text += f"\nK1:    {camera.k1[0].item():.2f}"
+
+        return text
+
+    def update_toggles(self):
+        """Update the toggles."""
+        key = cv2.waitKey(100) & 0xFF
+        if key == ord("h"):
+            self.hl_toggle = not self.hl_toggle
+        elif key == ord("u"):
+            self.up_toggle = not self.up_toggle
+        elif key == ord("l"):
+            self.lat_toggle = not self.lat_toggle
+        elif key == ord("c"):
+            self.conf_toggle = not self.conf_toggle
+        elif key == ord("d"):
+            self.undist_toggle = not self.undist_toggle
+        elif key == ord("g"):
+            self.grid_toggle = not self.grid_toggle
+        elif key == ord("b"):
+            self.box_toggle = not self.box_toggle
+
+        elif key == ord("1"):
+            self.camera_model = "pinhole"
+        elif key == ord("2"):
+            self.camera_model = "simple_radial"
+        elif key == ord("3"):
+            self.camera_model = "simple_divisional"
+
+        elif key == ord("q"):
+            return True
+
+        return False
+
+    def run(self):
+        """Run the interactive demo."""
+        while True:
+            start = time()
+            ret, frame = self.cap.read()
+
+            if not ret:
+                print("Error: Failed to retrieve frame.")
+                break
+
+            # create tensor from frame
+            img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            img = torch.tensor(img).permute(2, 0, 1) / 255.0
+
+            calibration = self.model.calibrate(img.to(self.device), camera_model=self.camera_model)
+
+            # render results to the frame
+            frame = self.render_frame(frame, calibration)
+            frame = add_text(frame, self.format_results(calibration))
+
+            end = time()
+            frame = add_text(
+                frame, f"FPS: {1 / (end - start):04.1f}", align_left=False, align_top=False
+            )
+
+            cv2.imshow("GeoCalib Demo", frame)
+
+            if self.update_toggles():
+                break
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--camera_id",
+        type=int,
+        default=None,
+        help="Camera ID to use. If none, will ask for ip of droidcam.",
+    )
+    args = parser.parse_args()
+
+    print(description)
+
+    device = get_device()
+    print(f"Running on: {device}")
+
+    # setup video capture
+    if args.camera_id is not None:
+        cap = VideoCapture(args.camera_id)
+    else:
+        ip = input("Enter the IP address of the camera: ")
+        cap = VideoCapture(f"http://{ip}:4747/video/force/1920x1080")
+
+    if not cap.isOpened():
+        raise ValueError("Error: Could not open camera.")
+
+    demo = InteractiveDemo(cap, device)
+    demo.run()
+
+
+if __name__ == "__main__":
+    main()

geocalib/lm_optimizer.py

@@ -0,0 +1,642 @@
+"""Implementation of the Levenberg-Marquardt optimizer for camera calibration."""
+
+import logging
+import time
+from types import SimpleNamespace
+from typing import Any, Callable, Dict, Tuple
+
+import torch
+import torch.nn as nn
+
+from geocalib.camera import BaseCamera, camera_models
+from geocalib.gravity import Gravity
+from geocalib.misc import J_focal2fov
+from geocalib.perspective_fields import J_perspective_field, get_perspective_field
+from geocalib.utils import focal2fov, rad2deg
+
+logger = logging.getLogger(__name__)
+
+
+def get_trivial_estimation(data: Dict[str, torch.Tensor], camera_model: BaseCamera) -> BaseCamera:
+    """Get initial camera for optimization with roll=0, pitch=0, vfov=0.7 * max(h, w).
+
+    Args:
+        data (Dict[str, torch.Tensor]): Input data dictionary.
+        camera_model (BaseCamera): Camera model to use.
+
+    Returns:
+        BaseCamera: Initial camera for optimization.
+    """
+    """Get initial camera for optimization with roll=0, pitch=0, vfov=0.7 * max(h, w)."""
+    ref = data.get("up_field", data["latitude_field"])
+    ref = ref.detach()
+
+    h, w = ref.shape[-2:]
+    batch_h, batch_w = (
+        ref.new_ones((ref.shape[0],)) * h,
+        ref.new_ones((ref.shape[0],)) * w,
+    )
+
+    init_r = ref.new_zeros((ref.shape[0],))
+    init_p = ref.new_zeros((ref.shape[0],))
+
+    focal = data.get("prior_focal", 0.7 * torch.max(batch_h, batch_w))
+    init_vfov = focal2fov(focal, h)
+
+    params = {"width": batch_w, "height": batch_h, "vfov": init_vfov}
+    params |= {"scales": data["scales"]} if "scales" in data else {}
+    params |= {"k1": data["prior_k1"]} if "prior_k1" in data else {}
+    camera = camera_model.from_dict(params)
+    camera = camera.float().to(ref.device)
+
+    gravity = Gravity.from_rp(init_r, init_p).float().to(ref.device)
+
+    if "prior_gravity" in data:
+        gravity = data["prior_gravity"].float().to(ref.device)
+        gravity = Gravity(gravity) if isinstance(gravity, torch.Tensor) else gravity
+
+    return camera, gravity
+
+
+def scaled_loss(
+    x: torch.Tensor, fn: Callable, a: float
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Apply a loss function to a tensor and pre- and post-scale it.
+
+    Args:
+        x: the data tensor, should already be squared: `x = y**2`.
+        fn: the loss function, with signature `fn(x) -> y`.
+        a: the scale parameter.
+
+    Returns:
+        The value of the loss, and its first and second derivatives.
+    """
+    a2 = a**2
+    loss, loss_d1, loss_d2 = fn(x / a2)
+    return loss * a2, loss_d1, loss_d2 / a2
+
+
+def huber_loss(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """The classical robust Huber loss, with first and second derivatives."""
+    mask = x <= 1
+    sx = torch.sqrt(x + 1e-8)  # avoid nan in backward pass
+    isx = torch.max(sx.new_tensor(torch.finfo(torch.float).eps), 1 / sx)
+    loss = torch.where(mask, x, 2 * sx - 1)
+    loss_d1 = torch.where(mask, torch.ones_like(x), isx)
+    loss_d2 = torch.where(mask, torch.zeros_like(x), -isx / (2 * x))
+    return loss, loss_d1, loss_d2
+
+
+def early_stop(new_cost: torch.Tensor, prev_cost: torch.Tensor, atol: float, rtol: float) -> bool:
+    """Early stopping criterion based on cost convergence."""
+    return torch.allclose(new_cost, prev_cost, atol=atol, rtol=rtol)
+
+
+def update_lambda(
+    lamb: torch.Tensor,
+    prev_cost: torch.Tensor,
+    new_cost: torch.Tensor,
+    lambda_min: float = 1e-6,
+    lambda_max: float = 1e2,
+) -> torch.Tensor:
+    """Update damping factor for Levenberg-Marquardt optimization."""
+    new_lamb = lamb.new_zeros(lamb.shape)
+    new_lamb = lamb * torch.where(new_cost > prev_cost, 10, 0.1)
+    lamb = torch.clamp(new_lamb, lambda_min, lambda_max)
+    return lamb
+
+
+def optimizer_step(
+    G: torch.Tensor, H: torch.Tensor, lambda_: torch.Tensor, eps: float = 1e-6
+) -> torch.Tensor:
+    """One optimization step with Gauss-Newton or Levenberg-Marquardt.
+
+    Args:
+        G (torch.Tensor): Batched gradient tensor of size (..., N).
+        H (torch.Tensor): Batched hessian tensor of size (..., N, N).
+        lambda_ (torch.Tensor): Damping factor for LM (use GN if lambda_=0) with shape (B,).
+        eps (float, optional): Epsilon for damping. Defaults to 1e-6.
+
+    Returns:
+        torch.Tensor: Batched update tensor of size (..., N).
+    """
+    diag = H.diagonal(dim1=-2, dim2=-1)
+    diag = diag * lambda_.unsqueeze(-1)  # (B, 3)
+
+    H = H + diag.clamp(min=eps).diag_embed()
+
+    H_, G_ = H.cpu(), G.cpu()
+    try:
+        U = torch.linalg.cholesky(H_)
+    except RuntimeError:
+        logger.warning("Cholesky decomposition failed. Stopping.")
+        delta = H.new_zeros((H.shape[0], H.shape[-1]))  # (B, 3)
+    else:
+        delta = torch.cholesky_solve(G_[..., None], U)[..., 0]
+
+    return delta.to(H.device)
+
+
+# mypy: ignore-errors
+class LMOptimizer(nn.Module):
+    """Levenberg-Marquardt optimizer for camera calibration."""
+
+    default_conf = {
+        # Camera model parameters
+        "camera_model": "pinhole",  # {"pinhole", "simple_radial", "simple_spherical"}
+        "shared_intrinsics": False,  # share focal length across all images in batch
+        # LM optimizer parameters
+        "num_steps": 30,
+        "lambda_": 0.1,
+        "fix_lambda": False,
+        "early_stop": True,
+        "atol": 1e-8,
+        "rtol": 1e-8,
+        "use_spherical_manifold": True,  # use spherical manifold for gravity optimization
+        "use_log_focal": True,  # use log focal length for optimization
+        # Loss function parameters
+        "up_loss_fn_scale": 1e-2,
+        "lat_loss_fn_scale": 1e-2,
+        # Misc
+        "verbose": False,
+    }
+
+    def __init__(self, conf: Dict[str, Any]):
+        """Initialize the LM optimizer."""
+        super().__init__()
+        self.conf = conf = SimpleNamespace(**{**self.default_conf, **conf})
+        self.num_steps = conf.num_steps
+
+        self.set_camera_model(conf.camera_model)
+        self.setup_optimization_and_priors(shared_intrinsics=conf.shared_intrinsics)
+
+    def set_camera_model(self, camera_model: str) -> None:
+        """Set the camera model to use for the optimization.
+
+        Args:
+            camera_model (str): Camera model to use.
+        """
+        assert (
+            camera_model in camera_models.keys()
+        ), f"Unknown camera model: {camera_model} not in {camera_models.keys()}"
+        self.camera_model = camera_models[camera_model]
+        self.camera_has_distortion = hasattr(self.camera_model, "dist")
+
+        logger.debug(
+            f"Using camera model: {camera_model} (with distortion: {self.camera_has_distortion})"
+        )
+
+    def setup_optimization_and_priors(
+        self, data: Dict[str, torch.Tensor] = None, shared_intrinsics: bool = False
+    ) -> None:
+        """Setup the optimization and priors for the LM optimizer.
+
+        Args:
+            data (Dict[str, torch.Tensor], optional): Dict potentially containing priors. Defaults
+            to None.
+            shared_intrinsics (bool, optional): Whether to share the intrinsics across the batch.
+            Defaults to False.
+        """
+        if data is None:
+            data = {}
+        self.shared_intrinsics = shared_intrinsics
+
+        if shared_intrinsics:  # si => must use pinhole
+            assert (
+                self.camera_model == camera_models["pinhole"]
+            ), f"Shared intrinsics only supported with pinhole camera model: {self.camera_model}"
+
+        self.estimate_gravity = True
+        if "prior_gravity" in data:
+            self.estimate_gravity = False
+            logger.debug("Using provided gravity as prior.")
+
+        self.estimate_focal = True
+        if "prior_focal" in data:
+            self.estimate_focal = False
+            logger.debug("Using provided focal as prior.")
+
+        self.estimate_k1 = True
+        if "prior_k1" in data:
+            self.estimate_k1 = False
+            logger.debug("Using provided k1 as prior.")
+
+        self.gravity_delta_dims = (0, 1) if self.estimate_gravity else (-1,)
+        self.focal_delta_dims = (
+            (max(self.gravity_delta_dims) + 1,) if self.estimate_focal else (-1,)
+        )
+        self.k1_delta_dims = (max(self.focal_delta_dims) + 1,) if self.estimate_k1 else (-1,)
+
+        logger.debug(f"Camera Model:       {self.camera_model}")
+        logger.debug(f"Optimizing gravity: {self.estimate_gravity} ({self.gravity_delta_dims})")
+        logger.debug(f"Optimizing focal:   {self.estimate_focal} ({self.focal_delta_dims})")
+        logger.debug(f"Optimizing k1:      {self.estimate_k1} ({self.k1_delta_dims})")
+
+        logger.debug(f"Shared intrinsics:  {self.shared_intrinsics}")
+
+    def calculate_residuals(
+        self, camera: BaseCamera, gravity: Gravity, data: Dict[str, torch.Tensor]
+    ) -> Dict[str, torch.Tensor]:
+        """Calculate the residuals for the optimization.
+
+        Args:
+            camera (BaseCamera): Optimized camera.
+            gravity (Gravity): Optimized gravity.
+            data (Dict[str, torch.Tensor]): Input data containing the up and latitude fields.
+
+        Returns:
+            Dict[str, torch.Tensor]: Residuals for the optimization.
+        """
+        perspective_up, perspective_lat = get_perspective_field(camera, gravity)
+        perspective_lat = torch.sin(perspective_lat)
+
+        residuals = {}
+        if "up_field" in data:
+            up_residual = (data["up_field"] - perspective_up).permute(0, 2, 3, 1)
+            residuals["up_residual"] = up_residual.reshape(up_residual.shape[0], -1, 2)
+
+        if "latitude_field" in data:
+            target_lat = torch.sin(data["latitude_field"])
+            lat_residual = (target_lat - perspective_lat).permute(0, 2, 3, 1)
+            residuals["latitude_residual"] = lat_residual.reshape(lat_residual.shape[0], -1, 1)
+
+        return residuals
+
+    def calculate_costs(
+        self, residuals: torch.Tensor, data: Dict[str, torch.Tensor]
+    ) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
+        """Calculate the costs and weights for the optimization.
+
+        Args:
+            residuals (torch.Tensor): Residuals for the optimization.
+            data (Dict[str, torch.Tensor]): Input data containing the up and latitude confidence.
+
+        Returns:
+            Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]: Costs and weights for the
+            optimization.
+        """
+        costs, weights = {}, {}
+
+        if "up_residual" in residuals:
+            up_cost = (residuals["up_residual"] ** 2).sum(dim=-1)
+            up_cost, up_weight, _ = scaled_loss(up_cost, huber_loss, self.conf.up_loss_fn_scale)
+
+            if "up_confidence" in data:
+                up_conf = data["up_confidence"].reshape(up_weight.shape[0], -1)
+                up_weight = up_weight * up_conf
+                up_cost = up_cost * up_conf
+
+            costs["up_cost"] = up_cost
+            weights["up_weights"] = up_weight
+
+        if "latitude_residual" in residuals:
+            lat_cost = (residuals["latitude_residual"] ** 2).sum(dim=-1)
+            lat_cost, lat_weight, _ = scaled_loss(lat_cost, huber_loss, self.conf.lat_loss_fn_scale)
+
+            if "latitude_confidence" in data:
+                lat_conf = data["latitude_confidence"].reshape(lat_weight.shape[0], -1)
+                lat_weight = lat_weight * lat_conf
+                lat_cost = lat_cost * lat_conf
+
+            costs["latitude_cost"] = lat_cost
+            weights["latitude_weights"] = lat_weight
+
+        return costs, weights
+
+    def calculate_gradient_and_hessian(
+        self,
+        J: torch.Tensor,
+        residuals: torch.Tensor,
+        weights: torch.Tensor,
+        shared_intrinsics: bool,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Calculate the gradient and Hessian for given the Jacobian, residuals, and weights.
+
+        Args:
+            J (torch.Tensor): Jacobian.
+            residuals (torch.Tensor): Residuals.
+            weights (torch.Tensor): Weights.
+            shared_intrinsics (bool): Whether to share the intrinsics across the batch.
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]: Gradient and Hessian.
+        """
+        dims = ()
+        if self.estimate_gravity:
+            dims = (0, 1)
+        if self.estimate_focal:
+            dims += (2,)
+        if self.camera_has_distortion and self.estimate_k1:
+            dims += (3,)
+        assert dims, "No parameters to optimize"
+
+        J = J[..., dims]
+
+        Grad = torch.einsum("...Njk,...Nj->...Nk", J, residuals)
+        Grad = weights[..., None] * Grad
+        Grad = Grad.sum(-2)  # (B, N_params)
+
+        if shared_intrinsics:
+            # reshape to (1, B * (N_params-1) + 1)
+            Grad_g = Grad[..., :2].reshape(1, -1)
+            Grad_f = Grad[..., 2].reshape(1, -1).sum(-1, keepdim=True)
+            Grad = torch.cat([Grad_g, Grad_f], dim=-1)
+
+        Hess = torch.einsum("...Njk,...Njl->...Nkl", J, J)
+        Hess = weights[..., None, None] * Hess
+        Hess = Hess.sum(-3)
+
+        if shared_intrinsics:
+            H_g = torch.block_diag(*list(Hess[..., :2, :2]))
+            J_fg = Hess[..., :2, 2].flatten()
+            J_gf = Hess[..., 2, :2].flatten()
+            J_f = Hess[..., 2, 2].sum()
+            dims = H_g.shape[-1] + 1
+            Hess = Hess.new_zeros((dims, dims), dtype=torch.float32)
+            Hess[:-1, :-1] = H_g
+            Hess[-1, :-1] = J_gf
+            Hess[:-1, -1] = J_fg
+            Hess[-1, -1] = J_f
+            Hess = Hess.unsqueeze(0)
+
+        return Grad, Hess
+
+    def setup_system(
+        self,
+        camera: BaseCamera,
+        gravity: Gravity,
+        residuals: Dict[str, torch.Tensor],
+        weights: Dict[str, torch.Tensor],
+        as_rpf: bool = False,
+        shared_intrinsics: bool = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Calculate the gradient and Hessian for the optimization.
+
+        Args:
+            camera (BaseCamera): Optimized camera.
+            gravity (Gravity): Optimized gravity.
+            residuals (Dict[str, torch.Tensor]): Residuals for the optimization.
+            weights (Dict[str, torch.Tensor]): Weights for the optimization.
+            as_rpf (bool, optional): Wether to calculate the gradient and Hessian with respect to
+            roll, pitch, and focal length. Defaults to False.
+            shared_intrinsics (bool, optional): Whether to share the intrinsics across the batch.
+            Defaults to False.
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]: Gradient and Hessian for the optimization.
+        """
+        J_up, J_lat = J_perspective_field(
+            camera,
+            gravity,
+            spherical=self.conf.use_spherical_manifold and not as_rpf,
+            log_focal=self.conf.use_log_focal and not as_rpf,
+        )
+
+        J_up = J_up.reshape(J_up.shape[0], -1, J_up.shape[-2], J_up.shape[-1])  # (B, N, 2, 3)
+        J_lat = J_lat.reshape(J_lat.shape[0], -1, J_lat.shape[-2], J_lat.shape[-1])  # (B, N, 1, 3)
+
+        n_params = (
+            2 * self.estimate_gravity
+            + self.estimate_focal
+            + (self.camera_has_distortion and self.estimate_k1)
+        )
+        Grad = J_up.new_zeros(J_up.shape[0], n_params)
+        Hess = J_up.new_zeros(J_up.shape[0], n_params, n_params)
+
+        if shared_intrinsics:
+            N_params = Grad.shape[0] * (n_params - 1) + 1
+            Grad = Grad.new_zeros(1, N_params)
+            Hess = Hess.new_zeros(1, N_params, N_params)
+
+        if "up_residual" in residuals:
+            Up_Grad, Up_Hess = self.calculate_gradient_and_hessian(
+                J_up, residuals["up_residual"], weights["up_weights"], shared_intrinsics
+            )
+
+            if self.conf.verbose:
+                logger.info(f"Up J:\n{Up_Grad.mean(0)}")
+
+            Grad = Grad + Up_Grad
+            Hess = Hess + Up_Hess
+
+        if "latitude_residual" in residuals:
+            Lat_Grad, Lat_Hess = self.calculate_gradient_and_hessian(
+                J_lat,
+                residuals["latitude_residual"],
+                weights["latitude_weights"],
+                shared_intrinsics,
+            )
+
+            if self.conf.verbose:
+                logger.info(f"Lat J:\n{Lat_Grad.mean(0)}")
+
+            Grad = Grad + Lat_Grad
+            Hess = Hess + Lat_Hess
+
+        return Grad, Hess
+
+    def estimate_uncertainty(
+        self,
+        camera_opt: BaseCamera,
+        gravity_opt: Gravity,
+        errors: Dict[str, torch.Tensor],
+        weights: Dict[str, torch.Tensor],
+    ) -> Dict[str, torch.Tensor]:
+        """Estimate the uncertainty of the optimized camera and gravity at the final step.
+
+        Args:
+            camera_opt (BaseCamera): Final optimized camera.
+            gravity_opt (Gravity): Final optimized gravity.
+            errors (Dict[str, torch.Tensor]): Costs for the optimization.
+            weights (Dict[str, torch.Tensor]): Weights for the optimization.
+
+        Returns:
+            Dict[str, torch.Tensor]: Uncertainty estimates for the optimized camera and gravity.
+        """
+        _, Hess = self.setup_system(
+            camera_opt, gravity_opt, errors, weights, as_rpf=True, shared_intrinsics=False
+        )
+        Cov = torch.inverse(Hess)
+
+        roll_uncertainty = Cov.new_zeros(Cov[..., 0, 0].shape)
+        pitch_uncertainty = Cov.new_zeros(Cov[..., 0, 0].shape)
+        gravity_uncertainty = Cov.new_zeros(Cov[..., 0, 0].shape)
+        if self.estimate_gravity:
+            roll_uncertainty = Cov[..., 0, 0]
+            pitch_uncertainty = Cov[..., 1, 1]
+
+            try:
+                delta_uncertainty = Cov[..., :2, :2]
+                eigenvalues = torch.linalg.eigvalsh(delta_uncertainty.cpu())
+                gravity_uncertainty = torch.max(eigenvalues, dim=-1).values.to(Cov.device)
+            except RuntimeError:
+                logger.warning("Could not calculate gravity uncertainty")
+                gravity_uncertainty = Cov.new_zeros(Cov.shape[0])
+
+        focal_uncertainty = Cov.new_zeros(Cov[..., 0, 0].shape)
+        fov_uncertainty = Cov.new_zeros(Cov[..., 0, 0].shape)
+        if self.estimate_focal:
+            focal_uncertainty = Cov[..., self.focal_delta_dims[0], self.focal_delta_dims[0]]
+            fov_uncertainty = (
+                J_focal2fov(camera_opt.f[..., 1], camera_opt.size[..., 1]) ** 2 * focal_uncertainty
+            )
+
+        return {
+            "covariance": Cov,
+            "roll_uncertainty": torch.sqrt(roll_uncertainty),
+            "pitch_uncertainty": torch.sqrt(pitch_uncertainty),
+            "gravity_uncertainty": torch.sqrt(gravity_uncertainty),
+            "focal_uncertainty": torch.sqrt(focal_uncertainty) / 2,
+            "vfov_uncertainty": torch.sqrt(fov_uncertainty / 2),
+        }
+
+    def update_estimate(
+        self, camera: BaseCamera, gravity: Gravity, delta: torch.Tensor
+    ) -> Tuple[BaseCamera, Gravity]:
+        """Update the camera and gravity estimates with the given delta.
+
+        Args:
+            camera (BaseCamera): Optimized camera.
+            gravity (Gravity): Optimized gravity.
+            delta (torch.Tensor): Delta to update the camera and gravity estimates.
+
+        Returns:
+            Tuple[BaseCamera, Gravity]: Updated camera and gravity estimates.
+        """
+        delta_gravity = (
+            delta[..., self.gravity_delta_dims]
+            if self.estimate_gravity
+            else delta.new_zeros(delta.shape[:-1] + (2,))
+        )
+        new_gravity = gravity.update(delta_gravity, spherical=self.conf.use_spherical_manifold)
+
+        delta_f = (
+            delta[..., self.focal_delta_dims]
+            if self.estimate_focal
+            else delta.new_zeros(delta.shape[:-1] + (1,))
+        )
+        new_camera = camera.update_focal(delta_f, as_log=self.conf.use_log_focal)
+
+        delta_dist = (
+            delta[..., self.k1_delta_dims]
+            if self.camera_has_distortion and self.estimate_k1
+            else delta.new_zeros(delta.shape[:-1] + (1,))
+        )
+        if self.camera_has_distortion:
+            new_camera = new_camera.update_dist(delta_dist)
+
+        return new_camera, new_gravity
+
+    def optimize(
+        self,
+        data: Dict[str, torch.Tensor],
+        camera_opt: BaseCamera,
+        gravity_opt: Gravity,
+    ) -> Tuple[BaseCamera, Gravity, Dict[str, torch.Tensor]]:
+        """Optimize the camera and gravity estimates.
+
+        Args:
+            data (Dict[str, torch.Tensor]): Input data.
+            camera_opt (BaseCamera): Optimized camera.
+            gravity_opt (Gravity): Optimized gravity.
+
+        Returns:
+            Tuple[BaseCamera, Gravity, Dict[str, torch.Tensor]]: Optimized camera, gravity
+            estimates and optimization information.
+        """
+        key = list(data.keys())[0]
+        B = data[key].shape[0]
+
+        lamb = data[key].new_ones(B) * self.conf.lambda_
+        if self.shared_intrinsics:
+            lamb = data[key].new_ones(1) * self.conf.lambda_
+
+        infos = {"stop_at": self.num_steps}
+        for i in range(self.num_steps):
+            if self.conf.verbose:
+                logger.info(f"Step {i+1}/{self.num_steps}")
+
+            errors = self.calculate_residuals(camera_opt, gravity_opt, data)
+            costs, weights = self.calculate_costs(errors, data)
+
+            if i == 0:
+                prev_cost = sum(c.mean(-1) for c in costs.values())
+                for k, c in costs.items():
+                    infos[f"initial_{k}"] = c.mean(-1)
+
+                infos["initial_cost"] = prev_cost
+
+            Grad, Hess = self.setup_system(
+                camera_opt,
+                gravity_opt,
+                errors,
+                weights,
+                shared_intrinsics=self.shared_intrinsics,
+            )
+            delta = optimizer_step(Grad, Hess, lamb)  # (B, N_params)
+
+            if self.shared_intrinsics:
+                delta_g = delta[..., :-1].reshape(B, 2)
+                delta_f = delta[..., -1].expand(B, 1)
+                delta = torch.cat([delta_g, delta_f], dim=-1)
+
+            # calculate new cost
+            camera_opt, gravity_opt = self.update_estimate(camera_opt, gravity_opt, delta)
+            new_cost, _ = self.calculate_costs(
+                self.calculate_residuals(camera_opt, gravity_opt, data), data
+            )
+            new_cost = sum(c.mean(-1) for c in new_cost.values())
+
+            if not self.conf.fix_lambda and not self.shared_intrinsics:
+                lamb = update_lambda(lamb, prev_cost, new_cost)
+
+            if self.conf.verbose:
+                logger.info(f"Cost:\nPrev: {prev_cost}\nNew:  {new_cost}")
+                logger.info(f"Camera:\n{camera_opt._data}")
+
+            if early_stop(new_cost, prev_cost, atol=self.conf.atol, rtol=self.conf.rtol):
+                infos["stop_at"] = min(i + 1, infos["stop_at"])
+
+                if self.conf.early_stop:
+                    if self.conf.verbose:
+                        logger.info(f"Early stopping at step {i+1}")
+                    break
+
+            prev_cost = new_cost
+
+            if i == self.num_steps - 1 and self.conf.early_stop:
+                logger.warning("Reached maximum number of steps without convergence.")
+
+        final_errors = self.calculate_residuals(camera_opt, gravity_opt, data)  # (B, N, 3)
+        final_cost, weights = self.calculate_costs(final_errors, data)  # (B, N)
+
+        if not self.training:
+            infos |= self.estimate_uncertainty(camera_opt, gravity_opt, final_errors, weights)
+
+        infos["stop_at"] = camera_opt.new_ones(camera_opt.shape[0]) * infos["stop_at"]
+        for k, c in final_cost.items():
+            infos[f"final_{k}"] = c.mean(-1)
+
+        infos["final_cost"] = sum(c.mean(-1) for c in final_cost.values())
+
+        return camera_opt, gravity_opt, infos
+
+    def forward(self, data: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        """Run the LM optimization."""
+        camera_init, gravity_init = get_trivial_estimation(data, self.camera_model)
+
+        self.setup_optimization_and_priors(data, shared_intrinsics=self.shared_intrinsics)
+
+        start = time.time()
+        camera_opt, gravity_opt, infos = self.optimize(data, camera_init, gravity_init)
+
+        if self.conf.verbose:
+            logger.info(f"Optimization took {(time.time() - start)*1000:.2f} ms")
+
+            logger.info(f"Initial camera:\n{rad2deg(camera_init.vfov)}")
+            logger.info(f"Optimized camera:\n{rad2deg(camera_opt.vfov)}")
+
+            logger.info(f"Initial gravity:\n{rad2deg(gravity_init.rp)}")
+            logger.info(f"Optimized gravity:\n{rad2deg(gravity_opt.rp)}")
+
+        return {"camera": camera_opt, "gravity": gravity_opt, **infos}

geocalib/misc.py

@@ -0,0 +1,318 @@
+"""Miscellaneous functions and classes for the geocalib_inference package."""
+
+import functools
+import inspect
+import logging
+from typing import Callable, List
+
+import numpy as np
+import torch
+
+logger = logging.getLogger(__name__)
+
+# mypy: ignore-errors
+
+
+def autocast(func: Callable) -> Callable:
+    """Cast the inputs of a TensorWrapper method to PyTorch tensors if they are numpy arrays.
+
+    Use the device and dtype of the wrapper.
+
+    Args:
+        func (Callable): Method of a TensorWrapper class.
+
+    Returns:
+        Callable: Wrapped method.
+    """
+
+    @functools.wraps(func)
+    def wrap(self, *args):
+        device = torch.device("cpu")
+        dtype = None
+        if isinstance(self, TensorWrapper):
+            if self._data is not None:
+                device = self.device
+                dtype = self.dtype
+        elif not inspect.isclass(self) or not issubclass(self, TensorWrapper):
+            raise ValueError(self)
+
+        cast_args = []
+        for arg in args:
+            if isinstance(arg, np.ndarray):
+                arg = torch.from_numpy(arg)
+                arg = arg.to(device=device, dtype=dtype)
+            cast_args.append(arg)
+        return func(self, *cast_args)
+
+    return wrap
+
+
+class TensorWrapper:
+    """Wrapper for PyTorch tensors."""
+
+    _data = None
+
+    @autocast
+    def __init__(self, data: torch.Tensor):
+        """Wrapper for PyTorch tensors."""
+        self._data = data
+
+    @property
+    def shape(self) -> torch.Size:
+        """Shape of the underlying tensor."""
+        return self._data.shape[:-1]
+
+    @property
+    def device(self) -> torch.device:
+        """Get the device of the underlying tensor."""
+        return self._data.device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        """Get the dtype of the underlying tensor."""
+        return self._data.dtype
+
+    def __getitem__(self, index) -> torch.Tensor:
+        """Get the underlying tensor."""
+        return self.__class__(self._data[index])
+
+    def __setitem__(self, index, item):
+        """Set the underlying tensor."""
+        self._data[index] = item.data
+
+    def to(self, *args, **kwargs):
+        """Move the underlying tensor to a new device."""
+        return self.__class__(self._data.to(*args, **kwargs))
+
+    def cpu(self):
+        """Move the underlying tensor to the CPU."""
+        return self.__class__(self._data.cpu())
+
+    def cuda(self):
+        """Move the underlying tensor to the GPU."""
+        return self.__class__(self._data.cuda())
+
+    def pin_memory(self):
+        """Pin the underlying tensor to memory."""
+        return self.__class__(self._data.pin_memory())
+
+    def float(self):
+        """Cast the underlying tensor to float."""
+        return self.__class__(self._data.float())
+
+    def double(self):
+        """Cast the underlying tensor to double."""
+        return self.__class__(self._data.double())
+
+    def detach(self):
+        """Detach the underlying tensor."""
+        return self.__class__(self._data.detach())
+
+    def numpy(self):
+        """Convert the underlying tensor to a numpy array."""
+        return self._data.detach().cpu().numpy()
+
+    def new_tensor(self, *args, **kwargs):
+        """Create a new tensor of the same type and device."""
+        return self._data.new_tensor(*args, **kwargs)
+
+    def new_zeros(self, *args, **kwargs):
+        """Create a new tensor of the same type and device."""
+        return self._data.new_zeros(*args, **kwargs)
+
+    def new_ones(self, *args, **kwargs):
+        """Create a new tensor of the same type and device."""
+        return self._data.new_ones(*args, **kwargs)
+
+    def new_full(self, *args, **kwargs):
+        """Create a new tensor of the same type and device."""
+        return self._data.new_full(*args, **kwargs)
+
+    def new_empty(self, *args, **kwargs):
+        """Create a new tensor of the same type and device."""
+        return self._data.new_empty(*args, **kwargs)
+
+    def unsqueeze(self, *args, **kwargs):
+        """Create a new tensor of the same type and device."""
+        return self.__class__(self._data.unsqueeze(*args, **kwargs))
+
+    def squeeze(self, *args, **kwargs):
+        """Create a new tensor of the same type and device."""
+        return self.__class__(self._data.squeeze(*args, **kwargs))
+
+    @classmethod
+    def stack(cls, objects: List, dim=0, *, out=None):
+        """Stack a list of objects with the same type and shape."""
+        data = torch.stack([obj._data for obj in objects], dim=dim, out=out)
+        return cls(data)
+
+    @classmethod
+    def __torch_function__(cls, func, types, args=(), kwargs=None):
+        """Support torch functions."""
+        if kwargs is None:
+            kwargs = {}
+        return cls.stack(*args, **kwargs) if func is torch.stack else NotImplemented
+
+
+class EuclideanManifold:
+    """Simple euclidean manifold."""
+
+    @staticmethod
+    def J_plus(x: torch.Tensor) -> torch.Tensor:
+        """Plus operator Jacobian."""
+        return torch.eye(x.shape[-1]).to(x)
+
+    @staticmethod
+    def plus(x: torch.Tensor, delta: torch.Tensor) -> torch.Tensor:
+        """Plus operator."""
+        return x + delta
+
+
+class SphericalManifold:
+    """Implementation of the spherical manifold.
+
+    Following the derivation from 'Integrating Generic Sensor Fusion Algorithms with Sound State
+    Representations through Encapsulation of Manifolds' by Hertzberg et al. (B.2, p. 25).
+
+    Householder transformation following Algorithm 5.1.1 (p. 210) from 'Matrix Computations' by
+    Golub et al.
+    """
+
+    @staticmethod
+    def householder_vector(x: torch.Tensor) -> torch.Tensor:
+        """Return the Householder vector and beta.
+
+        Algorithm 5.1.1 (p. 210) from 'Matrix Computations' by Golub et al. (Johns Hopkins Studies
+        in Mathematical Sciences) but using the nth element of the input vector as pivot instead of
+        first.
+
+        This computes the vector v with v(n) = 1 and beta such that H = I - beta * v * v^T is
+        orthogonal and H * x = ||x||_2 * e_n.
+
+        Args:
+            x (torch.Tensor): [..., n] tensor.
+
+        Returns:
+            torch.Tensor: v of shape [..., n]
+            torch.Tensor: beta of shape [...]
+        """
+        sigma = torch.sum(x[..., :-1] ** 2, -1)
+        xpiv = x[..., -1]
+        norm = torch.norm(x, dim=-1)
+        if torch.any(sigma < 1e-7):
+            sigma = torch.where(sigma < 1e-7, sigma + 1e-7, sigma)
+            logger.warning("sigma < 1e-7")
+
+        vpiv = torch.where(xpiv < 0, xpiv - norm, -sigma / (xpiv + norm))
+        beta = 2 * vpiv**2 / (sigma + vpiv**2)
+        v = torch.cat([x[..., :-1] / vpiv[..., None], torch.ones_like(vpiv)[..., None]], -1)
+        return v, beta
+
+    @staticmethod
+    def apply_householder(y: torch.Tensor, v: torch.Tensor, beta: torch.Tensor) -> torch.Tensor:
+        """Apply Householder transformation.
+
+        Args:
+            y (torch.Tensor): Vector to transform of shape [..., n].
+            v (torch.Tensor): Householder vector of shape [..., n].
+            beta (torch.Tensor): Householder beta of shape [...].
+
+        Returns:
+            torch.Tensor: Transformed vector of shape [..., n].
+        """
+        return y - v * (beta * torch.einsum("...i,...i->...", v, y))[..., None]
+
+    @classmethod
+    def J_plus(cls, x: torch.Tensor) -> torch.Tensor:
+        """Plus operator Jacobian."""
+        v, beta = cls.householder_vector(x)
+        H = -torch.einsum("..., ...k, ...l->...kl", beta, v, v)
+        H = H + torch.eye(H.shape[-1]).to(H)
+        return H[..., :-1]  # J
+
+    @classmethod
+    def plus(cls, x: torch.Tensor, delta: torch.Tensor) -> torch.Tensor:
+        """Plus operator.
+
+        Equation 109 (p. 25) from 'Integrating Generic Sensor Fusion Algorithms with Sound State
+        Representations through Encapsulation of Manifolds' by Hertzberg et al. but using the nth
+        element of the input vector as pivot instead of first.
+
+        Args:
+            x: point on the manifold
+            delta: tangent vector
+        """
+        eps = 1e-7
+        # keep norm is not equal to 1
+        nx = torch.norm(x, dim=-1, keepdim=True)
+        nd = torch.norm(delta, dim=-1, keepdim=True)
+
+        # make sure we don't divide by zero in backward as torch.where computes grad for both
+        # branches
+        nd_ = torch.where(nd < eps, nd + eps, nd)
+        sinc = torch.where(nd < eps, nd.new_ones(nd.shape), torch.sin(nd_) / nd_)
+
+        # cos is applied to last dim instead of first
+        exp_delta = torch.cat([sinc * delta, torch.cos(nd)], -1)
+
+        v, beta = cls.householder_vector(x)
+        return nx * cls.apply_householder(exp_delta, v, beta)
+
+
+@torch.jit.script
+def J_vecnorm(vec: torch.Tensor) -> torch.Tensor:
+    """Compute the jacobian of vec / norm2(vec).
+
+    Args:
+        vec (torch.Tensor): [..., D] tensor.
+
+    Returns:
+        torch.Tensor: [..., D, D] Jacobian.
+    """
+    D = vec.shape[-1]
+    norm_x = torch.norm(vec, dim=-1, keepdim=True).unsqueeze(-1)  # (..., 1, 1)
+
+    if (norm_x == 0).any():
+        norm_x = norm_x + 1e-6
+
+    xxT = torch.einsum("...i,...j->...ij", vec, vec)  # (..., D, D)
+    identity = torch.eye(D, device=vec.device, dtype=vec.dtype)  # (D, D)
+
+    return identity / norm_x - (xxT / norm_x**3)  # (..., D, D)
+
+
+@torch.jit.script
+def J_focal2fov(focal: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
+    """Compute the jacobian of the focal2fov function."""
+    return -4 * h / (4 * focal**2 + h**2)
+
+
+@torch.jit.script
+def J_up_projection(uv: torch.Tensor, abc: torch.Tensor, wrt: str = "uv") -> torch.Tensor:
+    """Compute the jacobian of the up-vector projection.
+
+    Args:
+        uv (torch.Tensor): Normalized image coordinates of shape (..., 2).
+        abc (torch.Tensor): Gravity vector of shape (..., 3).
+        wrt (str, optional): Parameter to differentiate with respect to. Defaults to "uv".
+
+    Raises:
+        ValueError: If the wrt parameter is unknown.
+
+    Returns:
+        torch.Tensor: Jacobian with respect to the parameter.
+    """
+    if wrt == "uv":
+        c = abc[..., 2][..., None, None, None]
+        return -c * torch.eye(2, device=uv.device, dtype=uv.dtype).expand(uv.shape[:-1] + (2, 2))
+
+    elif wrt == "abc":
+        J = uv.new_zeros(uv.shape[:-1] + (2, 3))
+        J[..., 0, 0] = 1
+        J[..., 1, 1] = 1
+        J[..., 0, 2] = -uv[..., 0]
+        J[..., 1, 2] = -uv[..., 1]
+        return J
+
+    else:
+        raise ValueError(f"Unknown wrt: {wrt}")

geocalib/modules.py

@@ -0,0 +1,575 @@
+"""Implementation of MSCAN from SegNeXt: Rethinking Convolutional Attention Design for Semantic 
+Segmentation (NeurIPS 2022) adapted from
+
+https://github.com/Visual-Attention-Network/SegNeXt/blob/main/mmseg/models/backbones/mscan.py
+
+
+Light Hamburger Decoder adapted from:
+
+https://github.com/Visual-Attention-Network/SegNeXt/blob/main/mmseg/models/decode_heads/ham_head.py
+"""
+
+from typing import Dict, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.modules.utils import _pair as to_2tuple
+
+# flake8: noqa: E266
+# mypy: ignore-errors
+
+
+class ConvModule(nn.Module):
+    """Replacement for mmcv.cnn.ConvModule to avoid mmcv dependency."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        padding: int = 0,
+        use_norm: bool = False,
+        bias: bool = True,
+    ):
+        """Simple convolution block.
+
+        Args:
+            in_channels (int): Input channels.
+            out_channels (int): Output channels.
+            kernel_size (int): Kernel size.
+            padding (int, optional): Padding. Defaults to 0.
+            use_norm (bool, optional): Whether to use normalization. Defaults to False.
+            bias (bool, optional): Whether to use bias. Defaults to True.
+        """
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding, bias=bias)
+        self.bn = nn.BatchNorm2d(out_channels) if use_norm else nn.Identity()
+        self.activate = nn.ReLU(inplace=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass."""
+        x = self.conv(x)
+        x = self.bn(x)
+        return self.activate(x)
+
+
+class ResidualConvUnit(nn.Module):
+    """Residual convolution module."""
+
+    def __init__(self, features):
+        """Simple residual convolution block.
+
+        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 = torch.nn.ReLU(inplace=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass."""
+        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: int, unit2only=False, upsample=True):
+        """Feature fusion block.
+
+        Args:
+            features (int): Number of features.
+            unit2only (bool, optional): Whether to use only the second unit. Defaults to False.
+            upsample (bool, optional): Whether to upsample. Defaults to True.
+        """
+        super().__init__()
+        self.upsample = upsample
+
+        if not unit2only:
+            self.resConfUnit1 = ResidualConvUnit(features)
+        self.resConfUnit2 = ResidualConvUnit(features)
+
+    def forward(self, *xs: torch.Tensor) -> torch.Tensor:
+        """Forward pass."""
+        output = xs[0]
+
+        if len(xs) == 2:
+            output = output + self.resConfUnit1(xs[1])
+
+        output = self.resConfUnit2(output)
+
+        if self.upsample:
+            output = F.interpolate(output, scale_factor=2, mode="bilinear", align_corners=False)
+
+        return output
+
+
+###################################################
+########### Light Hamburger Decoder ###############
+###################################################
+
+
+class NMF2D(nn.Module):
+    """Non-negative Matrix Factorization (NMF) for 2D data."""
+
+    def __init__(self):
+        """Non-negative Matrix Factorization (NMF) for 2D data."""
+        super().__init__()
+        self.S, self.D, self.R = 1, 512, 64
+        self.train_steps = 6
+        self.eval_steps = 7
+        self.inv_t = 1
+
+    def _build_bases(self, B: int, S: int, D: int, R: int, device: str = "cpu") -> torch.Tensor:
+        bases = torch.rand((B * S, D, R)).to(device)
+        return F.normalize(bases, dim=1)
+
+    def local_step(
+        self, x: torch.Tensor, bases: torch.Tensor, coef: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Update bases and coefficient."""
+        # (B * S, D, N)^T @ (B * S, D, R) -> (B * S, N, R)
+        numerator = torch.bmm(x.transpose(1, 2), bases)
+        # (B * S, N, R) @ [(B * S, D, R)^T @ (B * S, D, R)] -> (B * S, N, R)
+        denominator = coef.bmm(bases.transpose(1, 2).bmm(bases))
+        # Multiplicative Update
+        coef = coef * numerator / (denominator + 1e-6)
+        # (B * S, D, N) @ (B * S, N, R) -> (B * S, D, R)
+        numerator = torch.bmm(x, coef)
+        # (B * S, D, R) @ [(B * S, N, R)^T @ (B * S, N, R)] -> (B * S, D, R)
+        denominator = bases.bmm(coef.transpose(1, 2).bmm(coef))
+        # Multiplicative Update
+        bases = bases * numerator / (denominator + 1e-6)
+        return bases, coef
+
+    def compute_coef(
+        self, x: torch.Tensor, bases: torch.Tensor, coef: torch.Tensor
+    ) -> torch.Tensor:
+        """Compute coefficient."""
+        # (B * S, D, N)^T @ (B * S, D, R) -> (B * S, N, R)
+        numerator = torch.bmm(x.transpose(1, 2), bases)
+        # (B * S, N, R) @ (B * S, D, R)^T @ (B * S, D, R) -> (B * S, N, R)
+        denominator = coef.bmm(bases.transpose(1, 2).bmm(bases))
+        # multiplication update
+        return coef * numerator / (denominator + 1e-6)
+
+    def local_inference(
+        self, x: torch.Tensor, bases: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Local inference."""
+        # (B * S, D, N)^T @ (B * S, D, R) -> (B * S, N, R)
+        coef = torch.bmm(x.transpose(1, 2), bases)
+        coef = F.softmax(self.inv_t * coef, dim=-1)
+
+        steps = self.train_steps if self.training else self.eval_steps
+        for _ in range(steps):
+            bases, coef = self.local_step(x, bases, coef)
+
+        return bases, coef
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass."""
+        B, C, H, W = x.shape
+
+        # (B, C, H, W) -> (B * S, D, N)
+        D = C // self.S
+        N = H * W
+        x = x.view(B * self.S, D, N)
+
+        # (S, D, R) -> (B * S, D, R)
+        bases = self._build_bases(B, self.S, D, self.R, device=x.device)
+        bases, coef = self.local_inference(x, bases)
+        # (B * S, N, R)
+        coef = self.compute_coef(x, bases, coef)
+        # (B * S, D, R) @ (B * S, N, R)^T -> (B * S, D, N)
+        x = torch.bmm(bases, coef.transpose(1, 2))
+        # (B * S, D, N) -> (B, C, H, W)
+        x = x.view(B, C, H, W)
+        # (B * H, D, R) -> (B, H, N, D)
+        bases = bases.view(B, self.S, D, self.R)
+
+        return x
+
+
+class Hamburger(nn.Module):
+    """Hamburger Module."""
+
+    def __init__(self, ham_channels: int = 512):
+        """Hambuger Module.
+
+        Args:
+            ham_channels (int, optional): Number of channels in the hamburger module. Defaults to
+            512.
+        """
+        super().__init__()
+        self.ham_in = ConvModule(ham_channels, ham_channels, 1)
+        self.ham = NMF2D()
+        self.ham_out = ConvModule(ham_channels, ham_channels, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass."""
+        enjoy = self.ham_in(x)
+        enjoy = F.relu(enjoy, inplace=False)
+        enjoy = self.ham(enjoy)
+        enjoy = self.ham_out(enjoy)
+        ham = F.relu(x + enjoy, inplace=False)
+        return ham
+
+
+class LightHamHead(nn.Module):
+    """Is Attention Better Than Matrix Decomposition?
+
+    This head is the implementation of `HamNet <https://arxiv.org/abs/2109.04553>`.
+    """
+
+    def __init__(self):
+        """Light hamburger decoder head."""
+        super().__init__()
+        self.in_index = [0, 1, 2, 3]
+        self.in_channels = [64, 128, 320, 512]
+        self.out_channels = 64
+        self.ham_channels = 512
+        self.align_corners = False
+
+        self.squeeze = ConvModule(sum(self.in_channels), self.ham_channels, 1)
+
+        self.hamburger = Hamburger(self.ham_channels)
+
+        self.align = ConvModule(self.ham_channels, self.out_channels, 1)
+
+        self.linear_pred_uncertainty = nn.Sequential(
+            ConvModule(
+                in_channels=self.out_channels,
+                out_channels=self.out_channels,
+                kernel_size=3,
+                padding=1,
+                bias=False,
+            ),
+            nn.Conv2d(in_channels=self.out_channels, out_channels=1, kernel_size=1),
+        )
+
+        self.out_conv = ConvModule(self.out_channels, self.out_channels, 3, padding=1, bias=False)
+        self.ll_fusion = FeatureFusionBlock(self.out_channels, upsample=False)
+
+    def forward(self, features: Dict[str, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass."""
+        inputs = [features["hl"][i] for i in self.in_index]
+
+        inputs = [
+            F.interpolate(
+                level, size=inputs[0].shape[2:], mode="bilinear", align_corners=self.align_corners
+            )
+            for level in inputs
+        ]
+
+        inputs = torch.cat(inputs, dim=1)
+        x = self.squeeze(inputs)
+
+        x = self.hamburger(x)
+
+        feats = self.align(x)
+
+        assert "ll" in features, "Low-level features are required for this model"
+        feats = F.interpolate(feats, scale_factor=2, mode="bilinear", align_corners=False)
+        feats = self.out_conv(feats)
+        feats = F.interpolate(feats, scale_factor=2, mode="bilinear", align_corners=False)
+        feats = self.ll_fusion(feats, features["ll"].clone())
+
+        uncertainty = self.linear_pred_uncertainty(feats).squeeze(1)
+
+        return feats, uncertainty
+
+
+###################################################
+###########          MSCAN         ################
+###################################################
+
+
+class DWConv(nn.Module):
+    """Depthwise convolution."""
+
+    def __init__(self, dim: int = 768):
+        """Depthwise convolution.
+
+        Args:
+            dim (int, optional): Number of features. Defaults to 768.
+        """
+        super().__init__()
+        self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass."""
+        return self.dwconv(x)
+
+
+class Mlp(nn.Module):
+    """MLP module."""
+
+    def __init__(
+        self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.0
+    ):
+        """Initialize the MLP."""
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Conv2d(in_features, hidden_features, 1)
+        self.dwconv = DWConv(hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Conv2d(hidden_features, out_features, 1)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        """Forward pass."""
+        x = self.fc1(x)
+
+        x = self.dwconv(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+
+        return x
+
+
+class StemConv(nn.Module):
+    """Simple stem convolution module."""
+
+    def __init__(self, in_channels: int, out_channels: int):
+        """Simple stem convolution module.
+
+        Args:
+            in_channels (int): Input channels.
+            out_channels (int): Output channels.
+        """
+        super().__init__()
+
+        self.proj = nn.Sequential(
+            nn.Conv2d(
+                in_channels, out_channels // 2, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)
+            ),
+            nn.BatchNorm2d(out_channels // 2),
+            nn.GELU(),
+            nn.Conv2d(
+                out_channels // 2, out_channels, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)
+            ),
+            nn.BatchNorm2d(out_channels),
+        )
+
+    def forward(self, x):
+        """Forward pass."""
+        x = self.proj(x)
+        _, _, H, W = x.size()
+        x = x.flatten(2).transpose(1, 2)
+        return x, H, W
+
+
+class AttentionModule(nn.Module):
+    """Attention module."""
+
+    def __init__(self, dim: int):
+        """Attention module.
+
+        Args:
+            dim (int): Number of features.
+        """
+        super().__init__()
+        self.conv0 = nn.Conv2d(dim, dim, 5, padding=2, groups=dim)
+        self.conv0_1 = nn.Conv2d(dim, dim, (1, 7), padding=(0, 3), groups=dim)
+        self.conv0_2 = nn.Conv2d(dim, dim, (7, 1), padding=(3, 0), groups=dim)
+
+        self.conv1_1 = nn.Conv2d(dim, dim, (1, 11), padding=(0, 5), groups=dim)
+        self.conv1_2 = nn.Conv2d(dim, dim, (11, 1), padding=(5, 0), groups=dim)
+
+        self.conv2_1 = nn.Conv2d(dim, dim, (1, 21), padding=(0, 10), groups=dim)
+        self.conv2_2 = nn.Conv2d(dim, dim, (21, 1), padding=(10, 0), groups=dim)
+        self.conv3 = nn.Conv2d(dim, dim, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass."""
+        u = x.clone()
+        attn = self.conv0(x)
+
+        attn_0 = self.conv0_1(attn)
+        attn_0 = self.conv0_2(attn_0)
+
+        attn_1 = self.conv1_1(attn)
+        attn_1 = self.conv1_2(attn_1)
+
+        attn_2 = self.conv2_1(attn)
+        attn_2 = self.conv2_2(attn_2)
+        attn = attn + attn_0 + attn_1 + attn_2
+
+        attn = self.conv3(attn)
+        return attn * u
+
+
+class SpatialAttention(nn.Module):
+    """Spatial attention module."""
+
+    def __init__(self, dim: int):
+        """Spatial attention module.
+
+        Args:
+            dim (int): Number of features.
+        """
+        super().__init__()
+        self.d_model = dim
+        self.proj_1 = nn.Conv2d(dim, dim, 1)
+        self.activation = nn.GELU()
+        self.spatial_gating_unit = AttentionModule(dim)
+        self.proj_2 = nn.Conv2d(dim, dim, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass."""
+        shorcut = x.clone()
+        x = self.proj_1(x)
+        x = self.activation(x)
+        x = self.spatial_gating_unit(x)
+        x = self.proj_2(x)
+        x = x + shorcut
+        return x
+
+
+class Block(nn.Module):
+    """MSCAN block."""
+
+    def __init__(
+        self, dim: int, mlp_ratio: float = 4.0, drop: float = 0.0, act_layer: nn.Module = nn.GELU
+    ):
+        """MSCAN block.
+
+        Args:
+            dim (int): Number of features.
+            mlp_ratio (float, optional): Ratio of the hidden features in the MLP. Defaults to 4.0.
+            drop (float, optional): Dropout rate. Defaults to 0.0.
+            act_layer (nn.Module, optional): Activation layer. Defaults to nn.GELU.
+        """
+        super().__init__()
+        self.norm1 = nn.BatchNorm2d(dim)
+        self.attn = SpatialAttention(dim)
+        self.drop_path = nn.Identity()  # only used in training
+        self.norm2 = nn.BatchNorm2d(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(
+            in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop
+        )
+        layer_scale_init_value = 1e-2
+        self.layer_scale_1 = nn.Parameter(
+            layer_scale_init_value * torch.ones((dim)), requires_grad=True
+        )
+        self.layer_scale_2 = nn.Parameter(
+            layer_scale_init_value * torch.ones((dim)), requires_grad=True
+        )
+
+    def forward(self, x: torch.Tensor, H: int, W: int) -> torch.Tensor:
+        """Forward pass."""
+        B, N, C = x.shape
+        x = x.permute(0, 2, 1).view(B, C, H, W)
+        x = x + self.drop_path(self.layer_scale_1[..., None, None] * self.attn(self.norm1(x)))
+        x = x + self.drop_path(self.layer_scale_2[..., None, None] * self.mlp(self.norm2(x)))
+        return x.view(B, C, N).permute(0, 2, 1)
+
+
+class OverlapPatchEmbed(nn.Module):
+    """Image to Patch Embedding"""
+
+    def __init__(
+        self, patch_size: int = 7, stride: int = 4, in_chans: int = 3, embed_dim: int = 768
+    ):
+        """Image to Patch Embedding.
+
+        Args:
+            patch_size (int, optional): Image patch size. Defaults to 7.
+            stride (int, optional): Stride. Defaults to 4.
+            in_chans (int, optional): Number of input channels. Defaults to 3.
+            embed_dim (int, optional): Embedding dimension. Defaults to 768.
+        """
+        super().__init__()
+        patch_size = to_2tuple(patch_size)
+
+        self.proj = nn.Conv2d(
+            in_chans,
+            embed_dim,
+            kernel_size=patch_size,
+            stride=stride,
+            padding=(patch_size[0] // 2, patch_size[1] // 2),
+        )
+        self.norm = nn.BatchNorm2d(embed_dim)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, int, int]:
+        """Forward pass."""
+        x = self.proj(x)
+        _, _, H, W = x.shape
+        x = self.norm(x)
+        x = x.flatten(2).transpose(1, 2)
+        return x, H, W
+
+
+class MSCAN(nn.Module):
+    """Multi-scale convolutional attention network."""
+
+    def __init__(self):
+        """Multi-scale convolutional attention network."""
+        super().__init__()
+        self.in_channels = 3
+        self.embed_dims = [64, 128, 320, 512]
+        self.mlp_ratios = [8, 8, 4, 4]
+        self.drop_rate = 0.0
+        self.drop_path_rate = 0.1
+        self.depths = [3, 3, 12, 3]
+        self.num_stages = 4
+
+        for i in range(self.num_stages):
+            if i == 0:
+                patch_embed = StemConv(3, self.embed_dims[0])
+            else:
+                patch_embed = OverlapPatchEmbed(
+                    patch_size=7 if i == 0 else 3,
+                    stride=4 if i == 0 else 2,
+                    in_chans=self.in_chans if i == 0 else self.embed_dims[i - 1],
+                    embed_dim=self.embed_dims[i],
+                )
+
+            block = nn.ModuleList(
+                [
+                    Block(
+                        dim=self.embed_dims[i],
+                        mlp_ratio=self.mlp_ratios[i],
+                        drop=self.drop_rate,
+                    )
+                    for _ in range(self.depths[i])
+                ]
+            )
+            norm = nn.LayerNorm(self.embed_dims[i])
+
+            setattr(self, f"patch_embed{i + 1}", patch_embed)
+            setattr(self, f"block{i + 1}", block)
+            setattr(self, f"norm{i + 1}", norm)
+
+    def forward(self, data):
+        """Forward pass."""
+        # rgb -> bgr and from [0, 1] to [0, 255]
+        x = data["image"][:, [2, 1, 0], :, :] * 255.0
+        B = x.shape[0]
+
+        outs = []
+        for i in range(self.num_stages):
+            patch_embed = getattr(self, f"patch_embed{i + 1}")
+            block = getattr(self, f"block{i + 1}")
+            norm = getattr(self, f"norm{i + 1}")
+            x, H, W = patch_embed(x)
+            for blk in block:
+                x = blk(x, H, W)
+            x = norm(x)
+            x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
+            outs.append(x)
+
+        return {"features": outs}

geocalib/perspective_fields.py

@@ -0,0 +1,366 @@
+"""Implementation of perspective fields.
+
+Adapted from https://github.com/jinlinyi/PerspectiveFields/blob/main/perspective2d/utils/panocam.py
+"""
+
+from typing import Tuple
+
+import torch
+from torch.nn import functional as F
+
+from geocalib.camera import BaseCamera
+from geocalib.gravity import Gravity
+from geocalib.misc import J_up_projection, J_vecnorm, SphericalManifold
+
+# flake8: noqa: E266
+
+
+def get_horizon_line(camera: BaseCamera, gravity: Gravity, relative: bool = True) -> torch.Tensor:
+    """Get the horizon line from the camera parameters.
+
+    Args:
+        camera (Camera): Camera parameters.
+        gravity (Gravity): Gravity vector.
+        relative (bool, optional): Whether to normalize horizon line by img_h. Defaults to True.
+
+    Returns:
+        torch.Tensor: In image frame, fraction of image left/right border intersection with
+        respect to image height.
+    """
+    camera = camera.unsqueeze(0) if len(camera.shape) == 0 else camera
+    gravity = gravity.unsqueeze(0) if len(gravity.shape) == 0 else gravity
+
+    # project horizon midpoint to image plane
+    horizon_midpoint = camera.new_tensor([0, 0, 1])
+    horizon_midpoint = camera.K @ gravity.R @ horizon_midpoint
+    midpoint = horizon_midpoint[:2] / horizon_midpoint[2]
+
+    # compute left and right offset to borders
+    left_offset = midpoint[0] * torch.tan(gravity.roll)
+    right_offset = (camera.size[0] - midpoint[0]) * torch.tan(gravity.roll)
+    left, right = midpoint[1] + left_offset, midpoint[1] - right_offset
+
+    horizon = camera.new_tensor([left, right])
+    return horizon / camera.size[1] if relative else horizon
+
+
+def get_up_field(camera: BaseCamera, gravity: Gravity, normalize: bool = True) -> torch.Tensor:
+    """Get the up vector field from the camera parameters.
+
+    Args:
+        camera (Camera): Camera parameters.
+        normalize (bool, optional): Whether to normalize the up vector. Defaults to True.
+
+    Returns:
+        torch.Tensor: up vector field as tensor of shape (..., h, w, 2).
+    """
+    camera = camera.unsqueeze(0) if len(camera.shape) == 0 else camera
+    gravity = gravity.unsqueeze(0) if len(gravity.shape) == 0 else gravity
+
+    w, h = camera.size[0].unbind(-1)
+    h, w = h.round().to(int), w.round().to(int)
+
+    uv = camera.normalize(camera.pixel_coordinates())
+
+    # projected up is (a, b) - c * (u, v)
+    abc = gravity.vec3d
+    projected_up2d = abc[..., None, :2] - abc[..., 2, None, None] * uv  # (..., N, 2)
+
+    if hasattr(camera, "dist"):
+        d_uv = camera.distort(uv, return_scale=True)[0]  # (..., N, 1)
+        d_uv = torch.diag_embed(d_uv.expand(d_uv.shape[:-1] + (2,)))  # (..., N, 2, 2)
+        offset = camera.up_projection_offset(uv)  # (..., N, 2)
+        offset = torch.einsum("...i,...j->...ij", offset, uv)  # (..., N, 2, 2)
+
+        # (..., N, 2)
+        projected_up2d = torch.einsum("...Nij,...Nj->...Ni", d_uv + offset, projected_up2d)
+
+    if normalize:
+        projected_up2d = F.normalize(projected_up2d, dim=-1)  # (..., N, 2)
+
+    return projected_up2d.reshape(camera.shape[0], h, w, 2)
+
+
+def J_up_field(
+    camera: BaseCamera, gravity: Gravity, spherical: bool = False, log_focal: bool = False
+) -> torch.Tensor:
+    """Get the jacobian of the up field.
+
+    Args:
+        camera (Camera): Camera parameters.
+        gravity (Gravity): Gravity vector.
+        spherical (bool, optional): Whether to use spherical coordinates. Defaults to False.
+        log_focal (bool, optional): Whether to use log-focal length. Defaults to False.
+
+    Returns:
+        torch.Tensor: Jacobian of the up field as a tensor of shape (..., h, w, 2, 2, 3).
+    """
+    camera = camera.unsqueeze(0) if len(camera.shape) == 0 else camera
+    gravity = gravity.unsqueeze(0) if len(gravity.shape) == 0 else gravity
+
+    w, h = camera.size[0].unbind(-1)
+    h, w = h.round().to(int), w.round().to(int)
+
+    # Forward
+    xy = camera.pixel_coordinates()
+    uv = camera.normalize(xy)
+
+    projected_up2d = gravity.vec3d[..., None, :2] - gravity.vec3d[..., 2, None, None] * uv
+
+    # Backward
+    J = []
+
+    # (..., N, 2, 2)
+    J_norm2proj = J_vecnorm(
+        get_up_field(camera, gravity, normalize=False).reshape(camera.shape[0], -1, 2)
+    )
+
+    # distortion values
+    if hasattr(camera, "dist"):
+        d_uv = camera.distort(uv, return_scale=True)[0]  # (..., N, 1)
+        d_uv = torch.diag_embed(d_uv.expand(d_uv.shape[:-1] + (2,)))  # (..., N, 2, 2)
+        offset = camera.up_projection_offset(uv)  # (..., N, 2)
+        offset_uv = torch.einsum("...i,...j->...ij", offset, uv)  # (..., N, 2, 2)
+
+    ######################
+    ## Gravity Jacobian ##
+    ######################
+
+    J_proj2abc = J_up_projection(uv, gravity.vec3d, wrt="abc")  # (..., N, 2, 3)
+
+    if hasattr(camera, "dist"):
+        # (..., N, 2, 3)
+        J_proj2abc = torch.einsum("...Nij,...Njk->...Nik", d_uv + offset_uv, J_proj2abc)
+
+    J_abc2delta = SphericalManifold.J_plus(gravity.vec3d) if spherical else gravity.J_rp()
+    J_proj2delta = torch.einsum("...Nij,...jk->...Nik", J_proj2abc, J_abc2delta)
+    J_up2delta = torch.einsum("...Nij,...Njk->...Nik", J_norm2proj, J_proj2delta)
+    J.append(J_up2delta)
+
+    ######################
+    ### Focal Jacobian ###
+    ######################
+
+    J_proj2uv = J_up_projection(uv, gravity.vec3d, wrt="uv")  # (..., N, 2, 2)
+
+    if hasattr(camera, "dist"):
+        J_proj2up = torch.einsum("...Nij,...Njk->...Nik", d_uv + offset_uv, J_proj2uv)
+        J_proj2duv = torch.einsum("...i,...j->...ji", offset, projected_up2d)
+
+        inner = (uv * projected_up2d).sum(-1)[..., None, None]
+        J_proj2offset1 = inner * camera.J_up_projection_offset(uv, wrt="uv")
+        J_proj2offset2 = torch.einsum("...i,...j->...ij", offset, projected_up2d)  # (..., N, 2, 2)
+        J_proj2uv = (J_proj2duv + J_proj2offset1 + J_proj2offset2) + J_proj2up
+
+    J_uv2f = camera.J_normalize(xy)  # (..., N, 2, 2)
+
+    if log_focal:
+        J_uv2f = J_uv2f * camera.f[..., None, None, :]  # (..., N, 2, 2)
+
+    J_uv2f = J_uv2f.sum(-1)  # (..., N, 2)
+
+    J_proj2f = torch.einsum("...ij,...j->...i", J_proj2uv, J_uv2f)  # (..., N, 2)
+    J_up2f = torch.einsum("...Nij,...Nj->...Ni", J_norm2proj, J_proj2f)[..., None]  # (..., N, 2, 1)
+    J.append(J_up2f)
+
+    ######################
+    ##### K1 Jacobian ####
+    ######################
+
+    if hasattr(camera, "dist"):
+        J_duv = camera.J_distort(uv, wrt="scale2dist")
+        J_duv = torch.diag_embed(J_duv.expand(J_duv.shape[:-1] + (2,)))  # (..., N, 2, 2)
+        J_offset = torch.einsum(
+            "...i,...j->...ij", camera.J_up_projection_offset(uv, wrt="dist"), uv
+        )
+        J_proj2k1 = torch.einsum("...Nij,...Nj->...Ni", J_duv + J_offset, projected_up2d)
+        J_k1 = torch.einsum("...Nij,...Nj->...Ni", J_norm2proj, J_proj2k1)[..., None]
+        J.append(J_k1)
+
+    n_params = sum(j.shape[-1] for j in J)
+    return torch.cat(J, axis=-1).reshape(camera.shape[0], h, w, 2, n_params)
+
+
+def get_latitude_field(camera: BaseCamera, gravity: Gravity) -> torch.Tensor:
+    """Get the latitudes of the camera pixels in radians.
+
+    Latitudes are defined as the angle between the ray and the up vector.
+
+    Args:
+        camera (Camera): Camera parameters.
+        gravity (Gravity): Gravity vector.
+
+    Returns:
+        torch.Tensor: Latitudes in radians as a tensor of shape (..., h, w, 1).
+    """
+    camera = camera.unsqueeze(0) if len(camera.shape) == 0 else camera
+    gravity = gravity.unsqueeze(0) if len(gravity.shape) == 0 else gravity
+
+    w, h = camera.size[0].unbind(-1)
+    h, w = h.round().to(int), w.round().to(int)
+
+    uv1, _ = camera.image2world(camera.pixel_coordinates())
+    rays = camera.pixel_bearing_many(uv1)
+
+    lat = torch.einsum("...Nj,...j->...N", rays, gravity.vec3d)
+
+    eps = 1e-6
+    lat_asin = torch.asin(lat.clamp(min=-1 + eps, max=1 - eps))
+
+    return lat_asin.reshape(camera.shape[0], h, w, 1)
+
+
+def J_latitude_field(
+    camera: BaseCamera, gravity: Gravity, spherical: bool = False, log_focal: bool = False
+) -> torch.Tensor:
+    """Get the jacobian of the latitude field.
+
+    Args:
+        camera (Camera): Camera parameters.
+        gravity (Gravity): Gravity vector.
+        spherical (bool, optional): Whether to use spherical coordinates. Defaults to False.
+        log_focal (bool, optional): Whether to use log-focal length. Defaults to False.
+
+    Returns:
+        torch.Tensor: Jacobian of the latitude field as a tensor of shape (..., h, w, 1, 3).
+    """
+    camera = camera.unsqueeze(0) if len(camera.shape) == 0 else camera
+    gravity = gravity.unsqueeze(0) if len(gravity.shape) == 0 else gravity
+
+    w, h = camera.size[0].unbind(-1)
+    h, w = h.round().to(int), w.round().to(int)
+
+    # Forward
+    xy = camera.pixel_coordinates()
+    uv1, _ = camera.image2world(xy)
+    uv1_norm = camera.pixel_bearing_many(uv1)  # (..., N, 3)
+
+    # Backward
+    J = []
+    J_norm2w_to_img = J_vecnorm(uv1)[..., :2]  # (..., N, 2)
+
+    ######################
+    ## Gravity Jacobian ##
+    ######################
+
+    J_delta = SphericalManifold.J_plus(gravity.vec3d) if spherical else gravity.J_rp()
+    J_delta = torch.einsum("...Ni,...ij->...Nj", uv1_norm, J_delta)  # (..., N, 2)
+    J.append(J_delta)
+
+    ######################
+    ### Focal Jacobian ###
+    ######################
+
+    J_w_to_img2f = camera.J_image2world(xy, "f")  # (..., N, 2, 2)
+    if log_focal:
+        J_w_to_img2f = J_w_to_img2f * camera.f[..., None, None, :]
+    J_w_to_img2f = J_w_to_img2f.sum(-1)  # (..., N, 2)
+
+    J_norm2f = torch.einsum("...Nij,...Nj->...Ni", J_norm2w_to_img, J_w_to_img2f)  # (..., N, 3)
+    J_f = torch.einsum("...Ni,...i->...N", J_norm2f, gravity.vec3d).unsqueeze(-1)  # (..., N, 1)
+    J.append(J_f)
+
+    ######################
+    ##### K1 Jacobian ####
+    ######################
+
+    if hasattr(camera, "dist"):
+        J_w_to_img2k1 = camera.J_image2world(xy, "dist")  # (..., N, 2)
+        # (..., N, 2)
+        J_norm2k1 = torch.einsum("...Nij,...Nj->...Ni", J_norm2w_to_img, J_w_to_img2k1)
+        # (..., N, 1)
+        J_k1 = torch.einsum("...Ni,...i->...N", J_norm2k1, gravity.vec3d).unsqueeze(-1)
+        J.append(J_k1)
+
+    n_params = sum(j.shape[-1] for j in J)
+    return torch.cat(J, axis=-1).reshape(camera.shape[0], h, w, 1, n_params)
+
+
+def get_perspective_field(
+    camera: BaseCamera,
+    gravity: Gravity,
+    use_up: bool = True,
+    use_latitude: bool = True,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Get the perspective field from the camera parameters.
+
+    Args:
+        camera (Camera): Camera parameters.
+        gravity (Gravity): Gravity vector.
+        use_up (bool, optional): Whether to include the up vector field. Defaults to True.
+        use_latitude (bool, optional): Whether to include the latitude field. Defaults to True.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Up and latitude fields as tensors of shape
+        (..., 2, h, w) and (..., 1, h, w).
+    """
+    assert use_up or use_latitude, "At least one of use_up or use_latitude must be True."
+
+    camera = camera.unsqueeze(0) if len(camera.shape) == 0 else camera
+    gravity = gravity.unsqueeze(0) if len(gravity.shape) == 0 else gravity
+
+    w, h = camera.size[0].unbind(-1)
+    h, w = h.round().to(int), w.round().to(int)
+
+    if use_up:
+        permute = (0, 3, 1, 2)
+        # (..., 2, h, w)
+        up = get_up_field(camera, gravity).permute(permute)
+    else:
+        shape = (camera.shape[0], 2, h, w)
+        up = camera.new_zeros(shape)
+
+    if use_latitude:
+        permute = (0, 3, 1, 2)
+        # (..., 1, h, w)
+        lat = get_latitude_field(camera, gravity).permute(permute)
+    else:
+        shape = (camera.shape[0], 1, h, w)
+        lat = camera.new_zeros(shape)
+
+    return up, lat
+
+
+def J_perspective_field(
+    camera: BaseCamera,
+    gravity: Gravity,
+    use_up: bool = True,
+    use_latitude: bool = True,
+    spherical: bool = False,
+    log_focal: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Get the jacobian of the perspective field.
+
+    Args:
+        camera (Camera): Camera parameters.
+        gravity (Gravity): Gravity vector.
+        use_up (bool, optional): Whether to include the up vector field. Defaults to True.
+        use_latitude (bool, optional): Whether to include the latitude field. Defaults to True.
+        spherical (bool, optional): Whether to use spherical coordinates. Defaults to False.
+        log_focal (bool, optional): Whether to use log-focal length. Defaults to False.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Up and latitude jacobians as tensors of shape
+        (..., h, w, 2, 4) and (..., h, w, 1, 4).
+    """
+    assert use_up or use_latitude, "At least one of use_up or use_latitude must be True."
+
+    camera = camera.unsqueeze(0) if len(camera.shape) == 0 else camera
+    gravity = gravity.unsqueeze(0) if len(gravity.shape) == 0 else gravity
+
+    w, h = camera.size[0].unbind(-1)
+    h, w = h.round().to(int), w.round().to(int)
+
+    if use_up:
+        J_up = J_up_field(camera, gravity, spherical, log_focal)  # (..., h, w, 2, 4)
+    else:
+        shape = (camera.shape[0], h, w, 2, 4)
+        J_up = camera.new_zeros(shape)
+
+    if use_latitude:
+        J_lat = J_latitude_field(camera, gravity, spherical, log_focal)  # (..., h, w, 1, 4)
+    else:
+        shape = (camera.shape[0], h, w, 1, 4)
+        J_lat = camera.new_zeros(shape)
+
+    return J_up, J_lat

geocalib/utils.py

@@ -0,0 +1,325 @@
+"""Image loading and general conversion utilities."""
+
+import collections.abc as collections
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Dict, Optional, Tuple
+
+import cv2
+import kornia
+import numpy as np
+import torch
+import torchvision
+
+# mypy: ignore-errors
+
+
+def fit_to_multiple(x: torch.Tensor, multiple: int, mode: str = "center", crop: bool = False):
+    """Get padding to make the image size a multiple of the given number.
+
+    Args:
+        x (torch.Tensor): Input tensor.
+        multiple (int, optional): Multiple to fit to.
+        crop (bool, optional): Whether to crop or pad. Defaults to False.
+
+    Returns:
+        torch.Tensor: Padding.
+    """
+    h, w = x.shape[-2:]
+
+    if crop:
+        pad_w = (w // multiple) * multiple - w
+        pad_h = (h // multiple) * multiple - h
+    else:
+        pad_w = (multiple - w % multiple) % multiple
+        pad_h = (multiple - h % multiple) % multiple
+
+    if mode == "center":
+        pad_l = pad_w // 2
+        pad_r = pad_w - pad_l
+        pad_t = pad_h // 2
+        pad_b = pad_h - pad_t
+    elif mode == "left":
+        pad_l, pad_r = 0, pad_w
+        pad_t, pad_b = 0, pad_h
+    else:
+        raise ValueError(f"Unknown mode {mode}")
+
+    return (pad_l, pad_r, pad_t, pad_b)
+
+
+def fit_features_to_multiple(
+    features: torch.Tensor, multiple: int = 32, crop: bool = False
+) -> Tuple[torch.Tensor, Tuple[int, int]]:
+    """Pad or crop image to a multiple of the given number.
+
+    Args:
+        features (torch.Tensor): Input features.
+        multiple (int, optional): Multiple. Defaults to 32.
+        crop (bool, optional): Whether to crop or pad. Defaults to False.
+
+    Returns:
+        Tuple[torch.Tensor, Tuple[int, int]]: Padded features and padding.
+    """
+    pad = fit_to_multiple(features, multiple, crop=crop)
+    return torch.nn.functional.pad(features, pad, mode="reflect"), pad
+
+
+class ImagePreprocessor:
+    """Preprocess images for calibration."""
+
+    default_conf = {
+        "resize": 320,  # target edge length, None for no resizing
+        "edge_divisible_by": None,
+        "side": "short",
+        "interpolation": "bilinear",
+        "align_corners": None,
+        "antialias": True,
+        "square_crop": False,
+        "add_padding_mask": False,
+        "resize_backend": "kornia",  # torchvision, kornia
+    }
+
+    def __init__(self, conf) -> None:
+        """Initialize the image preprocessor."""
+        self.conf = {**self.default_conf, **conf}
+        self.conf = SimpleNamespace(**self.conf)
+
+    def __call__(self, img: torch.Tensor, interpolation: Optional[str] = None) -> dict:
+        """Resize and preprocess an image, return image and resize scale."""
+        h, w = img.shape[-2:]
+        size = h, w
+
+        if self.conf.square_crop:
+            min_size = min(h, w)
+            offset = (h - min_size) // 2, (w - min_size) // 2
+            img = img[:, offset[0] : offset[0] + min_size, offset[1] : offset[1] + min_size]
+            size = img.shape[-2:]
+
+        if self.conf.resize is not None:
+            if interpolation is None:
+                interpolation = self.conf.interpolation
+            size = self.get_new_image_size(h, w)
+            img = self.resize(img, size, interpolation)
+
+        scale = torch.Tensor([img.shape[-1] / w, img.shape[-2] / h]).to(img)
+        T = np.diag([scale[0].cpu(), scale[1].cpu(), 1])
+
+        data = {
+            "scales": scale,
+            "image_size": np.array(size[::-1]),
+            "transform": T,
+            "original_image_size": np.array([w, h]),
+        }
+
+        if self.conf.edge_divisible_by is not None:
+            # crop to make the edge divisible by a number
+            w_, h_ = img.shape[-1], img.shape[-2]
+            img, _ = fit_features_to_multiple(img, self.conf.edge_divisible_by, crop=True)
+            crop_pad = torch.Tensor([img.shape[-1] - w_, img.shape[-2] - h_]).to(img)
+            data["crop_pad"] = crop_pad
+            data["image_size"] = np.array([img.shape[-1], img.shape[-2]])
+
+        data["image"] = img
+        return data
+
+    def resize(self, img: torch.Tensor, size: Tuple[int, int], interpolation: str) -> torch.Tensor:
+        """Resize an image using the specified backend."""
+        if self.conf.resize_backend == "kornia":
+            return kornia.geometry.transform.resize(
+                img,
+                size,
+                side=self.conf.side,
+                antialias=self.conf.antialias,
+                align_corners=self.conf.align_corners,
+                interpolation=interpolation,
+            )
+        elif self.conf.resize_backend == "torchvision":
+            return torchvision.transforms.Resize(size, antialias=self.conf.antialias)(img)
+        else:
+            raise ValueError(f"{self.conf.resize_backend} not implemented.")
+
+    def load_image(self, image_path: Path) -> dict:
+        """Load an image from a path and preprocess it."""
+        return self(load_image(image_path))
+
+    def get_new_image_size(self, h: int, w: int) -> Tuple[int, int]:
+        """Get the new image size after resizing."""
+        side = self.conf.side
+        if isinstance(self.conf.resize, collections.Iterable):
+            assert len(self.conf.resize) == 2
+            return tuple(self.conf.resize)
+        side_size = self.conf.resize
+        aspect_ratio = w / h
+        if side not in ("short", "long", "vert", "horz"):
+            raise ValueError(
+                f"side can be one of 'short', 'long', 'vert', and 'horz'. Got '{side}'"
+            )
+        return (
+            (side_size, int(side_size * aspect_ratio))
+            if side == "vert" or (side != "horz" and (side == "short") ^ (aspect_ratio < 1.0))
+            else (int(side_size / aspect_ratio), side_size)
+        )
+
+
+def numpy_image_to_torch(image: np.ndarray) -> torch.Tensor:
+    """Normalize the image tensor and reorder the dimensions."""
+    if image.ndim == 3:
+        image = image.transpose((2, 0, 1))  # HxWxC to CxHxW
+    elif image.ndim == 2:
+        image = image[None]  # add channel axis
+    else:
+        raise ValueError(f"Not an image: {image.shape}")
+    return torch.tensor(image / 255.0, dtype=torch.float)
+
+
+def torch_image_to_numpy(image: torch.Tensor) -> np.ndarray:
+    """Normalize and reorder the dimensions of an image tensor."""
+    if image.ndim == 3:
+        image = image.permute((1, 2, 0))  # CxHxW to HxWxC
+    elif image.ndim == 2:
+        image = image[None]  # add channel axis
+    else:
+        raise ValueError(f"Not an image: {image.shape}")
+    return (image.cpu().detach().numpy() * 255).astype(np.uint8)
+
+
+def read_image(path: Path, grayscale: bool = False) -> np.ndarray:
+    """Read an image from path as RGB or grayscale."""
+    if not Path(path).exists():
+        raise FileNotFoundError(f"No image at path {path}.")
+    mode = cv2.IMREAD_GRAYSCALE if grayscale else cv2.IMREAD_COLOR
+    image = cv2.imread(str(path), mode)
+    if image is None:
+        raise IOError(f"Could not read image at {path}.")
+    if not grayscale:
+        image = image[..., ::-1]
+    return image
+
+
+def write_image(img: torch.Tensor, path: Path):
+    """Write an image tensor to a file."""
+    img = torch_image_to_numpy(img) if isinstance(img, torch.Tensor) else img
+    cv2.imwrite(str(path), img[..., ::-1])
+
+
+def load_image(path: Path, grayscale: bool = False, return_tensor: bool = True) -> torch.Tensor:
+    """Load an image from a path and return as a tensor."""
+    image = read_image(path, grayscale=grayscale)
+    if return_tensor:
+        return numpy_image_to_torch(image)
+
+    assert image.ndim in [2, 3], f"Not an image: {image.shape}"
+    image = image[None] if image.ndim == 2 else image
+    return torch.tensor(image.copy(), dtype=torch.uint8)
+
+
+def skew_symmetric(v: torch.Tensor) -> torch.Tensor:
+    """Create a skew-symmetric matrix from a (batched) vector of size (..., 3).
+
+    Args:
+        (torch.Tensor): Vector of size (..., 3).
+
+    Returns:
+        (torch.Tensor): Skew-symmetric matrix of size (..., 3, 3).
+    """
+    z = torch.zeros_like(v[..., 0])
+    return torch.stack(
+        [z, -v[..., 2], v[..., 1], v[..., 2], z, -v[..., 0], -v[..., 1], v[..., 0], z], dim=-1
+    ).reshape(v.shape[:-1] + (3, 3))
+
+
+def rad2rotmat(
+    roll: torch.Tensor, pitch: torch.Tensor, yaw: Optional[torch.Tensor] = None
+) -> torch.Tensor:
+    """Convert (batched) roll, pitch, yaw angles (in radians) to rotation matrix.
+
+    Args:
+        roll (torch.Tensor): Roll angle in radians.
+        pitch (torch.Tensor): Pitch angle in radians.
+        yaw (torch.Tensor, optional): Yaw angle in radians. Defaults to None.
+
+    Returns:
+        torch.Tensor: Rotation matrix of shape (..., 3, 3).
+    """
+    if yaw is None:
+        yaw = roll.new_zeros(roll.shape)
+
+    Rx = pitch.new_zeros(pitch.shape + (3, 3))
+    Rx[..., 0, 0] = 1
+    Rx[..., 1, 1] = torch.cos(pitch)
+    Rx[..., 1, 2] = torch.sin(pitch)
+    Rx[..., 2, 1] = -torch.sin(pitch)
+    Rx[..., 2, 2] = torch.cos(pitch)
+
+    Ry = yaw.new_zeros(yaw.shape + (3, 3))
+    Ry[..., 0, 0] = torch.cos(yaw)
+    Ry[..., 0, 2] = -torch.sin(yaw)
+    Ry[..., 1, 1] = 1
+    Ry[..., 2, 0] = torch.sin(yaw)
+    Ry[..., 2, 2] = torch.cos(yaw)
+
+    Rz = roll.new_zeros(roll.shape + (3, 3))
+    Rz[..., 0, 0] = torch.cos(roll)
+    Rz[..., 0, 1] = torch.sin(roll)
+    Rz[..., 1, 0] = -torch.sin(roll)
+    Rz[..., 1, 1] = torch.cos(roll)
+    Rz[..., 2, 2] = 1
+
+    return Rz @ Rx @ Ry
+
+
+def fov2focal(fov: torch.Tensor, size: torch.Tensor) -> torch.Tensor:
+    """Compute focal length from (vertical/horizontal) field of view."""
+    return size / 2 / torch.tan(fov / 2)
+
+
+def focal2fov(focal: torch.Tensor, size: torch.Tensor) -> torch.Tensor:
+    """Compute (vertical/horizontal) field of view from focal length."""
+    return 2 * torch.arctan(size / (2 * focal))
+
+
+def pitch2rho(pitch: torch.Tensor, f: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
+    """Compute the distance from principal point to the horizon."""
+    return torch.tan(pitch) * f / h
+
+
+def rho2pitch(rho: torch.Tensor, f: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
+    """Compute the pitch angle from the distance to the horizon."""
+    return torch.atan(rho * h / f)
+
+
+def rad2deg(rad: torch.Tensor) -> torch.Tensor:
+    """Convert radians to degrees."""
+    return rad / torch.pi * 180
+
+
+def deg2rad(deg: torch.Tensor) -> torch.Tensor:
+    """Convert degrees to radians."""
+    return deg / 180 * torch.pi
+
+
+def get_device() -> str:
+    """Get the device (cpu, cuda, mps) available."""
+    device = "cpu"
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif torch.backends.mps.is_available():
+        device = "mps"
+    return device
+
+
+def print_calibration(results: Dict[str, torch.Tensor]) -> None:
+    """Print the calibration results."""
+    camera, gravity = results["camera"], results["gravity"]
+    vfov = rad2deg(camera.vfov)
+    roll, pitch = rad2deg(gravity.rp).unbind(-1)
+
+    print("\nEstimated parameters (Pred):")
+    print(f"Roll:  {roll.item():.1f}° (± {rad2deg(results['roll_uncertainty']).item():.1f})°")
+    print(f"Pitch: {pitch.item():.1f}° (± {rad2deg(results['pitch_uncertainty']).item():.1f})°")
+    print(f"vFoV:  {vfov.item():.1f}° (± {rad2deg(results['vfov_uncertainty']).item():.1f})°")
+    print(f"Focal: {camera.f[0, 1].item():.1f} px (± {results['focal_uncertainty'].item():.1f} px)")
+
+    if hasattr(camera, "k1"):
+        print(f"K1:    {camera.k1.item():.1f}")

geocalib/viz2d.py

@@ -0,0 +1,502 @@
+"""2D visualization primitives based on Matplotlib.
+
+1) Plot images with `plot_images`.
+2) Call functions to plot heatmaps, vector fields, and horizon lines.
+3) Optionally: save a .png or .pdf plot (nice in papers!) with `save_plot`.
+"""
+
+import matplotlib.patheffects as path_effects
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+from geocalib.perspective_fields import get_perspective_field
+from geocalib.utils import rad2deg
+
+# mypy: ignore-errors
+
+
+def plot_images(imgs, titles=None, cmaps="gray", dpi=200, pad=0.5, adaptive=True):
+    """Plot a list of images.
+
+    Args:
+        imgs (List[np.ndarray]): List of images to plot.
+        titles (List[str], optional): Titles. Defaults to None.
+        cmaps (str, optional): Colormaps. Defaults to "gray".
+        dpi (int, optional): Dots per inch. Defaults to 200.
+        pad (float, optional): Padding. Defaults to 0.5.
+        adaptive (bool, optional): Whether to adapt the aspect ratio. Defaults to True.
+
+    Returns:
+        plt.Figure: Figure of the images.
+    """
+    n = len(imgs)
+    if not isinstance(cmaps, (list, tuple)):
+        cmaps = [cmaps] * n
+
+    ratios = [i.shape[1] / i.shape[0] for i in imgs] if adaptive else [4 / 3] * n
+    figsize = [sum(ratios) * 4.5, 4.5]
+    fig, axs = plt.subplots(1, n, figsize=figsize, dpi=dpi, gridspec_kw={"width_ratios": ratios})
+    if n == 1:
+        axs = [axs]
+    for i, (img, ax) in enumerate(zip(imgs, axs)):
+        ax.imshow(img, cmap=plt.get_cmap(cmaps[i]))
+        ax.set_axis_off()
+        if titles:
+            ax.set_title(titles[i])
+    fig.tight_layout(pad=pad)
+
+    return fig
+
+
+def plot_image_grid(
+    imgs,
+    titles=None,
+    cmaps="gray",
+    dpi=100,
+    pad=0.5,
+    fig=None,
+    adaptive=True,
+    figs=3.0,
+    return_fig=False,
+    set_lim=False,
+) -> plt.Figure:
+    """Plot a grid of images.
+
+    Args:
+        imgs (List[np.ndarray]): List of images to plot.
+        titles (List[str], optional): Titles. Defaults to None.
+        cmaps (str, optional): Colormaps. Defaults to "gray".
+        dpi (int, optional): Dots per inch. Defaults to 100.
+        pad (float, optional): Padding. Defaults to 0.5.
+        fig (_type_, optional): Figure to plot on. Defaults to None.
+        adaptive (bool, optional): Whether to adapt the aspect ratio. Defaults to True.
+        figs (float, optional): Figure size. Defaults to 3.0.
+        return_fig (bool, optional): Whether to return the figure. Defaults to False.
+        set_lim (bool, optional): Whether to set the limits. Defaults to False.
+
+    Returns:
+        plt.Figure: Figure and axes or just axes.
+    """
+    nr, n = len(imgs), len(imgs[0])
+    if not isinstance(cmaps, (list, tuple)):
+        cmaps = [cmaps] * n
+
+    if adaptive:
+        ratios = [i.shape[1] / i.shape[0] for i in imgs[0]]  # W / H
+    else:
+        ratios = [4 / 3] * n
+
+    figsize = [sum(ratios) * figs, nr * figs]
+    if fig is None:
+        fig, axs = plt.subplots(
+            nr, n, figsize=figsize, dpi=dpi, gridspec_kw={"width_ratios": ratios}
+        )
+    else:
+        axs = fig.subplots(nr, n, gridspec_kw={"width_ratios": ratios})
+        fig.figure.set_size_inches(figsize)
+
+    if nr == 1 and n == 1:
+        axs = [[axs]]
+    elif n == 1:
+        axs = axs[:, None]
+    elif nr == 1:
+        axs = [axs]
+
+    for j in range(nr):
+        for i in range(n):
+            ax = axs[j][i]
+            ax.imshow(imgs[j][i], cmap=plt.get_cmap(cmaps[i]))
+            ax.set_axis_off()
+            if set_lim:
+                ax.set_xlim([0, imgs[j][i].shape[1]])
+                ax.set_ylim([imgs[j][i].shape[0], 0])
+            if titles:
+                ax.set_title(titles[j][i])
+    if isinstance(fig, plt.Figure):
+        fig.tight_layout(pad=pad)
+    return (fig, axs) if return_fig else axs
+
+
+def add_text(
+    idx,
+    text,
+    pos=(0.01, 0.99),
+    fs=15,
+    color="w",
+    lcolor="k",
+    lwidth=4,
+    ha="left",
+    va="top",
+    axes=None,
+    **kwargs,
+):
+    """Add text to a plot.
+
+    Args:
+        idx (int): Index of the axes.
+        text (str): Text to add.
+        pos (tuple, optional): Text position. Defaults to (0.01, 0.99).
+        fs (int, optional): Font size. Defaults to 15.
+        color (str, optional): Text color. Defaults to "w".
+        lcolor (str, optional): Line color. Defaults to "k".
+        lwidth (int, optional): Line width. Defaults to 4.
+        ha (str, optional): Horizontal alignment. Defaults to "left".
+        va (str, optional): Vertical alignment. Defaults to "top".
+        axes (List[plt.Axes], optional): Axes to put text on. Defaults to None.
+
+    Returns:
+        plt.Text: Text object.
+    """
+    if axes is None:
+        axes = plt.gcf().axes
+
+    ax = axes[idx]
+
+    t = ax.text(
+        *pos,
+        text,
+        fontsize=fs,
+        ha=ha,
+        va=va,
+        color=color,
+        transform=ax.transAxes,
+        zorder=5,
+        **kwargs,
+    )
+    if lcolor is not None:
+        t.set_path_effects(
+            [
+                path_effects.Stroke(linewidth=lwidth, foreground=lcolor),
+                path_effects.Normal(),
+            ]
+        )
+    return t
+
+
+def plot_heatmaps(
+    heatmaps,
+    vmin=-1e-6,  # include negative zero
+    vmax=None,
+    cmap="Spectral",
+    a=0.5,
+    axes=None,
+    contours_every=None,
+    contour_style="solid",
+    colorbar=False,
+):
+    """Plot heatmaps with optional contours.
+
+    To plot latitude field, set vmin=-90, vmax=90 and contours_every=15.
+
+    Args:
+        heatmaps (List[np.ndarray | torch.Tensor]): List of 2D heatmaps.
+        vmin (float, optional): Min Value. Defaults to -1e-6.
+        vmax (float, optional): Max Value. Defaults to None.
+        cmap (str, optional): Colormap. Defaults to "Spectral".
+        a (float, optional): Alpha value. Defaults to 0.5.
+        axes (List[plt.Axes], optional): Axes to plot on. Defaults to None.
+        contours_every (int, optional): If not none, will draw contours. Defaults to None.
+        contour_style (str, optional): Style of the contours. Defaults to "solid".
+        colorbar (bool, optional): Whether to show colorbar. Defaults to False.
+
+    Returns:
+        List[plt.Artist]: List of artists.
+    """
+    if axes is None:
+        axes = plt.gcf().axes
+    artists = []
+
+    for i in range(len(axes)):
+        a_ = a if isinstance(a, float) else a[i]
+
+        if isinstance(heatmaps[i], torch.Tensor):
+            heatmaps[i] = heatmaps[i].cpu().numpy()
+
+        alpha = a_
+        # Plot the heatmap
+        art = axes[i].imshow(
+            heatmaps[i],
+            alpha=alpha,
+            vmin=vmin,
+            vmax=vmax,
+            cmap=cmap,
+        )
+        if colorbar:
+            cmax = vmax or np.percentile(heatmaps[i], 99)
+            art.set_clim(vmin, cmax)
+            cbar = plt.colorbar(art, ax=axes[i])
+            artists.append(cbar)
+
+        artists.append(art)
+
+        if contours_every is not None:
+            # Add contour lines to the heatmap
+            contour_data = np.arange(vmin, vmax + contours_every, contours_every)
+
+            # Get the colormap colors for contour lines
+            contour_colors = [
+                plt.colormaps.get_cmap(cmap)(plt.Normalize(vmin=vmin, vmax=vmax)(level))
+                for level in contour_data
+            ]
+            contours = axes[i].contour(
+                heatmaps[i],
+                levels=contour_data,
+                linewidths=2,
+                colors=contour_colors,
+                linestyles=contour_style,
+            )
+
+            contours.set_clim(vmin, vmax)
+
+            fmt = {
+                level: f"{label}°"
+                for level, label in zip(contour_data, contour_data.astype(int).astype(str))
+            }
+            t = axes[i].clabel(contours, inline=True, fmt=fmt, fontsize=16, colors="white")
+
+            for label in t:
+                label.set_path_effects(
+                    [
+                        path_effects.Stroke(linewidth=1, foreground="k"),
+                        path_effects.Normal(),
+                    ]
+                )
+            artists.append(contours)
+
+    return artists
+
+
+def plot_horizon_lines(
+    cameras, gravities, line_colors="orange", lw=2, styles="solid", alpha=1.0, ax=None
+):
+    """Plot horizon lines on the perspective field.
+
+    Args:
+        cameras (List[Camera]): List of cameras.
+        gravities (List[Gravity]): Gravities.
+        line_colors (str, optional): Line Colors. Defaults to "orange".
+        lw (int, optional): Line width. Defaults to 2.
+        styles (str, optional): Line styles. Defaults to "solid".
+        alpha (float, optional): Alphas. Defaults to 1.0.
+        ax (List[plt.Axes], optional): Axes to draw horizon line on. Defaults to None.
+    """
+    if not isinstance(line_colors, list):
+        line_colors = [line_colors] * len(cameras)
+
+    if not isinstance(styles, list):
+        styles = [styles] * len(cameras)
+
+    fig = plt.gcf()
+    ax = fig.gca() if ax is None else ax
+
+    if isinstance(ax, plt.Axes):
+        ax = [ax] * len(cameras)
+
+    assert len(ax) == len(cameras), f"{len(ax)}, {len(cameras)}"
+
+    for i in range(len(cameras)):
+        _, lat = get_perspective_field(cameras[i], gravities[i])
+        # horizon line is zero level of the latitude field
+        lat = lat[0, 0].cpu().numpy()
+        contours = ax[i].contour(lat, levels=[0], linewidths=lw, colors=line_colors[i])
+        for contour_line in contours.collections:
+            contour_line.set_linestyle(styles[i])
+
+
+def plot_vector_fields(
+    vector_fields,
+    cmap="lime",
+    subsample=15,
+    scale=None,
+    lw=None,
+    alphas=0.8,
+    axes=None,
+):
+    """Plot vector fields.
+
+    Args:
+        vector_fields (List[torch.Tensor]): List of vector fields of shape (2, H, W).
+        cmap (str, optional): Color of the vectors. Defaults to "lime".
+        subsample (int, optional): Subsample the vector field. Defaults to 15.
+        scale (float, optional): Scale of the vectors. Defaults to None.
+        lw (float, optional): Line width of the vectors. Defaults to None.
+        alphas (float | np.ndarray, optional): Alpha per vector or global. Defaults to 0.8.
+        axes (List[plt.Axes], optional): List of axes to draw on. Defaults to None.
+
+    Returns:
+        List[plt.Artist]: List of artists.
+    """
+    if axes is None:
+        axes = plt.gcf().axes
+
+    vector_fields = [v.cpu().numpy() if isinstance(v, torch.Tensor) else v for v in vector_fields]
+
+    artists = []
+
+    H, W = vector_fields[0].shape[-2:]
+    if scale is None:
+        scale = subsample / min(H, W)
+
+    if lw is None:
+        lw = 0.1 / subsample
+
+    if alphas is None:
+        alphas = np.ones_like(vector_fields[0][0])
+        alphas = np.stack([alphas] * len(vector_fields), 0)
+    elif isinstance(alphas, float):
+        alphas = np.ones_like(vector_fields[0][0]) * alphas
+        alphas = np.stack([alphas] * len(vector_fields), 0)
+    else:
+        alphas = np.array(alphas)
+
+    subsample = min(W, H) // subsample
+    offset_x = ((W % subsample) + subsample) // 2
+
+    samples_x = np.arange(offset_x, W, subsample)
+    samples_y = np.arange(int(subsample * 0.9), H, subsample)
+
+    x_grid, y_grid = np.meshgrid(samples_x, samples_y)
+
+    for i in range(len(axes)):
+        # vector field of shape (2, H, W) with vectors of norm == 1
+        vector_field = vector_fields[i]
+
+        a = alphas[i][samples_y][:, samples_x]
+        x, y = vector_field[:, samples_y][:, :, samples_x]
+
+        c = cmap
+        if not isinstance(cmap, str):
+            c = cmap[i][samples_y][:, samples_x].reshape(-1, 3)
+
+        s = scale * min(H, W)
+        arrows = axes[i].quiver(
+            x_grid,
+            y_grid,
+            x,
+            y,
+            scale=s,
+            scale_units="width" if H > W else "height",
+            units="width" if H > W else "height",
+            alpha=a,
+            color=c,
+            angles="xy",
+            antialiased=True,
+            width=lw,
+            headaxislength=3.5,
+            zorder=5,
+        )
+
+        artists.append(arrows)
+
+    return artists
+
+
+def plot_latitudes(
+    latitude,
+    is_radians=True,
+    vmin=-90,
+    vmax=90,
+    cmap="seismic",
+    contours_every=15,
+    alpha=0.4,
+    axes=None,
+    **kwargs,
+):
+    """Plot latitudes.
+
+    Args:
+        latitude (List[torch.Tensor]): List of latitudes.
+        is_radians (bool, optional): Whether the latitudes are in radians. Defaults to True.
+        vmin (int, optional): Min value to clip to. Defaults to -90.
+        vmax (int, optional): Max value to clip to. Defaults to 90.
+        cmap (str, optional): Colormap. Defaults to "seismic".
+        contours_every (int, optional): Contours every. Defaults to 15.
+        alpha (float, optional): Alpha value. Defaults to 0.4.
+        axes (List[plt.Axes], optional): Axes to plot on. Defaults to None.
+
+    Returns:
+        List[plt.Artist]: List of artists.
+    """
+    if axes is None:
+        axes = plt.gcf().axes
+
+    assert len(axes) == len(latitude), f"{len(axes)}, {len(latitude)}"
+    lat = [rad2deg(lat) for lat in latitude] if is_radians else latitude
+    return plot_heatmaps(
+        lat,
+        vmin=vmin,
+        vmax=vmax,
+        cmap=cmap,
+        a=alpha,
+        axes=axes,
+        contours_every=contours_every,
+        **kwargs,
+    )
+
+
+def plot_perspective_fields(cameras, gravities, axes=None, **kwargs):
+    """Plot perspective fields.
+
+    Args:
+        cameras (List[Camera]): List of cameras.
+        gravities (List[Gravity]): List of gravities.
+        axes (List[plt.Axes], optional): Axes to plot on. Defaults to None.
+
+    Returns:
+        List[plt.Artist]: List of artists.
+    """
+    if axes is None:
+        axes = plt.gcf().axes
+
+    assert len(axes) == len(cameras), f"{len(axes)}, {len(cameras)}"
+
+    artists = []
+    for i in range(len(axes)):
+        up, lat = get_perspective_field(cameras[i], gravities[i])
+        artists += plot_vector_fields([up[0]], axes=[axes[i]], **kwargs)
+        artists += plot_latitudes([lat[0, 0]], axes=[axes[i]], **kwargs)
+
+    return artists
+
+
+def plot_confidences(
+    confidence,
+    as_log=True,
+    vmin=-4,
+    vmax=0,
+    cmap="turbo",
+    alpha=0.4,
+    axes=None,
+    **kwargs,
+):
+    """Plot confidences.
+
+    Args:
+        confidence (List[torch.Tensor]): Confidence maps.
+        as_log (bool, optional): Whether to plot in log scale. Defaults to True.
+        vmin (int, optional): Min value to clip to. Defaults to -4.
+        vmax (int, optional): Max value to clip to. Defaults to 0.
+        cmap (str, optional): Colormap. Defaults to "turbo".
+        alpha (float, optional): Alpha value. Defaults to 0.4.
+        axes (List[plt.Axes], optional): Axes to plot on. Defaults to None.
+
+    Returns:
+        List[plt.Artist]: List of artists.
+    """
+    if axes is None:
+        axes = plt.gcf().axes
+
+    assert len(axes) == len(confidence), f"{len(axes)}, {len(confidence)}"
+
+    if as_log:
+        confidence = [torch.log10(c.clip(1e-5)).clip(vmin, vmax) for c in confidence]
+
+    # normalize to [0, 1]
+    confidence = [(c - c.min()) / (c.max() - c.min()) for c in confidence]
+    return plot_heatmaps(confidence, vmin=0, vmax=1, cmap=cmap, a=alpha, axes=axes, **kwargs)
+
+
+def save_plot(path, **kw):
+    """Save the current figure without any white margin."""
+    plt.savefig(path, bbox_inches="tight", pad_inches=0, **kw)

gradio_app.py

@@ -0,0 +1,228 @@
+"""Gradio app for GeoCalib inference."""
+
+from copy import deepcopy
+from time import time
+
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+
+from geocalib import viz2d
+from geocalib.camera import camera_models
+from geocalib.extractor import GeoCalib
+from geocalib.perspective_fields import get_perspective_field
+from geocalib.utils import rad2deg
+
+# flake8: noqa
+# mypy: ignore-errors
+
+description = """
+<p align="center">
+  <h1 align="center"><ins>GeoCalib</ins> 📸<br>Single-image Calibration with Geometric Optimization</h1>
+  <p align="center">
+    <a href="https://www.linkedin.com/in/alexander-veicht/">Alexander Veicht</a>
+    ·
+    <a href="https://psarlin.com/">Paul-Edouard&nbsp;Sarlin</a>
+    ·
+    <a href="https://www.linkedin.com/in/philipplindenberger/">Philipp Lindenberger</a>
+    ·
+    <a href="https://www.microsoft.com/en-us/research/people/mapoll/">Marc&nbsp;Pollefeys</a>
+  </p>
+  <h2 align="center">
+    <p>ECCV 2024</p>
+    <a href="" align="center">Paper</a> | <!--TODO: update link-->
+    <a href="https://github.com/cvg/GeoCalib" align="center">Code</a> |
+    <a href="https://colab.research.google.com/drive/1oMzgPGppAPAIQxe-s7SRd_q8r7dVfnqo#scrollTo=etdzQZQzoo-K" align="center">Colab</a>
+  </h2>
+</p>
+
+## Getting Started
+GeoCalib accurately estimates the camera intrinsics and gravity direction from a single image by 
+combining geometric optimization with deep learning.
+
+To get started, upload an image or select one of the examples below.
+You can choose between different camera models and visualize the calibration results.
+
+"""
+
+example_images = [
+    ["assets/pinhole-church.jpg"],
+    ["assets/pinhole-garden.jpg"],
+    ["assets/fisheye-skyline.jpg"],
+    ["assets/fisheye-dog-pool.jpg"],
+]
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = GeoCalib().to(device)
+
+
+def format_output(results):
+    camera, gravity = results["camera"], results["gravity"]
+    vfov = rad2deg(camera.vfov)
+    roll, pitch = rad2deg(gravity.rp).unbind(-1)
+
+    txt = "Estimated parameters:\n"
+    txt += f"Roll:  {roll.item():.2f}° (± {rad2deg(results['roll_uncertainty']).item():.2f})°\n"
+    txt += f"Pitch: {pitch.item():.2f}° (± {rad2deg(results['pitch_uncertainty']).item():.2f})°\n"
+    txt += f"vFoV:  {vfov.item():.2f}° (± {rad2deg(results['vfov_uncertainty']).item():.2f})°\n"
+    txt += (
+        f"Focal: {camera.f[0, 1].item():.2f} px (± {results['focal_uncertainty'].item():.2f} px)\n"
+    )
+    if hasattr(camera, "k1"):
+        txt += f"K1:    {camera.k1[0].item():.2f}\n"
+    return txt
+
+
+@spaces.GPU(duration=10)
+def inference(img, camera_model):
+    out = model.calibrate(img.to(device), camera_model=camera_model)
+    save_keys = ["camera", "gravity"] + [f"{k}_uncertainty" for k in ["roll", "pitch", "vfov"]]
+    res = {k: v.cpu() for k, v in out.items() if k in save_keys}
+    # not converting to numpy results in gpu abort
+    res["up_confidence"] = out["up_confidence"].cpu().numpy()
+    res["latitude_confidence"] = out["latitude_confidence"].cpu().numpy()
+    return res
+
+
+def process_results(
+    image_path,
+    camera_model,
+    plot_up,
+    plot_up_confidence,
+    plot_latitude,
+    plot_latitude_confidence,
+    plot_undistort,
+):
+    """Process the image and return the calibration results."""
+
+    if image_path is None:
+        raise gr.Error("Please upload an image first.")
+
+    img = model.load_image(image_path)
+    print("Running inference...")
+    start = time()
+    inference_result = inference(img, camera_model)
+    print(f"Done ({time() - start:.2f}s)")
+    inference_result["image"] = img.cpu()
+
+    if inference_result is None:
+        return ("", np.ones((128, 256, 3)), None)
+
+    plot_img = update_plot(
+        inference_result,
+        plot_up,
+        plot_up_confidence,
+        plot_latitude,
+        plot_latitude_confidence,
+        plot_undistort,
+    )
+
+    return format_output(inference_result), plot_img, inference_result
+
+
+def update_plot(
+    inference_result,
+    plot_up,
+    plot_up_confidence,
+    plot_latitude,
+    plot_latitude_confidence,
+    plot_undistort,
+):
+    """Update the plot based on the selected options."""
+    if inference_result is None:
+        gr.Error("Please calibrate an image first.")
+        return np.ones((128, 256, 3))
+
+    camera, gravity = inference_result["camera"], inference_result["gravity"]
+    img = inference_result["image"].permute(1, 2, 0).numpy()
+
+    if plot_undistort:
+        if not hasattr(camera, "k1"):
+            return img
+
+        return camera.undistort_image(inference_result["image"][None])[0].permute(1, 2, 0).numpy()
+
+    up, lat = get_perspective_field(camera, gravity)
+
+    fig = viz2d.plot_images([img], pad=0)
+    ax = fig.get_axes()
+
+    if plot_up:
+        viz2d.plot_vector_fields([up[0]], axes=[ax[0]])
+
+    if plot_latitude:
+        viz2d.plot_latitudes([lat[0, 0]], axes=[ax[0]])
+
+    if plot_up_confidence:
+        viz2d.plot_confidences([inference_result["up_confidence"][0]], axes=[ax[0]])
+
+    if plot_latitude_confidence:
+        viz2d.plot_confidences([inference_result["latitude_confidence"][0]], axes=[ax[0]])
+
+    fig.canvas.draw()
+    img = np.array(fig.canvas.renderer.buffer_rgba())
+
+    return img
+
+
+# Create the Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown(description)
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("""## Input Image""")
+            image_path = gr.Image(label="Upload image to calibrate", type="filepath")
+            choice_input = gr.Dropdown(
+                choices=list(camera_models.keys()), label="Choose a camera model.", value="pinhole"
+            )
+            submit_btn = gr.Button("Calibrate 📸")
+            gr.Examples(examples=example_images, inputs=[image_path, choice_input])
+
+        with gr.Column():
+            gr.Markdown("""## Results""")
+            image_output = gr.Image(label="Calibration Results")
+            gr.Markdown("### Plot Options")
+            plot_undistort = gr.Checkbox(
+                label="undistort",
+                value=False,
+                info="Undistorted image "
+                + "(this is only available for models with distortion "
+                + "parameters and will overwrite other options).",
+            )
+
+            with gr.Row():
+                plot_up = gr.Checkbox(label="up-vectors", value=True)
+                plot_up_confidence = gr.Checkbox(label="up confidence", value=False)
+                plot_latitude = gr.Checkbox(label="latitude", value=True)
+                plot_latitude_confidence = gr.Checkbox(label="latitude confidence", value=False)
+
+            gr.Markdown("### Calibration Results")
+            text_output = gr.Textbox(label="Estimated parameters", type="text", lines=5)
+
+    # Define the action when the button is clicked
+    inference_state = gr.State()
+    plot_inputs = [
+        inference_state,
+        plot_up,
+        plot_up_confidence,
+        plot_latitude,
+        plot_latitude_confidence,
+        plot_undistort,
+    ]
+    submit_btn.click(
+        fn=process_results,
+        inputs=[image_path, choice_input] + plot_inputs[1:],
+        outputs=[text_output, image_output, inference_state],
+    )
+
+    # Define the action when the plot checkboxes are clicked
+    plot_up.change(fn=update_plot, inputs=plot_inputs, outputs=image_output)
+    plot_up_confidence.change(fn=update_plot, inputs=plot_inputs, outputs=image_output)
+    plot_latitude.change(fn=update_plot, inputs=plot_inputs, outputs=image_output)
+    plot_latitude_confidence.change(fn=update_plot, inputs=plot_inputs, outputs=image_output)
+    plot_undistort.change(fn=update_plot, inputs=plot_inputs, outputs=image_output)
+
+
+# Launch the app
+demo.launch()

hubconf.py

@@ -0,0 +1,14 @@
+"""Entrypoint for torch hub."""
+
+dependencies = ["torch", "torchvision", "opencv-python", "kornia", "matplotlib"]
+
+from geocalib import GeoCalib
+
+
+def model(*args, **kwargs):
+    """Pre-trained Geocalib model.
+
+    Args:
+        weights (str): trained variant, "pinhole" (default) or "distorted".
+    """
+    return GeoCalib(*args, **kwargs)

pyproject.toml

@@ -0,0 +1,49 @@
+[build-system]
+requires = ["setuptools", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "geocalib"
+version = "1.0"
+description = "GeoCalib Inference Package"
+authors = [
+    { name = "Alexander Veicht" },
+    { name = "Paul-Edouard Sarlin" },
+    { name = "Philipp Lindenberger" },
+]
+readme = "README.md"
+requires-python = ">=3.9"
+license = { file = "LICENSE" }
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: Apache Software License",
+    "Operating System :: OS Independent",
+]
+urls = { Repository = "https://github.com/cvg/GeoCalib" }
+
+dynamic = ["dependencies"]
+
+[project.optional-dependencies]
+dev = ["black==23.9.1", "flake8", "isort==5.12.0"]
+
+[tool.setuptools]
+packages = ["geocalib"]
+
+[tool.setuptools.dynamic]
+dependencies = { file = ["requirements.txt"] }
+
+
+[tool.black]
+line-length = 100
+exclude = "(venv/|docs/|third_party/)"
+
+[tool.isort]
+profile = "black"
+line_length = 100
+atomic = true
+
+[tool.flake8]
+max-line-length = 100
+docstring-convention = "google"
+ignore = ["E203", "W503", "E402"]
+exclude = [".git", "__pycache__", "venv", "docs", "third_party", "scripts"]

requirements.txt

@@ -0,0 +1,5 @@
+torch
+torchvision
+opencv-python
+kornia
+matplotlib

siclib/LICENSE

@@ -0,0 +1,190 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   Copyright 2024 ETH Zurich
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

siclib/__init__.py

@@ -0,0 +1,15 @@
+import logging
+
+formatter = logging.Formatter(
+    fmt="[%(asctime)s %(name)s %(levelname)s] %(message)s", datefmt="%m/%d/%Y %H:%M:%S"
+)
+handler = logging.StreamHandler()
+handler.setFormatter(formatter)
+handler.setLevel(logging.INFO)
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+logger.addHandler(handler)
+logger.propagate = False
+
+__module_name__ = __name__

siclib/configs/deepcalib.yaml

@@ -0,0 +1,12 @@
+defaults:
+  - data: openpano-radial
+  - train: deepcalib
+  - model: deepcalib
+  - _self_
+
+data:
+  train_batch_size: 32
+  val_batch_size: 32
+  test_batch_size: 32
+  augmentations:
+    name: "deepcalib"

siclib/configs/geocalib-radial.yaml

@@ -0,0 +1,38 @@
+defaults:
+  - data: openpano-radial
+  - train: geocalib
+  - model: geocalib
+  - _self_
+
+data:
+  # smaller batch size since lm takes more memory
+  train_batch_size: 18
+  val_batch_size: 18
+  test_batch_size: 18
+
+model:
+  optimizer:
+    camera_model: simple_radial
+
+  weights: weights/geocalib.tar
+
+train:
+  lr: 1e-5 # smaller lr since we are fine-tuning
+  num_steps: 200_000 # adapt to see same number of samples as previous training
+
+  lr_schedule:
+    type: SequentialLR
+    on_epoch: false
+    options:
+      # adapt to see same number of samples as previous training
+      milestones: [5_000]
+      schedulers:
+        - type: LinearLR
+          options:
+            start_factor: 1e-3
+            total_iters: 5_000
+        - type: MultiStepLR
+          options:
+            gamma: 0.1
+            # adapt to see same number of samples as previous training
+            milestones: [110_000, 170_000]

siclib/configs/geocalib.yaml

@@ -0,0 +1,10 @@
+defaults:
+  - data: openpano
+  - train: geocalib
+  - model: geocalib
+  - _self_
+
+data:
+  train_batch_size: 24
+  val_batch_size: 24
+  test_batch_size: 24

siclib/configs/model/deepcalib.yaml

@@ -0,0 +1,7 @@
+name: networks.deepcalib
+bounds:
+  roll: [-45, 45]
+  # rho = torch.tan(pitch) / torch.tan(vfov / 2) / 2 -> rho in [-1/0.3526, 1/0.0872]
+  rho: [-2.83607487, 2.83607487]
+  vfov: [20, 105]
+  k1_hat: [-0.7, 0.7]

siclib/configs/model/geocalib.yaml

@@ -0,0 +1,31 @@
+name: networks.geocalib
+
+ll_enc:
+  name: encoders.low_level_encoder
+
+backbone:
+  name: encoders.mscan
+  weights: weights/mscan_b.pth
+
+perspective_decoder:
+  name: decoders.perspective_decoder
+
+  up_decoder:
+    name: decoders.up_decoder
+    loss_type: l1
+    use_uncertainty_loss: true
+    decoder:
+      name: decoders.light_hamburger
+      predict_uncertainty: true
+
+  latitude_decoder:
+    name: decoders.latitude_decoder
+    loss_type: l1
+    use_uncertainty_loss: true
+    decoder:
+      name: decoders.light_hamburger
+      predict_uncertainty: true
+
+optimizer:
+  name: optimization.lm_optimizer
+  camera_model: pinhole

siclib/configs/train/deepcalib.yaml

@@ -0,0 +1,22 @@
+seed: 0
+num_steps: 20_000
+log_every_iter: 500
+eval_every_iter: 3000
+test_every_epoch: 1
+writer: null
+lr: 1.0e-4
+clip_grad: 1.0
+lr_schedule:
+  type: null
+optimizer: adam
+submodules: []
+median_metrics:
+  - roll_error
+  - pitch_error
+  - vfov_error
+recall_metrics:
+  roll_error: [1, 5, 10]
+  pitch_error: [1, 5, 10]
+  vfov_error: [1, 5, 10]
+
+plot: [3, "siclib.visualization.visualize_batch.make_perspective_figures"]

siclib/configs/train/geocalib.yaml

@@ -0,0 +1,50 @@
+seed: 0
+num_steps: 150_000
+
+writer: null
+log_every_iter: 500
+eval_every_iter: 1000
+
+lr: 1e-4
+optimizer: adamw
+clip_grad: 1.0
+best_key: loss/param_total
+
+lr_schedule:
+  type: SequentialLR
+  on_epoch: false
+  options:
+    milestones: [4_000]
+    schedulers:
+      - type: LinearLR
+        options:
+          start_factor: 1e-3
+          total_iters: 4_000
+      - type: MultiStepLR
+        options:
+          gamma: 0.1
+          milestones: [80_000, 130_000]
+
+submodules: []
+
+median_metrics:
+  - roll_error
+  - pitch_error
+  - gravity_error
+  - vfov_error
+  - up_angle_error
+  - latitude_angle_error
+  - up_angle_recall@1
+  - up_angle_recall@5
+  - up_angle_recall@10
+  - latitude_angle_recall@1
+  - latitude_angle_recall@5
+  - latitude_angle_recall@10
+
+recall_metrics:
+  roll_error: [1, 3, 5, 10]
+  pitch_error: [1, 3, 5, 10]
+  gravity_error: [1, 3, 5, 10]
+  vfov_error: [1, 3, 5, 10]
+
+plot: [3, "siclib.visualization.visualize_batch.make_perspective_figures"]

siclib/datasets/__init__.py

@@ -0,0 +1,25 @@
+import importlib.util
+
+from siclib.datasets.base_dataset import BaseDataset
+from siclib.utils.tools import get_class
+
+
+def get_dataset(name):
+    import_paths = [name, f"{__name__}.{name}"]
+    for path in import_paths:
+        try:
+            spec = importlib.util.find_spec(path)
+        except ModuleNotFoundError:
+            spec = None
+        if spec is not None:
+            try:
+                return get_class(path, BaseDataset)
+            except AssertionError:
+                mod = __import__(path, fromlist=[""])
+                try:
+                    return mod.__main_dataset__
+                except AttributeError as exc:
+                    print(exc)
+                    continue
+
+    raise RuntimeError(f'Dataset {name} not found in any of [{" ".join(import_paths)}]')

siclib/datasets/augmentations.py

@@ -0,0 +1,359 @@
+from typing import Union
+
+import albumentations as A
+import cv2
+import numpy as np
+import torch
+from albumentations.pytorch.transforms import ToTensorV2
+from omegaconf import OmegaConf
+
+
+# flake8: noqa
+# mypy: ignore-errors
+class IdentityTransform(A.ImageOnlyTransform):
+    def apply(self, img, **params):
+        return img
+
+    def get_transform_init_args_names(self):
+        return ()
+
+
+class RandomAdditiveShade(A.ImageOnlyTransform):
+    def __init__(
+        self,
+        nb_ellipses=10,
+        transparency_limit=[-0.5, 0.8],
+        kernel_size_limit=[150, 350],
+        always_apply=False,
+        p=0.5,
+    ):
+        super().__init__(always_apply, p)
+        self.nb_ellipses = nb_ellipses
+        self.transparency_limit = transparency_limit
+        self.kernel_size_limit = kernel_size_limit
+
+    def apply(self, img, **params):
+        if img.dtype == np.float32:
+            shaded = self._py_additive_shade(img * 255.0)
+            shaded /= 255.0
+        elif img.dtype == np.uint8:
+            shaded = self._py_additive_shade(img.astype(np.float32))
+            shaded = shaded.astype(np.uint8)
+        else:
+            raise NotImplementedError(f"Data augmentation not available for type: {img.dtype}")
+        return shaded
+
+    def _py_additive_shade(self, img):
+        grayscale = len(img.shape) == 2
+        if grayscale:
+            img = img[None]
+        min_dim = min(img.shape[:2]) / 4
+        mask = np.zeros(img.shape[:2], img.dtype)
+        for i in range(self.nb_ellipses):
+            ax = int(max(np.random.rand() * min_dim, min_dim / 5))
+            ay = int(max(np.random.rand() * min_dim, min_dim / 5))
+            max_rad = max(ax, ay)
+            x = np.random.randint(max_rad, img.shape[1] - max_rad)  # center
+            y = np.random.randint(max_rad, img.shape[0] - max_rad)
+            angle = np.random.rand() * 90
+            cv2.ellipse(mask, (x, y), (ax, ay), angle, 0, 360, 255, -1)
+
+        transparency = np.random.uniform(*self.transparency_limit)
+        ks = np.random.randint(*self.kernel_size_limit)
+        if (ks % 2) == 0:  # kernel_size has to be odd
+            ks += 1
+        mask = cv2.GaussianBlur(mask.astype(np.float32), (ks, ks), 0)
+        shaded = img * (1 - transparency * mask[..., np.newaxis] / 255.0)
+        out = np.clip(shaded, 0, 255)
+        if grayscale:
+            out = out.squeeze(0)
+        return out
+
+    def get_transform_init_args_names(self):
+        return "transparency_limit", "kernel_size_limit", "nb_ellipses"
+
+
+def kw(entry: Union[float, dict], n=None, **default):
+    if not isinstance(entry, dict):
+        entry = {"p": entry}
+    entry = OmegaConf.create(entry)
+    if n is not None:
+        entry = default.get(n, entry)
+    return OmegaConf.merge(default, entry)
+
+
+def kwi(entry: Union[float, dict], n=None, **default):
+    conf = kw(entry, n=n, **default)
+    return {k: conf[k] for k in set(default.keys()).union(set(["p"]))}
+
+
+def replay_str(transforms, s="Replay:\n", log_inactive=True):
+    for t in transforms:
+        if "transforms" in t.keys():
+            s = replay_str(t["transforms"], s=s)
+        elif t["applied"] or log_inactive:
+            s += t["__class_fullname__"] + " " + str(t["applied"]) + "\n"
+    return s
+
+
+class BaseAugmentation(object):
+    base_default_conf = {
+        "name": "???",
+        "shuffle": False,
+        "p": 1.0,
+        "verbose": False,
+        "dtype": "uint8",  # (byte, float)
+    }
+
+    default_conf = {}
+
+    def __init__(self, conf={}):
+        """Perform some logic and call the _init method of the child model."""
+        default_conf = OmegaConf.merge(
+            OmegaConf.create(self.base_default_conf),
+            OmegaConf.create(self.default_conf),
+        )
+        OmegaConf.set_struct(default_conf, True)
+        if isinstance(conf, dict):
+            conf = OmegaConf.create(conf)
+        self.conf = OmegaConf.merge(default_conf, conf)
+        OmegaConf.set_readonly(self.conf, True)
+        self._init(self.conf)
+
+        self.conf = OmegaConf.merge(self.conf, conf)
+        if self.conf.verbose:
+            self.compose = A.ReplayCompose
+        else:
+            self.compose = A.Compose
+        if self.conf.dtype == "uint8":
+            self.dtype = np.uint8
+            self.preprocess = A.FromFloat(always_apply=True, dtype="uint8")
+            self.postprocess = A.ToFloat(always_apply=True)
+        elif self.conf.dtype == "float32":
+            self.dtype = np.float32
+            self.preprocess = A.ToFloat(always_apply=True)
+            self.postprocess = IdentityTransform()
+        else:
+            raise ValueError(f"Unsupported dtype {self.conf.dtype}")
+        self.to_tensor = ToTensorV2()
+
+    def _init(self, conf):
+        """Child class overwrites this, setting up a list of transforms"""
+        self.transforms = []
+
+    def __call__(self, image, return_tensor=False):
+        """image as HW or HWC"""
+        if isinstance(image, torch.Tensor):
+            image = image.cpu().numpy()
+        data = {"image": image}
+        if image.dtype != self.dtype:
+            data = self.preprocess(**data)
+        transforms = self.transforms
+        if self.conf.shuffle:
+            order = [i for i, _ in enumerate(transforms)]
+            np.random.shuffle(order)
+            transforms = [transforms[i] for i in order]
+        transformed = self.compose(transforms, p=self.conf.p)(**data)
+        if self.conf.verbose:
+            print(replay_str(transformed["replay"]["transforms"]))
+        transformed = self.postprocess(**transformed)
+        if return_tensor:
+            return self.to_tensor(**transformed)["image"]
+        else:
+            return transformed["image"]
+
+
+class IdentityAugmentation(BaseAugmentation):
+    default_conf = {}
+
+    def _init(self, conf):
+        self.transforms = [IdentityTransform(p=1.0)]
+
+
+class DarkAugmentation(BaseAugmentation):
+    default_conf = {"p": 0.75}
+
+    def _init(self, conf):
+        bright_contr = 0.5
+        blur = 0.1
+        random_gamma = 0.1
+        hue = 0.1
+        self.transforms = [
+            A.RandomRain(p=0.2),
+            A.RandomBrightnessContrast(
+                **kw(
+                    bright_contr,
+                    brightness_limit=(-0.4, 0.0),
+                    contrast_limit=(-0.3, 0.0),
+                )
+            ),
+            A.OneOf(
+                [
+                    A.Blur(**kwi(blur, p=0.1, blur_limit=(3, 9), n="blur")),
+                    A.MotionBlur(**kwi(blur, p=0.2, blur_limit=(3, 25), n="motion_blur")),
+                    A.ISONoise(),
+                    A.ImageCompression(),
+                ],
+                **kwi(blur, p=0.1),
+            ),
+            A.RandomGamma(**kw(random_gamma, gamma_limit=(15, 65))),
+            A.OneOf(
+                [
+                    A.Equalize(),
+                    A.CLAHE(p=0.2),
+                    A.ToGray(),
+                    A.ToSepia(p=0.1),
+                    A.HueSaturationValue(**kw(hue, val_shift_limit=(-100, -40))),
+                ],
+                p=0.5,
+            ),
+        ]
+
+
+class DefaultAugmentation(BaseAugmentation):
+    default_conf = {"p": 1.0}
+
+    def _init(self, conf):
+        self.transforms = [
+            A.RandomBrightnessContrast(p=0.2),
+            A.HueSaturationValue(p=0.2),
+            A.ToGray(p=0.2),
+            A.ImageCompression(quality_lower=30, quality_upper=100, p=0.5),
+            A.OneOf(
+                [
+                    A.MotionBlur(p=0.2),
+                    A.MedianBlur(blur_limit=3, p=0.1),
+                    A.Blur(blur_limit=3, p=0.1),
+                ],
+                p=0.2,
+            ),
+        ]
+
+
+class PerspectiveAugmentation(BaseAugmentation):
+    default_conf = {"p": 1.0}
+
+    def _init(self, conf):
+        self.transforms = [
+            A.RandomBrightnessContrast(p=0.2),
+            A.HueSaturationValue(p=0.2),
+            A.ToGray(p=0.2),
+            A.ImageCompression(quality_lower=30, quality_upper=100, p=0.5),
+            A.OneOf(
+                [
+                    A.MotionBlur(p=0.2),
+                    A.MedianBlur(blur_limit=3, p=0.1),
+                    A.Blur(blur_limit=3, p=0.1),
+                ],
+                p=0.2,
+            ),
+        ]
+
+
+class DeepCalibAugmentations(BaseAugmentation):
+    default_conf = {"p": 1.0}
+
+    def _init(self, conf):
+        self.transforms = [
+            A.RandomBrightnessContrast(p=0.5),
+            A.GaussNoise(var_limit=(5.0, 112.0), mean=0, per_channel=True, p=0.75),
+            A.Downscale(
+                scale_min=0.5,
+                scale_max=0.95,
+                interpolation=dict(downscale=cv2.INTER_AREA, upscale=cv2.INTER_LINEAR),
+                p=0.5,
+            ),
+            A.Downscale(scale_min=0.5, scale_max=0.95, interpolation=cv2.INTER_LINEAR, p=0.5),
+            A.ImageCompression(quality_lower=20, quality_upper=85, p=1, always_apply=True),
+            A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2, p=0.4),
+            A.Sharpen(alpha=(0.2, 0.5), lightness=(0.5, 1.0), p=0.5),
+            A.ToGray(always_apply=False, p=0.2),
+            A.GaussianBlur(blur_limit=(3, 5), sigma_limit=0, p=0.25),
+            A.MotionBlur(blur_limit=5, allow_shifted=True, p=0.25),
+            A.MultiplicativeNoise(multiplier=[0.85, 1.15], elementwise=True, p=0.5),
+        ]
+
+
+class GeoCalibAugmentations(BaseAugmentation):
+    default_conf = {"p": 1.0}
+
+    def _init(self, conf):
+        self.color_transforms = [
+            A.RandomGamma(gamma_limit=(80, 180), p=0.8),
+            A.RandomToneCurve(scale=0.1, p=0.5),
+            A.RandomBrightnessContrast(p=0.5),
+            A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2, p=0.4),
+            A.OneOf([A.ToGray(p=0.1), A.ToSepia(p=0.1), IdentityTransform(p=0.8)], p=1),
+        ]
+
+        self.noise_transforms = [
+            A.GaussNoise(var_limit=(5.0, 112.0), mean=0, per_channel=True, p=0.75),
+            A.ImageCompression(quality_lower=20, quality_upper=100, p=1),
+            A.ISONoise(color_shift=(0.01, 0.05), intensity=(0.1, 0.5), p=0.5),
+            A.OneOrOther(
+                first=A.Compose(
+                    [
+                        A.AdvancedBlur(
+                            p=1,
+                            blur_limit=(3, 7),
+                            sigmaX_limit=(0.2, 1.0),
+                            sigmaY_limit=(0.2, 1.0),
+                            rotate_limit=(-90, 90),
+                            beta_limit=(0.5, 8.0),
+                            noise_limit=(0.9, 1.1),
+                        ),
+                        A.Sharpen(p=0.5, alpha=(0.2, 0.5), lightness=(0.5, 1.0)),
+                    ]
+                ),
+                second=A.Compose(
+                    [
+                        A.Sharpen(p=0.5, alpha=(0.2, 0.5), lightness=(0.5, 1.0)),
+                        A.AdvancedBlur(
+                            p=1,
+                            blur_limit=(3, 7),
+                            sigmaX_limit=(0.2, 1.0),
+                            sigmaY_limit=(0.2, 1.0),
+                            rotate_limit=(-90, 90),
+                            beta_limit=(0.5, 8.0),
+                            noise_limit=(0.9, 1.1),
+                        ),
+                    ]
+                ),
+            ),
+        ]
+
+        self.image_transforms = [
+            A.OneOf(
+                [
+                    A.Downscale(
+                        scale_min=0.5,
+                        scale_max=0.99,
+                        interpolation=dict(downscale=down, upscale=up),
+                        p=1,
+                    )
+                    for down, up in [
+                        (cv2.INTER_AREA, cv2.INTER_LINEAR),
+                        (cv2.INTER_LINEAR, cv2.INTER_CUBIC),
+                        (cv2.INTER_CUBIC, cv2.INTER_LINEAR),
+                        (cv2.INTER_LINEAR, cv2.INTER_AREA),
+                    ]
+                ],
+                p=1,
+            )
+        ]
+
+        self.transforms = [
+            *self.color_transforms,
+            *self.noise_transforms,
+            *self.image_transforms,
+        ]
+
+
+augmentations = {
+    "default": DefaultAugmentation,
+    "dark": DarkAugmentation,
+    "perspective": PerspectiveAugmentation,
+    "deepcalib": DeepCalibAugmentations,
+    "geocalib": GeoCalibAugmentations,
+    "identity": IdentityAugmentation,
+}

siclib/datasets/base_dataset.py

@@ -0,0 +1,218 @@
+"""Base class for dataset.
+
+See mnist.py for an example of dataset.
+"""
+
+import collections
+import logging
+from abc import ABCMeta, abstractmethod
+
+import omegaconf
+import torch
+from omegaconf import OmegaConf
+from torch.utils.data import DataLoader, Sampler, get_worker_info
+from torch.utils.data._utils.collate import default_collate_err_msg_format, np_str_obj_array_pattern
+
+from siclib.utils.tensor import string_classes
+from siclib.utils.tools import set_num_threads, set_seed
+
+logger = logging.getLogger(__name__)
+
+# mypy: ignore-errors
+
+
+class LoopSampler(Sampler):
+    """Infinite sampler that loops over a given number of elements."""
+
+    def __init__(self, loop_size: int, total_size: int = None):
+        """Initialize the sampler.
+
+        Args:
+            loop_size (int): Number of elements to loop over.
+            total_size (int, optional): Total number of elements. Defaults to None.
+        """
+        self.loop_size = loop_size
+        self.total_size = total_size - (total_size % loop_size)
+
+    def __iter__(self):
+        """Return an iterator over the elements."""
+        return (i % self.loop_size for i in range(self.total_size))
+
+    def __len__(self):
+        """Return the number of elements."""
+        return self.total_size
+
+
+def worker_init_fn(i):
+    """Initialize the workers with a different seed."""
+    info = get_worker_info()
+    if hasattr(info.dataset, "conf"):
+        conf = info.dataset.conf
+        set_seed(info.id + conf.seed)
+        set_num_threads(conf.num_threads)
+    else:
+        set_num_threads(1)
+
+
+def collate(batch):
+    """Difference with PyTorch default_collate: it can stack of other objects."""
+    if not isinstance(batch, list):  # no batching
+        return batch
+    elem = batch[0]
+    elem_type = type(elem)
+    if isinstance(elem, torch.Tensor):
+        # out = None
+        if torch.utils.data.get_worker_info() is not None:
+            # If we're in a background process, concatenate directly into a
+            # shared memory tensor to avoid an extra copy
+            numel = sum([x.numel() for x in batch])
+            try:
+                _ = elem.untyped_storage()._new_shared(numel)
+            except AttributeError:
+                _ = elem.storage()._new_shared(numel)
+        return torch.stack(batch, dim=0)
+    elif (
+        elem_type.__module__ == "numpy"
+        and elem_type.__name__ != "str_"
+        and elem_type.__name__ != "string_"
+    ):
+        if elem_type.__name__ in ["ndarray", "memmap"]:
+            # array of string classes and object
+            if np_str_obj_array_pattern.search(elem.dtype.str) is not None:
+                raise TypeError(default_collate_err_msg_format.format(elem.dtype))
+            return collate([torch.as_tensor(b) for b in batch])
+        elif elem.shape == ():  # scalars
+            return torch.as_tensor(batch)
+    elif isinstance(elem, float):
+        return torch.tensor(batch, dtype=torch.float64)
+    elif isinstance(elem, int):
+        return torch.tensor(batch)
+    elif isinstance(elem, string_classes):
+        return batch
+    elif isinstance(elem, collections.abc.Mapping):
+        return {key: collate([d[key] for d in batch]) for key in elem}
+    elif isinstance(elem, tuple) and hasattr(elem, "_fields"):  # namedtuple
+        return elem_type(*(collate(samples) for samples in zip(*batch)))
+    elif isinstance(elem, collections.abc.Sequence):
+        # check to make sure that the elements in batch have consistent size
+        it = iter(batch)
+        elem_size = len(next(it))
+        if any(len(elem) != elem_size for elem in it):
+            raise RuntimeError("each element in list of batch should be of equal size")
+        transposed = zip(*batch)
+        return [collate(samples) for samples in transposed]
+    elif elem is None:
+        return elem
+    else:
+        # try to stack anyway in case the object implements stacking.
+        return torch.stack(batch, 0)
+
+
+class BaseDataset(metaclass=ABCMeta):
+    """Base class for dataset.
+
+    What the dataset model is expect to declare:
+    default_conf: dictionary of the default configuration of the dataset.
+    It overwrites base_default_conf in BaseModel, and it is overwritten by
+    the user-provided configuration passed to __init__.
+    Configurations can be nested.
+
+    _init(self, conf): initialization method, where conf is the final
+    configuration object (also accessible with `self.conf`). Accessing
+    unknown configuration entries will raise an error.
+
+    get_dataset(self, split): method that returns an instance of
+    torch.utils.data.Dataset corresponding to the requested split string,
+    which can be `'train'`, `'val'`, or `'test'`.
+    """
+
+    base_default_conf = {
+        "name": "???",
+        "num_workers": "???",
+        "train_batch_size": "???",
+        "val_batch_size": "???",
+        "test_batch_size": "???",
+        "shuffle_training": True,
+        "batch_size": 1,
+        "num_threads": 1,
+        "seed": 0,
+        "prefetch_factor": 2,
+    }
+    default_conf = {}
+
+    def __init__(self, conf):
+        """Perform some logic and call the _init method of the child model."""
+        default_conf = OmegaConf.merge(
+            OmegaConf.create(self.base_default_conf),
+            OmegaConf.create(self.default_conf),
+        )
+        OmegaConf.set_struct(default_conf, True)
+        if isinstance(conf, dict):
+            conf = OmegaConf.create(conf)
+        self.conf = OmegaConf.merge(default_conf, conf)
+        OmegaConf.set_readonly(self.conf, True)
+        logger.info(f"Creating dataset {self.__class__.__name__}")
+        self._init(self.conf)
+
+    @abstractmethod
+    def _init(self, conf):
+        """To be implemented by the child class."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_dataset(self, split):
+        """To be implemented by the child class."""
+        raise NotImplementedError
+
+    def get_data_loader(self, split, shuffle=None, pinned=False, distributed=False):
+        """Return a data loader for a given split."""
+        assert split in ["train", "val", "test"]
+        dataset = self.get_dataset(split)
+        try:
+            batch_size = self.conf[f"{split}_batch_size"]
+        except omegaconf.MissingMandatoryValue:
+            batch_size = self.conf.batch_size
+        num_workers = self.conf.get("num_workers", batch_size)
+        if distributed:
+            shuffle = False
+            sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+        else:
+            sampler = None
+            if shuffle is None:
+                shuffle = split == "train" and self.conf.shuffle_training
+        return DataLoader(
+            dataset,
+            batch_size=batch_size,
+            shuffle=shuffle,
+            sampler=sampler,
+            pin_memory=pinned,
+            collate_fn=collate,
+            num_workers=num_workers,
+            worker_init_fn=worker_init_fn,
+            prefetch_factor=self.conf.prefetch_factor,
+        )
+
+    def get_overfit_loader(self, split: str):
+        """Return an overfit data loader.
+
+        The training set is composed of a single duplicated batch, while
+        the validation and test sets contain a single copy of this same batch.
+        This is useful to debug a model and make sure that losses and metrics
+        correlate well.
+        """
+        assert split in {"train", "val", "test"}
+        dataset = self.get_dataset("train")
+        sampler = LoopSampler(
+            self.conf.batch_size,
+            len(dataset) if split == "train" else self.conf.batch_size,
+        )
+        num_workers = self.conf.get("num_workers", self.conf.batch_size)
+        return DataLoader(
+            dataset,
+            batch_size=self.conf.batch_size,
+            pin_memory=True,
+            num_workers=num_workers,
+            sampler=sampler,
+            worker_init_fn=worker_init_fn,
+            collate_fn=collate,
+        )

siclib/datasets/configs/openpano-radial.yaml

@@ -0,0 +1,41 @@
+name: openpano_radial
+base_dir: data/openpano
+pano_dir: "${.base_dir}/panoramas"
+images_per_pano: 16
+resize_factor: null
+n_workers: 1
+device: cpu
+overwrite: true
+parameter_dists:
+  roll:
+    type: uniform # uni[-45, 45]
+    options:
+      loc: -0.7853981633974483 # -45 degrees
+      scale: 1.5707963267948966 # 90 degrees
+  pitch:
+    type: uniform # uni[-45, 45]
+    options:
+      loc: -0.7853981633974483 # -45 degrees
+      scale: 1.5707963267948966 # 90 degrees
+  vfov:
+    type: uniform # uni[20, 105]
+    options:
+      loc: 0.3490658503988659 # 20 degrees
+      scale: 1.48352986419518 # 85 degrees
+  k1_hat:
+    type: truncnorm
+    options:
+      a: -4.285714285714286 # corresponds to -0.3
+      b: 4.285714285714286 # corresponds to 0.3
+      loc: 0
+      scale: 0.07
+  resize_factor:
+    type: uniform
+    options:
+      loc: 1.2
+      scale: 0.5
+  shape:
+    type: fix
+    value:
+      - 640
+      - 640

siclib/datasets/configs/openpano.yaml

@@ -0,0 +1,34 @@
+name: openpano
+base_dir: data/openpano
+pano_dir: "${.base_dir}/panoramas"
+images_per_pano: 16
+resize_factor: null
+n_workers: 1
+device: cpu
+overwrite: true
+parameter_dists:
+  roll:
+    type: uniform # uni[-45, 45]
+    options:
+      loc: -0.7853981633974483 # -45 degrees
+      scale: 1.5707963267948966 # 90 degrees
+  pitch:
+    type: uniform # uni[-45, 45]
+    options:
+      loc: -0.7853981633974483 # -45 degrees
+      scale: 1.5707963267948966 # 90 degrees
+  vfov:
+    type: uniform # uni[20, 105]
+    options:
+      loc: 0.3490658503988659 # 20 degrees
+      scale: 1.48352986419518 # 85 degrees
+  resize_factor:
+    type: uniform
+    options:
+      loc: 1.2
+      scale: 0.5
+  shape:
+    type: fix
+    value:
+      - 640
+      - 640

siclib/datasets/create_dataset_from_pano.py

@@ -0,0 +1,350 @@
+"""Script to create a dataset from panorama images."""
+
+import hashlib
+import logging
+from concurrent import futures
+from pathlib import Path
+
+import hydra
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import scipy
+import torch
+from omegaconf import DictConfig, OmegaConf
+from tqdm import tqdm
+
+from siclib.geometry.camera import camera_models
+from siclib.geometry.gravity import Gravity
+from siclib.utils.conversions import deg2rad, focal2fov, fov2focal, rad2deg
+from siclib.utils.image import load_image, write_image
+
+logger = logging.getLogger(__name__)
+
+
+# mypy: ignore-errors
+
+
+def max_radius(a, b):
+    """Compute the maximum radius of a Brown distortion model."""
+    discrim = a * a - 4 * b
+    # if torch.isfinite(discrim) and discrim >= 0.0:
+    #     discrim = np.sqrt(discrim) - a
+    #     if discrim > 0.0:
+    #         return 2.0 / discrim
+
+    valid = torch.isfinite(discrim) & (discrim >= 0.0)
+    discrim = torch.sqrt(discrim) - a
+    valid &= discrim > 0.0
+    return 2.0 / torch.where(valid, discrim, 0)
+
+
+def brown_max_radius(k1, k2):
+    """Compute the maximum radius of a Brown distortion model."""
+    # fold the constants from the derivative into a and b
+    a = k1 * 3
+    b = k2 * 5
+    return torch.sqrt(max_radius(a, b))
+
+
+class ParallelProcessor:
+    """Generic parallel processor class."""
+
+    def __init__(self, max_workers):
+        """Init processor and pbars."""
+        self.max_workers = max_workers
+        self.executor = futures.ProcessPoolExecutor(max_workers=self.max_workers)
+        self.pbars = {}
+
+    def update_pbar(self, pbar_key):
+        """Update progressbar."""
+        pbar = self.pbars.get(pbar_key)
+        pbar.update(1)
+
+    def submit_tasks(self, task_func, task_args, pbar_key):
+        """Submit tasks."""
+        pbar = tqdm(total=len(task_args), desc=f"Processing {pbar_key}", ncols=80)
+        self.pbars[pbar_key] = pbar
+
+        def update_pbar(future):
+            self.update_pbar(pbar_key)
+
+        futures = []
+        for args in task_args:
+            future = self.executor.submit(task_func, *args)
+            future.add_done_callback(update_pbar)
+            futures.append(future)
+
+        return futures
+
+    def wait_for_completion(self, futures):
+        """Wait for completion and return results."""
+        results = []
+        for f in futures:
+            results += f.result()
+
+        for key in self.pbars.keys():
+            self.pbars[key].close()
+
+        return results
+
+    def shutdown(self):
+        """Close the executer."""
+        self.executor.shutdown()
+
+
+class DatasetGenerator:
+    """Dataset generator class to create perspective datasets from panoramas."""
+
+    default_conf = {
+        "name": "???",
+        # paths
+        "base_dir": "???",
+        "pano_dir": "${.base_dir}/panoramas",
+        "pano_train": "${.pano_dir}/train",
+        "pano_val": "${.pano_dir}/val",
+        "pano_test": "${.pano_dir}/test",
+        "perspective_dir": "${.base_dir}/${.name}",
+        "perspective_train": "${.perspective_dir}/train",
+        "perspective_val": "${.perspective_dir}/val",
+        "perspective_test": "${.perspective_dir}/test",
+        "train_csv": "${.perspective_dir}/train.csv",
+        "val_csv": "${.perspective_dir}/val.csv",
+        "test_csv": "${.perspective_dir}/test.csv",
+        # data options
+        "camera_model": "pinhole",
+        "parameter_dists": {
+            "roll": {
+                "type": "uniform",
+                "options": {"loc": deg2rad(-45), "scale": deg2rad(90)},  # in [-45, 45]
+            },
+            "pitch": {
+                "type": "uniform",
+                "options": {"loc": deg2rad(-45), "scale": deg2rad(90)},  # in [-45, 45]
+            },
+            "vfov": {
+                "type": "uniform",
+                "options": {"loc": deg2rad(20), "scale": deg2rad(85)},  # in [20, 105]
+            },
+            "resize_factor": {
+                "type": "uniform",
+                "options": {"loc": 1.0, "scale": 1.0},  # factor in [1.0, 2.0]
+            },
+            "shape": {"type": "fix", "value": (640, 640)},
+        },
+        "images_per_pano": 16,
+        "n_workers": 10,
+        "device": "cpu",
+        "overwrite": False,
+    }
+
+    def __init__(self, conf):
+        """Init the class by merging and storing the config."""
+        self.conf = OmegaConf.merge(
+            OmegaConf.create(self.default_conf),
+            OmegaConf.create(conf),
+        )
+        logger.info(f"Config:\n{OmegaConf.to_yaml(self.conf)}")
+
+        self.infos = {}
+        self.device = self.conf.device
+
+        self.camera_model = camera_models[self.conf.camera_model]
+
+    def sample_value(self, parameter_name, seed=None):
+        """Sample a value from the specified distribution."""
+        param_conf = self.conf["parameter_dists"][parameter_name]
+
+        if param_conf.type == "fix":
+            return torch.tensor(param_conf.value)
+
+        # fix seed for reproducibility
+        generator = None
+        if seed:
+            if not isinstance(seed, (int, float)):
+                seed = int(hashlib.sha256(seed.encode()).hexdigest(), 16) % (2**32)
+            generator = np.random.default_rng(seed)
+
+        sampler = getattr(scipy.stats, param_conf.type)
+        return torch.tensor(sampler.rvs(random_state=generator, **param_conf.options))
+
+    def plot_distributions(self):
+        """Plot parameter distributions."""
+        fig, ax = plt.subplots(3, 3, figsize=(15, 10))
+        for i, split in enumerate(["train", "val", "test"]):
+            roll_vals = [rad2deg(row["roll"]) for row in self.infos[split]]
+            ax[i, 0].hist(roll_vals, bins=100)
+            ax[i, 0].set_xlabel("Roll (°)")
+            ax[i, 0].set_ylabel(f"Count {split}")
+
+            pitch_vals = [rad2deg(row["pitch"]) for row in self.infos[split]]
+            ax[i, 1].hist(pitch_vals, bins=100)
+            ax[i, 1].set_xlabel("Pitch (°)")
+            ax[i, 1].set_ylabel(f"Count {split}")
+
+            vfov_vals = [rad2deg(row["vfov"]) for row in self.infos[split]]
+            ax[i, 2].hist(vfov_vals, bins=100)
+            ax[i, 2].set_xlabel("vFoV (°)")
+            ax[i, 2].set_ylabel(f"Count {split}")
+
+        plt.tight_layout()
+        plt.savefig(Path(self.conf.perspective_dir) / "distributions.pdf")
+
+        fig, ax = plt.subplots(3, 3, figsize=(15, 10))
+        for i, k1 in enumerate(["roll", "pitch", "vfov"]):
+            for j, k2 in enumerate(["roll", "pitch", "vfov"]):
+                ax[i, j].scatter(
+                    [rad2deg(row[k1]) for row in self.infos["train"]],
+                    [rad2deg(row[k2]) for row in self.infos["train"]],
+                    s=1,
+                    label="train",
+                )
+
+                ax[i, j].scatter(
+                    [rad2deg(row[k1]) for row in self.infos["val"]],
+                    [rad2deg(row[k2]) for row in self.infos["val"]],
+                    s=1,
+                    label="val",
+                )
+
+                ax[i, j].scatter(
+                    [rad2deg(row[k1]) for row in self.infos["test"]],
+                    [rad2deg(row[k2]) for row in self.infos["test"]],
+                    s=1,
+                    label="test",
+                )
+
+                ax[i, j].set_xlabel(k1)
+                ax[i, j].set_ylabel(k2)
+                ax[i, j].legend()
+
+        plt.tight_layout()
+        plt.savefig(Path(self.conf.perspective_dir) / "distributions_scatter.pdf")
+
+    def generate_images_from_pano(self, pano_path: Path, out_dir: Path):
+        """Generate perspective images from a single panorama."""
+        infos = []
+
+        pano = load_image(pano_path).to(self.device)
+
+        yaws = np.linspace(0, 2 * np.pi, self.conf.images_per_pano, endpoint=False)
+        params = {
+            k: [self.sample_value(k, pano_path.stem + k + str(i)) for i in yaws]
+            for k in self.conf.parameter_dists
+            if k != "shape"
+        }
+        shapes = [self.sample_value("shape", pano_path.stem + "shape") for _ in yaws]
+        params |= {
+            "height": [shape[0] for shape in shapes],
+            "width": [shape[1] for shape in shapes],
+        }
+
+        if "k1_hat" in params:
+            height = torch.tensor(params["height"])
+            width = torch.tensor(params["width"])
+            k1_hat = torch.tensor(params["k1_hat"])
+            vfov = torch.tensor(params["vfov"])
+            focal = fov2focal(vfov, height)
+            focal = focal
+            rel_focal = focal / height
+            k1 = k1_hat * rel_focal
+
+            # distance to image corner
+            # r_max_im = f_px * r_max * (1 + k1*r_max**2)
+            # function of r_max_im: f_px = r_max_im / (r_max * (1 + k1*r_max**2))
+            min_permissible_rmax = torch.sqrt((height / 2) ** 2 + (width / 2) ** 2)
+            r_max = brown_max_radius(k1=k1, k2=0)
+            lowest_possible_f_px = min_permissible_rmax / (r_max * (1 + k1 * r_max**2))
+            valid = lowest_possible_f_px <= focal
+
+            f = torch.where(valid, focal, lowest_possible_f_px)
+            vfov = focal2fov(f, height)
+
+            params["vfov"] = vfov
+            params |= {"k1": k1}
+
+        cam = self.camera_model.from_dict(params).float().to(self.device)
+        gravity = Gravity.from_rp(params["roll"], params["pitch"]).float().to(self.device)
+
+        if (out_dir / f"{pano_path.stem}_0.jpg").exists() and not self.conf.overwrite:
+            for i in range(self.conf.images_per_pano):
+                perspective_name = f"{pano_path.stem}_{i}.jpg"
+                info = {"fname": perspective_name} | {k: v[i].item() for k, v in params.items()}
+                infos.append(info)
+
+            logger.info(f"Perspectives for {pano_path.stem} already exist.")
+
+            return infos
+
+        perspective_images = cam.get_img_from_pano(
+            pano_img=pano, gravity=gravity, yaws=yaws, resize_factor=params["resize_factor"]
+        )
+
+        for i, perspective_image in enumerate(perspective_images):
+            perspective_name = f"{pano_path.stem}_{i}.jpg"
+
+            n_pixels = perspective_image.shape[-2] * perspective_image.shape[-1]
+            valid = (torch.sum(perspective_image.sum(0) == 0) / n_pixels) < 0.01
+            if not valid:
+                logger.debug(f"Perspective {perspective_name} has too many black pixels.")
+                continue
+
+            write_image(perspective_image, out_dir / perspective_name)
+
+            info = {"fname": perspective_name} | {k: v[i].item() for k, v in params.items()}
+            infos.append(info)
+
+        return infos
+
+    def generate_split(self, split: str, parallel_processor: ParallelProcessor):
+        """Generate a single split of a dataset."""
+        self.infos[split] = []
+        panorama_paths = [
+            path
+            for path in Path(self.conf[f"pano_{split}"]).glob("*")
+            if not path.name.startswith(".")
+        ]
+
+        out_dir = Path(self.conf[f"perspective_{split}"])
+        logger.info(f"Writing perspective images to {str(out_dir)}")
+        if not out_dir.exists():
+            out_dir.mkdir(parents=True)
+
+        futures = parallel_processor.submit_tasks(
+            self.generate_images_from_pano, [(f, out_dir) for f in panorama_paths], split
+        )
+        self.infos[split] = parallel_processor.wait_for_completion(futures)
+        # parallel_processor.shutdown()
+
+        metadata = pd.DataFrame(data=self.infos[split])
+        metadata.to_csv(self.conf[f"{split}_csv"])
+
+    def generate_dataset(self):
+        """Generate all splits of a dataset."""
+        out_dir = Path(self.conf.perspective_dir)
+        if not out_dir.exists():
+            out_dir.mkdir(parents=True)
+
+        OmegaConf.save(self.conf, out_dir / "config.yaml")
+
+        processor = ParallelProcessor(self.conf.n_workers)
+        for split in ["train", "val", "test"]:
+            self.generate_split(split=split, parallel_processor=processor)
+
+        processor.shutdown()
+
+        for split in ["train", "val", "test"]:
+            logger.info(f"Generated {len(self.infos[split])} {split} images.")
+
+        self.plot_distributions()
+
+
+@hydra.main(version_base=None, config_path="configs", config_name="SUN360")
+def main(cfg: DictConfig) -> None:
+    """Run dataset generation."""
+    generator = DatasetGenerator(conf=cfg)
+    generator.generate_dataset()
+
+
+if __name__ == "__main__":
+    main()

siclib/datasets/simple_dataset.py

@@ -0,0 +1,237 @@
+"""Dataset for images created with 'create_dataset_from_pano.py'."""
+
+import logging
+from pathlib import Path
+from typing import Any, Dict, List, Tuple
+
+import pandas as pd
+import torch
+from omegaconf import DictConfig
+
+from siclib.datasets.augmentations import IdentityAugmentation, augmentations
+from siclib.datasets.base_dataset import BaseDataset
+from siclib.geometry.camera import SimpleRadial
+from siclib.geometry.gravity import Gravity
+from siclib.geometry.perspective_fields import get_perspective_field
+from siclib.utils.conversions import fov2focal
+from siclib.utils.image import ImagePreprocessor, load_image
+from siclib.utils.tools import fork_rng
+
+logger = logging.getLogger(__name__)
+
+# mypy: ignore-errors
+
+
+def load_csv(
+    csv_file: Path, img_root: Path
+) -> Tuple[List[Dict[str, Any]], torch.Tensor, torch.Tensor]:
+    """Load a CSV file containing image information.
+
+    Args:
+        csv_file (str): Path to the CSV file.
+        img_root (str): Path to the root directory containing the images.
+
+    Returns:
+        list: List of dictionaries containing the image paths and camera parameters.
+    """
+    df = pd.read_csv(csv_file)
+
+    infos, params, gravity = [], [], []
+    for _, row in df.iterrows():
+        h = row["height"]
+        w = row["width"]
+        px = row.get("px", w / 2)
+        py = row.get("py", h / 2)
+        vfov = row["vfov"]
+        f = fov2focal(torch.tensor(vfov), h)
+        k1 = row.get("k1", 0)
+        k2 = row.get("k2", 0)
+        params.append(torch.tensor([w, h, f, f, px, py, k1, k2]))
+
+        roll = row["roll"]
+        pitch = row["pitch"]
+        gravity.append(torch.tensor([roll, pitch]))
+
+        infos.append({"name": row["fname"], "file_name": str(img_root / row["fname"])})
+
+    params = torch.stack(params).float()
+    gravity = torch.stack(gravity).float()
+    return infos, params, gravity
+
+
+class SimpleDataset(BaseDataset):
+    """Dataset for images created with 'create_dataset_from_pano.py'."""
+
+    default_conf = {
+        # paths
+        "dataset_dir": "???",
+        "train_img_dir": "${.dataset_dir}/train",
+        "val_img_dir": "${.dataset_dir}/val",
+        "test_img_dir": "${.dataset_dir}/test",
+        "train_csv": "${.dataset_dir}/train.csv",
+        "val_csv": "${.dataset_dir}/val.csv",
+        "test_csv": "${.dataset_dir}/test.csv",
+        # data options
+        "use_up": True,
+        "use_latitude": True,
+        "use_prior_focal": False,
+        "use_prior_gravity": False,
+        "use_prior_k1": False,
+        # image options
+        "grayscale": False,
+        "preprocessing": ImagePreprocessor.default_conf,
+        "augmentations": {"name": "geocalib", "verbose": False},
+        "p_rotate": 0.0,  # probability to rotate image by +/- 90°
+        "reseed": False,
+        "seed": 0,
+        # data loader options
+        "num_workers": 8,
+        "prefetch_factor": 2,
+        "train_batch_size": 32,
+        "val_batch_size": 32,
+        "test_batch_size": 32,
+    }
+
+    def _init(self, conf):
+        pass
+
+    def get_dataset(self, split: str) -> torch.utils.data.Dataset:
+        """Return a dataset for a given split."""
+        return _SimpleDataset(self.conf, split)
+
+
+class _SimpleDataset(torch.utils.data.Dataset):
+    """Dataset for dataset for images created with 'create_dataset_from_pano.py'."""
+
+    def __init__(self, conf: DictConfig, split: str):
+        """Initialize the dataset."""
+        self.conf = conf
+        self.split = split
+        self.img_dir = Path(conf.get(f"{split}_img_dir"))
+
+        self.preprocessor = ImagePreprocessor(conf.preprocessing)
+
+        # load image information
+        assert f"{split}_csv" in conf, f"Missing {split}_csv in conf"
+        infos_path = self.conf.get(f"{split}_csv")
+        self.infos, self.parameters, self.gravity = load_csv(infos_path, self.img_dir)
+
+        # define augmentations
+        aug_name = conf.augmentations.name
+        assert (
+            aug_name in augmentations.keys()
+        ), f'{aug_name} not in {" ".join(augmentations.keys())}'
+
+        if self.split == "train":
+            self.augmentation = augmentations[aug_name](conf.augmentations)
+        else:
+            self.augmentation = IdentityAugmentation()
+
+    def __len__(self):
+        return len(self.infos)
+
+    def __getitem__(self, idx):
+        if not self.conf.reseed:
+            return self.getitem(idx)
+        with fork_rng(self.conf.seed + idx, False):
+            return self.getitem(idx)
+
+    def _read_image(
+        self, infos: Dict[str, Any], parameters: torch.Tensor, gravity: torch.Tensor
+    ) -> Dict[str, Any]:
+        path = Path(str(infos["file_name"]))
+
+        # load image as uint8 and HWC for augmentation
+        image = load_image(path, self.conf.grayscale, return_tensor=False)
+        image = self.augmentation(image, return_tensor=True)
+
+        # create radial camera -> same as pinhole if k1 = 0
+        camera = SimpleRadial(parameters[None]).float()
+
+        roll, pitch = gravity[None].unbind(-1)
+        gravity = Gravity.from_rp(roll, pitch)
+
+        # preprocess
+        data = self.preprocessor(image)
+        camera = camera.scale(data["scales"])
+        camera = camera.crop(data["crop_pad"]) if "crop_pad" in data else camera
+
+        priors = {"prior_gravity": gravity} if self.conf.use_prior_gravity else {}
+        priors |= {"prior_focal": camera.f[..., 1]} if self.conf.use_prior_focal else {}
+        priors |= {"prior_k1": camera.k1} if self.conf.use_prior_k1 else {}
+        return {
+            "name": infos["name"],
+            "path": str(path),
+            "camera": camera[0],
+            "gravity": gravity[0],
+            **priors,
+            **data,
+        }
+
+    def _get_perspective(self, data):
+        """Get perspective field."""
+        camera = data["camera"]
+        gravity = data["gravity"]
+
+        up_field, lat_field = get_perspective_field(
+            camera, gravity, use_up=self.conf.use_up, use_latitude=self.conf.use_latitude
+        )
+
+        out = {}
+        if self.conf.use_up:
+            out["up_field"] = up_field[0]
+        if self.conf.use_latitude:
+            out["latitude_field"] = lat_field[0]
+
+        return out
+
+    def getitem(self, idx: int):
+        """Return a sample from the dataset."""
+        infos = self.infos[idx]
+        parameters = self.parameters[idx]
+        gravity = self.gravity[idx]
+        data = self._read_image(infos, parameters, gravity)
+
+        if self.conf.use_up or self.conf.use_latitude:
+            data |= self._get_perspective(data)
+
+        return data
+
+
+if __name__ == "__main__":
+    # Create a dump of the dataset
+    import argparse
+
+    import matplotlib.pyplot as plt
+
+    from siclib.visualization.visualize_batch import make_perspective_figures
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--name", type=str, required=True)
+    parser.add_argument("--data_dir", type=str)
+    parser.add_argument("--split", type=str, default="train")
+    parser.add_argument("--shuffle", action="store_true")
+    parser.add_argument("--n_rows", type=int, default=4)
+    parser.add_argument("--dpi", type=int, default=100)
+    args = parser.parse_intermixed_args()
+
+    dconf = SimpleDataset.default_conf
+    dconf["name"] = args.name
+    dconf["num_workers"] = 0
+    dconf["prefetch_factor"] = None
+
+    dconf["dataset_dir"] = args.data_dir
+    dconf[f"{args.split}_batch_size"] = args.n_rows
+
+    torch.set_grad_enabled(False)
+
+    dataset = SimpleDataset(dconf)
+    loader = dataset.get_data_loader(args.split, args.shuffle)
+
+    with fork_rng(seed=42):
+        for data in loader:
+            pred = data
+            break
+        fig = make_perspective_figures(pred, data, n_pairs=args.n_rows)
+
+    plt.show()

siclib/datasets/utils/align_megadepth.py

@@ -0,0 +1,41 @@
+import argparse
+import subprocess
+from pathlib import Path
+
+# flake8: noqa
+# mypy: ignore-errors
+
+parser = argparse.ArgumentParser(description="Aligns a COLMAP model and plots the horizon lines")
+parser.add_argument(
+    "--base_dir", type=str, help="Path to the base directory of the MegaDepth dataset"
+)
+parser.add_argument("--out_dir", type=str, help="Path to the output directory")
+args = parser.parse_args()
+
+base_dir = Path(args.base_dir)
+out_dir = Path(args.out_dir)
+
+scenes = [d.name for d in base_dir.iterdir() if d.is_dir()]
+print(scenes[:3], len(scenes))
+
+# exit()
+
+for scene in scenes:
+    image_dir = base_dir / scene / "images"
+    sfm_dir = base_dir / scene / "sparse" / "manhattan" / "0"
+
+    # Align model
+    align_dir = out_dir / scene / "sparse" / "align"
+    align_dir.mkdir(exist_ok=True, parents=True)
+
+    print(f"image_dir ({image_dir.exists()}): {image_dir}")
+    print(f"sfm_dir   ({sfm_dir.exists()}): {sfm_dir}")
+    print(f"align_dir ({align_dir.exists()}): {align_dir}")
+
+    cmd = (
+        "colmap model_orientation_aligner "
+        + f"--image_path {image_dir} "
+        + f"--input_path {sfm_dir} "
+        + f"--output_path {str(align_dir)}"
+    )
+    subprocess.run(cmd, shell=True)

siclib/datasets/utils/download_openpano.py

@@ -0,0 +1,75 @@
+"""Helper script to download and extract OpenPano dataset."""
+
+import argparse
+import shutil
+from pathlib import Path
+
+import torch
+from tqdm import tqdm
+
+from siclib import logger
+
+PANO_URL = "https://cvg-data.inf.ethz.ch/GeoCalib_ECCV2024/openpano.zip"
+
+
+def download_and_extract_dataset(name: str, url: Path, output: Path) -> None:
+    """Download and extract a dataset from a URL."""
+    dataset_dir = output / name
+    if not output.exists():
+        output.mkdir(parents=True)
+
+    if dataset_dir.exists():
+        logger.info(f"Dataset {name} already exists at {dataset_dir}, skipping download.")
+        return
+
+    zip_file = output / f"{name}.zip"
+
+    if not zip_file.exists():
+        logger.info(f"Downloading dataset {name} to {zip_file} from {url}.")
+        torch.hub.download_url_to_file(url, zip_file)
+
+    logger.info(f"Extracting dataset {name} in {output}.")
+    shutil.unpack_archive(zip_file, output, format="zip")
+    zip_file.unlink()
+
+
+def main():
+    """Prepare the OpenPano dataset."""
+    parser = argparse.ArgumentParser(description="Download and extract OpenPano dataset.")
+    parser.add_argument("--name", type=str, default="openpano", help="Name of the dataset.")
+    parser.add_argument(
+        "--laval_dir", type=str, default="data/laval-tonemap", help="Path the Laval dataset."
+    )
+
+    args = parser.parse_args()
+
+    out_dir = Path("data")
+    download_and_extract_dataset(args.name, PANO_URL, out_dir)
+
+    pano_dir = out_dir / args.name / "panoramas"
+    for split in ["train", "test", "val"]:
+        with open(pano_dir / f"{split}_panos.txt", "r") as f:
+            pano_list = f.readlines()
+            pano_list = [fname.strip() for fname in pano_list]
+
+        for fname in tqdm(pano_list, ncols=80, desc=f"Copying {split} panoramas"):
+            laval_path = Path(args.laval_dir) / fname
+            target_path = pano_dir / split / fname
+
+            # pano either exists in laval or is in split
+            if target_path.exists():
+                continue
+
+            if laval_path.exists():
+                shutil.copy(laval_path, target_path)
+            else:  # not in laval and not in split
+                logger.warning(f"Panorama {fname} not found in {args.laval_dir} or {split} split.")
+
+    n_train = len(list(pano_dir.glob("train/*.jpg")))
+    n_test = len(list(pano_dir.glob("test/*.jpg")))
+    n_val = len(list(pano_dir.glob("val/*.jpg")))
+    logger.info(f"{args.name} contains {n_train}/{n_test}/{n_val} train/test/val panoramas.")
+
+
+if __name__ == "__main__":
+    main()

siclib/datasets/utils/tonemapping.py

@@ -0,0 +1,316 @@
+import os
+
+os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
+import argparse
+
+import cv2
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib import colors
+from tqdm import tqdm
+
+# flake8: noqa
+# mypy: ignore-errors
+
+
+class tonemap:
+    def __init__(self):
+        pass
+
+    def process(self, img):
+        return img
+
+    def inv_process(self, img):
+        return img
+
+
+# Log correction
+class log_tonemap(tonemap):
+    # Constructor
+    # Base of log
+    # Scale of tonemapped
+    # Offset
+    def __init__(self, base, scale=1, offset=1):
+        self.base = base
+        self.scale = scale
+        self.offset = offset
+
+    def process(self, img):
+        tonemapped = (np.log(img + self.offset) / np.log(self.base)) * self.scale
+        return tonemapped
+
+    def inv_process(self, img):
+        inverse_tonemapped = np.power(self.base, (img) / self.scale) - self.offset
+        return inverse_tonemapped
+
+
+class log_tonemap_clip(tonemap):
+    # Constructor
+    # Base of log
+    # Scale of tonemapped
+    # Offset
+    def __init__(self, base, scale=1, offset=1):
+        self.base = base
+        self.scale = scale
+        self.offset = offset
+
+    def process(self, img):
+        tonemapped = np.clip((np.log(img * self.scale + self.offset) / np.log(self.base)), 0, 2) - 1
+        return tonemapped
+
+    def inv_process(self, img):
+        inverse_tonemapped = (np.power(self.base, (img + 1)) - self.offset) / self.scale
+        return inverse_tonemapped
+
+
+# Gamma Tonemap
+class gamma_tonemap(tonemap):
+    def __init__(
+        self,
+        gamma,
+    ):
+        self.gamma = gamma
+
+    def process(self, img):
+        tonemapped = np.power(img, 1 / self.gamma)
+        return tonemapped
+
+    def inv_process(self, img):
+        inverse_tonemapped = np.power(img, self.gamma)
+        return inverse_tonemapped
+
+
+class linear_clip(tonemap):
+    def __init__(self, scale, mean):
+        self.scale = scale
+        self.mean = mean
+
+    def process(self, img):
+        tonemapped = np.clip((img - self.mean) / self.scale, -1, 1)
+        return tonemapped
+
+    def inv_process(self, img):
+        inverse_tonemapped = img * self.scale + self.mean
+        return inverse_tonemapped
+
+
+def make_tonemap_HDR(opt):
+    if opt.mode == "luminance":
+        res_tonemap = log_tonemap_clip(10, 1.0, 1.0)
+    else:  # temperature
+        res_tonemap = linear_clip(5000.0, 5000.0)
+    return res_tonemap
+
+
+class LDRfromHDR:
+    def __init__(
+        self, tonemap="none", orig_scale=False, clip=True, quantization=0, color_jitter=0, noise=0
+    ):
+        self.tonemap_str, val = tonemap
+        if tonemap[0] == "gamma":
+            self.tonemap = gamma_tonemap(val)
+        elif tonemap[0] == "log10":
+            self.tonemap = log_tonemap(val)
+        else:
+            print("Warning: No tonemap specified, using linear")
+
+        self.clip = clip
+        self.orig_scale = orig_scale
+        self.bits = quantization
+        self.jitter = color_jitter
+        self.noise = noise
+
+        self.wbModel = None
+
+    def process(self, HDR):
+        LDR, normalized_scale = self.rescale(HDR)
+        LDR = self.apply_clip(LDR)
+        LDR = self.apply_scale(LDR, normalized_scale)
+        LDR = self.apply_tonemap(LDR)
+        LDR = self.colorJitter(LDR)
+        LDR = self.gaussianNoise(LDR)
+        LDR = self.quantize(LDR)
+        LDR = self.apply_white_balance(LDR)
+        return LDR, normalized_scale
+
+    def rescale(self, img, percentile=90, max_mapping=0.8):
+        r_percentile = np.percentile(img, percentile)
+        alpha = max_mapping / (r_percentile + 1e-10)
+
+        img_reexposed = img * alpha
+
+        normalized_scale = normalizeScale(1 / alpha)
+
+        return img_reexposed, normalized_scale
+
+    def rescaleAlpha(self, img, percentile=90, max_mapping=0.8):
+        r_percentile = np.percentile(img, percentile)
+        alpha = max_mapping / (r_percentile + 1e-10)
+
+        return alpha
+
+    def apply_clip(self, img):
+        if self.clip:
+            img = np.clip(img, 0, 1)
+        return img
+
+    def apply_scale(self, img, scale):
+        if self.orig_scale:
+            scale = unNormalizeScale(scale)
+            img = img * scale
+        return img
+
+    def apply_tonemap(self, img):
+        if self.tonemap_str == "none":
+            return img
+        gammaed = self.tonemap.process(img)
+        return gammaed
+
+    def quantize(self, img):
+        if self.bits == 0:
+            return img
+        max_val = np.power(2, self.bits)
+        img = img * max_val
+        img = np.floor(img)
+        img = img / max_val
+        return img
+
+    def colorJitter(self, img):
+        if self.jitter == 0:
+            return img
+        hsv = colors.rgb_to_hsv(img)
+        hue_offset = np.random.normal(0, self.jitter, 1)
+        hsv[:, :, 0] = (hsv[:, :, 0] + hue_offset) % 1.0
+        rgb = colors.hsv_to_rgb(hsv)
+        return rgb
+
+    def gaussianNoise(self, img):
+        if self.noise == 0:
+            return img
+        noise_amount = np.random.uniform(0, self.noise, 1)
+        noise_img = np.random.normal(0, noise_amount, img.shape)
+        img = img + noise_img
+        img = np.clip(img, 0, 1).astype(np.float32)
+        return img
+
+    def apply_white_balance(self, img):
+        if self.wbModel is None:
+            return img
+        img = self.wbModel.correctImage(img)
+        return img.copy()
+
+
+def make_LDRfromHDR(opt):
+    LDR_from_HDR = LDRfromHDR(
+        opt.tonemap_LDR, opt.orig_scale, opt.clip, opt.quantization, opt.color_jitter, opt.noise
+    )
+    return LDR_from_HDR
+
+
+def torchnormalizeEV(EV, mean=5.12, scale=6, clip=True):
+    # Normalize based on the computed distribution between -1 1
+    EV -= mean
+    EV = EV / scale
+
+    if clip:
+        EV = torch.clip(EV, min=-1, max=1)
+
+    return EV
+
+
+def torchnormalizeEV0(EV, mean=5.12, scale=6, clip=True):
+    # Normalize based on the computed distribution between 0 1
+    EV -= mean
+    EV = EV / scale
+
+    if clip:
+        EV = torch.clip(EV, min=-1, max=1)
+
+    EV += 0.5
+    EV = EV / 2
+
+    return EV
+
+
+def normalizeScale(x, scale=4):
+    x = np.log10(x + 1)
+
+    x = x / (scale / 2)
+    x = x - 1
+
+    return x
+
+
+def unNormalizeScale(x, scale=4):
+    x = x + 1
+    x = x * (scale / 2)
+
+    x = np.power(10, x) - 1
+
+    return x
+
+
+def normalizeIlluminance(x, scale=5):
+    x = np.log10(x + 1)
+
+    x = x / (scale / 2)
+    x = x - 1
+
+    return x
+
+
+def unNormalizeIlluminance(x, scale=5):
+    x = x + 1
+    x = x * (scale / 2)
+
+    x = np.power(10, x) - 1
+
+    return x
+
+
+def main(args):
+    processor = LDRfromHDR(
+        # tonemap=("log10", 10),
+        tonemap=("gamma", args.gamma),
+        orig_scale=False,
+        clip=True,
+        quantization=0,
+        color_jitter=0,
+        noise=0,
+    )
+
+    img_list = list(os.listdir(args.hdr_dir))
+    img_list = [f for f in img_list if f.endswith(args.extension)]
+    img_list = [f for f in img_list if not f.startswith("._")]
+
+    if not os.path.exists(args.out_dir):
+        os.makedirs(args.out_dir)
+
+    for fname in tqdm(img_list):
+        fname_out = ".".join(fname.split(".")[:-1])
+        out = os.path.join(args.out_dir, f"{fname_out}.jpg")
+        if os.path.exists(out) and not args.overwrite:
+            continue
+
+        fpath = os.path.join(args.hdr_dir, fname)
+        img = cv2.imread(fpath, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+        ldr, scale = processor.process(img)
+
+        ldr = (ldr * 255).astype(np.uint8)
+        ldr = cv2.cvtColor(ldr, cv2.COLOR_RGB2BGR)
+        cv2.imwrite(out, ldr)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--hdr_dir", type=str, default="hdr")
+    parser.add_argument("--out_dir", type=str, default="ldr")
+    parser.add_argument("--extension", type=str, default=".exr")
+    parser.add_argument("--overwrite", action="store_true")
+    parser.add_argument("--gamma", type=float, default=2)
+    args = parser.parse_args()
+
+    main(args)

siclib/eval/__init__.py

@@ -0,0 +1,18 @@
+import torch
+
+from siclib.eval.eval_pipeline import EvalPipeline
+from siclib.utils.tools import get_class
+
+
+def get_benchmark(benchmark):
+    return get_class(f"{__name__}.{benchmark}", EvalPipeline)
+
+
+@torch.no_grad()
+def run_benchmark(benchmark, eval_conf, experiment_dir, model=None):
+    """This overwrites existing benchmarks"""
+    experiment_dir.mkdir(exist_ok=True, parents=True)
+    bm = get_benchmark(benchmark)
+
+    pipeline = bm(eval_conf)
+    return pipeline.run(experiment_dir, model=model, overwrite=True, overwrite_eval=True)

siclib/eval/configs/deepcalib.yaml

@@ -0,0 +1,3 @@
+model:
+  name: networks.deepcalib
+  weights: weights/deepcalib.tar

siclib/eval/configs/geocalib-pinhole.yaml

@@ -0,0 +1,2 @@
+model:
+  name: networks.geocalib_pretrained