[
{
"path": "AutoReconForDens3R/.gitattributes",
"content": "*.ipynb linguist-documentation"
},
{
"path": "AutoReconForDens3R/.github/ISSUE_TEMPLATE/bug-report.md",
"content": "---\nname: Bug report\nabout: Create a report to help us improve\ntitle: ''\nlabels: ''\nassignees: ''\n\n---\n\n**Describe the bug**\nA clear and concise description of what the bug is.\n\n**To Reproduce**\nSteps to reproduce the behavior:\n1. Go to '...'\n2. Click on '....'\n3. Scroll down to '....'\n4. See error\n\n**Expected behavior**\nA clear and concise description of what you expected to happen.\n\n**Screenshots**\nIf applicable, add screenshots to help explain your problem.\n\n**Additional context**\nAdd any other context about the problem here.\n"
},
{
"path": "AutoReconForDens3R/.github/ISSUE_TEMPLATE/feature_request.md",
"content": "---\nname: Feature request\nabout: Suggest an idea for this project\ntitle: ''\nlabels: ''\nassignees: ''\n\n---\n\n**Is your feature request related to a problem? Please describe.**\nA clear and concise description of what the problem is. Ex. I'm always frustrated when [...]\n\n**Describe the solution you'd like**\nA clear and concise description of what you want to happen.\n\n**Describe alternatives you've considered**\nA clear and concise description of any alternative solutions or features you've considered.\n\n**Additional context**\nAdd any other context or screenshots about the feature request here.\n"
},
{
"path": "AutoReconForDens3R/.github/workflows/core_code_checks.yml",
"content": "name: Core Tests.\n\non:\n push:\n branches: [main]\n pull_request:\n branches: [main]\n\npermissions:\n contents: read\n\njobs:\n build:\n runs-on: ubuntu-latest\n\n steps:\n - uses: actions/checkout@v3\n - name: Set up Python 3.7.13\n uses: actions/setup-python@v4\n with:\n python-version: '3.7.13'\n - uses: actions/cache@v2\n with:\n path: ${{ env.pythonLocation }}\n key: ${{ env.pythonLocation }}-${{ hashFiles('pyproject.toml') }}\n - name: Install dependencies\n run: |\n pip install --upgrade --upgrade-strategy eager -e .[dev]\n - name: Run license checks\n run: |\n ./scripts/licensing/license_headers.sh --check\n - name: Check notebook cell metadata\n run: |\n python ./scripts/docs/add_nb_tags.py --check\n - name: Run isort\n run: isort docs/ nerfstudio/ scripts/ tests/ --profile black --check\n - name: Run Black\n run: black docs/ nerfstudio/ scripts/ tests/ --check\n - name: Python Pylint\n run: |\n pylint nerfstudio tests scripts\n - name: Test with pytest\n run: |\n pytest\n"
},
{
"path": "AutoReconForDens3R/.github/workflows/publish.yml",
"content": "# This workflows will upload a Python Package using twine when a release is created\n# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries\n\nname: Upload Python Package\n\non:\n release:\n types: [created]\n\njobs:\n deploy:\n runs-on: ubuntu-latest\n\n steps:\n - uses: actions/checkout@v2\n - name: Set up Python\n uses: actions/setup-python@v1\n with:\n python-version: '3.8'\n - name: Install dependencies\n run: |\n python -m pip install build twine\n - name: Strip unsupported tags in README\n run: |\n sed -i '//,//d' README.md\n - name: Build and publish\n env:\n PYPI_TOKEN: ${{ secrets.PYPI_TOKEN }}\n run: |\n python -m build\n twine upload --username __token__ --password $PYPI_TOKEN dist/*\n"
},
{
"path": "AutoReconForDens3R/.github/workflows/viewer_build_deploy.yml",
"content": "name: Viewer Build and Deploy.\n\non:\n push:\n branches: [main]\n pull_request:\n branches: [main]\n\njobs:\n build:\n name: Build\n runs-on: ubuntu-latest\n defaults:\n run:\n working-directory: ./nerfstudio/viewer/app\n steps:\n - name: Checkout code\n uses: actions/checkout@master\n\n - name: Install Node.js\n uses: actions/setup-node@v3\n with:\n node-version: 17.8.0\n cache: 'yarn'\n cache-dependency-path: ./nerfstudio/viewer/app/yarn.lock\n\n - name: Install packages\n run: yarn install\n\n - name: Build project\n run: CI=false yarn build\n\n - name: Upload production-ready build files\n uses: actions/upload-artifact@v2\n with:\n name: production-files\n path: ./nerfstudio/viewer/app/build\n\n deploy:\n name: Deploy\n needs: build\n runs-on: ubuntu-latest\n\n env:\n SSH_KEY: ${{secrets.SSH_KEY}}\n\n steps:\n - uses: actions/checkout@v3\n - name: Set up Python 3.8.12\n uses: actions/setup-python@v4\n with:\n python-version: '3.8.12'\n\n - name: Install dependencies\n run: |\n pip install -r ./nerfstudio/viewer/app/requirements.txt\n\n - name: Download artifact\n uses: actions/download-artifact@v2\n with:\n name: production-files\n path: ./nerfstudio/viewer/app/build\n\n - name: Get branch name (merge)\n if: github.event_name != 'pull_request'\n shell: bash\n run: echo \"BRANCH_NAME=$(echo ${GITHUB_REF#refs/heads/} | tr / -)\" >> $GITHUB_ENV\n\n - name: Get branch name (pull request)\n if: github.event_name == 'pull_request'\n shell: bash\n run: echo \"BRANCH_NAME=$(echo ${GITHUB_HEAD_REF} | tr / -)\" >> $GITHUB_ENV\n\n # TODO: detect file or scheme changes of the viewer and only\n # increment the version.txt file when there is a change.\n # Update the version.txt code and push to master when things change.\n # https://github.com/marketplace/actions/changed-files\n # - name: Run changed-files with defaults on the dir1\n # id: changed-files-for-dir1\n # uses: tj-actions/changed-files@v29.0.3\n # with:\n # path: nerfstudio/viewer/app\n\n # - name: List all added files in dir1\n # run: |\n # for file in ${{ steps.changed-files-for-dir1.outputs.modified_files }}; do\n # echo \"$file was modified\"\n # done\n\n - run: |\n python ./nerfstudio/viewer/app/run_deploy.py \\\n --branch-name ${{ env.BRANCH_NAME }} \\\n --ssh-key-string \"$SSH_KEY\" \\\n --local-folder ./nerfstudio/viewer/app/build \\\n --package-json-filename ./nerfstudio/viewer/app/package.json \\\n --increment-version \"False\"\n - run: cat ~/.ssh/config\n"
},
{
"path": "AutoReconForDens3R/.gitignore",
"content": "*_debug_*\n\n# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\n!scripts/downloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\nshare/python-wheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.py,cover\n.hypothesis/\n.pytest_cache/\ncover/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\ndb.sqlite3-journal\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\n.pybuilder/\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# IPython\nprofile_default/\nipython_config.py\n\n# pyenv\n# For a library or package, you might want to ignore these files since the code is\n# intended to run in multiple environments; otherwise, check them in:\n# .python-version\n\n# pipenv\n# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.\n# However, in case of collaboration, if having platform-specific dependencies or dependencies\n# having no cross-platform support, pipenv may install dependencies that don't work, or not\n# install all needed dependencies.\n#Pipfile.lock\n\n# poetry\n# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.\n# This is especially recommended for binary packages to ensure reproducibility, and is more\n# commonly ignored for libraries.\n# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control\n#poetry.lock\n\n# pdm\n# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.\n#pdm.lock\n# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it\n# in version control.\n# https://pdm.fming.dev/#use-with-ide\n.pdm.toml\n\n# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm\n__pypackages__/\n\n# Celery stuff\ncelerybeat-schedule\ncelerybeat.pid\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.envrc\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n\n# pytype static type analyzer\n.pytype/\n\n# Cython debug symbols\ncython_debug/\n\n# PyCharm\n# JetBrains specific template is maintained in a separate JetBrains.gitignore that can\n# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore\n# and can be added to the global gitignore or merged into this file. For a more nuclear\n# option (not recommended) you can uncomment the following to ignore the entire idea folder.\n#.idea/\n\n# Experiments and outputs\noutputs/\nexports/\nrenders/\noutputs\nvis/\nmeshes/\ngradio_cached_examples/\n# tensorboard log files\nevents.out.*\n\n# Data\ndata\n!*/data\n\n# Misc\nold/\ntemp*\n!temporal*\n.nfs*\nexternal/\n__MACOSX/\nnode_modules/\nbash/\ncache/\npackage-lock.json\ncamera_paths/\n\n._.DS_Store\n.DS_Store\n\n.vscode/settings.json\nvis_results/\nW/\n*.ply\nnotebooks/\nmedia/\noutputs/\n/Dockerfile\n/.readthedocs.yaml\nvis_meshes/\nREADME_*.md"
},
{
"path": "AutoReconForDens3R/.prettierrc.js",
"content": "module.exports = {\n trailingComma: 'all',\n arrowParens: 'always',\n singleQuote: true,\n jsxSingleQuote: false,\n };"
},
{
"path": "AutoReconForDens3R/LICENSE",
"content": " Apache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright [yyyy] [name of copyright owner]\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n"
},
{
"path": "AutoReconForDens3R/README.md",
"content": "# AutoRecon: Automated 3D Object Discovery and Reconstruction\n\n### [Project Page](https://zju3dv.github.io/autorecon) | [Paper](https://zju3dv.github.io/autorecon/files/autorecon.pdf)\n\n\n\n> [AutoRecon: Automated 3D Object Discovery and Reconstruction](https://zju3dv.github.io/autorecon/files/autorecon.pdf) \n> Yuang Wang, Xingyi He, Sida Peng, Haotong Lin, Hujun Bao, Xiaowei Zhou \n> CVPR 2023\n\n# About\nThis is a refactored version of the [AutoRecon](https://zju3dv.github.io/autorecon) project based on the NeRFStudio and the SDFStudio codebase. We separate the project into two parts. The coarse decomposition part is implemented in the [AutoDecomp](https://github.com/zju3dv/AutoDecomp) repo, which can be used as a general tool for 3D object discovery and preprocessing casual captures for object reconstruction. The neural surface reconstruction part is implemented here.\n\n\n# Installation\nPlease refer to the [installation guide](docs/INSTALL.md) for detailed instructions.\n\n\n# Run the pipeline\n## Run the pipeline with your own data\nHere we take a demo data as an example. In this example, we assume only a sequential stream of images is available. You can easily adapt your data to use AutoRecon.\n1. Download the demo data from [Google Drive](https://drive.google.com/drive/folders/1IFbK9b7gzqwh9QkZe6zoLmcSFr_zW3rJ?usp=drive_link) and put it under `data` (-> `data/custom_data_example/...`)\n2. Run the pipeline with `exps/code-release/run_pipeline_demo_low-res.sh`\n\n> NOTE: In the demo script, we assume the images come from a sequential video stream and use sequential matching for SfM. If you have unordered images, the default setting might lead to inferior results and it is recommended to use exhaustive matching or vocab-tree instead.\n\n\n## Run the pipeline with annotated data in the IDR format\n1. Download the BlendedMVS data from [GoogleDrive](https://drive.google.com/drive/folders/1ZLQ0hap6o_Tjr7S6H_EAn17pz4_qFluW?usp=sharing) and put it under `data` (-> `data/BlendedMVS/...`)\n2. Run one of the scripts in `exps/code-release/bmvs`\n\n## Run the pipeline with CO3D data\n1. Download the CO3D data from [GoogleDrive](https://drive.google.com/drive/folders/1u-ugNhwFVtV6TKZ2J29iwcdQY2iwVnoi?usp=sharing) and put it under `data` (-> `data/CO3D_DEMO/...`)\n2. Run one of the scripts in `exps/code-release/co3d_demo`\n\n\n# Extract Mesh\nYou can take the following script as a reference to extract mesh:\n```bash\n# Extract mesh with MC\nLOG_DIR=\"path_to_log_dir\"\n\nMC_RES=512\nMESH_FN=\"extracted_mesh_res-${MC_RES}.ply\"\nMESH_PATH=\"${LOG_DIR}/${MESH_FN}\"\n\nns-extract-mesh \\\n\t--load-config $LOG_DIR/config.yml \\\n --load-dir $LOG_DIR/sdfstudio_models \\\n\t--output-path $MESH_PATH \\\n --chunk_size 25000 --store_float16 True \\\n --resolution $MC_RES \\\n --use_train_scene_box True \\\n --seg_aware_sdf False \\\n --remove_internal_geometry None \\\n --remove_non_maximum_connected_components True \\\n --close_holes False --simplify_mesh_final False\n```\n\n> NOTE: We postprocess the extracted mesh before evaluation by removing possible internal geometries with ambient occlusion. The postprocessing code depends on pymeshlab which only works on some machines. You can remove possible internal geometries manually with ambient occlusion using MeshLab.\n\n\n# Citation\nIf you find this code useful for your research, please use the following BibTeX entry.\n\n```bibtex\n@inproceedings{wang2023autorecon,\n title={AutoRecon: Automated 3D Object Discovery and Reconstruction},\n author={Wang, Yuang and He, Xingyi and Peng, Sida and Lin, Haotong and Bao, Hujun and Zhou, Xiaowei},\n booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},\n pages={21382--21391},\n year={2023}\n}\n```\n\n\n# Acknowledgement\nThis code is built upon the awesome projects including [nerfstudio](https://github.com/nerfstudio-project/nerfstudio/), [sdfstudio](https://github.com/autonomousvision/sdfstudio/blob/master/README.md), [nerfacc](https://github.com/KAIR-BAIR/nerfacc), [tyro](https://github.com/brentyi/tyro) and more. Thanks for these great projects!\n"
},
{
"path": "AutoReconForDens3R/colab/demo.ipynb",
"content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"colab_type\": \"text\",\n \"id\": \"view-in-github\"\n },\n \"source\": [\n \"\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"SiiXJ7K_fePG\"\n },\n \"source\": [\n \"
\\n\",\n \"\\n\",\n \"\\n\",\n \"# Nerfstudio: A collaboration friendly studio for NeRFs\\n\",\n \"\\n\",\n \"\\n\",\n \"\\n\",\n \"\\n\",\n \"This colab shows how to train and view NeRFs from Nerfstudio both on pre-made datasets or from your own videos/images.\\n\",\n \"\\n\",\n \"\\\\\\\\\\n\",\n \"\\n\",\n \"Credit to [NeX](https://nex-mpi.github.io/) for Google Colab format.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"Yyx5h6kz5ga7\"\n },\n \"source\": [\n \"## Frequently Asked Questions\\n\",\n \"\\n\",\n \"* **Downloading custom data is stalling (no output):**\\n\",\n \" * This is a bug in Colab. The data is processing, but may take a while to complete. You will know processing completed if `data/nerfstudio/custom_data/transforms.json` exists. Terminating the cell early will result in not being able to train.\\n\",\n \"* **Processing custom data is taking a long time:**\\n\",\n \" * The time it takes to process data depends on the number of images and its resolution. If processing is taking too long, try lowering the resolution of your custom data.\\n\",\n \"* **Error: Data processing did not complete:**\\n\",\n \" * This means that the data processing script did not fully complete. This could be because there were not enough images, or that the images were of low quality. We recommend images with little to no motion blur and lots of visual overlap of the scene to increase the chances of successful processing.\\n\",\n \"* **Training is not showing progress**:\\n\",\n \" * The lack of output is a bug in Colab. You can see the training progress from the viewer.\\n\",\n \"* **Viewer Quality is bad / Low resolution**:\\n\",\n \" * This may be because more GPU is being used on training that rendering the viewer. Try pausing training or decreasing training utilization.\\n\",\n \"* **WARNING: Running pip as the 'root' user...:**:\\n\",\n \" * This and other pip warnings or errors can be safely ignored.\\n\",\n \"* **Other problems?**\\n\",\n \" * Feel free to create an issue on our [GitHub repo](https://github.com/nerfstudio-project/nerfstudio).\\n\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"metadata\": {\n \"cellView\": \"form\",\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\"\n },\n \"id\": \"RGr33zHaHak0\",\n \"outputId\": \"2641dce6-cd57-4d31-d8b1-860bc1adf17f\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"⏬ Downloading https://github.com/jaimergp/miniforge/releases/latest/download/Mambaforge-colab-Linux-x86_64.sh...\\n\",\n \"📦 Installing...\\n\",\n \"📌 Adjusting configuration...\\n\",\n \"🩹 Patching environment...\\n\",\n \"⏲ Done in 0:00:27\\n\",\n \"🔁 Restarting kernel...\\n\"\n ]\n }\n ],\n \"source\": [\n \"#@markdown
Please re-run `Downloading and Processing Data`, or view the FAQ for more info.
'))\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"cellView\": \"form\",\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\"\n },\n \"id\": \"WGt8ukG6Htg3\",\n \"outputId\": \"fa946890-c7d8-4e46-a54e-7231bc5a2059\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\u001b[2;36m[19:48:48]\\u001b[0m\\u001b[2;36m \\u001b[0mSkipping \\u001b[1;36m0\\u001b[0m files in dataset split train. \\u001b]8;id=527413;file:///content/nerfstudio/nerfstudio/data/dataparsers/nerfstudio_dataparser.py\\u001b\\\\\\u001b[2mnerfstudio_dataparser.py\\u001b[0m\\u001b]8;;\\u001b\\\\\\u001b[2m:\\u001b[0m\\u001b]8;id=243595;file:///content/nerfstudio/nerfstudio/data/dataparsers/nerfstudio_dataparser.py#91\\u001b\\\\\\u001b[2m91\\u001b[0m\\u001b]8;;\\u001b\\\\\\n\",\n \"\\u001b[2;36m \\u001b[0m\\u001b[2;36m \\u001b[0mSkipping \\u001b[1;36m0\\u001b[0m files in dataset split test. \\u001b]8;id=109270;file:///content/nerfstudio/nerfstudio/data/dataparsers/nerfstudio_dataparser.py\\u001b\\\\\\u001b[2mnerfstudio_dataparser.py\\u001b[0m\\u001b]8;;\\u001b\\\\\\u001b[2m:\\u001b[0m\\u001b]8;id=464675;file:///content/nerfstudio/nerfstudio/data/dataparsers/nerfstudio_dataparser.py#91\\u001b\\\\\\u001b[2m91\\u001b[0m\\u001b]8;;\\u001b\\\\\\n\",\n \"\\u001b[2KLoading data batch \\u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\\u001b[0m \\u001b[35m100%\\u001b[0m \\u001b[36m0:00:00\\u001b[0m\\n\",\n \"\\u001b[?25h/usr/local/lib/python3.7/site-packages/torch/utils/data/dataloader.py:566: UserWarning: This DataLoader will create 4 worker processes in total. Our suggested max number of worker in current system is 2, which is smaller than what this DataLoader is going to create. Please be aware that excessive worker creation might get DataLoader running slow or even freeze, lower the worker number to avoid potential slowness/freeze if necessary.\\n\",\n \" cpuset_checked))\\n\",\n \"\\u001b[2KLoading data batch \\u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\\u001b[0m \\u001b[35m100%\\u001b[0m \\u001b[36m0:00:00\\u001b[0m\\n\",\n \"\\u001b[?25h/usr/local/lib/python3.7/site-packages/torchvision/models/_utils.py:209: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and will be removed in 0.15, please use 'weights' instead.\\n\",\n \" f\\\"The parameter '{pretrained_param}' is deprecated since 0.13 and will be removed in 0.15, \\\"\\n\",\n \"/usr/local/lib/python3.7/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and will be removed in 0.15. The current behavior is equivalent to passing `weights=AlexNet_Weights.IMAGENET1K_V1`. You can also use `weights=AlexNet_Weights.DEFAULT` to get the most up-to-date weights.\\n\",\n \" warnings.warn(msg)\\n\",\n \"Loading latest checkpoint from load_dir\\n\",\n \"✅ Done loading checkpoint from \\n\",\n \"outputs/data-nerfstudio-poster/nerfacto/\\u001b[1;36m2022\\u001b[0m-\\u001b[1;36m10\\u001b[0m-29_192844/nerfstudio_models/step-\\u001b[1;36m000014000.\\u001b[0mckpt\\n\",\n \"\\u001b[1;32mCreating trajectory video\\u001b[0m\\n\",\n \"\\u001b[2K🎥 Rendering 🎥 \\u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\\u001b[0m \\u001b[35m100%\\u001b[0m \\u001b[31m0.14 fps\\u001b[0m \\u001b[33m11:47\\u001b[0m\\n\",\n \"\\u001b[2K\\u001b[32m( ● )\\u001b[0m \\u001b[33mSaving video\\u001b[0m\\n\",\n \"\\u001b[1A\\u001b[2K\\u001b[92m────────────────────────────────────────────── \\u001b[0m\\u001b[32m 🎉 🎉 🎉 Success 🎉 🎉 🎉\\u001b[0m\\u001b[92m ──────────────────────────────────────────────\\u001b[0m\\n\",\n \" \\u001b[32mSaved video to renders/output.mp4\\u001b[0m \\n\",\n \"\\u001b[0m\"\n ]\n }\n ],\n \"source\": [\n \"#@title # Render Video { vertical-output: true }\\n\",\n \"#@markdown
Export the camera path from within the viewer, then run this cell.
\\n\",\n \"#@markdown
The rendered video should be at renders/output.mp4!
'))\\n\",\n \"%cd $training_run_dir\\n\",\n \"uploaded = files.upload()\\n\",\n \"uploaded_camera_path_filename = list(uploaded.keys())[0]\\n\",\n \"\\n\",\n \"config_filename = training_run_dir + \\\"/config.yml\\\"\\n\",\n \"camera_path_filename = training_run_dir + \\\"/\\\" + uploaded_camera_path_filename\\n\",\n \"camera_path_filename = camera_path_filename.replace(\\\" \\\", \\\"\\\\\\\\ \\\").replace(\\\"(\\\", \\\"\\\\\\\\(\\\").replace(\\\")\\\", \\\"\\\\\\\\)\\\")\\n\",\n \"\\n\",\n \"%cd /content/\\n\",\n \"!ns-render --load-config $config_filename --traj filename --camera-path-filename $camera_path_filename --output-path renders/output.mp4\"\n ]\n }\n ],\n \"metadata\": {\n \"accelerator\": \"GPU\",\n \"colab\": {\n \"collapsed_sections\": [],\n \"include_colab_link\": true,\n \"provenance\": []\n },\n \"gpuClass\": \"standard\",\n \"kernelspec\": {\n \"display_name\": \"Python 3.8.13 ('nerfstudio')\",\n \"language\": \"python\",\n \"name\": \"python3\"\n },\n \"language_info\": {\n \"name\": \"python\",\n \"version\": \"3.8.13\"\n },\n \"vscode\": {\n \"interpreter\": {\n \"hash\": \"c59f626636933ef1dc834fb3684b382f705301c5306cf8436d2da634c2289783\"\n }\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 0\n}\n"
},
{
"path": "AutoReconForDens3R/docs/INSTALL.md",
"content": "# Installation\n\n## Create a conda environment\n```bash\nconda create --name auto_recon -y python=3.9\nconda activate auto_recon\npython -m pip install --upgrade pip setuptools\n```\n\n## Clone the repo\n```shell\ngit clone --recurse-submodules git@github.com:zju3dv/AutoRecon.git\n```\n\n## Install AutoDecomp\n```bash\ncd third_party/AutoDecomp\n```\nThen, please install AutoDecomp based on its [installation guide](https://github.com/zju3dv/AutoDecomp/blob/main/docs/INSTALL.md).\nWe expect that it is installed in `third_party/AutoDecomp`.\n\n## Install other dependencies\nInstall pytorch with CUDA (this repo has been tested with CUDA 11.7) and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn)\n```bash\npip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch\n```\n\nInstall faiss following the [official guide](https://github.com/facebookresearch/faiss/blob/main/INSTALL.md)\n```bash\n# for example\nconda install -c conda-forge faiss-gpu\n```\n\n## Install AutoRecon\n```bash\ncd path/to/AutoRecon\npip install -e .\n# install tab completion\nns-install-cli\n```\n"
},
{
"path": "AutoReconForDens3R/exps/code-release/bmvs/scan1.sh",
"content": "\nscan_name=\"scan1\"\n\nDATA_ROOT=data\nINST_REL_DIR=BlendedMVS/$scan_name\nFORCE_RERUN=True\n\n# Coarse decomposition\npython third_party/AutoDecomp/auto_decomp/cli/inference_transformer.py --config-name=cvpr_idr \\\n data_root=$DATA_ROOT \\\n inst_rel_dir=$INST_REL_DIR \\\n triangulation.force_rerun=$FORCE_RERUN \\\n triangulation.n_feature_workers=1 triangulation.n_recon_workers=1 \\\n dino_feature.force_extract=$FORCE_RERUN dino_feature.n_workers=1\n\n# Reconstruction\nns-train neus-facto-wbg-reg_sep-plane-nerf \\\n --experiment-name neusfacto-wbg-reg_sep-plane-nerf_60k_plane-h-ratio-0.3_bmvs-${scan_name}_cvpr \\\n --vis tensorboard \\\n --trainer.steps_per_eval_image 2500 \\\n --trainer.steps_per_eval_batch 2500 \\\n --trainer.max_num_iterations 60001 \\\n --trainer.steps_per_save 60000 \\\n --pipeline.model.cos_anneal_end 10000 \\\n --pipeline.model.plane_height_ratio 0.3 \\\n --optimizers.fields.scheduler.max_steps 60001 \\\n --pipeline.datamanager.eval_camera_res_scale_factor 0.25 \\\n --pipeline.model.sdf_field.hash_grid_progressive_training False \\\n autorecon-data --data $DATA_ROOT/$INST_REL_DIR \\\n --anno_dirname 'triangulate_loftr-720000_pairs-from-poses_np-10/auto-deocomp_sfm-transformer_cvpr' \\\n --camera_filename cameras_cameras_norm-obj-side-2.0.npz \\\n --object_filename objects_cameras_norm-obj-side-2.0.npz \\\n --parse_images_from_camera_dict True \\\n --image_dirname \"images\" --image_extension \".jpg\" \\\n --include_image_features False --include_coarse_features False \\\n --use_accurate_scene_box True \\\n --collider_type 'box_near_far' \\\n --near_far 0.1 100.0 \\\n --compute_fg_bbox_mask True \\\n --force_recompute_fg_bbox_mask False \\\n --decomposition_mode regularization \\\n --downsample_ptcd -1\n\n# LOG_DIR=\"outputs/\"\n\n# MC_RES=512 # 1024 (default)\n# MESH_FN=\"extracted_mesh_res-${MC_RES}.ply\"\n# MESH_PATH=\"${LOG_DIR}/${MESH_FN}\"\n\n# ns-extract-mesh \\\n# \t--load-config $LOG_DIR/config.yml \\\n# --load-dir $LOG_DIR/sdfstudio_models \\\n# \t--output-path $MESH_PATH \\\n# --chunk_size 25000 --store_float16 True \\\n# --resolution $MC_RES \\\n# --use_train_scene_box True \\\n# --seg_aware_sdf False \\\n# --remove_internal_geometry None \\\n# --remove_non_maximum_connected_components True \\\n# --close_holes False --simplify_mesh_final False\n"
},
{
"path": "AutoReconForDens3R/exps/code-release/bmvs/scan2.sh",
"content": "\nscan_name=\"scan2\"\n\nDATA_ROOT=data\nINST_REL_DIR=BlendedMVS/$scan_name\nFORCE_RERUN=True\n\n# Coarse decomposition\npython third_party/AutoDecomp/auto_decomp/cli/inference_transformer.py --config-name=cvpr_idr \\\n data_root=$DATA_ROOT \\\n inst_rel_dir=$INST_REL_DIR \\\n triangulation.force_rerun=$FORCE_RERUN \\\n triangulation.n_feature_workers=1 triangulation.n_recon_workers=1 \\\n dino_feature.force_extract=$FORCE_RERUN dino_feature.n_workers=1\n\n# Reconstruction\nns-train neus-facto-wbg-reg_sep-plane-nerf \\\n --experiment-name neusfacto-wbg-reg_sep-plane-nerf_60k_plane-h-ratio-0.3_bmvs-${scan_name}_cvpr \\\n --vis tensorboard \\\n --trainer.steps_per_eval_image 2500 \\\n --trainer.steps_per_eval_batch 2500 \\\n --trainer.max_num_iterations 60001 \\\n --trainer.steps_per_save 60000 \\\n --pipeline.model.cos_anneal_end 10000 \\\n --pipeline.model.plane_height_ratio 0.3 \\\n --optimizers.fields.scheduler.max_steps 60001 \\\n --pipeline.datamanager.eval_camera_res_scale_factor 0.25 \\\n --pipeline.model.sdf_field.hash_grid_progressive_training False \\\n autorecon-data --data $DATA_ROOT/$INST_REL_DIR \\\n --anno_dirname 'triangulate_loftr-720000_pairs-from-poses_np-10/auto-deocomp_sfm-transformer_cvpr' \\\n --camera_filename cameras_cameras_norm-obj-side-2.0.npz \\\n --object_filename objects_cameras_norm-obj-side-2.0.npz \\\n --parse_images_from_camera_dict True \\\n --image_dirname \"images\" --image_extension \".jpg\" \\\n --include_image_features False --include_coarse_features False \\\n --use_accurate_scene_box True \\\n --collider_type 'box_near_far' \\\n --near_far 0.1 100.0 \\\n --compute_fg_bbox_mask True \\\n --force_recompute_fg_bbox_mask False \\\n --decomposition_mode regularization \\\n --downsample_ptcd -1\n\n# LOG_DIR=\"outputs/\"\n\n# MC_RES=512 # 1024 (default)\n# MESH_FN=\"extracted_mesh_res-${MC_RES}.ply\"\n# MESH_PATH=\"${LOG_DIR}/${MESH_FN}\"\n\n# ns-extract-mesh \\\n# \t--load-config $LOG_DIR/config.yml \\\n# --load-dir $LOG_DIR/sdfstudio_models \\\n# \t--output-path $MESH_PATH \\\n# --chunk_size 25000 --store_float16 True \\\n# --resolution $MC_RES \\\n# --use_train_scene_box True \\\n# --seg_aware_sdf False \\\n# --remove_internal_geometry None \\\n# --remove_non_maximum_connected_components False \\\n# --close_holes False --simplify_mesh_final False\n"
},
{
"path": "AutoReconForDens3R/exps/code-release/bmvs/scan3.sh",
"content": "\nscan_name=\"scan3\"\n\nDATA_ROOT=data\nINST_REL_DIR=BlendedMVS/$scan_name\nFORCE_RERUN=True\n\n# Coarse decomposition\npython third_party/AutoDecomp/auto_decomp/cli/inference_transformer.py --config-name=cvpr_idr \\\n data_root=$DATA_ROOT \\\n inst_rel_dir=$INST_REL_DIR \\\n triangulation.force_rerun=$FORCE_RERUN \\\n triangulation.n_feature_workers=1 triangulation.n_recon_workers=1 \\\n dino_feature.force_extract=$FORCE_RERUN dino_feature.n_workers=1\n\n# Reconstruction\nns-train neus-facto-wbg-reg_sep-plane-nerf \\\n --experiment-name neusfacto-wbg-reg_sep-plane-nerf_60k_plane-h-ratio-0.3_bmvs-${scan_name}_cvpr \\\n --vis tensorboard \\\n --trainer.steps_per_eval_image 2500 \\\n --trainer.steps_per_eval_batch 2500 \\\n --trainer.max_num_iterations 60001 \\\n --trainer.steps_per_save 60000 \\\n --pipeline.model.cos_anneal_end 10000 \\\n --pipeline.model.plane_height_ratio 0.3 \\\n --optimizers.fields.scheduler.max_steps 60001 \\\n --pipeline.datamanager.eval_camera_res_scale_factor 0.25 \\\n --pipeline.model.sdf_field.hash_grid_progressive_training False \\\n autorecon-data --data $DATA_ROOT/$INST_REL_DIR \\\n --anno_dirname 'triangulate_loftr-720000_pairs-from-poses_np-10/auto-deocomp_sfm-transformer_cvpr' \\\n --camera_filename cameras_cameras_norm-obj-side-2.0.npz \\\n --object_filename objects_cameras_norm-obj-side-2.0.npz \\\n --parse_images_from_camera_dict True \\\n --image_dirname \"images\" --image_extension \".jpg\" \\\n --include_image_features False --include_coarse_features False \\\n --use_accurate_scene_box True \\\n --collider_type 'box_near_far' \\\n --near_far 0.1 100.0 \\\n --compute_fg_bbox_mask True \\\n --force_recompute_fg_bbox_mask False \\\n --decomposition_mode regularization \\\n --downsample_ptcd -1\n\n# LOG_DIR=\"outputs/\"\n\n# MC_RES=512 # 1024 (default)\n# MESH_FN=\"extracted_mesh_res-${MC_RES}.ply\"\n# MESH_PATH=\"${LOG_DIR}/${MESH_FN}\"\n\n# ns-extract-mesh \\\n# \t--load-config $LOG_DIR/config.yml \\\n# --load-dir $LOG_DIR/sdfstudio_models \\\n# \t--output-path $MESH_PATH \\\n# --chunk_size 25000 --store_float16 True \\\n# --resolution $MC_RES \\\n# --use_train_scene_box True \\\n# --seg_aware_sdf False \\\n# --remove_internal_geometry None \\\n# --remove_non_maximum_connected_components False \\\n# --close_holes False --simplify_mesh_final False\n"
},
{
"path": "AutoReconForDens3R/exps/code-release/bmvs/scan4.sh",
"content": "\nscan_name=\"scan4\"\n\nDATA_ROOT=data\nINST_REL_DIR=BlendedMVS/$scan_name\nFORCE_RERUN=True\n\n# Coarse decomposition\npython third_party/AutoDecomp/auto_decomp/cli/inference_transformer.py --config-name=cvpr_idr \\\n data_root=$DATA_ROOT \\\n inst_rel_dir=$INST_REL_DIR \\\n triangulation.force_rerun=$FORCE_RERUN \\\n triangulation.n_feature_workers=1 triangulation.n_recon_workers=1 \\\n dino_feature.force_extract=$FORCE_RERUN dino_feature.n_workers=1\n\n# Reconstruction\nns-train neus-facto-wbg-reg_sep-plane-nerf \\\n --experiment-name neusfacto-wbg-reg_sep-plane-nerf_60k_plane-h-ratio-0.3_bmvs-${scan_name}_cvpr \\\n --vis tensorboard \\\n --trainer.steps_per_eval_image 2500 \\\n --trainer.steps_per_eval_batch 2500 \\\n --trainer.max_num_iterations 60001 \\\n --trainer.steps_per_save 60000 \\\n --pipeline.model.cos_anneal_end 10000 \\\n --pipeline.model.plane_height_ratio 0.3 \\\n --optimizers.fields.scheduler.max_steps 60001 \\\n --pipeline.datamanager.eval_camera_res_scale_factor 0.25 \\\n --pipeline.model.sdf_field.hash_grid_progressive_training False \\\n autorecon-data --data $DATA_ROOT/$INST_REL_DIR \\\n --anno_dirname 'triangulate_loftr-720000_pairs-from-poses_np-10/auto-deocomp_sfm-transformer_cvpr' \\\n --camera_filename cameras_cameras_norm-obj-side-2.0.npz \\\n --object_filename objects_cameras_norm-obj-side-2.0.npz \\\n --parse_images_from_camera_dict True \\\n --image_dirname \"images\" --image_extension \".jpg\" \\\n --include_image_features False --include_coarse_features False \\\n --use_accurate_scene_box True \\\n --collider_type 'box_near_far' \\\n --near_far 0.1 100.0 \\\n --compute_fg_bbox_mask True \\\n --force_recompute_fg_bbox_mask False \\\n --decomposition_mode regularization \\\n --downsample_ptcd -1\n\n# LOG_DIR=\"outputs/\"\n\n# MC_RES=512 # 1024 (default)\n# MESH_FN=\"extracted_mesh_res-${MC_RES}.ply\"\n# MESH_PATH=\"${LOG_DIR}/${MESH_FN}\"\n\n# ns-extract-mesh \\\n# \t--load-config $LOG_DIR/config.yml \\\n# --load-dir $LOG_DIR/sdfstudio_models \\\n# \t--output-path $MESH_PATH \\\n# --chunk_size 25000 --store_float16 True \\\n# --resolution $MC_RES \\\n# --use_train_scene_box True \\\n# --seg_aware_sdf False \\\n# --remove_internal_geometry None \\\n# --remove_non_maximum_connected_components False \\\n# --close_holes False --simplify_mesh_final False\n"
},
{
"path": "AutoReconForDens3R/exps/code-release/bmvs/scan5.sh",
"content": "\nscan_name=\"scan5\"\n\nDATA_ROOT=data\nINST_REL_DIR=BlendedMVS/$scan_name\nFORCE_RERUN=True\n\n# Coarse decomposition\npython third_party/AutoDecomp/auto_decomp/cli/inference_transformer.py --config-name=cvpr_idr \\\n data_root=$DATA_ROOT \\\n inst_rel_dir=$INST_REL_DIR \\\n triangulation.force_rerun=$FORCE_RERUN \\\n triangulation.n_feature_workers=1 triangulation.n_recon_workers=1 \\\n dino_feature.force_extract=$FORCE_RERUN dino_feature.n_workers=1\n\n# Reconstruction\nns-train neus-facto-wbg-reg_sep-plane-nerf \\\n --experiment-name neusfacto-wbg-reg_sep-plane-nerf_60k_plane-h-ratio-0.3_bmvs-${scan_name}_cvpr \\\n --vis tensorboard \\\n --trainer.steps_per_eval_image 2500 \\\n --trainer.steps_per_eval_batch 2500 \\\n --trainer.max_num_iterations 60001 \\\n --trainer.steps_per_save 60000 \\\n --pipeline.model.cos_anneal_end 10000 \\\n --pipeline.model.plane_height_ratio 0.3 \\\n --optimizers.fields.scheduler.max_steps 60001 \\\n --pipeline.datamanager.eval_camera_res_scale_factor 0.25 \\\n --pipeline.model.sdf_field.hash_grid_progressive_training False \\\n autorecon-data --data $DATA_ROOT/$INST_REL_DIR \\\n --anno_dirname 'triangulate_loftr-720000_pairs-from-poses_np-10/auto-deocomp_sfm-transformer_cvpr' \\\n --camera_filename cameras_cameras_norm-obj-side-2.0.npz \\\n --object_filename objects_cameras_norm-obj-side-2.0.npz \\\n --parse_images_from_camera_dict True \\\n --image_dirname \"images\" --image_extension \".jpg\" \\\n --include_image_features False --include_coarse_features False \\\n --use_accurate_scene_box True \\\n --collider_type 'box_near_far' \\\n --near_far 0.1 100.0 \\\n --compute_fg_bbox_mask True \\\n --force_recompute_fg_bbox_mask False \\\n --decomposition_mode regularization \\\n --downsample_ptcd -1\n\n# LOG_DIR=\"outputs/\"\n\n# MC_RES=512 # 1024 (default)\n# MESH_FN=\"extracted_mesh_res-${MC_RES}.ply\"\n# MESH_PATH=\"${LOG_DIR}/${MESH_FN}\"\n\n# ns-extract-mesh \\\n# \t--load-config $LOG_DIR/config.yml \\\n# --load-dir $LOG_DIR/sdfstudio_models \\\n# \t--output-path $MESH_PATH \\\n# --chunk_size 25000 --store_float16 True \\\n# --resolution $MC_RES \\\n# --use_train_scene_box True \\\n# --seg_aware_sdf False \\\n# --remove_internal_geometry None \\\n# --remove_non_maximum_connected_components False \\\n# --close_holes False --simplify_mesh_final False\n"
},
{
"path": "AutoReconForDens3R/exps/code-release/bmvs/scan6.sh",
"content": "\nscan_name=\"scan6\"\n\nDATA_ROOT=data\nINST_REL_DIR=BlendedMVS/$scan_name\nFORCE_RERUN=True\n\n# Coarse decomposition\npython third_party/AutoDecomp/auto_decomp/cli/inference_transformer.py --config-name=cvpr_idr \\\n data_root=$DATA_ROOT \\\n inst_rel_dir=$INST_REL_DIR \\\n triangulation.force_rerun=$FORCE_RERUN \\\n triangulation.n_feature_workers=1 triangulation.n_recon_workers=1 \\\n dino_feature.force_extract=$FORCE_RERUN dino_feature.n_workers=1\n\n# Reconstruction\nns-train neus-facto-wbg-reg_sep-plane-nerf \\\n --experiment-name neusfacto-wbg-reg_sep-plane-nerf_60k_plane-h-ratio-0.3_bmvs-${scan_name}_cvpr \\\n --vis tensorboard \\\n --trainer.steps_per_eval_image 2500 \\\n --trainer.steps_per_eval_batch 2500 \\\n --trainer.max_num_iterations 60001 \\\n --trainer.steps_per_save 60000 \\\n --pipeline.model.cos_anneal_end 10000 \\\n --pipeline.model.plane_height_ratio 0.3 \\\n --optimizers.fields.scheduler.max_steps 60001 \\\n --pipeline.datamanager.eval_camera_res_scale_factor 0.25 \\\n --pipeline.model.sdf_field.hash_grid_progressive_training False \\\n autorecon-data --data $DATA_ROOT/$INST_REL_DIR \\\n --anno_dirname 'triangulate_loftr-720000_pairs-from-poses_np-10/auto-deocomp_sfm-transformer_cvpr' \\\n --camera_filename cameras_cameras_norm-obj-side-2.0.npz \\\n --object_filename objects_cameras_norm-obj-side-2.0.npz \\\n --parse_images_from_camera_dict True \\\n --image_dirname \"images\" --image_extension \".jpg\" \\\n --include_image_features False --include_coarse_features False \\\n --use_accurate_scene_box True \\\n --collider_type 'box_near_far' \\\n --near_far 0.1 100.0 \\\n --compute_fg_bbox_mask True \\\n --force_recompute_fg_bbox_mask False \\\n --decomposition_mode regularization \\\n --downsample_ptcd -1\n\n# LOG_DIR=\"outputs/\"\n\n# MC_RES=512 # 1024 (default)\n# MESH_FN=\"extracted_mesh_res-${MC_RES}.ply\"\n# MESH_PATH=\"${LOG_DIR}/${MESH_FN}\"\n\n# ns-extract-mesh \\\n# \t--load-config $LOG_DIR/config.yml \\\n# --load-dir $LOG_DIR/sdfstudio_models \\\n# \t--output-path $MESH_PATH \\\n# --chunk_size 25000 --store_float16 True \\\n# --resolution $MC_RES \\\n# --use_train_scene_box True \\\n# --seg_aware_sdf False \\\n# --remove_internal_geometry None \\\n# --remove_non_maximum_connected_components False \\\n# --close_holes False --simplify_mesh_final False\n"
},
{
"path": "AutoReconForDens3R/exps/code-release/co3d_demo/scan1.sh",
"content": "\nscan_name=scan1\n\nDATA_ROOT=data\nINST_REL_DIR=CO3D_DEMO/$scan_name\nFORCE_RERUN=True\n\npython third_party/AutoDecomp/auto_decomp/cli/inference_transformer.py --config-name=cvpr \\\n data_root=$DATA_ROOT \\\n inst_rel_dir=$INST_REL_DIR \\\n sparse_recon.n_images=40 \\\n sparse_recon.force_rerun=$FORCE_RERUN \\\n sparse_recon.n_feature_workers=1 sparse_recon.n_recon_workers=1 \\\n triangulation.force_rerun=$FORCE_RERUN \\\n triangulation.n_feature_workers=1 triangulation.n_recon_workers=1 \\\n dino_feature.force_extract=$FORCE_RERUN dino_feature.n_workers=1 \n\nns-train neus-facto-wbg-reg_sep-plane-nerf \\\n --experiment-name neusfacto-wbg-reg_sep-plane-nerf_60k_plane-h-ratio-0.3_co3d-${scan_name}_cvpr \\\n --vis tensorboard \\\n --trainer.steps_per_eval_image 2500 \\\n --trainer.steps_per_eval_batch 2500 \\\n --trainer.max_num_iterations 60001 \\\n --trainer.steps_per_save 60000 \\\n --pipeline.model.cos_anneal_end 10000 \\\n --pipeline.model.plane_height_ratio 0.3 \\\n --optimizers.fields.scheduler.max_steps 60001 \\\n --pipeline.datamanager.eval_camera_res_scale_factor 0.25 \\\n --pipeline.model.sdf_field.hash_grid_progressive_training False \\\n autorecon-data --data $DATA_ROOT/$INST_REL_DIR \\\n --anno_dirname 'triangulate_loftr-720000_sequential_np-10/auto-deocomp_sfm-transformer_cvpr' \\\n --camera_filename cameras_cameras_norm-obj-side-2.0.npz \\\n --object_filename objects_cameras_norm-obj-side-2.0.npz \\\n --parse_images_from_camera_dict True \\\n --image_dirname \"images\" --image_extension \".jpg\" \\\n --include_image_features False --include_coarse_features False \\\n --use_accurate_scene_box True \\\n --collider_type 'box_near_far' \\\n --near_far 0.1 100.0 \\\n --compute_fg_bbox_mask True \\\n --force_recompute_fg_bbox_mask False \\\n --decomposition_mode regularization \\\n --downsample_ptcd -1\n\n# LOG_DIR=\"outputs/\"\n\n# MC_RES=512 # 1024 (default)\n# MESH_FN=\"extracted_mesh_res-${MC_RES}.ply\"\n# MESH_PATH=\"${LOG_DIR}/${MESH_FN}\"\n\n# ns-extract-mesh \\\n# \t--load-config $LOG_DIR/config.yml \\\n# --load-dir $LOG_DIR/sdfstudio_models \\\n# \t--output-path $MESH_PATH \\\n# --chunk_size 25000 --store_float16 True \\\n# --resolution $MC_RES \\\n# --use_train_scene_box True \\\n# --seg_aware_sdf False \\\n# --remove_internal_geometry None \\\n# --remove_non_maximum_connected_components True \\\n# --close_holes False --simplify_mesh_final False\n"
},
{
"path": "AutoReconForDens3R/exps/code-release/co3d_demo/scan2.sh",
"content": "\nscan_name=scan2\n\nDATA_ROOT=data\nINST_REL_DIR=CO3D_DEMO/$scan_name\nFORCE_RERUN=True\n\npython third_party/AutoDecomp/auto_decomp/cli/inference_transformer.py --config-name=cvpr \\\n data_root=$DATA_ROOT \\\n inst_rel_dir=$INST_REL_DIR \\\n sparse_recon.n_images=40 \\\n sparse_recon.force_rerun=$FORCE_RERUN \\\n sparse_recon.n_feature_workers=1 sparse_recon.n_recon_workers=1 \\\n triangulation.force_rerun=$FORCE_RERUN \\\n triangulation.n_feature_workers=1 triangulation.n_recon_workers=1 \\\n dino_feature.force_extract=$FORCE_RERUN dino_feature.n_workers=1 \n\nns-train neus-facto-wbg-reg_sep-plane-nerf \\\n --experiment-name neusfacto-wbg-reg_sep-plane-nerf_60k_plane-h-ratio-0.3_co3d-${scan_name}_cvpr \\\n --vis tensorboard \\\n --trainer.steps_per_eval_image 2500 \\\n --trainer.steps_per_eval_batch 2500 \\\n --trainer.max_num_iterations 60001 \\\n --trainer.steps_per_save 60000 \\\n --pipeline.model.cos_anneal_end 10000 \\\n --pipeline.model.plane_height_ratio 0.3 \\\n --optimizers.fields.scheduler.max_steps 60001 \\\n --pipeline.datamanager.eval_camera_res_scale_factor 0.25 \\\n --pipeline.model.sdf_field.hash_grid_progressive_training False \\\n autorecon-data --data $DATA_ROOT/$INST_REL_DIR \\\n --anno_dirname 'triangulate_loftr-720000_sequential_np-10/auto-deocomp_sfm-transformer_cvpr' \\\n --camera_filename cameras_cameras_norm-obj-side-2.0.npz \\\n --object_filename objects_cameras_norm-obj-side-2.0.npz \\\n --parse_images_from_camera_dict True \\\n --image_dirname \"images\" --image_extension \".jpg\" \\\n --include_image_features False --include_coarse_features False \\\n --use_accurate_scene_box True \\\n --collider_type 'box_near_far' \\\n --near_far 0.1 100.0 \\\n --compute_fg_bbox_mask True \\\n --force_recompute_fg_bbox_mask False \\\n --decomposition_mode regularization \\\n --downsample_ptcd -1\n\n# LOG_DIR=\"outputs/\"\n\n# MC_RES=512 # 1024 (default)\n# MESH_FN=\"extracted_mesh_res-${MC_RES}.ply\"\n# MESH_PATH=\"${LOG_DIR}/${MESH_FN}\"\n\n# ns-extract-mesh \\\n# \t--load-config $LOG_DIR/config.yml \\\n# --load-dir $LOG_DIR/sdfstudio_models \\\n# \t--output-path $MESH_PATH \\\n# --chunk_size 25000 --store_float16 True \\\n# --resolution $MC_RES \\\n# --use_train_scene_box True \\\n# --seg_aware_sdf False \\\n# --remove_internal_geometry None \\\n# --remove_non_maximum_connected_components True \\\n# --close_holes False --simplify_mesh_final False\n"
},
{
"path": "AutoReconForDens3R/exps/code-release/co3d_demo/scan3.sh",
"content": "\nscan_name=scan3\n\nDATA_ROOT=data\nINST_REL_DIR=CO3D_DEMO/$scan_name\nFORCE_RERUN=True\n\npython third_party/AutoDecomp/auto_decomp/cli/inference_transformer.py --config-name=cvpr \\\n data_root=$DATA_ROOT \\\n inst_rel_dir=$INST_REL_DIR \\\n sparse_recon.n_images=40 \\\n sparse_recon.force_rerun=$FORCE_RERUN \\\n sparse_recon.n_feature_workers=1 sparse_recon.n_recon_workers=1 \\\n triangulation.force_rerun=$FORCE_RERUN \\\n triangulation.n_feature_workers=1 triangulation.n_recon_workers=1 \\\n dino_feature.force_extract=$FORCE_RERUN dino_feature.n_workers=1 \n\nns-train neus-facto-wbg-reg_sep-plane-nerf \\\n --experiment-name neusfacto-wbg-reg_sep-plane-nerf_60k_plane-h-ratio-0.3_co3d-${scan_name}_cvpr \\\n --vis tensorboard \\\n --trainer.steps_per_eval_image 2500 \\\n --trainer.steps_per_eval_batch 2500 \\\n --trainer.max_num_iterations 60001 \\\n --trainer.steps_per_save 60000 \\\n --pipeline.model.cos_anneal_end 10000 \\\n --pipeline.model.plane_height_ratio 0.3 \\\n --optimizers.fields.scheduler.max_steps 60001 \\\n --pipeline.datamanager.eval_camera_res_scale_factor 0.25 \\\n --pipeline.model.sdf_field.hash_grid_progressive_training False \\\n autorecon-data --data $DATA_ROOT/$INST_REL_DIR \\\n --anno_dirname 'triangulate_loftr-720000_sequential_np-10/auto-deocomp_sfm-transformer_cvpr' \\\n --camera_filename cameras_cameras_norm-obj-side-2.0.npz \\\n --object_filename objects_cameras_norm-obj-side-2.0.npz \\\n --parse_images_from_camera_dict True \\\n --image_dirname \"images\" --image_extension \".jpg\" \\\n --include_image_features False --include_coarse_features False \\\n --use_accurate_scene_box True \\\n --collider_type 'box_near_far' \\\n --near_far 0.1 100.0 \\\n --compute_fg_bbox_mask True \\\n --force_recompute_fg_bbox_mask False \\\n --decomposition_mode regularization \\\n --downsample_ptcd -1\n\n# LOG_DIR=\"outputs/\"\n\n# MC_RES=512 # 1024 (default)\n# MESH_FN=\"extracted_mesh_res-${MC_RES}.ply\"\n# MESH_PATH=\"${LOG_DIR}/${MESH_FN}\"\n\n# ns-extract-mesh \\\n# \t--load-config $LOG_DIR/config.yml \\\n# --load-dir $LOG_DIR/sdfstudio_models \\\n# \t--output-path $MESH_PATH \\\n# --chunk_size 25000 --store_float16 True \\\n# --resolution $MC_RES \\\n# --use_train_scene_box True \\\n# --seg_aware_sdf False \\\n# --remove_internal_geometry None \\\n# --remove_non_maximum_connected_components True \\\n# --close_holes False --simplify_mesh_final False\n"
},
{
"path": "AutoReconForDens3R/exps/code-release/co3d_demo/scan4.sh",
"content": "\nscan_name=scan4\n\nDATA_ROOT=data\nINST_REL_DIR=CO3D_DEMO/$scan_name\nFORCE_RERUN=True\n\npython third_party/AutoDecomp/auto_decomp/cli/inference_transformer.py --config-name=cvpr \\\n data_root=$DATA_ROOT \\\n inst_rel_dir=$INST_REL_DIR \\\n sparse_recon.n_images=40 \\\n sparse_recon.force_rerun=$FORCE_RERUN \\\n sparse_recon.n_feature_workers=1 sparse_recon.n_recon_workers=1 \\\n triangulation.force_rerun=$FORCE_RERUN \\\n triangulation.n_feature_workers=1 triangulation.n_recon_workers=1 \\\n dino_feature.force_extract=$FORCE_RERUN dino_feature.n_workers=1 \n\nns-train neus-facto-wbg-reg_sep-plane-nerf \\\n --experiment-name neusfacto-wbg-reg_sep-plane-nerf_60k_plane-h-ratio-0.3_co3d-${scan_name}_cvpr \\\n --vis tensorboard \\\n --trainer.steps_per_eval_image 2500 \\\n --trainer.steps_per_eval_batch 2500 \\\n --trainer.max_num_iterations 60001 \\\n --trainer.steps_per_save 60000 \\\n --pipeline.model.cos_anneal_end 10000 \\\n --pipeline.model.plane_height_ratio 0.3 \\\n --optimizers.fields.scheduler.max_steps 60001 \\\n --pipeline.datamanager.eval_camera_res_scale_factor 0.25 \\\n --pipeline.model.sdf_field.hash_grid_progressive_training False \\\n autorecon-data --data $DATA_ROOT/$INST_REL_DIR \\\n --anno_dirname 'triangulate_loftr-720000_sequential_np-10/auto-deocomp_sfm-transformer_cvpr' \\\n --camera_filename cameras_cameras_norm-obj-side-2.0.npz \\\n --object_filename objects_cameras_norm-obj-side-2.0.npz \\\n --parse_images_from_camera_dict True \\\n --image_dirname \"images\" --image_extension \".jpg\" \\\n --include_image_features False --include_coarse_features False \\\n --use_accurate_scene_box True \\\n --collider_type 'box_near_far' \\\n --near_far 0.1 100.0 \\\n --compute_fg_bbox_mask True \\\n --force_recompute_fg_bbox_mask False \\\n --decomposition_mode regularization \\\n --downsample_ptcd -1\n\n# LOG_DIR=\"outputs/\"\n\n# MC_RES=512 # 1024 (default)\n# MESH_FN=\"extracted_mesh_res-${MC_RES}.ply\"\n# MESH_PATH=\"${LOG_DIR}/${MESH_FN}\"\n\n# ns-extract-mesh \\\n# \t--load-config $LOG_DIR/config.yml \\\n# --load-dir $LOG_DIR/sdfstudio_models \\\n# \t--output-path $MESH_PATH \\\n# --chunk_size 25000 --store_float16 True \\\n# --resolution $MC_RES \\\n# --use_train_scene_box True \\\n# --seg_aware_sdf False \\\n# --remove_internal_geometry None \\\n# --remove_non_maximum_connected_components True \\\n# --close_holes False --simplify_mesh_final False\n"
},
{
"path": "AutoReconForDens3R/exps/code-release/co3d_demo/scan5.sh",
"content": "\nscan_name=scan5\n\nDATA_ROOT=data\nINST_REL_DIR=CO3D_DEMO/$scan_name\nFORCE_RERUN=True\n\npython third_party/AutoDecomp/auto_decomp/cli/inference_transformer.py --config-name=cvpr \\\n data_root=$DATA_ROOT \\\n inst_rel_dir=$INST_REL_DIR \\\n sparse_recon.n_images=40 \\\n sparse_recon.force_rerun=$FORCE_RERUN \\\n sparse_recon.n_feature_workers=1 sparse_recon.n_recon_workers=1 \\\n triangulation.force_rerun=$FORCE_RERUN \\\n triangulation.n_feature_workers=1 triangulation.n_recon_workers=1 \\\n dino_feature.force_extract=$FORCE_RERUN dino_feature.n_workers=1 \n\nns-train neus-facto-wbg-reg_sep-plane-nerf \\\n --experiment-name neusfacto-wbg-reg_sep-plane-nerf_60k_plane-h-ratio-0.3_co3d-${scan_name}_cvpr \\\n --vis tensorboard \\\n --trainer.steps_per_eval_image 2500 \\\n --trainer.steps_per_eval_batch 2500 \\\n --trainer.max_num_iterations 60001 \\\n --trainer.steps_per_save 60000 \\\n --pipeline.model.cos_anneal_end 10000 \\\n --pipeline.model.plane_height_ratio 0.3 \\\n --optimizers.fields.scheduler.max_steps 60001 \\\n --pipeline.datamanager.eval_camera_res_scale_factor 0.25 \\\n --pipeline.model.sdf_field.hash_grid_progressive_training False \\\n autorecon-data --data $DATA_ROOT/$INST_REL_DIR \\\n --anno_dirname 'triangulate_loftr-720000_sequential_np-10/auto-deocomp_sfm-transformer_cvpr' \\\n --camera_filename cameras_cameras_norm-obj-side-2.0.npz \\\n --object_filename objects_cameras_norm-obj-side-2.0.npz \\\n --parse_images_from_camera_dict True \\\n --image_dirname \"images\" --image_extension \".jpg\" \\\n --include_image_features False --include_coarse_features False \\\n --use_accurate_scene_box True \\\n --collider_type 'box_near_far' \\\n --near_far 0.1 100.0 \\\n --compute_fg_bbox_mask True \\\n --force_recompute_fg_bbox_mask False \\\n --decomposition_mode regularization \\\n --downsample_ptcd -1\n\n# LOG_DIR=\"outputs/\"\n\n# MC_RES=512 # 1024 (default)\n# MESH_FN=\"extracted_mesh_res-${MC_RES}.ply\"\n# MESH_PATH=\"${LOG_DIR}/${MESH_FN}\"\n\n# ns-extract-mesh \\\n# \t--load-config $LOG_DIR/config.yml \\\n# --load-dir $LOG_DIR/sdfstudio_models \\\n# \t--output-path $MESH_PATH \\\n# --chunk_size 25000 --store_float16 True \\\n# --resolution $MC_RES \\\n# --use_train_scene_box True \\\n# --seg_aware_sdf False \\\n# --remove_internal_geometry None \\\n# --remove_non_maximum_connected_components True \\\n# --close_holes False --simplify_mesh_final False\n"
},
{
"path": "AutoReconForDens3R/exps/code-release/run_dens3r_recon.sh",
"content": "\n#/bin/bash\n\nDATA_ROOT=/path/to/your/data\nINST_REL_DIR=rel_path/to/your/data\nDENS3R_MODEL_PATH=/path/to/pth/file/of/model\nN_IMAGES=40\nNEUS_TRAIN_RESO_DOWN_FACTOR=0.5\n\nFORCE_RERUN=True\nEXP_NAME=neusfacto-wbg-reg_sep-plane-nerf_60k_plane-h-ratio-0.3_demo_cvpr\nNERF_MODE=neus-facto-wbg-reg_sep-plane-nerf\n\n\n# reconstruct scene with dens3r\n# python third_party/AutoDecomp/auto_decomp/cli/inference_transformer.py \\\n# --config-name=cvpr \\\n# data_root=$DATA_ROOT \\\n# inst_rel_dir=$INST_REL_DIR \\\n# sparse_recon.n_images=$N_IMAGES \\\n# sparse_recon.force_rerun=$FORCE_RERUN \\\n# sparse_recon.n_feature_workers=1 sparse_recon.n_recon_workers=1 \\\n# triangulation.force_rerun=$FORCE_RERUN \\\n# triangulation.n_feature_workers=1 triangulation.n_recon_workers=1 \\\n# dino_feature.force_extract=$FORCE_RERUN dino_feature.n_workers=1 \\\n# sfm_mode=dens3r \\\n# sparse_recon.dens3r_model_ckpt=$DENS3R_MODEL_PATH\n\n\n# reconstruct object with AutoRecon\n# python scripts/train.py \\\n# $NERF_MODE \\\n# --experiment-name $EXP_NAME \\\n# --vis tensorboard \\\n# --trainer.steps_per_eval_image 2500 \\\n# --trainer.steps_per_eval_batch 2500 \\\n# --trainer.max_num_iterations 60001 \\\n# --trainer.steps_per_save 60000 \\\n# --pipeline.model.cos_anneal_end 10000 \\\n# --pipeline.model.plane_height_ratio 0.3 \\\n# --optimizers.fields.scheduler.max_steps 60001 \\\n# --pipeline.datamanager.eval_camera_res_scale_factor 0.25 \\\n# --pipeline.datamanager.camera_res_scale_factor $NEUS_TRAIN_RESO_DOWN_FACTOR \\\n# --pipeline.model.sdf_field.hash_grid_progressive_training False \\\n# autorecon-data --data $DATA_ROOT/$INST_REL_DIR \\\n# --anno_dirname 'dens3r_40_sequential/auto-deocomp_sfm-transformer_cvpr' \\\n# --camera_filename cameras_cameras_norm-obj-side-2.0.npz \\\n# --object_filename objects_cameras_norm-obj-side-2.0.npz \\\n# --parse_images_from_camera_dict True \\\n# --image_dirname \"images\" --image_extension \".jpg\" \\\n# --include_image_features False --include_coarse_features False \\\n# --use_accurate_scene_box True \\\n# --collider_type 'box_near_far' \\\n# --near_far 0.1 100.0 \\\n# --compute_fg_bbox_mask True \\\n# --force_recompute_fg_bbox_mask False \\\n# --decomposition_mode regularization \\\n# --downsample_ptcd -1\n\n\n# Extract mesh with MC\n\n# LOG_DIR=\"outputs/$EXP_NAME/$NERF_MODE/data\" # default is outputs/$EXP_NAME/$NERF_MODE/date\n# MC_RES=512\n# MESH_FN=\"extracted_mesh_res-${MC_RES}.ply\"\n# MESH_PATH=\"${LOG_DIR}/sdfstudio_models/${MESH_FN}\"\n\n# python scripts/extract_mesh.py \\\n# --load-config $LOG_DIR/config.yml \\\n# --load-dir $LOG_DIR/sdfstudio_models \\\n# --output-path $MESH_PATH \\\n# --chunk_size 25000 --store_float16 True \\\n# --resolution $MC_RES \\\n# --use_train_scene_box True \\\n# --seg_aware_sdf False \\\n# --remove_internal_geometry None \\\n# --remove_non_maximum_connected_components True \\\n# --close_holes False --simplify_mesh_final False"
},
{
"path": "AutoReconForDens3R/exps/code-release/run_pipeline_demo_low-res.sh",
"content": "\nDATA_ROOT=data\nINST_REL_DIR=custom_data_example/co3d_chair\nFORCE_RERUN=True\n\n# Coarse decomposition\npython third_party/AutoDecomp/auto_decomp/cli/inference_transformer.py --config-name=cvpr \\\n data_root=$DATA_ROOT \\\n inst_rel_dir=$INST_REL_DIR \\\n sparse_recon.n_images=40 \\\n sparse_recon.force_rerun=$FORCE_RERUN \\\n sparse_recon.n_feature_workers=1 sparse_recon.n_recon_workers=1 \\\n triangulation.force_rerun=$FORCE_RERUN \\\n triangulation.n_feature_workers=1 triangulation.n_recon_workers=1 \\\n dino_feature.force_extract=$FORCE_RERUN dino_feature.n_workers=1 \n\n# reconstruction\n# TODO: parse anno_dirname & object_filename from cache\nns-train neus-facto-wbg-reg_sep-plane-nerf \\\n --experiment-name neusfacto-wbg-reg_sep-plane-nerf_60k_plane-h-ratio-0.3_demo_cvpr \\\n --vis tensorboard \\\n --trainer.steps_per_eval_image 2500 \\\n --trainer.steps_per_eval_batch 2500 \\\n --trainer.max_num_iterations 60001 \\\n --trainer.steps_per_save 60000 \\\n --pipeline.model.cos_anneal_end 10000 \\\n --pipeline.model.plane_height_ratio 0.3 \\\n --optimizers.fields.scheduler.max_steps 60001 \\\n --pipeline.datamanager.eval_camera_res_scale_factor 0.25 \\\n --pipeline.model.sdf_field.hash_grid_progressive_training False \\\n autorecon-data --data $DATA_ROOT/$INST_REL_DIR \\\n --anno_dirname 'triangulate_loftr-720000_sequential_np-10/auto-deocomp_sfm-transformer_cvpr' \\\n --camera_filename cameras_cameras_norm-obj-side-2.0.npz \\\n --object_filename objects_cameras_norm-obj-side-2.0.npz \\\n --parse_images_from_camera_dict True \\\n --image_dirname \"images\" --image_extension \".jpg\" \\\n --include_image_features False --include_coarse_features False \\\n --use_accurate_scene_box True \\\n --collider_type 'box_near_far' \\\n --near_far 0.1 100.0 \\\n --compute_fg_bbox_mask True \\\n --force_recompute_fg_bbox_mask False \\\n --decomposition_mode regularization \\\n --downsample_ptcd -1\n\n# Extract mesh with MC\n# LOG_DIR=\"path_to_log_dir\"\n\n# MC_RES=512\n# MESH_FN=\"extracted_mesh_res-${MC_RES}.ply\"\n# MESH_PATH=\"${LOG_DIR}/${MESH_FN}\"\n\n# ns-extract-mesh \\\n# \t--load-config $LOG_DIR/config.yml \\\n# --load-dir $LOG_DIR/sdfstudio_models \\\n# \t--output-path $MESH_PATH \\\n# --chunk_size 25000 --store_float16 True \\\n# --resolution $MC_RES \\\n# --use_train_scene_box True \\\n# --seg_aware_sdf False \\\n# --remove_internal_geometry None \\\n# --remove_non_maximum_connected_components True \\\n# --close_holes False --simplify_mesh_final False\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/__init__.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/cameras/__init__.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/cameras/camera_optimizers.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nPose and Intrinsics Optimizers\n\"\"\"\n\nfrom __future__ import annotations\n\nimport functools\nfrom dataclasses import dataclass, field\nfrom typing import Type, Union\n\nimport torch\nimport tyro\nfrom torch import nn\nfrom torchtyping import TensorType\nfrom typing_extensions import Literal, assert_never\n\nfrom nerfstudio.cameras.lie_groups import exp_map_SE3, exp_map_SO3xR3\nfrom nerfstudio.configs import base_config as cfg\nfrom nerfstudio.engine.optimizers import AdamOptimizerConfig\nfrom nerfstudio.engine.schedulers import SchedulerConfig\nfrom nerfstudio.utils import poses as pose_utils\n\n\n@dataclass\nclass CameraOptimizerConfig(cfg.InstantiateConfig):\n \"\"\"Configuration of optimization for camera poses.\"\"\"\n\n _target: Type = field(default_factory=lambda: CameraOptimizer)\n\n mode: Literal[\"off\", \"SO3xR3\", \"SE3\"] = \"off\"\n \"\"\"Pose optimization strategy to use. If enabled, we recommend SO3xR3.\"\"\"\n\n position_noise_std: float = 0.0\n \"\"\"Noise to add to initial positions. Useful for debugging.\"\"\"\n\n orientation_noise_std: float = 0.0\n \"\"\"Noise to add to initial orientations. Useful for debugging.\"\"\"\n\n optimizer: AdamOptimizerConfig = AdamOptimizerConfig(lr=6e-4, eps=1e-15)\n \"\"\"ADAM parameters for camera optimization.\"\"\"\n\n scheduler: SchedulerConfig = SchedulerConfig(max_steps=10000)\n \"\"\"Learning rate scheduler for camera optimizer..\"\"\"\n\n param_group: tyro.conf.Suppress[str] = \"camera_opt\"\n \"\"\"Name of the parameter group used for pose optimization. Can be any string that doesn't conflict with other\n groups.\"\"\"\n\n\nclass CameraOptimizer(nn.Module):\n \"\"\"Layer that modifies camera poses to be optimized as well as the field during training.\"\"\"\n\n config: CameraOptimizerConfig\n\n def __init__(\n self,\n config: CameraOptimizerConfig,\n num_cameras: int,\n device: Union[torch.device, str],\n **kwargs, # pylint: disable=unused-argument\n ) -> None:\n super().__init__()\n self.config = config\n self.num_cameras = num_cameras\n self.device = device\n\n # Initialize learnable parameters.\n if self.config.mode == \"off\":\n pass\n elif self.config.mode in (\"SO3xR3\", \"SE3\"):\n self.pose_adjustment = torch.nn.Parameter(torch.zeros((num_cameras, 6), device=device))\n else:\n assert_never(self.config.mode)\n\n # Initialize pose noise; useful for debugging.\n if config.position_noise_std != 0.0 or config.orientation_noise_std != 0.0:\n assert config.position_noise_std >= 0.0 and config.orientation_noise_std >= 0.0\n std_vector = torch.tensor(\n [config.position_noise_std] * 3 + [config.orientation_noise_std] * 3, device=device\n )\n self.pose_noise = exp_map_SE3(torch.normal(torch.zeros((num_cameras, 6), device=device), std_vector))\n else:\n self.pose_noise = None\n\n def forward(\n self,\n indices: TensorType[\"num_cameras\"],\n ) -> TensorType[\"num_cameras\", 3, 4]:\n \"\"\"Indexing into camera adjustments.\n Args:\n indices: indices of Cameras to optimize.\n Returns:\n Tranformation matrices from optimized camera coordinates coordinates\n to given camera coordinates.\n \"\"\"\n outputs = []\n\n # Apply learned transformation delta.\n if self.config.mode == \"off\":\n pass\n elif self.config.mode == \"SO3xR3\":\n outputs.append(exp_map_SO3xR3(self.pose_adjustment[indices, :]))\n elif self.config.mode == \"SE3\":\n outputs.append(exp_map_SE3(self.pose_adjustment[indices, :]))\n else:\n assert_never(self.config.mode)\n\n # Apply initial pose noise.\n if self.pose_noise is not None:\n outputs.append(self.pose_noise[indices, :, :])\n\n # Return: identity if no transforms are needed, otherwise multiply transforms together.\n if len(outputs) == 0:\n # Note that using repeat() instead of tile() here would result in unnecessary copies.\n return torch.eye(4, device=self.device)[None, :3, :4].tile(indices.shape[0], 1, 1)\n return functools.reduce(pose_utils.multiply, outputs)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/cameras/camera_paths.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nCode for camera paths.\n\"\"\"\n\nfrom typing import Any, Dict, Optional, Tuple\n\nimport torch\n\nimport nerfstudio.utils.poses as pose_utils\nfrom nerfstudio.cameras import camera_utils\nfrom nerfstudio.cameras.camera_utils import get_interpolated_poses_many\nfrom nerfstudio.cameras.cameras import Cameras\nfrom nerfstudio.viewer.server.utils import three_js_perspective_camera_focal_length\n\n\ndef get_interpolated_camera_path(cameras: Cameras, steps: int) -> Cameras:\n \"\"\"Generate a camera path between two cameras.\n\n Args:\n cameras: Cameras object containing intrinsics of all cameras.\n steps: The number of steps to interpolate between the two cameras.\n\n Returns:\n A new set of cameras along a path.\n \"\"\"\n Ks = cameras.get_intrinsics_matrices().cpu().numpy()\n poses = cameras.camera_to_worlds().cpu().numpy()\n poses, Ks = get_interpolated_poses_many(poses, Ks, steps_per_transition=steps)\n\n cameras = Cameras(fx=Ks[:, 0, 0], fy=Ks[:, 1, 1], cx=Ks[0, 0, 2], cy=Ks[0, 1, 2], camera_to_worlds=poses)\n return cameras\n\n\ndef get_spiral_path(\n camera: Cameras,\n steps: int = 30,\n radius: Optional[float] = None,\n radiuses: Optional[Tuple[float]] = None,\n rots: int = 2,\n zrate: float = 0.5,\n) -> Cameras:\n \"\"\"\n Returns a list of camera in a sprial trajectory.\n\n Args:\n camera: The camera to start the spiral from.\n steps: The number of cameras in the generated path.\n radius: The radius of the spiral for all xyz directions.\n radiuses: The list of radii for the spiral in xyz directions.\n rots: The number of rotations to apply to the camera.\n zrate: How much to change the z position of the camera.\n\n Returns:\n A spiral camera path.\n \"\"\"\n\n assert radius is not None or radiuses is not None, \"Either radius or radiuses must be specified.\"\n assert camera.ndim == 1, \"We assume only one batch dim here\"\n if radius is not None and radiuses is None:\n rad = torch.tensor([radius] * 3, device=camera.device)\n elif radiuses is not None and radius is None:\n rad = torch.tensor(radiuses, device=camera.device)\n else:\n raise ValueError(\"Only one of radius or radiuses must be specified.\")\n\n camera = camera.flatten()\n\n up = camera.camera_to_worlds[0, :3, 2] # scene is z up\n focal = torch.min(camera.fx[0], camera.fy[0])\n target = torch.tensor([0, 0, -focal], device=camera.device) # camera looking in -z direction\n\n c2w = camera.camera_to_worlds[0]\n c2wh_global = pose_utils.to4x4(c2w)\n\n local_c2whs = []\n for theta in torch.linspace(0.0, 2.0 * torch.pi * rots, steps + 1)[:-1]:\n center = (\n torch.tensor([torch.cos(theta), -torch.sin(theta), -torch.sin(theta * zrate)], device=camera.device) * rad\n )\n lookat = center - target\n c2w = camera_utils.viewmatrix(lookat, up, center)\n c2wh = pose_utils.to4x4(c2w)\n local_c2whs.append(c2wh)\n\n new_c2ws = []\n for local_c2wh in local_c2whs:\n c2wh = torch.matmul(c2wh_global, local_c2wh)\n new_c2ws.append(c2wh[:3, :4])\n new_c2ws = torch.stack(new_c2ws, dim=0)\n\n return Cameras(\n fx=camera.fx[0],\n fy=camera.fy[0],\n cx=camera.cx[0],\n cy=camera.cy[0],\n camera_to_worlds=new_c2ws,\n )\n\n\ndef get_path_from_json(camera_path: Dict[str, Any]) -> Cameras:\n \"\"\"Takes a camera path dictionary and returns a trajectory as a Camera instance.\n\n Args:\n camera_path: A dictionary of the camera path information coming from the viewer.\n\n Returns:\n A Cameras instance with the camera path.\n \"\"\"\n\n image_height = camera_path[\"render_height\"]\n image_width = camera_path[\"render_width\"]\n\n c2ws = []\n fxs = []\n fys = []\n for camera in camera_path[\"camera_path\"]:\n # pose\n c2w = torch.tensor(camera[\"camera_to_world\"]).view(4, 4)[:3]\n c2ws.append(c2w)\n # field of view\n fov = camera[\"fov\"]\n focal_length = three_js_perspective_camera_focal_length(fov, image_height)\n fxs.append(focal_length)\n fys.append(focal_length)\n\n camera_to_worlds = torch.stack(c2ws, dim=0)\n fx = torch.tensor(fxs)\n fy = torch.tensor(fys)\n return Cameras(\n fx=fx,\n fy=fy,\n cx=image_width / 2,\n cy=image_height / 2,\n camera_to_worlds=camera_to_worlds,\n )\n\n\ndef get_path_from_npz(camera_path: Dict[str, Any]) -> Cameras:\n image_height, image_width = int(camera_path[\"height\"]), int(camera_path[\"width\"])\n \n K33_shared = torch.tensor(camera_path['K33'], dtype=torch.float32)\n T44_c2w_all = torch.tensor(camera_path['T44_c2w'], dtype=torch.float32) # (n, 4, 4)\n K33 = K33_shared[None].repeat(T44_c2w_all.shape[0], 1, 1)\n \n return Cameras(\n T44_c2w_all[:, :3],\n fx=K33[:, 0, 0], fy=K33[:, 1, 1], cx=K33[:, 0, 2], cy=K33[:, 1, 2],\n width=image_width, height=image_height,\n )\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/cameras/camera_utils.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nCamera transformation helper code.\n\"\"\"\n\nimport math\nfrom typing import List, Optional, Tuple\n\nimport numpy as np\nimport torch\nfrom torchtyping import TensorType\nfrom typing_extensions import Literal\n\n_EPS = np.finfo(float).eps * 4.0\n\n\ndef unit_vector(data, axis: Optional[int] = None) -> np.ndarray:\n \"\"\"Return ndarray normalized by length, i.e. Euclidean norm, along axis.\n\n Args:\n axis: the axis along which to normalize into unit vector\n out: where to write out the data to. If None, returns a new np ndarray\n \"\"\"\n data = np.array(data, dtype=np.float64, copy=True)\n if data.ndim == 1:\n data /= math.sqrt(np.dot(data, data))\n return data\n length = np.atleast_1d(np.sum(data * data, axis))\n np.sqrt(length, length)\n if axis is not None:\n length = np.expand_dims(length, axis)\n data /= length\n return data\n\n\ndef quaternion_from_matrix(matrix, isprecise: bool = False) -> np.ndarray:\n \"\"\"Return quaternion from rotation matrix.\n\n Args:\n matrix: rotation matrix to obtain quaternion\n isprecise: if True, input matrix is assumed to be precise rotation matrix and a faster algorithm is used.\n \"\"\"\n M = np.array(matrix, dtype=np.float64, copy=False)[:4, :4]\n if isprecise:\n q = np.empty((4,))\n t = np.trace(M)\n if t > M[3, 3]:\n q[0] = t\n q[3] = M[1, 0] - M[0, 1]\n q[2] = M[0, 2] - M[2, 0]\n q[1] = M[2, 1] - M[1, 2]\n else:\n i, j, k = 1, 2, 3\n if M[1, 1] > M[0, 0]:\n i, j, k = 2, 3, 1\n if M[2, 2] > M[i, i]:\n i, j, k = 3, 1, 2\n t = M[i, i] - (M[j, j] + M[k, k]) + M[3, 3]\n q[i] = t\n q[j] = M[i, j] + M[j, i]\n q[k] = M[k, i] + M[i, k]\n q[3] = M[k, j] - M[j, k]\n q *= 0.5 / math.sqrt(t * M[3, 3])\n else:\n m00 = M[0, 0]\n m01 = M[0, 1]\n m02 = M[0, 2]\n m10 = M[1, 0]\n m11 = M[1, 1]\n m12 = M[1, 2]\n m20 = M[2, 0]\n m21 = M[2, 1]\n m22 = M[2, 2]\n # symmetric matrix K\n K = np.array(\n [\n [m00 - m11 - m22, 0.0, 0.0, 0.0],\n [m01 + m10, m11 - m00 - m22, 0.0, 0.0],\n [m02 + m20, m12 + m21, m22 - m00 - m11, 0.0],\n [m21 - m12, m02 - m20, m10 - m01, m00 + m11 + m22],\n ]\n )\n K /= 3.0\n # quaternion is eigenvector of K that corresponds to largest eigenvalue\n w, V = np.linalg.eigh(K)\n q = V[np.array([3, 0, 1, 2]), np.argmax(w)]\n if q[0] < 0.0:\n np.negative(q, q)\n return q\n\n\ndef quaternion_slerp(quat0, quat1, fraction: float, spin: int = 0, shortestpath: bool = True) -> np.ndarray:\n \"\"\"Return spherical linear interpolation between two quaternions.\n Args:\n quat0: first quaternion\n quat1: second quaternion\n fraction: how much to interpolate between quat0 vs quat1 (if 0, closer to quat0; if 1, closer to quat1)\n spin: how much of an additional spin to place on the interpolation\n shortestpath: whether to return the short or long path to rotation\n \"\"\"\n q0 = unit_vector(quat0[:4])\n q1 = unit_vector(quat1[:4])\n if q0 is None or q1 is None:\n raise ValueError(\"Input quaternions invalid.\")\n if fraction == 0.0:\n return q0\n if fraction == 1.0:\n return q1\n d = np.dot(q0, q1)\n if abs(abs(d) - 1.0) < _EPS:\n return q0\n if shortestpath and d < 0.0:\n # invert rotation\n d = -d\n np.negative(q1, q1)\n angle = math.acos(d) + spin * math.pi\n if abs(angle) < _EPS:\n return q0\n isin = 1.0 / math.sin(angle)\n q0 *= math.sin((1.0 - fraction) * angle) * isin\n q1 *= math.sin(fraction * angle) * isin\n q0 += q1\n return q0\n\n\ndef quaternion_matrix(quaternion) -> np.ndarray:\n \"\"\"Return homogeneous rotation matrix from quaternion.\n\n Args:\n quaternion: value to convert to matrix\n \"\"\"\n q = np.array(quaternion, dtype=np.float64, copy=True)\n n = np.dot(q, q)\n if n < _EPS:\n return np.identity(4)\n q *= math.sqrt(2.0 / n)\n q = np.outer(q, q)\n return np.array(\n [\n [1.0 - q[2, 2] - q[3, 3], q[1, 2] - q[3, 0], q[1, 3] + q[2, 0], 0.0],\n [q[1, 2] + q[3, 0], 1.0 - q[1, 1] - q[3, 3], q[2, 3] - q[1, 0], 0.0],\n [q[1, 3] - q[2, 0], q[2, 3] + q[1, 0], 1.0 - q[1, 1] - q[2, 2], 0.0],\n [0.0, 0.0, 0.0, 1.0],\n ]\n )\n\n\ndef get_interpolated_poses(pose_a, pose_b, steps: int = 10) -> List[float]:\n \"\"\"Return interpolation of poses with specified number of steps.\n Args:\n poseA: first pose\n poseB: second pose\n steps: number of steps the interpolated pose path should contain\n \"\"\"\n\n quat_a = quaternion_from_matrix(pose_a[:3, :3])\n quat_b = quaternion_from_matrix(pose_b[:3, :3])\n\n ts = np.linspace(0, 1, steps)\n quats = [quaternion_slerp(quat_a, quat_b, t) for t in ts]\n trans = [(1 - t) * pose_a[:3, 3] + t * pose_b[:3, 3] for t in ts]\n\n poses_ab = []\n for quat, tran in zip(quats, trans):\n pose = np.identity(4)\n pose[:3, :3] = quaternion_matrix(quat)[:3, :3]\n pose[:3, 3] = tran\n poses_ab.append(pose)\n return poses_ab\n\n\ndef get_interpolated_k(k_a, k_b, steps: int = 10) -> TensorType[3, 4]:\n \"\"\"\n Returns interpolated path between two camera poses with specified number of steps.\n\n Args:\n KA: camera matrix 1\n KB: camera matrix 2\n steps: number of steps the interpolated pose path should contain\n \"\"\"\n Ks = []\n ts = np.linspace(0, 1, steps)\n for t in ts:\n new_k = k_a * (1.0 - t) + k_b * t\n Ks.append(new_k)\n return Ks\n\n\ndef get_interpolated_poses_many(\n poses: TensorType[\"num_poses\", 3, 4],\n Ks: TensorType[\"num_poses\", 3, 3],\n steps_per_transition=10,\n) -> Tuple[TensorType[\"num_poses\", 3, 4], TensorType[\"num_poses\", 3, 3]]:\n \"\"\"Return interpolated poses for many camera poses.\n\n Args:\n poses: list of camera poses\n Ks: list of camera intrinsics\n steps_per_transition: number of steps per transition\n\n Returns:\n tuple of new poses and intrinsics\n \"\"\"\n traj = []\n Ks = []\n for idx in range(poses.shape[0] - 1):\n pose_a = poses[idx]\n pose_b = poses[idx + 1]\n poses_ab = get_interpolated_poses(pose_a, pose_b, steps=steps_per_transition)\n traj += poses_ab\n Ks += get_interpolated_k(Ks[idx], Ks[idx + 1], steps_per_transition)\n return torch.stack(traj, dim=0), torch.stack(Ks, dim=0)\n\n\ndef normalize(x) -> TensorType[...]:\n \"\"\"Returns a normalized vector.\"\"\"\n return x / torch.linalg.norm(x)\n\n\ndef viewmatrix(lookat, up, pos) -> TensorType[...]:\n \"\"\"Returns a camera transformation matrix.\n\n Args:\n lookat: The direction the camera is looking.\n up: The upward direction of the camera.\n pos: The position of the camera.\n\n Returns:\n A camera transformation matrix.\n \"\"\"\n vec2 = normalize(lookat)\n vec1_avg = normalize(up)\n vec0 = normalize(torch.cross(vec1_avg, vec2))\n vec1 = normalize(torch.cross(vec2, vec0))\n m = torch.stack([vec0, vec1, vec2, pos], 1)\n return m\n\n\ndef get_distortion_params(\n k1: float = 0.0,\n k2: float = 0.0,\n k3: float = 0.0,\n k4: float = 0.0,\n p1: float = 0.0,\n p2: float = 0.0,\n) -> TensorType[...]:\n \"\"\"Returns a distortion parameters matrix.\n\n Args:\n k1: The first radial distortion parameter.\n k2: The second radial distortion parameter.\n k3: The third radial distortion parameter.\n k4: The fourth radial distortion parameter.\n p1: The first tangential distortion parameter.\n p2: The second tangential distortion parameter.\n Returns:\n torch.Tensor: A distortion parameters matrix.\n \"\"\"\n return torch.Tensor([k1, k2, k3, k4, p1, p2])\n\n\n@torch.jit.script\ndef _compute_residual_and_jacobian(\n x: torch.Tensor,\n y: torch.Tensor,\n xd: torch.Tensor,\n yd: torch.Tensor,\n distortion_params: torch.Tensor,\n) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor,]:\n \"\"\"Auxiliary function of radial_and_tangential_undistort() that computes residuals and jacobians.\n Adapted from MultiNeRF:\n https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/camera_utils.py#L427-L474\n\n Args:\n x: The updated x coordinates.\n y: The updated y coordinates.\n xd: The distorted x coordinates.\n yd: The distorted y coordinates.\n distortion_params: The distortion parameters [k1, k2, k3, k4, p1, p2].\n\n Returns:\n The residuals (fx, fy) and jacobians (fx_x, fx_y, fy_x, fy_y).\n \"\"\"\n\n k1 = distortion_params[..., 0]\n k2 = distortion_params[..., 1]\n k3 = distortion_params[..., 2]\n k4 = distortion_params[..., 3]\n p1 = distortion_params[..., 4]\n p2 = distortion_params[..., 5]\n\n # let r(x, y) = x^2 + y^2;\n # d(x, y) = 1 + k1 * r(x, y) + k2 * r(x, y) ^2 + k3 * r(x, y)^3 +\n # k4 * r(x, y)^4;\n r = x * x + y * y\n d = 1.0 + r * (k1 + r * (k2 + r * (k3 + r * k4)))\n\n # The perfect projection is:\n # xd = x * d(x, y) + 2 * p1 * x * y + p2 * (r(x, y) + 2 * x^2);\n # yd = y * d(x, y) + 2 * p2 * x * y + p1 * (r(x, y) + 2 * y^2);\n #\n # Let's define\n #\n # fx(x, y) = x * d(x, y) + 2 * p1 * x * y + p2 * (r(x, y) + 2 * x^2) - xd;\n # fy(x, y) = y * d(x, y) + 2 * p2 * x * y + p1 * (r(x, y) + 2 * y^2) - yd;\n #\n # We are looking for a solution that satisfies\n # fx(x, y) = fy(x, y) = 0;\n fx = d * x + 2 * p1 * x * y + p2 * (r + 2 * x * x) - xd\n fy = d * y + 2 * p2 * x * y + p1 * (r + 2 * y * y) - yd\n\n # Compute derivative of d over [x, y]\n d_r = k1 + r * (2.0 * k2 + r * (3.0 * k3 + r * 4.0 * k4))\n d_x = 2.0 * x * d_r\n d_y = 2.0 * y * d_r\n\n # Compute derivative of fx over x and y.\n fx_x = d + d_x * x + 2.0 * p1 * y + 6.0 * p2 * x\n fx_y = d_y * x + 2.0 * p1 * x + 2.0 * p2 * y\n\n # Compute derivative of fy over x and y.\n fy_x = d_x * y + 2.0 * p2 * y + 2.0 * p1 * x\n fy_y = d + d_y * y + 2.0 * p2 * x + 6.0 * p1 * y\n\n return fx, fy, fx_x, fx_y, fy_x, fy_y\n\n\n@torch.jit.script\ndef radial_and_tangential_undistort(\n coords: torch.Tensor,\n distortion_params: torch.Tensor,\n eps: float = 1e-3,\n max_iterations: int = 10,\n) -> torch.Tensor:\n \"\"\"Computes undistorted coords given opencv distortion parameters.\n Addapted from MultiNeRF\n https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/camera_utils.py#L477-L509\n\n Args:\n coords: The distorted coordinates.\n distortion_params: The distortion parameters [k1, k2, k3, k4, p1, p2].\n eps: The epsilon for the convergence.\n max_iterations: The maximum number of iterations to perform.\n\n Returns:\n The undistorted coordinates.\n \"\"\"\n\n # Initialize from the distorted point.\n x = coords[..., 0]\n y = coords[..., 1]\n\n for _ in range(max_iterations):\n fx, fy, fx_x, fx_y, fy_x, fy_y = _compute_residual_and_jacobian(\n x=x, y=y, xd=coords[..., 0], yd=coords[..., 1], distortion_params=distortion_params\n )\n denominator = fy_x * fx_y - fx_x * fy_y\n x_numerator = fx * fy_y - fy * fx_y\n y_numerator = fy * fx_x - fx * fy_x\n step_x = torch.where(torch.abs(denominator) > eps, x_numerator / denominator, torch.zeros_like(denominator))\n step_y = torch.where(torch.abs(denominator) > eps, y_numerator / denominator, torch.zeros_like(denominator))\n\n x = x + step_x\n y = y + step_y\n\n return torch.stack([x, y], dim=-1)\n\n\ndef rotation_matrix(a: TensorType[3], b: TensorType[3]) -> TensorType[3, 3]:\n \"\"\"Compute the rotation matrix that rotates vector a to vector b.\n\n Args:\n a: The vector to rotate.\n b: The vector to rotate to.\n Returns:\n The rotation matrix.\n \"\"\"\n a = a / torch.linalg.norm(a)\n b = b / torch.linalg.norm(b)\n v = torch.cross(a, b)\n c = torch.dot(a, b)\n # If vectors are exactly opposite, we add a little noise to one of them\n if c < -1 + 1e-8:\n eps = (torch.rand(3) - 0.5) * 0.01\n return rotation_matrix(a + eps, b)\n s = torch.linalg.norm(v)\n skew_sym_mat = torch.Tensor(\n [\n [0, -v[2], v[1]],\n [v[2], 0, -v[0]],\n [-v[1], v[0], 0],\n ]\n )\n return torch.eye(3) + skew_sym_mat + skew_sym_mat @ skew_sym_mat * ((1 - c) / (s**2 + 1e-8))\n\n\ndef auto_orient_and_center_poses(\n poses: TensorType[\"num_poses\":..., 4, 4], method: Literal[\"pca\", \"up\", \"none\"] = \"up\", center_poses: bool = True\n) -> TensorType[\"num_poses\":..., 3, 4]:\n \"\"\"Orients and centers the poses. We provide two methods for orientation: pca and up.\n\n pca: Orient the poses so that the principal component of the points is aligned with the axes.\n This method works well when all of the cameras are in the same plane.\n up: Orient the poses so that the average up vector is aligned with the z axis.\n This method works well when images are not at arbitrary angles.\n\n\n Args:\n poses: The poses to orient.\n method: The method to use for orientation.\n center_poses: If True, the poses are centered around the origin.\n\n Returns:\n The oriented poses.\n \"\"\"\n\n translation = poses[..., :3, 3]\n\n mean_translation = torch.mean(translation, dim=0)\n translation_diff = translation - mean_translation\n\n if center_poses:\n translation = mean_translation\n else:\n translation = torch.zeros_like(mean_translation)\n\n if method == \"pca\":\n _, eigvec = torch.linalg.eigh(translation_diff.T @ translation_diff)\n eigvec = torch.flip(eigvec, dims=(-1,))\n\n if torch.linalg.det(eigvec) < 0:\n eigvec[:, 2] = -eigvec[:, 2]\n\n transform = torch.cat([eigvec, eigvec @ -translation[..., None]], dim=-1)\n oriented_poses = transform @ poses\n\n if oriented_poses.mean(axis=0)[2, 1] < 0:\n oriented_poses[:, 1:3] = -1 * oriented_poses[:, 1:3]\n elif method == \"up\":\n up = torch.mean(poses[:, :3, 1], dim=0)\n up = up / torch.linalg.norm(up)\n\n rotation = rotation_matrix(up, torch.Tensor([0, 0, 1]))\n transform = torch.cat([rotation, rotation @ -translation[..., None]], dim=-1)\n oriented_poses = transform @ poses\n elif method == \"none\":\n transform = torch.eye(4)\n transform[:3, 3] = -translation\n transform = transform[:3, :]\n oriented_poses = transform @ poses\n\n return oriented_poses, transform\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/cameras/cameras.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nCamera Models\n\"\"\"\nimport base64\nimport math\nfrom dataclasses import dataclass\nfrom enum import Enum, auto\nfrom typing import Dict, List, Optional, Tuple, Union\n\nimport cv2\nimport torch\nimport torchvision\nfrom torch.nn.functional import normalize\nfrom torchtyping import TensorType\n\nimport nerfstudio.utils.poses as pose_utils\nfrom nerfstudio.cameras import camera_utils\nfrom nerfstudio.cameras.rays import RayBundle\nfrom nerfstudio.utils.tensor_dataclass import TensorDataclass\n\n\nclass CameraType(Enum):\n \"\"\"Supported camera types.\"\"\"\n\n PERSPECTIVE = auto()\n FISHEYE = auto()\n EQUIRECTANGULAR = auto()\n\n\nCAMERA_MODEL_TO_TYPE = {\n \"SIMPLE_PINHOLE\": CameraType.PERSPECTIVE,\n \"PINHOLE\": CameraType.PERSPECTIVE,\n \"SIMPLE_RADIAL\": CameraType.PERSPECTIVE,\n \"RADIAL\": CameraType.PERSPECTIVE,\n \"OPENCV\": CameraType.PERSPECTIVE,\n \"OPENCV_FISHEYE\": CameraType.FISHEYE,\n \"EQUIRECTANGULAR\": CameraType.EQUIRECTANGULAR,\n}\n\n\n@dataclass(init=False)\nclass Cameras(TensorDataclass):\n \"\"\"Dataparser outputs for the image dataset and the ray generator.\n\n Note: currently only supports cameras with the same principal points and types. The reason we type\n the focal lengths, principal points, and image sizes as tensors is to allow for batched cameras\n down the line in cases where your batches of camera data don't come from the same cameras.\n\n If a single value is provided, it is broadcasted to all cameras.\n\n Args:\n camera_to_worlds: Camera to world matrices. Tensor of per-image c2w matrices, in [R | t] format\n fx: Focal length x\n fy: Focal length y\n cx: Principal point x\n cy: Principal point y\n width: Image width\n height: Image height\n distortion_params: OpenCV 6 radial distortion coefficients\n camera_type: Type of camera model. This will be an int corresponding to the CameraType enum.\n times: Timestamps for each camera\n \"\"\"\n\n camera_to_worlds: TensorType[\"num_cameras\":..., 3, 4]\n fx: TensorType[\"num_cameras\":..., 1]\n fy: TensorType[\"num_cameras\":..., 1]\n cx: TensorType[\"num_cameras\":..., 1]\n cy: TensorType[\"num_cameras\":..., 1]\n width: TensorType[\"num_cameras\":..., 1]\n height: TensorType[\"num_cameras\":..., 1]\n distortion_params: Optional[TensorType[\"num_cameras\":..., 6]]\n camera_type: TensorType[\"num_cameras\":..., 1]\n times: Optional[TensorType[\"num_cameras\", 1]]\n\n def __init__(\n self,\n camera_to_worlds: TensorType[\"batch_c2ws\":..., 3, 4],\n fx: Union[TensorType[\"batch_fxs\":..., 1], float],\n fy: Union[TensorType[\"batch_fys\":..., 1], float],\n cx: Union[TensorType[\"batch_cxs\":..., 1], float],\n cy: Union[TensorType[\"batch_cys\":..., 1], float],\n width: Optional[Union[TensorType[\"batch_ws\":..., 1], int]] = None,\n height: Optional[Union[TensorType[\"batch_hs\":..., 1], int]] = None,\n distortion_params: Optional[TensorType[\"batch_dist_params\":..., 6]] = None,\n camera_type: Optional[\n Union[\n TensorType[\"batch_cam_types\":..., 1],\n int,\n List[CameraType],\n CameraType,\n ]\n ] = CameraType.PERSPECTIVE,\n times: Optional[TensorType[\"num_cameras\"]] = None,\n ):\n \"\"\"Initializes the Cameras object.\n\n Note on Input Tensor Dimensions: All of these tensors have items of dimensions TensorType[3, 4]\n (in the case of the c2w matrices), TensorType[6] (in the case of distortion params), or\n TensorType[1] (in the case of the rest of the elements). The dimensions before that are\n considered the batch dimension of that tensor (batch_c2ws, batch_fxs, etc.). We will broadcast\n all the tensors to be the same batch dimension. This means you can use any combination of the\n input types in the function signature and it won't break. Your batch size for all tensors\n must be broadcastable to the same size, and the resulting number of batch dimensions will be\n the batch dimension with the largest number of dimensions.\n \"\"\"\n\n # This will notify the tensordataclass that we have a field with more than 1 dimension\n self._field_custom_dimensions = {\"camera_to_worlds\": 2}\n\n self.camera_to_worlds = camera_to_worlds\n\n # fx fy calculation\n self.fx = self._init_get_fc_xy(fx, \"fx\") # @dataclass's post_init will take care of broadcasting\n self.fy = self._init_get_fc_xy(fy, \"fy\") # @dataclass's post_init will take care of broadcasting\n\n # cx cy calculation\n self.cx = self._init_get_fc_xy(cx, \"cx\") # @dataclass's post_init will take care of broadcasting\n self.cy = self._init_get_fc_xy(cy, \"cy\") # @dataclass's post_init will take care of broadcasting\n\n # Distortion Params Calculation:\n self.distortion_params = distortion_params # @dataclass's post_init will take care of broadcasting\n\n # @dataclass's post_init will take care of broadcasting\n self.height = self._init_get_height_width(height, self.cy)\n self.width = self._init_get_height_width(width, self.cx)\n self.camera_type = self._init_get_camera_type(camera_type)\n self.times = self._init_get_times(times)\n\n self.__post_init__() # This will do the dataclass post_init and broadcast all the tensors\n\n def _init_get_fc_xy(self, fc_xy, name):\n \"\"\"\n Parses the input focal length / principle point x or y and returns a tensor of the correct shape\n\n Only needs to make sure that we a 1 in the last dimension if it is a tensor. If it is a float, we\n just need to make it into a tensor and it will be broadcasted later in the __post_init__ function.\n\n Args:\n fc_xy: The focal length / principle point x or y\n name: The name of the variable. Used for error messages\n \"\"\"\n if isinstance(fc_xy, float):\n fc_xy = torch.Tensor([fc_xy], device=self.device)\n elif isinstance(fc_xy, torch.Tensor):\n if fc_xy.ndim == 0 or fc_xy.shape[-1] != 1:\n fc_xy = fc_xy.unsqueeze(-1)\n fc_xy = fc_xy.to(self.device)\n else:\n raise ValueError(f\"{name} must be a float or tensor, got {type(fc_xy)}\")\n return fc_xy\n\n def _init_get_camera_type(\n self,\n camera_type: Union[\n TensorType[\"batch_cam_types\":..., 1], TensorType[\"batch_cam_types\":...], int, List[CameraType], CameraType\n ],\n ) -> TensorType[\"num_cameras\":..., 1]:\n \"\"\"\n Parses the __init__() argument camera_type\n\n Camera Type Calculation:\n If CameraType, convert to int and then to tensor, then broadcast to all cameras\n If List of CameraTypes, convert to ints and then to tensor, then broadcast to all cameras\n If int, first go to tensor and then broadcast to all cameras\n If tensor, broadcast to all cameras\n\n Args:\n camera_type: camera_type argument from __init__()\n \"\"\"\n if isinstance(camera_type, CameraType):\n camera_type = torch.tensor([camera_type.value], device=self.device)\n elif isinstance(camera_type, List) and isinstance(camera_type[0], CameraType):\n camera_type = torch.tensor([[c.value] for c in camera_type], device=self.device)\n elif isinstance(camera_type, int):\n camera_type = torch.tensor([camera_type], device=self.device)\n elif isinstance(camera_type, torch.Tensor):\n assert not torch.is_floating_point(\n camera_type\n ), f\"camera_type tensor must be of type int, not: {camera_type.dtype}\"\n camera_type = camera_type.to(self.device)\n if camera_type.ndim == 0 or camera_type.shape[-1] != 1:\n camera_type = camera_type.unsqueeze(-1)\n # assert torch.all(\n # camera_type.view(-1)[0] == camera_type\n # ), \"Batched cameras of different camera_types will be allowed in the future.\"\n else:\n raise ValueError(\n 'Invalid camera_type. Must be CameraType, List[CameraType], int, or torch.Tensor[\"num_cameras\"]. \\\n Received: '\n + str(type(camera_type))\n )\n return camera_type\n\n def _init_get_height_width(\n self,\n h_w: Union[TensorType[\"batch_hws\":..., 1], TensorType[\"batch_hws\":...], int, None],\n c_x_y: TensorType[\"batch_cxys\":...],\n ) -> TensorType[\"num_cameras\":..., 1]:\n \"\"\"\n Parses the __init__() argument for height or width\n\n Height/Width Calculation:\n If int, first go to tensor and then broadcast to all cameras\n If tensor, broadcast to all cameras\n If none, use cx or cy * 2\n Else raise error\n\n Args:\n h_w: height or width argument from __init__()\n c_x_y: cx or cy for when h_w == None\n \"\"\"\n if isinstance(h_w, int):\n h_w = torch.Tensor([h_w]).to(torch.int64).to(self.device)\n elif isinstance(h_w, torch.Tensor):\n assert not torch.is_floating_point(h_w), f\"height and width tensor must be of type int, not: {h_w.dtype}\"\n h_w = h_w.to(torch.int64).to(self.device)\n if h_w.ndim == 0 or h_w.shape[-1] != 1:\n h_w = h_w.unsqueeze(-1)\n # assert torch.all(h_w == h_w.view(-1)[0]), \"Batched cameras of different h, w will be allowed in the future.\"\n elif h_w is None:\n h_w = torch.Tensor((c_x_y * 2).to(torch.int64).to(self.device))\n else:\n raise ValueError(\"Height must be an int, tensor, or None, received: \" + str(type(h_w)))\n return h_w\n\n def _init_get_times(self, times):\n if times is None:\n times = None\n elif isinstance(times, torch.Tensor):\n if times.ndim == 0 or times.shape[-1] != 1:\n times = times.unsqueeze(-1).to(self.device)\n else:\n raise ValueError(f\"times must be None or a tensor, got {type(times)}\")\n\n return times\n\n @property\n def device(self):\n \"\"\"Returns the device that the camera is on.\"\"\"\n return self.camera_to_worlds.device\n\n @property\n def image_height(self) -> TensorType[\"num_cameras\":..., 1]:\n \"\"\"Returns the height of the images.\"\"\"\n return self.height\n\n @property\n def image_width(self) -> TensorType[\"num_cameras\":..., 1]:\n \"\"\"Returns the height of the images.\"\"\"\n return self.width\n\n @property\n def is_jagged(self):\n \"\"\"\n Returns whether or not the cameras are \"jagged\" (i.e. the height and widths are different, meaning that\n you cannot concatenate the image coordinate maps together)\n \"\"\"\n h_jagged = not torch.all(self.height == self.height.view(-1)[0])\n w_jagged = not torch.all(self.width == self.width.view(-1)[0])\n return h_jagged or w_jagged\n\n def get_image_coords(\n self, pixel_offset: float = 0.5, index: Optional[Tuple] = None\n ) -> TensorType[\"height\", \"width\", 2]:\n \"\"\"This gets the image coordinates of one of the cameras in this object.\n\n If no index is specified, it will return the maximum possible sized height / width image coordinate map,\n by looking at the maximum height and width of all the cameras in this object.\n\n Args:\n pixel_offset: Offset for each pixel. Defaults to center of pixel (0.5)\n index: Tuple of indices into the batch dimensions of the camera. Defaults to None, which returns the 0th\n flattened camera\n\n Returns:\n Grid of image coordinates.\n \"\"\"\n if index is None:\n image_height = torch.max(self.image_height.view(-1))\n image_width = torch.max(self.image_width.view(-1))\n image_coords = torch.meshgrid(torch.arange(image_height), torch.arange(image_width), indexing=\"ij\")\n image_coords = torch.stack(image_coords, dim=-1) + pixel_offset # stored as (y, x) coordinates\n else:\n image_height = self.image_height[index].item()\n image_width = self.image_width[index].item()\n image_coords = torch.meshgrid(torch.arange(image_height), torch.arange(image_width), indexing=\"ij\")\n image_coords = torch.stack(image_coords, dim=-1) + pixel_offset # stored as (y, x) coordinates\n return image_coords\n\n def generate_rays( # pylint: disable=too-many-statements\n self,\n camera_indices: Union[TensorType[\"num_rays\":..., \"num_cameras_batch_dims\"], int],\n coords: Optional[TensorType[\"num_rays\":..., 2]] = None,\n camera_opt_to_camera: Optional[TensorType[\"num_rays\":..., 3, 4]] = None,\n distortion_params_delta: Optional[TensorType[\"num_rays\":..., 6]] = None,\n keep_shape: Optional[bool] = None,\n disable_distortion: bool = False,\n ) -> RayBundle:\n \"\"\"Generates rays for the given camera indices.\n\n This function will standardize the input arguments and then call the _generate_rays_from_coords function\n to generate the rays. Our goal is to parse the arguments and then get them into the right shape:\n - camera_indices: (num_rays:..., num_cameras_batch_dims)\n - coords: (num_rays:..., 2)\n - camera_opt_to_camera: (num_rays:..., 3, 4) or None\n - distortion_params_delta: (num_rays:..., 6) or None\n\n Read the docstring for _generate_rays_from_coords for more information on how we generate the rays\n after we have standardized the arguments.\n\n We are only concerned about different combinations of camera_indices and coords matrices, and the following\n are the 4 cases we have to deal with:\n 1. isinstance(camera_indices, int) and coords == None\n - In this case we broadcast our camera_indices / coords shape (h, w, 1 / 2 respectively)\n 2. isinstance(camera_indices, int) and coords != None\n - In this case, we broadcast camera_indices to the same batch dim as coords\n 3. not isinstance(camera_indices, int) and coords == None\n - In this case, we will need to set coords so that it is of shape (h, w, num_rays, 2), and broadcast\n all our other args to match the new definition of num_rays := (h, w) + num_rays\n 4. not isinstance(camera_indices, int) and coords != None\n - In this case, we have nothing to do, only check that the arguments are of the correct shape\n\n There is one more edge case we need to be careful with: when we have \"jagged cameras\" (ie: different heights\n and widths for each camera). This isn't problematic when we specify coords, since coords is already a tensor.\n When coords == None (ie: when we render out the whole image associated with this camera), we run into problems\n since there's no way to stack each coordinate map as all coordinate maps are all different shapes. In this case,\n we will need to flatten each individual coordinate map and concatenate them, giving us only one batch dimension,\n regaurdless of the number of prepended extra batch dimensions in the camera_indices tensor.\n\n\n Args:\n camera_indices: Camera indices of the flattened cameras object to generate rays for.\n coords: Coordinates of the pixels to generate rays for. If None, the full image will be rendered.\n camera_opt_to_camera: Optional transform for the camera to world matrices.\n distortion_params_delta: Optional delta for the distortion parameters.\n keep_shape: If None, then we default to the regular behavior of flattening if cameras is jagged, otherwise\n keeping dimensions. If False, we flatten at the end. If True, then we keep the shape of the\n camera_indices and coords tensors (if we can).\n disable_distortion: If True, disables distortion.\n\n Returns:\n Rays for the given camera indices and coords.\n \"\"\"\n # Check the argument types to make sure they're valid and all shaped correctly\n assert isinstance(camera_indices, (torch.Tensor, int)), \"camera_indices must be a tensor or int\"\n assert coords is None or isinstance(coords, torch.Tensor), \"coords must be a tensor or None\"\n assert camera_opt_to_camera is None or isinstance(camera_opt_to_camera, torch.Tensor)\n assert distortion_params_delta is None or isinstance(distortion_params_delta, torch.Tensor)\n if isinstance(camera_indices, torch.Tensor) and isinstance(coords, torch.Tensor):\n num_rays_shape = camera_indices.shape[:-1]\n errormsg = \"Batch dims of inputs must match when inputs are all tensors\"\n assert coords.shape[:-1] == num_rays_shape, errormsg\n assert camera_opt_to_camera is None or camera_opt_to_camera.shape[:-2] == num_rays_shape, errormsg\n assert distortion_params_delta is None or distortion_params_delta.shape[:-1] == num_rays_shape, errormsg\n\n # If zero dimensional, we need to unsqueeze to get a batch dimension and then squeeze later\n if not self.shape:\n cameras = self.reshape((1,))\n assert torch.all(\n torch.tensor(camera_indices == 0) if isinstance(camera_indices, int) else camera_indices == 0\n ), \"Can only index into single camera with no batch dimensions if index is zero\"\n else:\n cameras = self\n\n # If the camera indices are an int, then we need to make sure that the camera batch is 1D\n if isinstance(camera_indices, int):\n assert (\n len(cameras.shape) == 1\n ), \"camera_indices must be a tensor if cameras are batched with more than 1 batch dimension\"\n camera_indices = torch.tensor([camera_indices], device=cameras.device)\n\n assert camera_indices.shape[-1] == len(\n cameras.shape\n ), \"camera_indices must have shape (num_rays:..., num_cameras_batch_dims)\"\n\n # If keep_shape is True, then we need to make sure that the camera indices in question\n # are all the same height and width and can actually be batched while maintaining the image\n # shape\n if keep_shape is True:\n assert torch.all(cameras.height[camera_indices] == cameras.height[camera_indices[0]]) and torch.all(\n cameras.width[camera_indices] == cameras.width[camera_indices[0]]\n ), \"Can only keep shape if all cameras have the same height and width\"\n\n # If the cameras don't all have same height / width, if coords is not none, we will need to generate\n # a flat list of coords for each camera and then concatenate otherwise our rays will be jagged.\n # Camera indices, camera_opt, and distortion will also need to be broadcasted accordingly which is non-trivial\n if cameras.is_jagged and coords is None and (keep_shape is None or keep_shape is False):\n index_dim = camera_indices.shape[-1]\n camera_indices = camera_indices.reshape(-1, index_dim)\n _coords = [cameras.get_image_coords(index=tuple(index)).reshape(-1, 2) for index in camera_indices]\n camera_indices = torch.cat(\n [index.unsqueeze(0).repeat(coords.shape[0], 1) for index, coords in zip(camera_indices, _coords)],\n )\n coords = torch.cat(_coords, dim=0)\n assert coords.shape[0] == camera_indices.shape[0]\n # Need to get the coords of each indexed camera and flatten all coordinate maps and concatenate them\n\n # The case where we aren't jagged && keep_shape (since otherwise coords is already set) and coords\n # is None. In this case we append (h, w) to the num_rays dimensions for all tensors. In this case,\n # each image in camera_indices has to have the same shape since otherwise we would have error'd when\n # we checked keep_shape is valid or we aren't jagged.\n if coords is None:\n index_dim = camera_indices.shape[-1]\n index = camera_indices.reshape(-1, index_dim)[0]\n coords: torch.Tensor = cameras.get_image_coords(index=tuple(index)) # (h, w, 2)\n coords = coords.reshape(coords.shape[:2] + (1,) * len(camera_indices.shape[:-1]) + (2,)) # (h, w, 1..., 2)\n coords = coords.expand(coords.shape[:2] + camera_indices.shape[:-1] + (2,)) # (h, w, num_rays, 2)\n camera_opt_to_camera = ( # (h, w, num_rays, 3, 4) or None\n camera_opt_to_camera.broadcast_to(coords.shape[:-1] + (3, 4))\n if camera_opt_to_camera is not None\n else None\n )\n distortion_params_delta = ( # (h, w, num_rays, 6) or None\n distortion_params_delta.broadcast_to(coords.shape[:-1] + (6,))\n if distortion_params_delta is not None\n else None\n )\n\n # If camera indices was an int or coords was none, we need to broadcast our indices along batch dims\n camera_indices = camera_indices.broadcast_to(coords.shape[:-1] + (len(cameras.shape),)).to(torch.long)\n\n # Checking our tensors have been standardized\n assert isinstance(coords, torch.Tensor) and isinstance(camera_indices, torch.Tensor)\n assert camera_indices.shape[-1] == len(cameras.shape)\n assert camera_opt_to_camera is None or camera_opt_to_camera.shape[:-2] == coords.shape[:-1]\n assert distortion_params_delta is None or distortion_params_delta.shape[:-1] == coords.shape[:-1]\n\n # This will do the actual work of generating the rays now that we have standardized the inputs\n # raybundle.shape == (num_rays) when done\n # pylint: disable=protected-access\n raybundle = cameras._generate_rays_from_coords(\n camera_indices, coords, camera_opt_to_camera, distortion_params_delta, disable_distortion=disable_distortion\n )\n\n # If we have mandated that we don't keep the shape, then we flatten\n if keep_shape is False:\n raybundle = raybundle.flatten()\n\n # TODO: We should have to squeeze the last dimension here if we started with zero batch dims, but never have to,\n # so there might be a rogue squeeze happening somewhere, and this may cause some unintended behaviour\n # that we haven't caught yet with tests\n return raybundle\n\n # pylint: disable=too-many-statements\n def _generate_rays_from_coords(\n self,\n camera_indices: TensorType[\"num_rays\":..., \"num_cameras_batch_dims\"],\n coords: TensorType[\"num_rays\":..., 2],\n camera_opt_to_camera: Optional[TensorType[\"num_rays\":..., 3, 4]] = None,\n distortion_params_delta: Optional[TensorType[\"num_rays\":..., 6]] = None,\n disable_distortion: bool = False,\n ) -> RayBundle:\n \"\"\"Generates rays for the given camera indices and coords where self isn't jagged\n\n This is a fairly complex function, so let's break this down slowly.\n\n Shapes involved:\n - num_rays: This is your output raybundle shape. It dictates the number and shape of the rays generated\n - num_cameras_batch_dims: This is the number of dimensions of our camera\n\n Args:\n camera_indices: Camera indices of the flattened cameras object to generate rays for.\n The shape of this is such that indexing into camera_indices[\"num_rays\":...] will return the\n index into each batch dimension of the camera in order to get the correct camera specified by\n \"num_rays\".\n Example:\n >>> cameras = Cameras(...)\n >>> cameras.shape\n (2, 3, 4)\n >>> camera_indices = torch.tensor([0, 0, 0]) # We need an axis of length 3 since cameras.ndim == 3\n >>> camera_indices.shape\n (3,)\n >>> coords = torch.tensor([1,1])\n >>> coords.shape\n (2,)\n >>> out_rays = cameras.generate_rays(camera_indices=camera_indices, coords = coords)\n # This will generate a RayBundle with a single ray for the\n # camera at cameras[0,0,0] at image coordinates (1,1), so out_rays.shape == ()\n >>> out_rays.shape\n ()\n >>> camera_indices = torch.tensor([[0,0,0]])\n >>> camera_indices.shape\n (1, 3)\n >>> coords = torch.tensor([[1,1]])\n >>> coords.shape\n (1, 2)\n >>> out_rays = cameras.generate_rays(camera_indices=camera_indices, coords = coords)\n # This will generate a RayBundle with a single ray for the\n # camera at cameras[0,0,0] at point (1,1), so out_rays.shape == (1,)\n # since we added an extra dimension in front of camera_indices\n >>> out_rays.shape\n (1,)\n\n If you want more examples, check tests/cameras/test_cameras and the function check_generate_rays_shape\n\n The bottom line is that for camera_indices: (num_rays:..., num_cameras_batch_dims), num_rays is the\n output shape and if you index into the output RayBundle with some indices [i:...], if you index into\n camera_indices with camera_indices[i:...] as well, you will get a 1D tensor containing the batch\n indices into the original cameras object corresponding to that ray (ie: you will get the camera\n from our batched cameras corresponding to the ray at RayBundle[i:...]).\n\n coords: Coordinates of the pixels to generate rays for. If None, the full image will be rendered, meaning\n height and width get prepended to the num_rays dimensions. Indexing into coords with [i:...] will\n get you the image coordinates [x, y] of that specific ray located at output RayBundle[i:...].\n\n camera_opt_to_camera: Optional transform for the camera to world matrices.\n In terms of shape, it follows the same rules as coords, but indexing into it with [i:...] gets you\n the 2D camera to world transform matrix for the camera optimization at RayBundle[i:...].\n\n distortion_params_delta: Optional delta for the distortion parameters.\n In terms of shape, it follows the same rules as coords, but indexing into it with [i:...] gets you\n the 1D tensor with the 6 distortion parameters for the camera optimization at RayBundle[i:...].\n\n disable_distortion: If True, disables distortion.\n\n Returns:\n Rays for the given camera indices and coords. RayBundle.shape == num_rays\n \"\"\"\n # Make sure we're on the right devices\n camera_indices = camera_indices.to(self.device)\n coords = coords.to(self.device)\n\n # Checking to make sure everything is of the right shape and type\n num_rays_shape = camera_indices.shape[:-1]\n assert camera_indices.shape == num_rays_shape + (self.ndim,)\n assert coords.shape == num_rays_shape + (2,)\n assert coords.shape[-1] == 2\n assert camera_opt_to_camera is None or camera_opt_to_camera.shape == num_rays_shape + (3, 4)\n assert distortion_params_delta is None or distortion_params_delta.shape == num_rays_shape + (6,)\n\n # Here, we've broken our indices down along the num_cameras_batch_dims dimension allowing us to index by all\n # of our output rays at each dimension of our cameras object\n true_indices = [camera_indices[..., i] for i in range(camera_indices.shape[-1])]\n\n # Get all our focal lengths, principal points and make sure they are the right shapes\n y = coords[..., 0] # (num_rays,) get rid of the last dimension\n x = coords[..., 1] # (num_rays,) get rid of the last dimension\n fx, fy = self.fx[true_indices].squeeze(-1), self.fy[true_indices].squeeze(-1) # (num_rays,)\n cx, cy = self.cx[true_indices].squeeze(-1), self.cy[true_indices].squeeze(-1) # (num_rays,)\n assert (\n y.shape == num_rays_shape\n and x.shape == num_rays_shape\n and fx.shape == num_rays_shape\n and fy.shape == num_rays_shape\n and cx.shape == num_rays_shape\n and cy.shape == num_rays_shape\n ), (\n str(num_rays_shape)\n + str(y.shape)\n + str(x.shape)\n + str(fx.shape)\n + str(fy.shape)\n + str(cx.shape)\n + str(cy.shape)\n )\n\n # Get our image coordinates and image coordinates offset by 1 (offsets used for dx, dy calculations)\n # Also make sure the shapes are correct\n coord = torch.stack([(x - cx) / fx, -(y - cy) / fy], -1) # (num_rays, 2)\n coord_x_offset = torch.stack([(x - cx + 1) / fx, -(y - cy) / fy], -1) # (num_rays, 2)\n coord_y_offset = torch.stack([(x - cx) / fx, -(y - cy + 1) / fy], -1) # (num_rays, 2)\n assert (\n coord.shape == num_rays_shape + (2,)\n and coord_x_offset.shape == num_rays_shape + (2,)\n and coord_y_offset.shape == num_rays_shape + (2,)\n )\n\n # Stack image coordinates and image coordinates offset by 1, check shapes too\n coord_stack = torch.stack([coord, coord_x_offset, coord_y_offset], dim=0) # (3, num_rays, 2)\n assert coord_stack.shape == (3,) + num_rays_shape + (2,)\n\n # Undistorts our images according to our distortion parameters\n if not disable_distortion:\n distortion_params = None\n if self.distortion_params is not None:\n distortion_params = self.distortion_params[true_indices]\n if distortion_params_delta is not None:\n distortion_params = distortion_params + distortion_params_delta\n elif distortion_params_delta is not None:\n distortion_params = distortion_params_delta\n\n # Do not apply distortion for equirectangular images\n if distortion_params is not None:\n mask = (self.camera_type[true_indices] != CameraType.EQUIRECTANGULAR.value).squeeze(-1) # (num_rays)\n coord_mask = torch.stack([mask, mask, mask], dim=0)\n if mask.any():\n coord_stack[coord_mask, :] = camera_utils.radial_and_tangential_undistort(\n coord_stack[coord_mask, :].reshape(3, -1, 2),\n distortion_params[mask, :],\n ).reshape(-1, 2)\n\n # Make sure after we have undistorted our images, the shapes are still correct\n assert coord_stack.shape == (3,) + num_rays_shape + (2,)\n\n # Gets our directions for all our rays in camera coordinates and checks shapes at the end\n # Here, directions_stack is of shape (3, num_rays, 3)\n # directions_stack[0] is the direction for ray in camera coordinates\n # directions_stack[1] is the direction for ray in camera coordinates offset by 1 in x\n # directions_stack[2] is the direction for ray in camera coordinates offset by 1 in y\n cam_types = torch.unique(self.camera_type, sorted=False)\n directions_stack = torch.empty((3,) + num_rays_shape + (3,), device=self.device)\n if CameraType.PERSPECTIVE.value in cam_types:\n mask = (self.camera_type[true_indices] == CameraType.PERSPECTIVE.value).squeeze(-1) # (num_rays)\n mask = torch.stack([mask, mask, mask], dim=0)\n directions_stack[..., 0][mask] = torch.masked_select(coord_stack[..., 0], mask).float()\n directions_stack[..., 1][mask] = torch.masked_select(coord_stack[..., 1], mask).float()\n directions_stack[..., 2][mask] = -1.0\n\n if CameraType.FISHEYE.value in cam_types:\n mask = (self.camera_type[true_indices] == CameraType.FISHEYE.value).squeeze(-1) # (num_rays)\n mask = torch.stack([mask, mask, mask], dim=0)\n\n theta = torch.sqrt(torch.sum(coord_stack**2, dim=-1))\n theta = torch.clip(theta, 0.0, math.pi)\n\n sin_theta = torch.sin(theta)\n\n directions_stack[..., 0][mask] = torch.masked_select(coord_stack[..., 0] * sin_theta / theta, mask).float()\n directions_stack[..., 1][mask] = torch.masked_select(coord_stack[..., 1] * sin_theta / theta, mask).float()\n directions_stack[..., 2][mask] = -torch.masked_select(torch.cos(theta), mask)\n\n if CameraType.EQUIRECTANGULAR.value in cam_types:\n mask = (self.camera_type[true_indices] == CameraType.EQUIRECTANGULAR.value).squeeze(-1) # (num_rays)\n mask = torch.stack([mask, mask, mask], dim=0)\n\n # For equirect, fx = fy = height = width/2\n # Then coord[..., 0] goes from -1 to 1 and coord[..., 1] goes from -1/2 to 1/2\n theta = -torch.pi * coord_stack[..., 0] # minus sign for right-handed\n phi = torch.pi * (0.5 - coord_stack[..., 1])\n # use spherical in local camera coordinates (+y up, x=0 and z<0 is theta=0)\n directions_stack[..., 0][mask] = torch.masked_select(-torch.sin(theta) * torch.sin(phi), mask).float()\n directions_stack[..., 1][mask] = torch.masked_select(torch.cos(phi), mask).float()\n directions_stack[..., 2][mask] = torch.masked_select(-torch.cos(theta) * torch.sin(phi), mask).float()\n\n for value in cam_types:\n if value not in [CameraType.PERSPECTIVE.value, CameraType.FISHEYE.value, CameraType.EQUIRECTANGULAR.value]:\n raise ValueError(f\"Camera type {value} not supported.\")\n\n assert directions_stack.shape == (3,) + num_rays_shape + (3,)\n\n c2w = self.camera_to_worlds[true_indices]\n assert c2w.shape == num_rays_shape + (3, 4)\n\n if camera_opt_to_camera is not None:\n c2w = pose_utils.multiply(c2w, camera_opt_to_camera)\n rotation = c2w[..., :3, :3] # (..., 3, 3)\n assert rotation.shape == num_rays_shape + (3, 3)\n\n directions_stack = torch.sum(\n directions_stack[..., None, :] * rotation, dim=-1\n ) # (..., 1, 3) * (..., 3, 3) -> (..., 3)\n\n directions_norm = torch.norm(directions_stack, dim=-1, keepdim=True)\n directions_norm = directions_norm[0]\n\n directions_stack = normalize(directions_stack, dim=-1)\n assert directions_stack.shape == (3,) + num_rays_shape + (3,)\n\n origins = c2w[..., :3, 3] # (..., 3)\n assert origins.shape == num_rays_shape + (3,)\n\n directions = directions_stack[0]\n assert directions.shape == num_rays_shape + (3,)\n\n # norms of the vector going between adjacent coords, giving us dx and dy per output ray\n dx = torch.sqrt(torch.sum((directions - directions_stack[1]) ** 2, dim=-1)) # (\"num_rays\":...,)\n dy = torch.sqrt(torch.sum((directions - directions_stack[2]) ** 2, dim=-1)) # (\"num_rays\":...,)\n assert dx.shape == num_rays_shape and dy.shape == num_rays_shape\n\n pixel_area = (dx * dy)[..., None] # (\"num_rays\":..., 1)\n assert pixel_area.shape == num_rays_shape + (1,)\n\n times = self.times[camera_indices, 0] if self.times is not None else None\n\n return RayBundle(\n origins=origins,\n directions=directions,\n pixel_area=pixel_area,\n camera_indices=camera_indices,\n directions_norm=directions_norm,\n times=times,\n )\n\n def to_json(\n self, camera_idx: int, image: Optional[TensorType[\"height\", \"width\", 2]] = None, max_size: Optional[int] = None\n ) -> Dict:\n \"\"\"Convert a camera to a json dictionary.\n\n Args:\n camera_idx: Index of the camera to convert.\n image: An image in range [0, 1] that is encoded to a base64 string.\n max_size: Max size to resize the image to if present.\n\n Returns:\n A JSON representation of the camera\n \"\"\"\n flattened = self.flatten()\n json_ = {\n \"type\": \"PinholeCamera\",\n \"cx\": flattened[camera_idx].cx.item(),\n \"cy\": flattened[camera_idx].cy.item(),\n \"fx\": flattened[camera_idx].fx.item(),\n \"fy\": flattened[camera_idx].fy.item(),\n \"camera_to_world\": self.camera_to_worlds[camera_idx].tolist(),\n \"camera_index\": camera_idx,\n \"times\": flattened[camera_idx].times.item() if self.times is not None else None,\n }\n if image is not None:\n image_uint8 = (image * 255).detach().type(torch.uint8)\n if max_size is not None:\n image_uint8 = image_uint8.permute(2, 0, 1)\n image_uint8 = torchvision.transforms.functional.resize(image_uint8, max_size) # type: ignore\n image_uint8 = image_uint8.permute(1, 2, 0)\n image_uint8 = image_uint8.cpu().numpy()\n data = cv2.imencode(\".jpg\", image_uint8)[1].tobytes()\n json_[\"image\"] = str(\"data:image/jpeg;base64,\" + base64.b64encode(data).decode(\"ascii\"))\n return json_\n\n def get_intrinsics_matrices(self) -> TensorType[\"num_cameras\":..., 3, 3]:\n \"\"\"Returns the intrinsic matrices for each camera.\n\n Returns:\n Pinhole camera intrinsics matrices\n \"\"\"\n K = torch.zeros((*self.shape, 3, 3), dtype=torch.float32)\n K[..., 0, 0] = self.fx.squeeze(-1)\n K[..., 1, 1] = self.fy.squeeze(-1)\n K[..., 0, 2] = self.cx.squeeze(-1)\n K[..., 1, 2] = self.cy.squeeze(-1)\n K[..., 2, 2] = 1.0\n return K\n\n def rescale_output_resolution(\n self, scaling_factor: Union[TensorType[\"num_cameras\":...], TensorType[\"num_cameras\":..., 1], float, int]\n ) -> None:\n \"\"\"Rescale the output resolution of the cameras.\n\n Args:\n scaling_factor: Scaling factor to apply to the output resolution.\n \"\"\"\n if isinstance(scaling_factor, (float, int)):\n scaling_factor = torch.tensor([scaling_factor]).to(self.device).broadcast_to((self.cx.shape))\n elif isinstance(scaling_factor, torch.Tensor) and scaling_factor.shape == self.shape:\n scaling_factor = scaling_factor.unsqueeze(-1)\n elif isinstance(scaling_factor, torch.Tensor) and scaling_factor.shape == (*self.shape, 1):\n pass\n else:\n raise ValueError(\n f\"Scaling factor must be a float, int, or a tensor of shape {self.shape} or {(*self.shape, 1)}.\"\n )\n\n self.fx = self.fx * scaling_factor\n self.fy = self.fy * scaling_factor\n self.cx = self.cx * scaling_factor\n self.cy = self.cy * scaling_factor\n self.height = (self.height * scaling_factor).to(torch.int64)\n self.width = (self.width * scaling_factor).to(torch.int64)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/cameras/lie_groups.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nHelper for Lie group operations. Currently only used for pose optimization.\n\"\"\"\nimport torch\nfrom torchtyping import TensorType\n\n\n# We make an exception on snake case conventions because SO3 != so3.\ndef exp_map_SO3xR3(tangent_vector: TensorType[\"b\", 6]) -> TensorType[\"b\", 3, 4]: # pylint: disable=invalid-name\n \"\"\"Compute the exponential map of the direct product group `SO(3) x R^3`.\n\n This can be used for learning pose deltas on SE(3), and is generally faster than `exp_map_SE3`.\n\n Args:\n tangent_vector: Tangent vector; length-3 translations, followed by an `so(3)` tangent vector.\n Returns:\n [R|t] tranformation matrices.\n \"\"\"\n # code for SO3 map grabbed from pytorch3d and stripped down to bare-bones\n log_rot = tangent_vector[:, 3:]\n nrms = (log_rot * log_rot).sum(1)\n rot_angles = torch.clamp(nrms, 1e-4).sqrt()\n rot_angles_inv = 1.0 / rot_angles\n fac1 = rot_angles_inv * rot_angles.sin()\n fac2 = rot_angles_inv * rot_angles_inv * (1.0 - rot_angles.cos())\n skews = torch.zeros((log_rot.shape[0], 3, 3), dtype=log_rot.dtype, device=log_rot.device)\n skews[:, 0, 1] = -log_rot[:, 2]\n skews[:, 0, 2] = log_rot[:, 1]\n skews[:, 1, 0] = log_rot[:, 2]\n skews[:, 1, 2] = -log_rot[:, 0]\n skews[:, 2, 0] = -log_rot[:, 1]\n skews[:, 2, 1] = log_rot[:, 0]\n skews_square = torch.bmm(skews, skews)\n\n ret = torch.zeros(tangent_vector.shape[0], 3, 4, dtype=tangent_vector.dtype, device=tangent_vector.device)\n ret[:, :3, :3] = (\n fac1[:, None, None] * skews\n + fac2[:, None, None] * skews_square\n + torch.eye(3, dtype=log_rot.dtype, device=log_rot.device)[None]\n )\n\n # Compute the translation\n ret[:, :3, 3] = tangent_vector[:, :3]\n return ret\n\n\ndef exp_map_SE3(tangent_vector: TensorType[\"b\", 6]) -> TensorType[\"b\", 3, 4]: # pylint: disable=invalid-name\n \"\"\"Compute the exponential map `se(3) -> SE(3)`.\n\n This can be used for learning pose deltas on `SE(3)`.\n\n Args:\n tangent_vector: A tangent vector from `se(3)`.\n\n Returns:\n [R|t] tranformation matrices.\n \"\"\"\n\n tangent_vector_lin = tangent_vector[:, :3].view(-1, 3, 1)\n tangent_vector_ang = tangent_vector[:, 3:].view(-1, 3, 1)\n\n theta = torch.linalg.norm(tangent_vector_ang, dim=1).unsqueeze(1)\n theta2 = theta**2\n theta3 = theta**3\n\n near_zero = theta < 1e-2\n non_zero = torch.ones(1, dtype=tangent_vector.dtype, device=tangent_vector.device)\n theta_nz = torch.where(near_zero, non_zero, theta)\n theta2_nz = torch.where(near_zero, non_zero, theta2)\n theta3_nz = torch.where(near_zero, non_zero, theta3)\n\n # Compute the rotation\n sine = theta.sin()\n cosine = torch.where(near_zero, 8 / (4 + theta2) - 1, theta.cos())\n sine_by_theta = torch.where(near_zero, 0.5 * cosine + 0.5, sine / theta_nz)\n one_minus_cosine_by_theta2 = torch.where(near_zero, 0.5 * sine_by_theta, (1 - cosine) / theta2_nz)\n ret = torch.zeros(tangent_vector.shape[0], 3, 4).to(dtype=tangent_vector.dtype, device=tangent_vector.device)\n ret[:, :3, :3] = one_minus_cosine_by_theta2 * tangent_vector_ang @ tangent_vector_ang.transpose(1, 2)\n\n ret[:, 0, 0] += cosine.view(-1)\n ret[:, 1, 1] += cosine.view(-1)\n ret[:, 2, 2] += cosine.view(-1)\n temp = sine_by_theta.view(-1, 1) * tangent_vector_ang.view(-1, 3)\n ret[:, 0, 1] -= temp[:, 2]\n ret[:, 1, 0] += temp[:, 2]\n ret[:, 0, 2] += temp[:, 1]\n ret[:, 2, 0] -= temp[:, 1]\n ret[:, 1, 2] -= temp[:, 0]\n ret[:, 2, 1] += temp[:, 0]\n\n # Compute the translation\n sine_by_theta = torch.where(near_zero, 1 - theta2 / 6, sine_by_theta)\n one_minus_cosine_by_theta2 = torch.where(near_zero, 0.5 - theta2 / 24, one_minus_cosine_by_theta2)\n theta_minus_sine_by_theta3_t = torch.where(near_zero, 1.0 / 6 - theta2 / 120, (theta - sine) / theta3_nz)\n\n ret[:, :, 3:] = sine_by_theta * tangent_vector_lin\n ret[:, :, 3:] += one_minus_cosine_by_theta2 * torch.cross(tangent_vector_ang, tangent_vector_lin, dim=1)\n ret[:, :, 3:] += theta_minus_sine_by_theta3_t * (\n tangent_vector_ang @ (tangent_vector_ang.transpose(1, 2) @ tangent_vector_lin)\n )\n return ret\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/cameras/rays.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nSome ray datastructures.\n\"\"\"\nimport random\nimport warnings\nfrom dataclasses import dataclass\nfrom typing import Callable, Dict, Optional, Tuple\n\nimport torch\nfrom torchtyping import TensorType\nfrom rich.console import Console\n\nfrom nerfstudio.utils.math import Gaussians, conical_frustum_to_gaussian\nfrom nerfstudio.utils.tensor_dataclass import TensorDataclass\n\nCONSOLE = Console(width=120)\n\nNEGATIVE_MIN_DELTA_WARNING = (\"[bold yellow]RaySamples.deltas includes negative values, \"\n \"which might be caused by nears > fars / bad weights (e.g., all zeros) \"\n \"/ numerical issues caused by nears being too closed with fars. \"\n \"This will lead to negative weights and NaN interlevel loss.\")\n_WARNED = False\n\n@dataclass\nclass Frustums(TensorDataclass):\n \"\"\"Describes region of space as a frustum.\"\"\"\n\n origins: TensorType[\"bs\":..., 3]\n \"\"\"xyz coordinate for ray origin.\"\"\"\n directions: TensorType[\"bs\":..., 3]\n \"\"\"Direction of ray.\"\"\"\n starts: TensorType[\"bs\":..., 1]\n \"\"\"Where the frustum starts along a ray.\"\"\"\n ends: TensorType[\"bs\":..., 1]\n \"\"\"Where the frustum ends along a ray.\"\"\"\n pixel_area: TensorType[\"bs\":..., 1]\n \"\"\"Projected area of pixel a distance 1 away from origin.\"\"\"\n offsets: Optional[TensorType[\"bs\":..., 3]] = None\n \"\"\"Offsets for each sample position\"\"\"\n\n def get_positions(self) -> TensorType[..., 3]:\n \"\"\"Calulates \"center\" position of frustum. Not weighted by mass.\n\n Returns:\n xyz positions.\n \"\"\"\n pos = self.origins + self.directions * (self.starts + self.ends) / 2\n if self.offsets is not None:\n pos = pos + self.offsets\n return pos\n\n def set_offsets(self, offsets):\n \"\"\"Sets offsets for this frustum for computing positions\"\"\"\n self.offsets = offsets\n\n def get_start_positions(self) -> TensorType[..., 3]:\n \"\"\"Calulates \"start\" position of frustum. We use start positions for MonoSDF\n because when we use error bounded sampling, we need to upsample many times.\n It's hard to merge two set of ray samples while keeping the mid points fixed.\n Every time we up sample the points the mid points will change and\n therefore we need to evaluate all points again which is 3 times slower.\n But we can skip the evaluation of sdf value if we use start position instead of mid position\n because after we merge the points, the starting point is the same and only the delta is changed.\n\n Returns:\n xyz positions.\n \"\"\"\n return self.origins + self.directions * self.starts\n\n def get_gaussian_blob(self) -> Gaussians:\n \"\"\"Calculates guassian approximation of conical frustum.\n\n Resturns:\n Conical frustums approximated by gaussian distribution.\n \"\"\"\n # Cone radius is set such that the square pixel_area matches the cone area.\n cone_radius = torch.sqrt(self.pixel_area) / 1.7724538509055159 # r = sqrt(pixel_area / pi)\n if self.offsets is not None:\n raise NotImplementedError()\n return conical_frustum_to_gaussian(\n origins=self.origins,\n directions=self.directions,\n starts=self.starts,\n ends=self.ends,\n radius=cone_radius,\n )\n\n @classmethod\n def get_mock_frustum(cls, device=\"cpu\") -> \"Frustums\":\n \"\"\"Helper function to generate a placeholder frustum.\n\n Returns:\n A size 1 frustum with meaningless values.\n \"\"\"\n return Frustums(\n origins=torch.ones((1, 3)).to(device),\n directions=torch.ones((1, 3)).to(device),\n starts=torch.ones((1, 1)).to(device),\n ends=torch.ones((1, 1)).to(device),\n pixel_area=torch.ones((1, 1)).to(device),\n )\n\n\n@dataclass\nclass RaySamples(TensorDataclass):\n \"\"\"Samples along a ray\"\"\"\n\n frustums: Frustums\n \"\"\"Frustums along ray.\"\"\"\n camera_indices: Optional[TensorType[\"bs\":..., 1]] = None\n \"\"\"Camera index.\"\"\"\n deltas: Optional[TensorType[\"bs\":..., 1]] = None\n \"\"\"\"width\" of each sample.\"\"\"\n spacing_starts: Optional[TensorType[\"bs\":..., \"num_samples\", 1]] = None\n \"\"\"Start of normalized bin edges along ray [0,1], before warping is applied, ie. linear in disparity sampling.\"\"\"\n spacing_ends: Optional[TensorType[\"bs\":..., \"num_samples\", 1]] = None\n \"\"\"Start of normalized bin edges along ray [0,1], before warping is applied, ie. linear in disparity sampling.\"\"\"\n spacing_to_euclidean_fn: Optional[Callable] = None\n \"\"\"Function to convert bins to euclidean distance.\"\"\"\n metadata: Optional[Dict[str, TensorType[\"bs\":..., \"latent_dims\"]]] = None\n \"\"\"addtional information relevant to generating ray samples\"\"\"\n\n times: Optional[TensorType[..., 1]] = None\n \"\"\"Times at which rays are sampled\"\"\"\n\n def __post_init__(self) -> None:\n super().__post_init__()\n if self.deltas is not None:\n negative_deltas_mask = self.deltas < 0\n if negative_deltas_mask.any():\n n_negative_deltas = negative_deltas_mask.sum()\n self.deltas[negative_deltas_mask] = 0.\n global _WARNED # TODO: use a context manager to count the number of time this warning is triggered\n if not _WARNED:\n _WARNED = True\n CONSOLE.print(NEGATIVE_MIN_DELTA_WARNING)\n CONSOLE.print(f\"[bold yellow]Truncate negative deltas to 0! ({n_negative_deltas} in total)\")\n \n def get_weights(self, densities: TensorType[..., \"num_samples\", 1]) -> TensorType[..., \"num_samples\", 1]:\n \"\"\"Return weights based on predicted densities\n\n Args:\n densities: Predicted densities for samples along ray\n\n Returns:\n Weights for each sample\n \"\"\"\n delta_density = self.deltas * densities\n alphas = 1 - torch.exp(-delta_density)\n\n transmittance = torch.cumsum(delta_density[..., :-1, :], dim=-2)\n transmittance = torch.cat(\n [torch.zeros((*transmittance.shape[:1], 1, 1), device=densities.device), transmittance], dim=-2\n )\n transmittance = torch.exp(-transmittance) # [..., \"num_samples\"]\n\n weights = alphas * transmittance # [..., \"num_samples\"]\n weights = torch.nan_to_num(weights) # nan weights might be caused by inf densities\n \n if torch.isnan(weights).any():\n __import__('ipdb').set_trace()\n\n return weights\n\n def get_weights_and_transmittance(\n self, densities: TensorType[..., \"num_samples\", 1]\n ) -> Tuple[TensorType[..., \"num_samples\", 1], TensorType[..., \"num_samples\", 1]]:\n \"\"\"Return weights and transmittance based on predicted densities\n\n Args:\n densities: Predicted densities for samples along ray\n\n Returns:\n Weights and transmittance for each sample\n \"\"\"\n\n delta_density = self.deltas * densities\n alphas = 1 - torch.exp(-delta_density)\n\n transmittance = torch.cumsum(delta_density[..., :-1, :], dim=-2)\n transmittance = torch.cat(\n [torch.zeros((*transmittance.shape[:1], 1, 1), device=densities.device), transmittance], dim=-2\n )\n transmittance = torch.exp(-transmittance) # [..., \"num_samples\"]\n\n weights = alphas * transmittance # [..., \"num_samples\"]\n weights = torch.nan_to_num(weights) # nan weights might be caused by inf densities\n\n return weights, transmittance\n\n def get_weights_from_alphas(self, alphas: TensorType[..., \"num_samples\", 1]) -> TensorType[..., \"num_samples\", 1]:\n \"\"\"Return weights based on predicted alphas\n\n Args:\n alphas: Predicted alphas (maybe from sdf) for samples along ray\n\n Returns:\n Weights for each sample\n \"\"\"\n\n transmittance = torch.cumprod(\n torch.cat([torch.ones((*alphas.shape[:1], 1, 1), device=alphas.device), 1.0 - alphas + 1e-7], 1), 1\n ) # [..., \"num_samples\"]\n\n weights = alphas * transmittance[:, :-1, :] # [..., \"num_samples\"]\n weights = torch.nan_to_num(weights)\n\n return weights\n\n def get_weights_and_transmittance_from_alphas(\n self, alphas: TensorType[..., \"num_samples\", 1]\n ) -> TensorType[..., \"num_samples\", 1]:\n \"\"\"Return weights based on predicted alphas\n\n Args:\n alphas: Predicted alphas (maybe from sdf) for samples along ray\n\n Returns:\n Weights for each sample\n \"\"\"\n\n transmittance = torch.cumprod(\n torch.cat([torch.ones((*alphas.shape[:1], 1, 1), device=alphas.device), 1.0 - alphas + 1e-7], 1), 1\n ) # [..., \"num_samples\"]\n\n weights = alphas * transmittance[:, :-1, :] # [..., \"num_samples\"]\n weights = torch.nan_to_num(weights)\n \n if torch.isnan(weights).any():\n __import__('ipdb').set_trace()\n\n return weights, transmittance\n\n\n@dataclass\nclass RayBundle(TensorDataclass):\n \"\"\"A bundle of ray parameters.\"\"\"\n\n # TODO(ethan): make sure the sizes with ... are correct\n origins: TensorType[..., 3]\n \"\"\"Ray origins (XYZ)\"\"\"\n directions: TensorType[..., 3]\n \"\"\"Unit ray direction vector\"\"\"\n pixel_area: TensorType[..., 1]\n \"\"\"Projected area of pixel a distance 1 away from origin\"\"\"\n directions_norm: Optional[TensorType[..., 1]] = None\n \"\"\"Norm of ray direction vector before normalization\"\"\"\n camera_indices: Optional[TensorType[..., 1]] = None\n \"\"\"Camera indices\"\"\"\n nears: Optional[TensorType[..., 1]] = None\n \"\"\"Distance along ray to start sampling\"\"\"\n fars: Optional[TensorType[..., 1]] = None\n \"\"\"Rays Distance along ray to stop sampling\"\"\"\n # metadata: Optional[Dict[str, TensorType[\"num_rays\", \"latent_dims\"]]] = None\n metadata: Optional[Dict[str, TensorType[\"bs\":..., \"latent_dims\"]]] = None\n \"\"\"Additional metadata or data needed for interpolation, will mimic shape of rays\"\"\"\n times: Optional[TensorType[..., 1]] = None\n \"\"\"Times at which rays are sampled\"\"\"\n\n def set_camera_indices(self, camera_index: int) -> None:\n \"\"\"Sets all of the the camera indices to a specific camera index.\n\n Args:\n camera_index: Camera index.\n \"\"\"\n self.camera_indices = torch.ones_like(self.origins[..., 0:1]).long() * camera_index\n\n def __len__(self):\n num_rays = torch.numel(self.origins) // self.origins.shape[-1]\n return num_rays\n\n def sample(self, num_rays: int) -> \"RayBundle\":\n \"\"\"Returns a RayBundle as a subset of rays.\n\n Args:\n num_rays: Number of rays in output RayBundle\n\n Returns:\n RayBundle with subset of rays.\n \"\"\"\n assert num_rays <= len(self)\n indices = random.sample(range(len(self)), k=num_rays)\n return self[indices]\n\n def get_row_major_sliced_ray_bundle(self, start_idx: int, end_idx: int) -> \"RayBundle\":\n \"\"\"Flattens RayBundle and extracts chunk given start and end indicies.\n\n Args:\n start_idx: Start index of RayBundle chunk.\n end_idx: End index of RayBundle chunk.\n\n Returns:\n Flattened RayBundle with end_idx-start_idx rays.\n\n \"\"\"\n return self.flatten()[start_idx:end_idx]\n\n def get_ray_samples(\n self,\n bin_starts: TensorType[\"bs\":..., \"num_samples\", 1],\n bin_ends: TensorType[\"bs\":..., \"num_samples\", 1],\n spacing_starts: Optional[TensorType[\"bs\":..., \"num_samples\", 1]] = None,\n spacing_ends: Optional[TensorType[\"bs\":..., \"num_samples\", 1]] = None,\n spacing_to_euclidean_fn: Optional[Callable] = None,\n ) -> RaySamples:\n \"\"\"Produces samples for each ray by projection points along the ray direction. Currently samples uniformly.\n\n Args:\n bin_starts: Distance from origin to start of bin.\n bin_ends: Distance from origin to end of bin.\n\n Returns:\n Samples projected along ray.\n \"\"\"\n deltas = bin_ends - bin_starts\n # if deltas.min() < 0:\n # CONSOLE.log(NEGATIVE_MIN_DELTA_WARNING)\n # # if nears and fars are too close to each other (e.g., the ray is intersecting the edge region of an aabb),\n # # the spacing between ray samples might not be handled properly using float32.\n # __import__('ipdb').set_trace()\n if self.camera_indices is not None:\n camera_indices = self.camera_indices[..., None]\n else:\n camera_indices = None\n\n shaped_raybundle_fields = self[..., None]\n\n frustums = Frustums(\n origins=shaped_raybundle_fields.origins, # [..., 1, 3]\n directions=shaped_raybundle_fields.directions, # [..., 1, 3]\n starts=bin_starts, # [..., num_samples, 1]\n ends=bin_ends, # [..., num_samples, 1]\n pixel_area=shaped_raybundle_fields.pixel_area, # [..., 1, 1]\n )\n\n ray_samples = RaySamples(\n frustums=frustums,\n camera_indices=camera_indices, # [..., 1, 1]\n deltas=deltas, # [..., num_samples, 1]\n spacing_starts=spacing_starts, # [..., num_samples, 1]\n spacing_ends=spacing_ends, # [..., num_samples, 1]\n spacing_to_euclidean_fn=spacing_to_euclidean_fn,\n metadata=shaped_raybundle_fields.metadata,\n times=None if self.times is None else self.times[..., None], # [..., 1, 1]\n )\n\n return ray_samples\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/configs/__init__.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/configs/autorecon/__init__.py",
"content": "from __future__ import annotations\n\nfrom typing import Dict, Sequence\nfrom functools import partial\n\nimport tyro\nfrom sklearn.decomposition import PCA\n\nfrom nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig\nfrom nerfstudio.configs.base_config import (\n Config,\n SchedulerConfig,\n TrainerConfig,\n ViewerConfig,\n)\nfrom nerfstudio.data.datamanagers.base_datamanager import (\n FlexibleDataManagerConfig,\n VanillaDataManagerConfig,\n)\nfrom nerfstudio.data.datamanagers.semantic_datamanager import SemanticDataManagerConfig\nfrom nerfstudio.data.datamanagers.autorecon_datamanager import AutoReconDataManagerConfig\nfrom nerfstudio.data.datamanagers.variable_res_datamanager import (\n VariableResDataManagerConfig,\n)\nfrom nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig\nfrom nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig\nfrom nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConfig\nfrom nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig\nfrom nerfstudio.data.dataparsers.phototourism_dataparser import (\n PhototourismDataParserConfig,\n)\nfrom nerfstudio.data.dataparsers.sdfstudio_dataparser import SDFStudioDataParserConfig\nfrom nerfstudio.data.dataparsers.autorecon_dataparser import AutoReconDataParserConfig\nfrom nerfstudio.data.utils.nerfstudio_collate import nerfstudio_collate\nfrom nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig, AdamWOptimizerConfig\nfrom nerfstudio.engine.schedulers import (\n ExponentialSchedulerConfig,\n MultiStepSchedulerConfig,\n NeuSSchedulerConfig,\n AutoReconSchedulerConfig\n)\n\nfrom nerfstudio.field_components.temporal_distortions import TemporalDistortionKind\nfrom nerfstudio.fields.sdf_field import SDFFieldConfig\nfrom nerfstudio.fields.feature_field import FeatureFieldConfig, FeatureSegFieldConfig\nfrom nerfstudio.models.dto import DtoOModelConfig\nfrom nerfstudio.models.instant_ngp import InstantNGPModelConfig\nfrom nerfstudio.models.mipnerf import MipNerfModel\nfrom nerfstudio.models.nerfacto import NerfactoModelConfig\nfrom nerfstudio.models.neuralreconW import NeuralReconWModelConfig\nfrom nerfstudio.models.neus import NeuSModelConfig\nfrom nerfstudio.models.neus_acc import NeuSAccModelConfig\nfrom nerfstudio.models.neus_facto import NeuSFactoModelConfig\nfrom nerfstudio.models.neus_facto_dff import NeuSFactoDFFModelConfig\nfrom nerfstudio.models.neus_facto_reg import NeuSFactoRegModelConfig\nfrom nerfstudio.models.distilled_neus_facto import DistilledNeuSFactoModelConfig\nfrom nerfstudio.models.semantic_nerfw import SemanticNerfWModelConfig\nfrom nerfstudio.models.tensorf import TensoRFModelConfig\nfrom nerfstudio.models.unisurf import UniSurfModelConfig\nfrom nerfstudio.models.vanilla_nerf import NeRFModel, VanillaModelConfig\nfrom nerfstudio.models.volsdf import VolSDFModelConfig\nfrom nerfstudio.pipelines.base_pipeline import (\n FlexibleInputPipelineConfig,\n VanillaPipelineConfig,\n)\nfrom nerfstudio.pipelines.dynamic_batch import DynamicBatchPipelineConfig\nfrom nerfstudio.utils.func_utils import get_first_element\nfrom nerfstudio.utils.pointclouds import BasicPointClouds\n\nfrom .baseline import neus_wbg, neus_wbg_mlp\nfrom .common import (\n neusfacto_autorecon, neus_facto_wbg, neus_facto_wbg_BgProbNet,\n neus_facto_wbg_medium, neus_facto_wbg_tiny, neus_facto_wbg_tiny_long_schedule\n)\nfrom .distilled_neusfacto import distilled_neus_facto_wbg, debug_distilled_neus_facto_wbg\nfrom .feature_field import neus_facto_wbg_fast_dff, neus_facto_dff_wbg\nfrom .neusfacto_fast import neus_facto_wbg_fast\nfrom .regularization import neus_facto_wbg_reg, neus_facto_wbg_reg_sep_plane, neus_facto_wbg_reg_sep_plane_nerf\nfrom .semantic_nerf import autorecon_semantic_nerf\n\n\nmethod_configs: Dict[str, Config] = {\n \"neus-wbg\": neus_wbg,\n \"neus-wbg_mlp\": neus_wbg_mlp,\n \"neus-facto-autorecon\": neusfacto_autorecon,\n \"neus-facto-wbg\": neus_facto_wbg,\n \"neus-facto-wbg_bg-prob-net\": neus_facto_wbg_BgProbNet,\n \"neus-facto-wbg_medium\": neus_facto_wbg_medium,\n \"neus-facto-wbg_tiny\": neus_facto_wbg_tiny,\n \"neus-facto-wbg_tiny_long-schedule\": neus_facto_wbg_tiny_long_schedule,\n \"distilled-neus-facto-wbg\": distilled_neus_facto_wbg,\n \"debug_distilled-neus-facto-wbg\": debug_distilled_neus_facto_wbg,\n \"neus-facto-wbg-fast_dff\": neus_facto_wbg_fast_dff,\n \"neus-facto-dff-wbg\": neus_facto_dff_wbg,\n \"neus-facto-wbg-fast\": neus_facto_wbg_fast,\n \"neus-facto-wbg-reg\": neus_facto_wbg_reg,\n \"neus-facto-wbg-reg_sep-plane\": neus_facto_wbg_reg_sep_plane,\n \"neus-facto-wbg-reg_sep-plane-nerf\": neus_facto_wbg_reg_sep_plane_nerf,\n \"autorecon_semantic-nerfw\": autorecon_semantic_nerf\n}\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/configs/autorecon/baseline.py",
"content": "from . import *\n\n# TODO: neus w/o hash grid + NeRF++ scene parameterization\n\n# neus (MLP-based) w/ a BG Model (l2 scene contraction)\nneus_wbg_mlp = Config(\n method_name=\"neus-wbg_mlp\",\n trainer=TrainerConfig(\n steps_per_eval_image=2500,\n steps_per_eval_batch=2500,\n steps_per_save=50000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=300000,\n mixed_precision=False,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=SDFStudioDataParserConfig(),\n train_num_rays_per_batch=512,\n eval_num_rays_per_batch=512,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n ),\n model=NeuSModelConfig(\n sdf_field=SDFFieldConfig(\n use_position_encoding=True,\n bias=0.2,\n beta_init=0.3,\n inside_outside=False,\n ),\n scene_contraction_order=\"none\",\n bg_sampler_type=\"lin_disp\", # TODO: ablate \"uniform_lin_disp\"\n background_model=\"mlp\",\n use_fg_aware_scene_contraction=False, # TODO: True might lead to better results\n bg_use_appearance_embedding=False,\n use_average_appearance_embedding=False,\n eval_num_rays_per_chunk=1024\n ),\n ),\n optimizers={\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000),\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000),\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\n\n# (deprecated) neus with a BG Model\nneus_wbg = Config(\n method_name=\"neus-wbg\",\n trainer=TrainerConfig(\n steps_per_eval_image=1000,\n steps_per_eval_batch=1000,\n steps_per_save=20000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=60001, # TODO: 20001 is enough?\n mixed_precision=False,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=SDFStudioDataParserConfig(),\n train_num_rays_per_batch=1024,\n eval_num_rays_per_batch=1024,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n ),\n model=NeuSModelConfig(\n sdf_field=SDFFieldConfig(\n use_grid_feature=True,\n num_layers=2,\n num_layers_color=2,\n hidden_dim=256,\n bias=0.5,\n beta_init=0.3,\n use_appearance_embedding=False, # TODO: improve impl efficiency (currently zero embeddings are used -> disable completely!)\n inside_outside=False, # only consider object-centric scenes\n ),\n cos_anneal_end=10000,\n bg_use_appearance_embedding=False,\n use_average_appearance_embedding=False,\n background_model=\"grid\",\n # TODO: check bg model size\n eval_num_rays_per_chunk=1024,\n ),\n ),\n optimizers={ # TODO: use higher lr for grid feature\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=1500, learning_rate_alpha=0.05, max_steps=60001),\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=1500, learning_rate_alpha=0.05, max_steps=60001),\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/configs/autorecon/common.py",
"content": "from . import *\n\n\n# assume the fg object is bounded by a unit cube, and the entire scene is\n# contracted to [-2, 2] as done in MipNeRF-360\nneusfacto_autorecon = Config(\n method_name=\"neus-facto-autorecon\",\n trainer=TrainerConfig(\n steps_per_eval_image=1000,\n steps_per_eval_batch=1000,\n steps_per_save=20000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=60001, # 20k iters is not enough for high quality reconstruction\n mixed_precision=True,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=AutoReconDataParserConfig(),\n train_num_rays_per_batch=2048,\n eval_num_rays_per_batch=1024,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n eval_camera_res_scale_factor=1.0,\n collate_fn=partial(nerfstudio_collate,\n extra_mappings={PCA: get_first_element,\n BasicPointClouds: get_first_element})\n ),\n model=NeuSFactoModelConfig(\n sdf_field=SDFFieldConfig(\n use_position_encoding=True,\n use_grid_feature=True,\n num_layers=2,\n num_layers_color=2,\n hidden_dim=256,\n bias=0.5,\n beta_init=0.3,\n direction_encoding_type=\"sh\",\n use_appearance_embedding=False,\n inside_outside=False,\n spatial_normalization_region='full', # [-2, 2] -> [0, 1] for feature grid interpolation\n ),\n cos_anneal_end=10000,\n background_model=\"none\",\n eval_num_rays_per_chunk=1024,\n proposal_use_uniform_sampler=False,\n proposal_net_spatial_normalization_region=\"full\" # [-2, 2] -> [0, 1] for feature grid interpolation\n ),\n ),\n optimizers={\n \"proposal_networks\": { # fixed gamma=0.3 and n_milestones=3\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": MultiStepSchedulerConfig(max_steps=60001),\n },\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=2e-3, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=1500, learning_rate_alpha=0.05, max_steps=60001),\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\n\n# neus-facto with a BG Model\n# This model has inferior fg reconstructions comparing to the neus-facto model.\nneus_facto_wbg = Config(\n method_name=\"neus-facto-wbg\",\n trainer=TrainerConfig(\n steps_per_eval_image=2500,\n steps_per_eval_batch=2500,\n steps_per_save=60000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=60001,\n mixed_precision=True,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n # TODO: set dataparser args\n dataparser=AutoReconDataParserConfig(), # NOTE: cannot set init args here, which would not be recognised.\n train_num_rays_per_batch=2048,\n eval_num_rays_per_batch=1024,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n eval_camera_res_scale_factor=1.0,\n collate_fn=partial(nerfstudio_collate,\n extra_mappings={PCA: get_first_element,\n BasicPointClouds: get_first_element})\n ),\n model=NeuSFactoModelConfig(\n sdf_field=SDFFieldConfig(\n use_position_encoding=True,\n use_grid_feature=True,\n num_layers=2,\n num_layers_color=2,\n hidden_dim=256,\n bias=0.5,\n beta_init=0.3,\n direction_encoding_type=\"sh\",\n use_appearance_embedding=False, # TODO: improve impl efficiency (currently zero embeddings are used -> disable completely!)\n inside_outside=False, # only consider object-centric scenes\n ),\n cos_anneal_end=10000,\n proposal_net_spatial_normalization_region=\"fg\",\n eval_num_rays_per_chunk=1024,\n use_average_appearance_embedding=False,\n background_model=\"grid\",\n bg_use_appearance_embedding=False,\n use_fg_aware_scene_contraction=True,\n fg_aware_scene_contraction_alpha=15.0, # TODO: 15.0 -> 5.0\n bg_sampler_type=\"uniform_lin_disp\",\n num_samples_outside=64,\n ),\n ),\n optimizers={\n \"proposal_networks\": { # fixed gamma=0.3 and n_milestones=3\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": MultiStepSchedulerConfig(max_steps=60001),\n },\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=2e-3, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=1500, learning_rate_alpha=0.05, max_steps=60001),\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=2e-3, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=1500, learning_rate_alpha=0.05, max_steps=60001),\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\n\n# neus-facto with a BG Model (w/ proposal-net sampler)\n# This model has inferior fg reconstructions comparing to the neus-facto model.\nneus_facto_wbg_BgProbNet = Config(\n method_name=\"neus-facto-wbg_bg-prop-net\",\n trainer=TrainerConfig(\n steps_per_eval_image=2500,\n steps_per_eval_batch=2500,\n steps_per_save=60000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=60001,\n mixed_precision=True,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n # TODO: set dataparser args\n dataparser=AutoReconDataParserConfig(), # NOTE: cannot set init args here, which would not be recognised.\n train_num_rays_per_batch=2048,\n eval_num_rays_per_batch=2048,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n eval_camera_res_scale_factor=1.0,\n collate_fn=partial(nerfstudio_collate,\n extra_mappings={PCA: get_first_element,\n BasicPointClouds: get_first_element})\n ),\n model=NeuSFactoModelConfig(\n sdf_field=SDFFieldConfig(\n use_position_encoding=True,\n use_grid_feature=True,\n num_layers=2,\n num_layers_color=2,\n hidden_dim=256,\n bias=0.1, # 0.5 -> 0.1\n beta_init=0.3,\n direction_encoding_type=\"sh\",\n use_appearance_embedding=False, # TODO: improve impl efficiency (currently zero embeddings are used -> disable completely!)\n inside_outside=False, # only consider object-centric scenes\n hash_grid_progressive_training=False,\n hash_grid_progressive_training_iters=0, # TODO: run ablation\n spatial_normalization_region=\"aabb\", # to better utilize the feature grid\n weight_norm=False, # weight_norm might lead to NaNs during training\n ),\n cos_anneal_end=0, # FIXME: 5000\n proposal_net_spatial_normalization_region=\"fg\",\n eval_num_rays_per_chunk=1024,\n use_average_appearance_embedding=False,\n background_model=\"grid\",\n bg_use_appearance_embedding=False,\n use_fg_aware_scene_contraction=True,\n fg_aware_scene_contraction_alpha=5.0,\n bg_sampler_type=\"proposal_network\",\n num_samples_outside=16, # TODO: run ablation (use a smaller number to avoid the fg modeling affected by bg)\n num_bg_proposal_samples_per_ray=(96, 64)\n ),\n ),\n optimizers={\n \"proposal_networks\": { # fixed gamma=0.3 and n_milestones=3\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n # \"scheduler\": MultiStepSchedulerConfig(max_steps=60001),\n \"scheduler\": None,\n },\n \"proposal_networks_bg\": { # fixed gamma=0.3 and n_milestones=3\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n # \"scheduler\": MultiStepSchedulerConfig(max_steps=30001),\n },\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=2e-3, eps=1e-15),\n # \"scheduler\": NeuSSchedulerConfig(warm_up_end=1500, learning_rate_alpha=0.05, max_steps=60001),\n \"scheduler\": MultiStepSchedulerConfig(max_steps=30001),\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\n\n# medium-sized model (deprecated)\nneus_facto_wbg_medium = Config(\n method_name=\"neus-facto-wbg_medium\",\n trainer=TrainerConfig(\n steps_per_eval_image=2000,\n steps_per_eval_batch=2000,\n steps_per_save=40000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=40001,\n mixed_precision=True,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n # TODO: set dataparser args\n dataparser=AutoReconDataParserConfig(), # NOTE: cannot set init args here, which would not be recognised.\n train_num_rays_per_batch=2048,\n eval_num_rays_per_batch=1024,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n eval_camera_res_scale_factor=1.0,\n collate_fn=partial(nerfstudio_collate,\n extra_mappings={PCA: get_first_element,\n BasicPointClouds: get_first_element})\n ),\n model=NeuSFactoModelConfig(\n sdf_field=SDFFieldConfig(\n use_position_encoding=False,\n use_grid_feature=True,\n hash_grid_num_levels=16,\n num_layers=2,\n num_layers_color=2,\n hidden_dim=64, # 256 -> 64\n geo_feat_dim=64, # 256 -> 64\n direction_encoding_type=\"sh\",\n hidden_dim_color=64, # 256 -> 64\n color_network_include_sdf=False,\n bias=0.5,\n beta_init=0.3,\n use_appearance_embedding=False, # TODO: improve impl efficiency (currently zero embeddings are used -> disable completely!)\n inside_outside=False, # only consider object-centric scenes\n ),\n bg_use_appearance_embedding=False,\n use_average_appearance_embedding=False,\n background_model=\"grid\",\n bg_hash_grid_num_levels=16,\n eval_num_rays_per_chunk=1024,\n proposal_net_spatial_normalization_region=\"fg\"\n ),\n ),\n optimizers={\n \"proposal_networks\": { # fixed gamma=0.3 and n_milestones=3\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": MultiStepSchedulerConfig(max_steps=40001),\n },\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=2e-3, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=1000, learning_rate_alpha=0.05, max_steps=40001),\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=2e-3, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=1000, learning_rate_alpha=0.05, max_steps=40001),\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\n\n# tiny model (deprecated) - same model size as the one in NeuS2\nneus_facto_wbg_tiny = Config(\n method_name=\"neus-facto-wbg_tiny\",\n trainer=TrainerConfig(\n steps_per_eval_image=2000,\n steps_per_eval_batch=2000,\n steps_per_save=20000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=20001,\n mixed_precision=True,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n # TODO: set dataparser args\n dataparser=AutoReconDataParserConfig(), # NOTE: cannot set init args here, which would not be recognised.\n train_num_rays_per_batch=2048,\n eval_num_rays_per_batch=1024,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n eval_camera_res_scale_factor=1.0,\n collate_fn=partial(nerfstudio_collate,\n extra_mappings={PCA: get_first_element,\n BasicPointClouds: get_first_element})\n ),\n model=NeuSFactoModelConfig(\n sdf_field=SDFFieldConfig(\n use_position_encoding=False,\n use_grid_feature=True,\n hash_grid_num_levels=14, # NOTE: following NeuS2\n num_layers=1, # NOTE: following NeuS2\n num_layers_color=2,\n hidden_dim=64, # NOTE: following NeuS2\n geo_feat_dim=15, # NOTE: following NeuS2\n direction_encoding_type=\"sh\", # NOTE: following NeuS2\n hidden_dim_color=64, # NOTE: following NeuS2\n color_network_include_sdf=False, # TODO: NeuS2 uses True\n bias=0.5,\n beta_init=0.3,\n use_appearance_embedding=False, # TODO: improve impl efficiency (currently zero embeddings are used -> disable completely!)\n inside_outside=False, # only consider object-centric scenes\n ),\n bg_use_appearance_embedding=False,\n use_average_appearance_embedding=False,\n background_model=\"grid\",\n bg_hash_grid_num_levels=14, # NOTE: following NeuS2\n eval_num_rays_per_chunk=1024,\n proposal_net_spatial_normalization_region=\"fg\"\n ),\n ),\n optimizers={\n \"proposal_networks\": { # fixed gamma=0.3 and n_milestones=3\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": MultiStepSchedulerConfig(max_steps=20001),\n },\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=500, learning_rate_alpha=0.05, max_steps=20001),\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=500, learning_rate_alpha=0.05, max_steps=20001),\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\n\n# tiny model (deprecated) - same model size as the one in NeuS2 with a longer training schedule\nneus_facto_wbg_tiny_long_schedule = Config(\n method_name=\"neus-facto-wbg_tiny_long-schedule\",\n trainer=TrainerConfig(\n steps_per_eval_image=2500,\n steps_per_eval_batch=2500,\n steps_per_save=30000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=30001,\n mixed_precision=True,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=AutoReconDataParserConfig(), # NOTE: cannot set init args here, which would not be recognised.\n train_num_rays_per_batch=2048,\n eval_num_rays_per_batch=1024,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n eval_camera_res_scale_factor=1.0,\n collate_fn=partial(nerfstudio_collate,\n extra_mappings={PCA: get_first_element,\n BasicPointClouds: get_first_element})\n ),\n model=NeuSFactoModelConfig(\n sdf_field=SDFFieldConfig(\n use_position_encoding=True,\n use_grid_feature=True,\n hash_grid_num_levels=14, # NOTE: following NeuS2\n num_layers=1, # NOTE: following NeuS2\n num_layers_color=2,\n hidden_dim=64, # NOTE: following NeuS2\n geo_feat_dim=15, # NOTE: following NeuS2\n direction_encoding_type=\"sh\", # NOTE: following NeuS2\n hidden_dim_color=64, # NOTE: following NeuS2\n color_network_include_sdf=False, # TODO: NeuS2 uses True\n bias=0.5,\n beta_init=0.3,\n use_appearance_embedding=False, # TODO: improve impl efficiency (currently zero embeddings are used -> disable completely!)\n inside_outside=False, # only consider object-centric scenes\n ),\n cos_anneal_end=5000,\n use_fg_aware_scene_contraction=True,\n fg_aware_scene_contraction_alpha=15.0,\n bg_sampler_type=\"uniform_lin_disp\", # TODO: ablation: \"proposal_network\"\n num_samples_outside=64,\n bg_use_appearance_embedding=False,\n use_average_appearance_embedding=False,\n background_model=\"grid\",\n bg_hash_grid_num_levels=14, # NOTE: following NeuS2\n eval_num_rays_per_chunk=1024,\n proposal_net_spatial_normalization_region=\"fg\"\n ),\n ),\n optimizers={\n \"proposal_networks\": { # fixed gamma=0.3 and n_milestones=3\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": MultiStepSchedulerConfig(max_steps=30001),\n },\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=2e-3, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=750, learning_rate_alpha=0.05, max_steps=30001),\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=2e-3, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=750, learning_rate_alpha=0.05, max_steps=30001),\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/configs/autorecon/distilled_neusfacto.py",
"content": "from . import *\n\n\n# 2-stages training of neus-facto:\n# 1. train nerfacto; (nerfacto is very robust to the scene condition)\n# 2. replace fg nerfacto_field with fg sdf_field (optionally supervised by pre-trained nerfacto)\ndistilled_neus_facto_wbg = Config(\n method_name=\"distilled-neus-facto-wbg\",\n trainer=TrainerConfig(\n steps_per_eval_image=2500,\n steps_per_eval_batch=2500,\n steps_per_save=30000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=30001,\n mixed_precision=True,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=AutoReconDataParserConfig(), # NOTE: cannot set init args here, which would not be recognised.\n train_num_rays_per_batch=4096, # 2048,\n eval_num_rays_per_batch=4096, # 1024,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n eval_camera_res_scale_factor=1.0,\n collate_fn=partial(nerfstudio_collate,\n extra_mappings={PCA: get_first_element,\n BasicPointClouds: get_first_element})\n ),\n model=DistilledNeuSFactoModelConfig(\n nerfacto_field_enabled=True, # for the 1st-stage training\n nerfacto_spatial_normalization_region=\"fg\", # bg nerfacto always uses full\n proposal_net_use_separate_contraction=True,\n proposal_net_contraction_scale_factor=0.25,\n proposal_net_spatial_normalization_region=\"full\",\n proposal_use_uniform_sampler=False,\n use_fg_aware_scene_contraction=True,\n fg_aware_scene_contraction_alpha=15.0,\n sdf_field=SDFFieldConfig( # NeuSFacto Medium\n use_position_encoding=True,\n use_grid_feature=True,\n num_layers=2,\n num_layers_color=2,\n hidden_dim=64, # 256 -> 64\n geo_feat_dim=64, # 256 -> 64 -> 15\n bias=0.5,\n beta_init=0.3,\n direction_encoding_type=\"sh\",\n use_appearance_embedding=False, # TODO: improve impl efficiency (currently zero embeddings are used -> disable completely!)\n inside_outside=False, # only consider object-centric scenes\n ),\n cos_anneal_end=0, # TODO: set 5000 while training NeuS\n eval_num_rays_per_chunk=1024,\n use_average_appearance_embedding=False,\n background_model=\"grid\",\n bg_use_appearance_embedding=False,\n bg_sampler_type=\"none\",\n num_samples_outside=0,\n ),\n ),\n optimizers={\n \"proposal_networks\": { # TODO: disable this optimizer during 2nd-stage training\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n \"fields\": { # for NerfactoField / SDFField\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n \"field_background\": { # TODO: disable this optimizer during 2nd-stage training\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\n# try to replicate the result of nerfacto (using the nerfacto field in fg & bg)\ndebug_distilled_neus_facto_wbg = Config(\n method_name=\"debug_distilled-neus-facto-wbg\",\n trainer=TrainerConfig(\n steps_per_eval_image=2500,\n steps_per_eval_batch=2500,\n steps_per_save=30000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=30001,\n mixed_precision=True,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=AutoReconDataParserConfig(), # NOTE: cannot set init args here, which would not be recognised.\n train_num_rays_per_batch=4096, # 2048,\n eval_num_rays_per_batch=4096, # 1024,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n eval_camera_res_scale_factor=1.0,\n collate_fn=partial(nerfstudio_collate,\n extra_mappings={PCA: get_first_element,\n BasicPointClouds: get_first_element})\n ),\n model=DistilledNeuSFactoModelConfig(\n nerfacto_field_enabled=True, # for the 1st-stage training\n nerfacto_spatial_normalization_region=\"full\", # bg nerfacto always uses full\n single_field_fg_bg=True,\n proposal_net_use_separate_contraction=True, # proposal net scene contraction\n proposal_net_contraction_scale_factor=1.0,\n proposal_net_spatial_normalization_region=\"full\",\n proposal_use_uniform_sampler=False,\n use_fg_aware_scene_contraction=False, # main field scene contraction\n scene_contraction_scale_factor=1.0,\n sdf_field=SDFFieldConfig( # NeuSFacto Medium\n use_position_encoding=True,\n use_grid_feature=True,\n num_layers=2,\n num_layers_color=2,\n hidden_dim=64, # 256 -> 64\n geo_feat_dim=64, # 256 -> 64 -> 15\n bias=0.5,\n beta_init=0.3,\n direction_encoding_type=\"sh\",\n use_appearance_embedding=False, # TODO: improve impl efficiency (currently zero embeddings are used -> disable completely!)\n inside_outside=False, # only consider object-centric scenes\n ),\n cos_anneal_end=0, # TODO: set 5000 while training NeuS\n eval_num_rays_per_chunk=1024,\n use_average_appearance_embedding=False,\n background_model=\"grid\",\n bg_use_appearance_embedding=False,\n bg_sampler_type=\"none\",\n num_samples_outside=0,\n ),\n ),\n optimizers={\n \"proposal_networks\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n \"fields\": { # for NerfactoField\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n \"field_background\": { # NOTE: comment out this part if single_field_fg_bg is True\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/configs/autorecon/feature_field.py",
"content": "from . import *\n\n\nneus_facto_wbg_fast_dff = Config(\n method_name=\"neus-facto-wbg-fast_dff\",\n trainer=TrainerConfig(\n steps_per_eval_image=1000,\n steps_per_eval_batch=1000,\n steps_per_save=5000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=5001,\n mixed_precision=True,\n load_pipeline_ckpt_only=True,\n load_pipeline_ckpt_strict=False\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=AutoReconDataParserConfig(),\n train_num_rays_per_batch=2048,\n eval_num_rays_per_batch=1024,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n eval_camera_res_scale_factor=1.0,\n collate_fn=partial(nerfstudio_collate,\n extra_mappings={PCA: get_first_element,\n BasicPointClouds: get_first_element}),\n train_mask_sample_start=2000, # TODO: can be more aggresive e.g., 1\n train_mask_sample_ratio=0.8 # quality of regions out of the fg bbox mask is bad, so learning good feature field on bad geometry is meaningless\n ),\n model=NeuSFactoDFFModelConfig(\n cos_anneal_end=0, # must set to 0, as the sdf_field is already trained!\n use_proposal_weight_anneal=False, # as the proposal net is already trained\n use_proposal_weight_anneal_bg=False, # as the proposal net is already trained\n sdf_field=SDFFieldConfig(\n use_position_encoding=True,\n use_grid_feature=True,\n num_layers=2,\n num_layers_color=2,\n hidden_dim=256,\n bias=0.1, # 0.5 -> 0.1\n beta_init=0.3,\n direction_encoding_type=\"sh\",\n use_appearance_embedding=False,\n inside_outside=False, # only consider object-centric scenes\n hash_grid_progressive_training=True,\n hash_grid_progressive_training_iters=-1, # use weights=1\n spatial_normalization_region=\"aabb\", # to better utilize the feature grid\n weight_norm=False # weight_norm might lead to NaNs during training\n ),\n handle_bottom_intersection_rays=True, # this is vital to force the fg model to reconstruct the ground plane\n proposal_net_spatial_normalization_region=\"fg\",\n eval_num_rays_per_chunk=1024,\n use_average_appearance_embedding=False,\n background_model=\"grid\",\n bg_use_appearance_embedding=False,\n use_fg_aware_scene_contraction=True,\n fg_aware_scene_contraction_alpha=5.0, # NOTE: the reconstruction result is sensitive to this parameter, better set adaptively according to the coarse sfm reconstruction\n bg_sampler_type=\"proposal_network\",\n num_samples_outside=16,\n num_bg_proposal_samples_per_ray=(96, 64),\n use_surface_rendering=True,\n feature_field=FeatureFieldConfig(\n hash_grid_max_res=128,\n use_position_encoding=True, # smooth encoding\n use_fused_encoding_mlp=True\n ),\n feature_field_scene_contraction_scale_factor=0.25, # TODO: tune scale factor\n feature_loss_mult=1.0,\n render_fg_seg=True,\n fg_seg_binary_thr=0.5,\n feature_seg_field=FeatureSegFieldConfig(\n knn=10,\n contrast_fg_bg=True\n )\n ),\n ),\n optimizers={\n \"feature_field\": {\n \"optimizer\": AdamOptimizerConfig(lr=2e-3, eps=1e-15),\n \"scheduler\": MultiStepSchedulerConfig(max_steps=8000),\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\n\n# nesu-facto-wgb with Feature Field\nneus_facto_dff_wbg = Config(\n method_name=\"neus-facto-dff-wbg\",\n trainer=TrainerConfig(\n steps_per_eval_image=2500,\n steps_per_eval_batch=2500,\n steps_per_save=10000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=20001,\n mixed_precision=True,\n load_pipeline_ckpt_only=True,\n load_pipeline_ckpt_strict=False\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=AutoReconDataParserConfig(), # NOTE: cannot set init args here, which would not be recognised.\n train_num_rays_per_batch=2048,\n eval_num_rays_per_batch=1024,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n eval_camera_res_scale_factor=1.0,\n collate_fn=partial(nerfstudio_collate,\n extra_mappings={PCA: get_first_element,\n BasicPointClouds: get_first_element})\n ),\n model=NeuSFactoDFFModelConfig(\n sdf_field=SDFFieldConfig(\n # TODO: use SH instead of sine\n use_position_encoding=False,\n use_grid_feature=True,\n hash_grid_precision=\"float16\",\n num_layers=2,\n num_layers_color=2,\n hidden_dim=256,\n bias=0.5,\n beta_init=0.3,\n use_appearance_embedding=False, # TODO: improve impl efficiency (currently zero embeddings are used -> disable completely!)\n inside_outside=False, # only consider object-centric scenes\n ),\n bg_use_appearance_embedding=False,\n use_average_appearance_embedding=False,\n background_model=\"grid\",\n eval_num_rays_per_chunk=1024,\n proposal_net_spatial_normalization_region=\"fg\",\n feature_field=FeatureFieldConfig(\n hash_grid_max_res=128,\n use_position_encoding=True, # smooth encoding\n use_fused_encoding_mlp=True\n ),\n feature_loss_mult=1.0,\n render_fg_seg=True,\n fg_seg_binary_thr=0.5,\n feature_seg_field=FeatureSegFieldConfig(\n knn=10,\n contrast_fg_bg=True\n )\n ),\n ),\n optimizers={\n \"feature_field\": {\n \"optimizer\": AdamOptimizerConfig(lr=2e-3, eps=1e-15),\n \"scheduler\": None,\n }\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/configs/autorecon/neusfacto_fast.py",
"content": "from . import *\n\nneus_facto_wbg_fast = Config(\n method_name=\"neus-facto-wbg-fast\",\n trainer=TrainerConfig(\n steps_per_eval_image=2500,\n steps_per_eval_batch=2500,\n steps_per_save=30000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=30001,\n mixed_precision=True,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=AutoReconDataParserConfig(),\n train_num_rays_per_batch=2048,\n eval_num_rays_per_batch=2048,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n eval_camera_res_scale_factor=1.0,\n collate_fn=partial(\n nerfstudio_collate, extra_mappings={PCA: get_first_element, BasicPointClouds: get_first_element}\n ),\n train_mask_sample_start=10000,\n train_mask_sample_ratio=0.8,\n ),\n model=NeuSFactoModelConfig(\n sdf_field=SDFFieldConfig(\n use_position_encoding=True,\n use_grid_feature=True,\n num_layers=2,\n num_layers_color=2,\n hidden_dim=256,\n bias=0.1, # 0.5 -> 0.1\n beta_init=0.3,\n direction_encoding_type=\"sh\",\n use_appearance_embedding=False,\n inside_outside=False, # only consider object-centric scenes\n hash_grid_progressive_training=True,\n hash_grid_progressive_training_iters=10000,\n spatial_normalization_region=\"aabb\", # to better utilize the feature grid\n weight_norm=False, # weight_norm might lead to NaNs during training\n ),\n handle_bottom_intersection_rays=True, # this is vital to force the fg model to reconstruct the ground plane\n cos_anneal_end=5000,\n proposal_net_spatial_normalization_region=\"fg\",\n eval_num_rays_per_chunk=1024,\n use_average_appearance_embedding=False,\n background_model=\"grid\",\n bg_use_appearance_embedding=False,\n use_fg_aware_scene_contraction=True,\n fg_aware_scene_contraction_alpha=5.0, # NOTE: the reconstruction result is sensitive to this parameter, better set adaptively according to the coarse sfm reconstruction\n bg_sampler_type=\"proposal_network\",\n num_samples_outside=16,\n num_bg_proposal_samples_per_ray=(96, 64),\n ),\n ),\n optimizers={\n \"proposal_networks\": { # fixed gamma=0.3 and n_milestones=3\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n \"proposal_networks_bg\": { # fixed gamma=0.3 and n_milestones=3\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=2e-3, eps=1e-15, weight_decay=1e-6),\n \"scheduler\": MultiStepSchedulerConfig(max_steps=30001),\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/configs/autorecon/regularization.py",
"content": "from . import *\n\n# nesu-facto w/ explicit pointcloud regularization + separate plane nerf field modeling\nneus_facto_wbg_reg_sep_plane_nerf = Config(\n method_name=\"neus-facto-wbg-reg_sep-plane-nerf\",\n trainer=TrainerConfig(\n steps_per_eval_image=2500,\n steps_per_eval_batch=2500,\n steps_per_save=30000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=30001,\n mixed_precision=True,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=AutoReconDataParserConfig(),\n train_num_rays_per_batch=2048,\n eval_num_rays_per_batch=2048,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n eval_camera_res_scale_factor=1.0,\n collate_fn=partial(\n nerfstudio_collate, extra_mappings={PCA: get_first_element, BasicPointClouds: get_first_element}\n ),\n train_mask_sample_start=10000,\n train_mask_sample_ratio=0.8,\n ),\n model=NeuSFactoRegModelConfig(\n sdf_field=SDFFieldConfig(\n use_position_encoding=True,\n use_grid_feature=True,\n num_layers=2,\n num_layers_color=2,\n hidden_dim=256,\n bias=0.1, # 0.5 -> 0.1\n beta_init=0.3,\n direction_encoding_type=\"sh\",\n use_appearance_embedding=False,\n inside_outside=False, # only consider object-centric scenes\n hash_grid_progressive_training=False,\n hash_grid_progressive_training_iters=0,\n spatial_normalization_region=\"aabb\", # to better utilize the feature grid\n weight_norm=False, # weight_norm might lead to NaNs during training\n ),\n # plane_sdf_field\n plane_field_type=\"nerf\",\n plane_height_ratio=0.3,\n num_samples_plane=16,\n ptcd_reg_plane_field_mult=0.0,\n # explicit regularization\n ptcd_reg_fg_mult=0.1,\n ptcd_reg_plane_mult=0.1,\n ptcd_reg_n_iters=15000,\n ptcd_reg_weight_anneal_fn=\"cosine\",\n mask_beta_prior_mult=0.0,\n handle_bottom_intersection_rays=True,\n cos_anneal_end=5000,\n proposal_net_spatial_normalization_region=\"fg\",\n eval_num_rays_per_chunk=1024,\n use_average_appearance_embedding=False,\n background_model=\"grid\",\n bg_use_appearance_embedding=False,\n use_fg_aware_scene_contraction=True,\n fg_aware_scene_contraction_alpha=5.0, # NOTE: the reconstruction result is sensitive to this parameter, better set adaptively according to the coarse sfm reconstruction\n bg_sampler_type=\"proposal_network\",\n num_samples_outside=16,\n num_bg_proposal_samples_per_ray=(96, 64),\n # additional eikonal samples\n eikonal_loss_mult=0.05,\n num_proposal_samples_per_ray_for_eikonal_loss=64,\n ),\n ),\n optimizers={\n \"proposal_networks\": { # fixed gamma=0.3 and n_milestones=3\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n \"proposal_networks_bg\": { # fixed gamma=0.3 and n_milestones=3\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=2e-3, eps=1e-15, weight_decay=1e-6),\n \"scheduler\": MultiStepSchedulerConfig(max_steps=30001),\n },\n \"field_plane\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\n\n# nesu-facto w/ explicit pointcloud regularization + separate plane sdf field modeling\nneus_facto_wbg_reg_sep_plane = Config(\n method_name=\"neus-facto-wbg-reg_sep-plane\",\n trainer=TrainerConfig(\n steps_per_eval_image=2500,\n steps_per_eval_batch=2500,\n steps_per_save=30000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=30001,\n mixed_precision=True,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=AutoReconDataParserConfig(),\n train_num_rays_per_batch=2048,\n eval_num_rays_per_batch=2048,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n eval_camera_res_scale_factor=1.0,\n collate_fn=partial(\n nerfstudio_collate, extra_mappings={PCA: get_first_element, BasicPointClouds: get_first_element}\n ),\n train_mask_sample_start=10000,\n train_mask_sample_ratio=0.8,\n ),\n model=NeuSFactoRegModelConfig(\n sdf_field=SDFFieldConfig(\n use_position_encoding=True,\n use_grid_feature=True,\n num_layers=2,\n num_layers_color=2,\n hidden_dim=256,\n bias=0.1, # 0.5 -> 0.1\n beta_init=0.3,\n direction_encoding_type=\"sh\",\n use_appearance_embedding=False,\n inside_outside=False, # only consider object-centric scenes\n hash_grid_progressive_training=False,\n hash_grid_progressive_training_iters=0,\n spatial_normalization_region=\"aabb\", # to better utilize the feature grid\n weight_norm=False, # weight_norm might lead to NaNs during training\n ),\n # plane_sdf_field\n plane_height_ratio=0.4,\n num_samples_plane=16,\n ptcd_reg_plane_field_mult=0.0,\n # explicit regularization\n ptcd_reg_fg_mult=0.1,\n ptcd_reg_plane_mult=0.1,\n ptcd_reg_n_iters=15000,\n ptcd_reg_weight_anneal_fn=\"cosine\",\n mask_beta_prior_mult=0.0,\n handle_bottom_intersection_rays=True,\n cos_anneal_end=0,\n proposal_net_spatial_normalization_region=\"fg\",\n eval_num_rays_per_chunk=1024,\n use_average_appearance_embedding=False,\n background_model=\"grid\",\n bg_use_appearance_embedding=False,\n use_fg_aware_scene_contraction=True,\n fg_aware_scene_contraction_alpha=5.0, # NOTE: the reconstruction result is sensitive to this parameter, better set adaptively according to the coarse sfm reconstruction\n bg_sampler_type=\"proposal_network\",\n num_samples_outside=16,\n num_bg_proposal_samples_per_ray=(96, 64),\n ),\n ),\n optimizers={\n \"proposal_networks\": { # fixed gamma=0.3 and n_milestones=3\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n \"proposal_networks_bg\": { # fixed gamma=0.3 and n_milestones=3\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=2e-3, eps=1e-15, weight_decay=1e-6),\n \"scheduler\": MultiStepSchedulerConfig(max_steps=30001),\n },\n \"field_plane\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-3, eps=1e-15, weight_decay=1e-6),\n \"scheduler\": MultiStepSchedulerConfig(max_steps=30001),\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\n\n# nesu-facto w/ explicit pointcloud regularization\nneus_facto_wbg_reg = Config(\n method_name=\"neus-facto-wbg-reg\",\n trainer=TrainerConfig(\n steps_per_eval_image=2500,\n steps_per_eval_batch=2500,\n steps_per_save=30000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=30001,\n mixed_precision=True,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=AutoReconDataParserConfig(),\n train_num_rays_per_batch=2048,\n eval_num_rays_per_batch=2048,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n eval_camera_res_scale_factor=1.0,\n collate_fn=partial(\n nerfstudio_collate, extra_mappings={PCA: get_first_element, BasicPointClouds: get_first_element}\n ),\n train_mask_sample_start=10000,\n train_mask_sample_ratio=0.8,\n ),\n model=NeuSFactoModelConfig(\n sdf_field=SDFFieldConfig(\n use_position_encoding=True,\n use_grid_feature=True,\n num_layers=2,\n num_layers_color=2,\n hidden_dim=256,\n bias=0.1, # 0.5 -> 0.1\n beta_init=0.3,\n direction_encoding_type=\"sh\",\n use_appearance_embedding=False,\n inside_outside=False, # only consider object-centric scenes\n hash_grid_progressive_training=True,\n hash_grid_progressive_training_iters=10000,\n spatial_normalization_region=\"aabb\", # to better utilize the feature grid\n weight_norm=False, # weight_norm might lead to NaNs during training\n ),\n # explicit regularization\n ptcd_reg_fg_mult=0.1,\n ptcd_reg_plane_mult=0.1,\n ptcd_reg_n_iters=15000,\n ptcd_reg_weight_anneal_fn=\"cosine\",\n mask_beta_prior_mult=0.0,\n handle_bottom_intersection_rays=False,\n cos_anneal_end=5000,\n proposal_net_spatial_normalization_region=\"fg\",\n eval_num_rays_per_chunk=1024,\n use_average_appearance_embedding=False,\n background_model=\"grid\",\n bg_use_appearance_embedding=False,\n use_fg_aware_scene_contraction=True,\n fg_aware_scene_contraction_alpha=5.0,\n bg_sampler_type=\"proposal_network\",\n num_samples_outside=16,\n num_bg_proposal_samples_per_ray=(96, 64),\n ),\n ),\n optimizers={\n \"proposal_networks\": { # fixed gamma=0.3 and n_milestones=3\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n \"proposal_networks_bg\": { # fixed gamma=0.3 and n_milestones=3\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=2e-3, eps=1e-15, weight_decay=1e-6),\n \"scheduler\": MultiStepSchedulerConfig(max_steps=30001),\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/configs/autorecon/semantic_nerf.py",
"content": "from . import *\n\n\nautorecon_semantic_nerf = Config(\n method_name=\"autorecon_semantic-nerfw\",\n trainer=TrainerConfig(\n steps_per_eval_batch=500, steps_per_eval_image=500,\n steps_per_save=30000, max_num_iterations=30001, mixed_precision=True\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=AutoReconDataManagerConfig(\n dataparser=AutoReconDataParserConfig(),\n train_num_rays_per_batch=4096,\n eval_num_rays_per_batch=4096\n ),\n model=SemanticNerfWModelConfig(\n eval_num_rays_per_chunk=1 << 16,\n semantic_mult=1.0,\n use_transient_embedding=False,\n scene_contraction_scale_factor=0.25,\n ),\n ),\n optimizers={\n \"proposal_networks\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 16),\n vis=\"viewer\",\n)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/configs/base_config.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Base Configs\"\"\"\n\n# pylint: disable=wrong-import-position\n\nfrom __future__ import annotations\n\nimport warnings\nfrom dataclasses import dataclass\nfrom datetime import datetime\nfrom pathlib import Path\nfrom typing import Any, Dict, List, Optional, Tuple, Type\n\nimport yaml\nfrom rich.console import Console\nfrom typing_extensions import Literal\n\nfrom nerfstudio.configs.config_utils import to_immutable_dict\n\n# model instances\nfrom nerfstudio.utils import writer\n\nwarnings.filterwarnings(\"ignore\", module=\"torchvision\")\n\nCONSOLE = Console(width=120)\n\nCFG = None # global config accessible by all modules\n\n\n# Pretty printing class\nclass PrintableConfig: # pylint: disable=too-few-public-methods\n \"\"\"Printable Config defining str function\"\"\"\n\n def __str__(self):\n lines = [self.__class__.__name__ + \":\"]\n for key, val in vars(self).items():\n if isinstance(val, Tuple):\n flattened_val = \"[\"\n for item in val:\n flattened_val += str(item) + \"\\n\"\n flattened_val = flattened_val.rstrip(\"\\n\")\n val = flattened_val + \"]\"\n lines += f\"{key}: {str(val)}\".split(\"\\n\")\n return \"\\n \".join(lines)\n\n\n# Base instantiate configs\n@dataclass\nclass InstantiateConfig(PrintableConfig): # pylint: disable=too-few-public-methods\n \"\"\"Config class for instantiating an the class specified in the _target attribute.\"\"\"\n\n _target: Type\n\n def setup(self, **kwargs) -> Any:\n \"\"\"Returns the instantiated object using the config.\"\"\"\n return self._target(self, **kwargs)\n\n\n# Machine related configs\n@dataclass\nclass MachineConfig(PrintableConfig):\n \"\"\"Configuration of machine setup\"\"\"\n\n seed: int = 42\n \"\"\"random seed initilization\"\"\"\n num_gpus: int = 1\n \"\"\"total number of gpus available for train/eval\"\"\"\n num_machines: int = 1\n \"\"\"total number of distributed machines available (for DDP)\"\"\"\n machine_rank: int = 0\n \"\"\"current machine's rank (for DDP)\"\"\"\n dist_url: str = \"auto\"\n \"\"\"distributed connection point (for DDP)\"\"\"\n\n\n@dataclass\nclass LocalWriterConfig(InstantiateConfig):\n \"\"\"Local Writer config\"\"\"\n\n _target: Type = writer.LocalWriter\n \"\"\"target class to instantiate\"\"\"\n enable: bool = False\n \"\"\"if True enables local logging, else disables\"\"\"\n stats_to_track: Tuple[writer.EventName, ...] = (\n writer.EventName.ITER_TRAIN_TIME,\n writer.EventName.TRAIN_RAYS_PER_SEC,\n writer.EventName.CURR_TEST_PSNR,\n writer.EventName.VIS_RAYS_PER_SEC,\n writer.EventName.TEST_RAYS_PER_SEC,\n )\n \"\"\"specifies which stats will be logged/printed to terminal\"\"\"\n max_log_size: int = 10\n \"\"\"maximum number of rows to print before wrapping. if 0, will print everything.\"\"\"\n\n def setup(self, banner_messages: Optional[List[str]] = None, **kwargs) -> Any:\n \"\"\"Instantiate local writer\n\n Args:\n banner_messages: List of strings that always print at the bottom of screen.\n \"\"\"\n return self._target(self, banner_messages=banner_messages, **kwargs)\n\n\n@dataclass\nclass LoggingConfig(PrintableConfig):\n \"\"\"Configuration of loggers and profilers\"\"\"\n\n relative_log_dir: Path = Path(\"./\")\n \"\"\"relative path to save all logged events\"\"\"\n steps_per_log: int = 10\n \"\"\"number of steps between logging stats\"\"\"\n max_buffer_size: int = 20\n \"\"\"maximum history size to keep for computing running averages of stats.\n e.g. if 20, averages will be computed over past 20 occurances.\"\"\"\n local_writer: LocalWriterConfig = LocalWriterConfig(enable=True)\n \"\"\"if provided, will print stats locally. if None, will disable printing\"\"\"\n enable_profiler: bool = True\n \"\"\"whether to enable profiling code; prints speed of functions at the end of a program.\n profiler logs run times of functions and prints at end of training\"\"\"\n\n\n# Trainer related configs\n@dataclass\nclass TrainerConfig(PrintableConfig):\n \"\"\"Configuration for training regimen\"\"\"\n\n steps_per_save: int = 1000\n \"\"\"Number of steps between saves.\"\"\"\n steps_per_eval_batch: int = 500\n \"\"\"Number of steps between randomly sampled batches of rays.\"\"\"\n steps_per_eval_image: int = 500\n \"\"\"Number of steps between single eval images.\"\"\"\n steps_per_eval_all_images: int = 25000\n \"\"\"Number of steps between eval all images.\"\"\"\n max_num_iterations: int = 1000000\n \"\"\"Maximum number of iterations to run.\"\"\"\n mixed_precision: bool = False\n \"\"\"Whether or not to use mixed precision for training.\"\"\"\n relative_model_dir: Path = Path(\"sdfstudio_models/\")\n \"\"\"Relative path to save all checkpoints.\"\"\"\n save_only_latest_checkpoint: bool = True\n \"\"\"Whether to only save the latest checkpoint or all checkpoints.\"\"\"\n # optional parameters if we want to resume training\n load_dir: Optional[Path] = None\n \"\"\"Optionally specify a pre-trained model directory to load from.\"\"\"\n load_step: Optional[int] = None\n \"\"\"Optionally specify model step to load from; if none, will find most recent model in load_dir.\"\"\"\n load_pipeline_ckpt_only: bool = False\n \"\"\"Only load pipeline checkpoints. Ignore optimizer, grad_sacler, ...\"\"\"\n load_pipeline_ckpt_strict: bool = True\n \"\"\"Whether or not to load pipeline ckpt strictly.\"\"\"\n load_config: Optional[Path] = None\n \"\"\"Optionally specify model config to load from; if none, will use the default config?\"\"\"\n load_scheduler: bool = True\n \"\"\"Whether to load the lr scheduler state_dict if exists\"\"\"\n\n\n# Viewer related configs\n@dataclass\nclass ViewerConfig(PrintableConfig):\n \"\"\"Configuration for viewer instantiation\"\"\"\n\n relative_log_filename: str = \"viewer_log_filename.txt\"\n \"\"\"Filename to use for the log file.\"\"\"\n start_train: bool = True\n \"\"\"whether to immediately start training upon loading viewer\n if False, will just visualize dataset but you can toggle training in viewer\"\"\"\n zmq_port: Optional[int] = None\n \"\"\"The zmq port to connect to for communication. If None, find an available port.\"\"\"\n launch_bridge_server: bool = True\n \"\"\"whether or not to launch the bridge server\"\"\"\n websocket_port: Optional[int] = 7007\n \"\"\"the default websocket port to connect to\"\"\"\n ip_address: str = \"127.0.0.1\"\n \"\"\"the ip address where the bridge server is running\"\"\"\n num_rays_per_chunk: int = 32768\n \"\"\"number of rays per chunk to render with viewer\"\"\"\n max_num_display_images: int = 512\n \"\"\"Maximum number of training images to display in the viewer, to avoid lag. This does not change which images are\n actually used in training/evaluation. If -1, display all.\"\"\"\n quit_on_train_completion: bool = False\n \"\"\"Whether to kill the training job when it has completed. Note this will stop rendering in the viewer.\"\"\"\n skip_openrelay: bool = False\n \"\"\"Avoid using openrelay to communicate with the viewer. Try disabling if you have trouble\n connecting to the viewer\"\"\"\n\n\nfrom nerfstudio.engine.optimizers import OptimizerConfig\nfrom nerfstudio.engine.schedulers import SchedulerConfig\nfrom nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig\n\n\n@dataclass\nclass Config(PrintableConfig):\n \"\"\"Full config contents\"\"\"\n\n output_dir: Path = Path(\"outputs\")\n \"\"\"relative or absolute output directory to save all checkpoints and logging\"\"\"\n method_name: Optional[str] = None\n \"\"\"Method name. Required to set in python or via cli\"\"\"\n experiment_name: Optional[str] = None\n \"\"\"Experiment name. If None, will automatically be set to dataset name\"\"\"\n timestamp: str = \"{timestamp}\"\n \"\"\"Experiment timestamp.\"\"\"\n machine: MachineConfig = MachineConfig()\n \"\"\"Machine configuration\"\"\"\n logging: LoggingConfig = LoggingConfig()\n \"\"\"Logging configuration\"\"\"\n viewer: ViewerConfig = ViewerConfig()\n \"\"\"Viewer configuration\"\"\"\n trainer: TrainerConfig = TrainerConfig()\n \"\"\"Trainer configuration\"\"\"\n pipeline: VanillaPipelineConfig = VanillaPipelineConfig()\n \"\"\"Pipeline configuration\"\"\"\n optimizers: Dict[str, Any] = to_immutable_dict(\n {\n \"fields\": {\n \"optimizer\": OptimizerConfig(),\n \"scheduler\": SchedulerConfig(),\n }\n }\n )\n \"\"\"Dictionary of optimizer groups and their schedulers\"\"\"\n vis: Literal[\"viewer\", \"wandb\", \"tensorboard\"] = \"wandb\"\n \"\"\"Which visualizer to use.\"\"\"\n data: Optional[Path] = None\n \"\"\"Alias for --pipeline.datamanager.dataparser.data\"\"\"\n\n def set_global(self) -> None:\n \"\"\"\n Example:\n from nerfstudio.configs import base_config\n base_config.CFG.xxx()\n FIXME: this is useless if the original CFG is already imported by other module\n TODO: make the global config read-only\n \"\"\"\n global CFG\n CFG = self\n\n def is_viewer_enabled(self) -> bool:\n \"\"\"Checks if a viewer is enabled.\"\"\"\n return \"viewer\" == self.vis\n\n def is_wandb_enabled(self) -> bool:\n \"\"\"Checks if wandb is enabled.\"\"\"\n return \"wandb\" == self.vis\n\n def is_tensorboard_enabled(self) -> bool:\n \"\"\"Checks if tensorboard is enabled.\"\"\"\n return \"tensorboard\" == self.vis\n\n def set_timestamp(self) -> None:\n \"\"\"Dynamically set the experiment timestamp\"\"\"\n if self.timestamp == \"{timestamp}\":\n self.timestamp = datetime.now().strftime(\"%Y-%m-%d_%H%M%S\")\n\n def set_experiment_name(self) -> None:\n \"\"\"Dynamically set the experiment name\"\"\"\n if self.experiment_name is None:\n self.experiment_name = str(self.pipeline.datamanager.dataparser.data).replace(\"../\", \"\").replace(\"/\", \"-\")\n\n def get_base_dir(self) -> Path:\n \"\"\"Retrieve the base directory to set relative paths\"\"\"\n # check the experiment and method names\n assert self.method_name is not None, \"Please set method name in config or via the cli\"\n self.set_experiment_name()\n return Path(f\"{self.output_dir}/{self.experiment_name}/{self.method_name}/{self.timestamp}\")\n\n def get_checkpoint_dir(self) -> Path:\n \"\"\"Retrieve the checkpoint directory\"\"\"\n return Path(self.get_base_dir() / self.trainer.relative_model_dir)\n\n def print_to_terminal(self) -> None:\n \"\"\"Helper to pretty print config to terminal\"\"\"\n CONSOLE.rule(\"Config\")\n CONSOLE.print(self)\n CONSOLE.rule(\"\")\n\n def save_config(self) -> None:\n \"\"\"Save config to base directory\"\"\"\n base_dir = self.get_base_dir()\n assert base_dir is not None\n base_dir.mkdir(parents=True, exist_ok=True)\n config_yaml_path = base_dir / \"config.yml\"\n CONSOLE.log(f\"Saving config to: {config_yaml_path}\")\n config_yaml_path.write_text(yaml.dump(self), \"utf8\")\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/configs/config_utils.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nSome utility code for configs.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom dataclasses import field\nfrom typing import Any, Dict\n\nfrom rich.console import Console\n\nCONSOLE = Console()\n# pylint: disable=import-outside-toplevel\n\n# cannot use mutable types directly within dataclass; abstracting default factory calls\ndef to_immutable_dict(d: Dict[str, Any]):\n \"\"\"Method to convert mutable dict to default factory dict\n\n Args:\n d: dictionary to convert into default factory dict for dataclass\n \"\"\"\n return field(default_factory=lambda: dict(d))\n\n\ndef convert_markup_to_ansi(markup_string: str) -> str:\n \"\"\"Convert rich-style markup to ANSI sequences for command-line formatting.\n\n Args:\n markup_string: Text with rich-style markup.\n\n Returns:\n Text formatted via ANSI sequences.\n \"\"\"\n with CONSOLE.capture() as out:\n CONSOLE.print(markup_string, soft_wrap=True)\n return out.get()\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/configs/method_configs.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nPut all the method implementations in one location.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom typing import Dict\n\nimport tyro\n\nfrom nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig\nfrom nerfstudio.configs.autorecon import method_configs as autorecon_method_configs\nfrom nerfstudio.configs.base_config import (\n Config,\n SchedulerConfig,\n TrainerConfig,\n ViewerConfig,\n)\nfrom nerfstudio.data.datamanagers.base_datamanager import (\n FlexibleDataManagerConfig,\n VanillaDataManagerConfig,\n)\nfrom nerfstudio.data.datamanagers.semantic_datamanager import SemanticDataManagerConfig\nfrom nerfstudio.data.datamanagers.variable_res_datamanager import (\n VariableResDataManagerConfig,\n)\nfrom nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig\nfrom nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig\nfrom nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConfig\nfrom nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig\nfrom nerfstudio.data.dataparsers.phototourism_dataparser import (\n PhototourismDataParserConfig,\n)\nfrom nerfstudio.data.dataparsers.sdfstudio_dataparser import SDFStudioDataParserConfig\nfrom nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig\nfrom nerfstudio.engine.schedulers import (\n ExponentialSchedulerConfig,\n MultiStepSchedulerConfig,\n NeuSSchedulerConfig,\n)\nfrom nerfstudio.field_components.temporal_distortions import TemporalDistortionKind\nfrom nerfstudio.fields.sdf_field import SDFFieldConfig\nfrom nerfstudio.models.dto import DtoOModelConfig\nfrom nerfstudio.models.instant_ngp import InstantNGPModelConfig\nfrom nerfstudio.models.mipnerf import MipNerfModel\nfrom nerfstudio.models.nerfacto import NerfactoModelConfig\nfrom nerfstudio.models.neuralreconW import NeuralReconWModelConfig\nfrom nerfstudio.models.neus import NeuSModelConfig\nfrom nerfstudio.models.neus_acc import NeuSAccModelConfig\nfrom nerfstudio.models.neus_facto import NeuSFactoModelConfig\nfrom nerfstudio.models.semantic_nerfw import SemanticNerfWModelConfig\nfrom nerfstudio.models.tensorf import TensoRFModelConfig\nfrom nerfstudio.models.unisurf import UniSurfModelConfig\nfrom nerfstudio.models.vanilla_nerf import NeRFModel, VanillaModelConfig\nfrom nerfstudio.models.volsdf import VolSDFModelConfig\nfrom nerfstudio.pipelines.base_pipeline import (\n FlexibleInputPipelineConfig,\n VanillaPipelineConfig,\n)\nfrom nerfstudio.pipelines.dynamic_batch import DynamicBatchPipelineConfig\n\nmethod_configs: Dict[str, Config] = {}\nmethod_configs.update(autorecon_method_configs)\n\ndescriptions = {\n \"nerfacto\": \"Recommended real-time model tuned for real captures. This model will be continually updated.\",\n \"instant-ngp\": \"Implementation of Instant-NGP. Recommended real-time model for bounded synthetic data.\",\n \"mipnerf\": \"High quality model for bounded scenes. (slow)\",\n \"semantic-nerfw\": \"Predicts semantic segmentations and filters out transient objects.\",\n \"vanilla-nerf\": \"Original NeRF model. (slow)\",\n \"volsdf\": \"Implementation of VolSDF.\",\n \"monosdf\": \"Implementation of MonoSDF.\",\n \"geo-volsdf\": \"Implementation of patch warping from GeoNeuS with VolSDF.\",\n \"neus\": \"Implementation of NeuS.\",\n \"mono-neus\": \"Implementation of MonoSDF with NeuS rendering formulation.\",\n \"geo-neus\": \"Implementation of patch warping from GeoNeuS with NeuS.\",\n \"unisurf\": \"Implementation of UniSurf.\",\n \"mono-unisurf\": \"Implementation of MonoSDF with unisurf rendering formulation.\",\n \"geo-unisurf\": \"Implementation of patch warping from GeoNeuS with UniSurf.\",\n \"tensorf\": \"tensorf\",\n \"dnerf\": \"Dynamic-NeRF model. (slow)\",\n \"phototourism\": \"Uses the Phototourism data.\",\n \"dto\": \"Occupancy field with density guided sampling\",\n \"neusW\": \"Implementation of Neural Reconstruction in the wild\",\n \"neus-acc\": \"Implementation of NeuS with empty space skipping.\",\n \"neus-facto\": \"Implementation of NeuS similar to nerfacto where proposal sampler is used.\",\n \"neus-facto-bigmlp\": \"NeuS-facto with big MLP, it is used in training heritage data with 8 gpus\",\n}\n\n# assume the entire scene is bounded by a unit cube\nmethod_configs[\"neus-facto\"] = Config(\n method_name=\"neus-facto\",\n trainer=TrainerConfig(\n steps_per_eval_image=5000,\n steps_per_eval_batch=5000,\n steps_per_save=20000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=20001,\n mixed_precision=False,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=SDFStudioDataParserConfig(),\n train_num_rays_per_batch=2048,\n eval_num_rays_per_batch=1024,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n ),\n model=NeuSFactoModelConfig(\n sdf_field=SDFFieldConfig(\n use_grid_feature=True,\n num_layers=2,\n num_layers_color=2,\n hidden_dim=256,\n bias=0.5,\n beta_init=0.3,\n use_appearance_embedding=False,\n ),\n background_model=\"none\", # FIXME\n eval_num_rays_per_chunk=1024,\n ),\n ),\n optimizers={\n \"proposal_networks\": { # fixed gamma=0.3 and n_milestones=3\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": MultiStepSchedulerConfig(max_steps=20000),\n },\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=500, learning_rate_alpha=0.05, max_steps=20000),\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=500, learning_rate_alpha=0.05, max_steps=20000),\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\nmethod_configs[\"neus-facto-bigmlp\"] = Config(\n method_name=\"neus-facto-bigmlp\",\n trainer=TrainerConfig(\n steps_per_eval_image=5000,\n steps_per_eval_batch=5000,\n steps_per_save=20000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=100001,\n mixed_precision=False,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=SDFStudioDataParserConfig(),\n train_num_rays_per_batch=2048,\n eval_num_rays_per_batch=1024,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n ),\n model=NeuSFactoModelConfig(\n sdf_field=SDFFieldConfig(num_layers=8, hidden_dim=512, num_layers_color=4), eval_num_rays_per_chunk=1024\n ),\n ),\n optimizers={\n \"proposal_networks\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": MultiStepSchedulerConfig(max_steps=100000),\n },\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-3, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=500, learning_rate_alpha=0.05, max_steps=100000),\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=500, learning_rate_alpha=0.05, max_steps=100000),\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\nmethod_configs[\"geo-volsdf\"] = Config(\n method_name=\"geo-volsdf\",\n trainer=TrainerConfig(\n steps_per_eval_image=5000,\n steps_per_eval_batch=5000,\n steps_per_save=20000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=200001,\n mixed_precision=False,\n ),\n pipeline=FlexibleInputPipelineConfig(\n datamanager=FlexibleDataManagerConfig(\n dataparser=SDFStudioDataParserConfig(),\n train_num_rays_per_batch=1024,\n eval_num_rays_per_batch=1024,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n ),\n model=VolSDFModelConfig(patch_warp_loss_mult=0.1, eval_num_rays_per_chunk=1024),\n ),\n optimizers={\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": MultiStepSchedulerConfig(\n max_steps=1000000\n ), # set max_steps to a large value so it never step and we will use the last_lr form the pretrained model\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": ExponentialSchedulerConfig(decay_rate=0.1, max_steps=200000),\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\nmethod_configs[\"monosdf\"] = Config(\n method_name=\"monosdf\",\n trainer=TrainerConfig(\n steps_per_eval_image=5000,\n steps_per_eval_batch=5000,\n steps_per_save=20000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=200000,\n mixed_precision=False,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=SDFStudioDataParserConfig(),\n train_num_rays_per_batch=1024,\n eval_num_rays_per_batch=1024,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n ),\n model=VolSDFModelConfig(mono_depth_loss_mult=0.1, mono_normal_loss_mult=0.05, eval_num_rays_per_chunk=1024),\n ),\n optimizers={\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": ExponentialSchedulerConfig(decay_rate=0.1, max_steps=200000),\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": ExponentialSchedulerConfig(decay_rate=0.1, max_steps=200000),\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\nmethod_configs[\"volsdf\"] = Config(\n method_name=\"volsdf\",\n trainer=TrainerConfig(\n steps_per_eval_image=5000,\n steps_per_eval_batch=5000,\n steps_per_save=20000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=100000,\n mixed_precision=False,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=SDFStudioDataParserConfig(),\n train_num_rays_per_batch=1024,\n eval_num_rays_per_batch=1024,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n ),\n model=VolSDFModelConfig(eval_num_rays_per_chunk=1024),\n ),\n optimizers={\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": ExponentialSchedulerConfig(decay_rate=0.1, max_steps=100000),\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": ExponentialSchedulerConfig(decay_rate=0.1, max_steps=100000),\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\nmethod_configs[\"geo-neus\"] = Config(\n method_name=\"geo-neus\",\n trainer=TrainerConfig(\n steps_per_eval_image=500,\n steps_per_eval_batch=5000,\n steps_per_save=20000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=200000,\n mixed_precision=False,\n ),\n pipeline=FlexibleInputPipelineConfig(\n datamanager=FlexibleDataManagerConfig(\n dataparser=SDFStudioDataParserConfig(),\n train_num_rays_per_batch=1024,\n eval_num_rays_per_batch=1024,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n ),\n model=NeuSModelConfig(patch_warp_loss_mult=0.1, eval_num_rays_per_chunk=1024),\n ),\n optimizers={\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000),\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000),\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\nmethod_configs[\"mono-neus\"] = Config(\n method_name=\"mono-neus\",\n trainer=TrainerConfig(\n steps_per_eval_image=500,\n steps_per_eval_batch=5000,\n steps_per_save=20000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=100000,\n mixed_precision=False,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=SDFStudioDataParserConfig(),\n train_num_rays_per_batch=1024,\n eval_num_rays_per_batch=1024,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n ),\n model=NeuSModelConfig(mono_depth_loss_mult=0.1, mono_normal_loss_mult=0.05, eval_num_rays_per_chunk=1024),\n ),\n optimizers={\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000),\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000),\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\nmethod_configs[\"neus\"] = Config(\n method_name=\"neus\",\n trainer=TrainerConfig(\n steps_per_eval_image=500,\n steps_per_eval_batch=5000,\n steps_per_save=20000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=100000, # FIXME: 300000?\n mixed_precision=False,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=SDFStudioDataParserConfig(),\n train_num_rays_per_batch=1024,\n eval_num_rays_per_batch=1024,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n ),\n model=NeuSModelConfig(eval_num_rays_per_chunk=1024),\n ),\n optimizers={\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000),\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000),\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\nmethod_configs[\"unisurf\"] = Config(\n method_name=\"unisurf\",\n trainer=TrainerConfig(\n steps_per_eval_image=500,\n steps_per_eval_batch=5000,\n steps_per_save=20000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=100000,\n mixed_precision=False,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=SDFStudioDataParserConfig(),\n train_num_rays_per_batch=1024,\n eval_num_rays_per_batch=1024,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n ),\n model=UniSurfModelConfig(eval_num_rays_per_chunk=1024),\n ),\n optimizers={\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000),\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000),\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\nmethod_configs[\"mono-unisurf\"] = Config(\n method_name=\"mono-unisurf\",\n trainer=TrainerConfig(\n steps_per_eval_image=500,\n steps_per_eval_batch=5000,\n steps_per_save=20000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=100000,\n mixed_precision=False,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=SDFStudioDataParserConfig(),\n train_num_rays_per_batch=1024,\n eval_num_rays_per_batch=1024,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n ),\n model=UniSurfModelConfig(mono_depth_loss_mult=0.1, mono_normal_loss_mult=0.05, eval_num_rays_per_chunk=1024),\n ),\n optimizers={\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000),\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000),\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\nmethod_configs[\"geo-unisurf\"] = Config(\n method_name=\"geo-unisurf\",\n trainer=TrainerConfig(\n steps_per_eval_image=500,\n steps_per_eval_batch=5000,\n steps_per_save=20000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=100000,\n mixed_precision=False,\n ),\n pipeline=FlexibleInputPipelineConfig(\n datamanager=FlexibleDataManagerConfig(\n dataparser=SDFStudioDataParserConfig(),\n train_num_rays_per_batch=1024,\n eval_num_rays_per_batch=1024,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n ),\n model=UniSurfModelConfig(patch_warp_loss_mult=0.1, eval_num_rays_per_chunk=1024),\n ),\n optimizers={\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000),\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000),\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\nmethod_configs[\"dto\"] = Config(\n method_name=\"dto\",\n trainer=TrainerConfig(\n steps_per_eval_image=2000,\n steps_per_eval_batch=5000,\n steps_per_save=5000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=100000,\n mixed_precision=False,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=SDFStudioDataParserConfig(),\n train_num_rays_per_batch=2048,\n eval_num_rays_per_batch=2048,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n ),\n model=DtoOModelConfig(eval_num_rays_per_chunk=1 << 10),\n ),\n optimizers={\n \"proposal_networks\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": MultiStepSchedulerConfig(max_steps=300000),\n },\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": MultiStepSchedulerConfig(max_steps=300000),\n },\n \"occupancy_field\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=500, learning_rate_alpha=0.05, max_steps=300000),\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\nmethod_configs[\"neusW\"] = Config(\n method_name=\"neusW\",\n trainer=TrainerConfig(\n steps_per_eval_image=5000,\n steps_per_eval_batch=5000,\n steps_per_save=5000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=100000,\n mixed_precision=False,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=SDFStudioDataParserConfig(),\n train_num_rays_per_batch=2048,\n eval_num_rays_per_batch=2048,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n ),\n model=NeuralReconWModelConfig(\n background_model=\"grid\", num_samples_outside=4, eikonal_loss_mult=0.0001, eval_num_rays_per_chunk=1 << 10\n ),\n ),\n optimizers={\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-3, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=500, learning_rate_alpha=0.05, max_steps=300000),\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": MultiStepSchedulerConfig(max_steps=300000),\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"wandb\",\n)\n\nmethod_configs[\"neus-acc\"] = Config(\n method_name=\"neus-acc\",\n trainer=TrainerConfig(\n steps_per_eval_image=5000,\n steps_per_eval_batch=5000,\n steps_per_save=20000,\n steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images\n max_num_iterations=20000,\n mixed_precision=False,\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=SDFStudioDataParserConfig(),\n train_num_rays_per_batch=2048,\n eval_num_rays_per_batch=1024,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"off\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n ),\n model=NeuSAccModelConfig(eval_num_rays_per_chunk=1024),\n ),\n optimizers={\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=500, learning_rate_alpha=0.05, max_steps=20000),\n },\n \"field_background\": {\n \"optimizer\": AdamOptimizerConfig(lr=5e-4, eps=1e-15),\n \"scheduler\": NeuSSchedulerConfig(warm_up_end=500, learning_rate_alpha=0.05, max_steps=20000),\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\nmethod_configs[\"nerfacto\"] = Config(\n method_name=\"nerfacto\",\n trainer=TrainerConfig(\n steps_per_eval_batch=5000, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=NerfstudioDataParserConfig(),\n train_num_rays_per_batch=4096,\n eval_num_rays_per_batch=4096,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"SO3xR3\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n ),\n model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15),\n ),\n optimizers={\n \"proposal_networks\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": MultiStepSchedulerConfig(max_steps=300000), # FIXME: 30000?\n },\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": MultiStepSchedulerConfig(max_steps=300000), # FIXME: 30000?\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\nmethod_configs[\"instant-ngp\"] = Config(\n method_name=\"instant-ngp\",\n trainer=TrainerConfig(\n steps_per_eval_batch=5000,\n steps_per_eval_image=5000,\n steps_per_save=20000,\n max_num_iterations=20001,\n mixed_precision=True,\n steps_per_eval_all_images=20000,\n ),\n pipeline=DynamicBatchPipelineConfig(\n datamanager=VanillaDataManagerConfig(dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=8192),\n model=InstantNGPModelConfig(render_step_size=0.005, eval_num_rays_per_chunk=8192),\n ),\n optimizers={\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": MultiStepSchedulerConfig(max_steps=20000),\n }\n },\n viewer=ViewerConfig(num_rays_per_chunk=64000),\n vis=\"viewer\",\n)\n\nmethod_configs[\"mipnerf\"] = Config(\n method_name=\"mipnerf\",\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=1024),\n model=VanillaModelConfig(\n _target=MipNerfModel,\n loss_coefficients={\"rgb_loss_coarse\": 0.1, \"rgb_loss_fine\": 1.0},\n num_coarse_samples=128,\n num_importance_samples=128,\n eval_num_rays_per_chunk=1024,\n ),\n ),\n optimizers={\n \"fields\": {\n \"optimizer\": RAdamOptimizerConfig(lr=5e-4, eps=1e-08),\n \"scheduler\": None,\n }\n },\n)\n\nmethod_configs[\"semantic-nerfw\"] = Config(\n method_name=\"semantic-nerfw\",\n trainer=TrainerConfig(\n steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=SemanticDataManagerConfig(\n dataparser=FriendsDataParserConfig(), train_num_rays_per_batch=4096, eval_num_rays_per_batch=8192\n ),\n model=SemanticNerfWModelConfig(eval_num_rays_per_chunk=1 << 16),\n ),\n optimizers={\n \"proposal_networks\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 16),\n vis=\"viewer\",\n)\n\nmethod_configs[\"vanilla-nerf\"] = Config(\n method_name=\"vanilla-nerf\",\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=BlenderDataParserConfig(),\n ),\n model=VanillaModelConfig(_target=NeRFModel),\n ),\n optimizers={\n \"fields\": {\n \"optimizer\": RAdamOptimizerConfig(lr=5e-4, eps=1e-08),\n \"scheduler\": None,\n },\n \"temporal_distortion\": {\n \"optimizer\": RAdamOptimizerConfig(lr=5e-4, eps=1e-08),\n \"scheduler\": None,\n },\n },\n)\n\nmethod_configs[\"tensorf\"] = Config(\n method_name=\"tensorf\",\n trainer=TrainerConfig(mixed_precision=False),\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(\n dataparser=BlenderDataParserConfig(),\n ),\n model=TensoRFModelConfig(),\n ),\n optimizers={\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=0.001),\n \"scheduler\": SchedulerConfig(lr_final=0.0001, max_steps=30000),\n },\n \"encodings\": {\n \"optimizer\": AdamOptimizerConfig(lr=0.02),\n \"scheduler\": SchedulerConfig(lr_final=0.002, max_steps=30000),\n },\n },\n)\n\nmethod_configs[\"dnerf\"] = Config(\n method_name=\"dnerf\",\n pipeline=VanillaPipelineConfig(\n datamanager=VanillaDataManagerConfig(dataparser=DNeRFDataParserConfig()),\n model=VanillaModelConfig(\n _target=NeRFModel,\n enable_temporal_distortion=True,\n temporal_distortion_params={\"kind\": TemporalDistortionKind.DNERF},\n ),\n ),\n optimizers={\n \"fields\": {\n \"optimizer\": RAdamOptimizerConfig(lr=5e-4, eps=1e-08),\n \"scheduler\": None,\n },\n \"temporal_distortion\": {\n \"optimizer\": RAdamOptimizerConfig(lr=5e-4, eps=1e-08),\n \"scheduler\": None,\n },\n },\n)\n\nmethod_configs[\"phototourism\"] = Config(\n method_name=\"phototourism\",\n trainer=TrainerConfig(\n steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True\n ),\n pipeline=VanillaPipelineConfig(\n datamanager=VariableResDataManagerConfig( # NOTE: one of the only differences with nerfacto\n dataparser=PhototourismDataParserConfig(), # NOTE: one of the only differences with nerfacto\n train_num_rays_per_batch=4096,\n eval_num_rays_per_batch=4096,\n camera_optimizer=CameraOptimizerConfig(\n mode=\"SO3xR3\", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)\n ),\n ),\n model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15),\n ),\n optimizers={\n \"proposal_networks\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n \"fields\": {\n \"optimizer\": AdamOptimizerConfig(lr=1e-2, eps=1e-15),\n \"scheduler\": None,\n },\n },\n viewer=ViewerConfig(num_rays_per_chunk=1 << 15),\n vis=\"viewer\",\n)\n\nAnnotatedBaseConfigUnion = tyro.conf.SuppressFixed[ # Don't show unparseable (fixed) arguments in helptext.\n tyro.conf.FlagConversionOff[\n tyro.extras.subcommand_type_from_defaults(defaults=method_configs, descriptions=descriptions)\n ]\n]\n\"\"\"Union[] type over config types, annotated with default instances for use with\ntyro.cli(). Allows the user to pick between one of several base configurations, and\nthen override values in it.\"\"\"\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/engine/__init__.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/engine/callbacks.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nCallback code used for training iterations\n\"\"\"\n\nfrom dataclasses import InitVar, dataclass\nfrom enum import Enum, auto\nfrom inspect import signature\nfrom typing import Callable, Dict, List, Optional, Tuple\n\nfrom nerfstudio.configs.base_config import TrainerConfig\n\n\n@dataclass\nclass TrainingCallbackAttributes:\n \"\"\"Attributes that can be used to configure training callbacks.\n The callbacks can be specified in the Dataloader or Model implementations.\n Instead of providing access to the entire Trainer object, we only provide these attributes.\n This should be least prone to errors and fairly clean from a user perspective.\"\"\"\n\n # TODO(ethan): type this without circular imports\n optimizers: Optional[InitVar]\n \"\"\"optimizers for training\"\"\"\n grad_scaler: Optional[InitVar]\n \"\"\"gradient scalers\"\"\"\n pipeline: Optional[InitVar]\n \"\"\"reference to training pipeline\"\"\"\n config: TrainerConfig\n \"\"\"the trainer config\"\"\"\n\n\nclass TrainingCallbackLocation(Enum):\n \"\"\"Enum for specifying where the training callback should be run.\"\"\"\n\n BEFORE_TRAIN_ITERATION = auto()\n AFTER_TRAIN_ITERATION = auto()\n\n\nclass TrainingCallback:\n \"\"\"Callback class used during training.\n The function 'func' with 'args' and 'kwargs' will be called every 'update_every_num_iters' training iterations,\n including at iteration 0. The function is called after the training iteration.\n\n Args:\n where_to_run: List of locations for when to run callbak (before/after iteration)\n func: The function that will be called.\n update_every_num_iters: How often to call the function `func`.\n iters: Tuple of iteration steps to perform callback\n args: args for the function 'func'.\n kwargs: kwargs for the function 'func'.\n \"\"\"\n\n def __init__(\n self,\n where_to_run: List[TrainingCallbackLocation],\n func: Callable,\n update_every_num_iters: Optional[int] = None,\n iters: Optional[Tuple[int, ...]] = None,\n args: Optional[List] = None,\n kwargs: Optional[Dict] = None,\n ):\n assert (\n \"step\" in signature(func).parameters.keys()\n ), f\"'step: int' must be an argument in the callback function 'func': {func.__name__}\"\n self.where_to_run = where_to_run\n self.update_every_num_iters = update_every_num_iters\n self.iters = iters\n self.func = func\n self.args = args if args is not None else []\n self.kwargs = kwargs if kwargs is not None else {}\n\n def run_callback(self, step: int):\n \"\"\"Callback to run after training step\n\n Args:\n step: current iteration step\n \"\"\"\n if self.update_every_num_iters is not None:\n if step % self.update_every_num_iters == 0:\n self.func(*self.args, **self.kwargs, step=step)\n elif self.iters is not None:\n if step in self.iters:\n self.func(*self.args, **self.kwargs, step=step)\n\n def run_callback_at_location(self, step: int, location: TrainingCallbackLocation):\n \"\"\"Runs the callback if it's supposed to be run at the given location.\n\n Args:\n step: current iteration step\n location: when to run callback (before/after iteration)\n \"\"\"\n if location in self.where_to_run:\n self.run_callback(step=step)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/engine/optimizers.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nOptimizers class.\n\"\"\"\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass\nfrom typing import Any, Dict, List, Optional, Type\nfrom rich.console import Console\n\nimport torch\nfrom torch.cuda.amp.grad_scaler import GradScaler\nfrom torch.nn.parameter import Parameter\n\nfrom nerfstudio.configs import base_config\nfrom nerfstudio.utils import writer\n\nCONSOLE = Console(width=128)\n\n\n# Optimizer related configs\n@dataclass\nclass OptimizerConfig(base_config.PrintableConfig):\n \"\"\"Basic optimizer config with RAdam\"\"\"\n\n _target: Type = torch.optim.Adam\n \"\"\"The optimizer class to use.\"\"\"\n lr: float = 0.0005\n \"\"\"The learning rate to use.\"\"\"\n eps: float = 1e-08\n \"\"\"The epsilon value to use.\"\"\"\n max_norm: Optional[float] = None\n \"\"\"The max norm to use for gradient clipping.\"\"\"\n\n # TODO: somehow make this more generic. i dont like the idea of overriding the setup function\n # but also not sure how to go about passing things into predefined torch objects.\n def setup(self, params) -> Any:\n \"\"\"Returns the instantiated object using the config.\"\"\"\n kwargs = vars(self).copy()\n kwargs.pop(\"_target\")\n kwargs.pop(\"max_norm\")\n return self._target(params, **kwargs)\n\n\n@dataclass\nclass AdamOptimizerConfig(OptimizerConfig):\n \"\"\"Basic optimizer config with Adam\"\"\"\n\n _target: Type = torch.optim.Adam\n weight_decay: float = 0\n\n\n@dataclass\nclass RAdamOptimizerConfig(OptimizerConfig):\n \"\"\"Basic optimizer config with RAdam\"\"\"\n\n _target: Type = torch.optim.RAdam\n\n\n@dataclass\nclass AdamWOptimizerConfig(OptimizerConfig):\n \"\"\"Basic optimizer config with AdamW\"\"\"\n\n _target: Type = torch.optim.AdamW\n weight_decay: float = 1e-2\n\n\ndef setup_optimizers(config: base_config.Config, param_groups: Dict[str, List[Parameter]]) -> \"Optimizers\":\n \"\"\"Helper to set up the optimizers\n\n Args:\n config: The trainer configuration object.\n param_groups: A dictionary of parameter groups to optimize.\n\n Returns:\n The optimizers object.\n \"\"\"\n optimizer_config = config.optimizers.copy()\n\n # Add the camera optimizer if enabled.\n camera_optimizer_config = config.pipeline.datamanager.camera_optimizer\n if camera_optimizer_config.mode != \"off\":\n assert camera_optimizer_config.param_group not in optimizer_config\n optimizer_config[camera_optimizer_config.param_group] = {\n \"optimizer\": config.pipeline.datamanager.camera_optimizer.optimizer,\n \"scheduler\": config.pipeline.datamanager.camera_optimizer.scheduler,\n }\n return Optimizers(optimizer_config, param_groups)\n\n\nclass Optimizers:\n \"\"\"A set of optimizers.\n\n Args:\n config: The optimizer configuration object.\n param_groups: A dictionary of parameter groups to optimize.\n \"\"\"\n\n def __init__(self, config: Dict[str, Any], param_groups: Dict[str, List[Parameter]]):\n self.config = config\n self.optimizers = {}\n self.schedulers = {}\n self.parameters = {}\n for param_group_name, params in param_groups.items():\n if param_group_name not in config:\n CONSOLE.log(f\"[bold red] No optimizer config for param group: '{param_group_name}', params will not be optimized!\")\n continue\n lr_init = config[param_group_name][\"optimizer\"].lr\n self.optimizers[param_group_name] = config[param_group_name][\"optimizer\"].setup(params=params)\n self.parameters[param_group_name] = params\n if config[param_group_name][\"scheduler\"]:\n self.schedulers[param_group_name] = config[param_group_name][\"scheduler\"].setup(\n optimizer=self.optimizers[param_group_name], lr_init=lr_init\n )\n\n def optimizer_step(self, param_group_name: str) -> None:\n \"\"\"Fetch and step corresponding optimizer.\n\n Args:\n param_group_name: name of optimizer to step forward\n \"\"\"\n self.optimizers[param_group_name].step()\n\n def scheduler_step(self, param_group_name: str) -> None:\n \"\"\"Fetch and step corresponding scheduler.\n\n Args:\n param_group_name: name of scheduler to step forward\n \"\"\"\n if self.config.param_group_name.scheduler: # type: ignore\n self.schedulers[param_group_name].step()\n\n def zero_grad_all(self) -> None:\n \"\"\"Zero the gradients for all optimizer parameters.\"\"\"\n for _, optimizer in self.optimizers.items():\n optimizer.zero_grad()\n\n def optimizer_scaler_step_all(self, grad_scaler: GradScaler) -> None:\n \"\"\"Take an optimizer step using a grad scaler.\n\n Args:\n grad_scaler: GradScaler to use\n \"\"\"\n for param_group, optimizer in self.optimizers.items():\n has_grads = any(\n p.grad is not None\n for group in optimizer.param_groups\n for p in group[\"params\"]\n )\n if not has_grads:\n continue\n max_norm = self.config[param_group][\"optimizer\"].max_norm\n if max_norm is not None:\n grad_scaler.unscale_(optimizer)\n torch.nn.utils.clip_grad_norm_(self.parameters[param_group], max_norm)\n grad_scaler.step(optimizer)\n\n def optimizer_step_all(self):\n \"\"\"Run step for all optimizers.\"\"\"\n for param_group, optimizer in self.optimizers.items():\n # note that they key is the parameter name\n max_norm = self.config[param_group][\"optimizer\"].max_norm\n if max_norm is not None:\n torch.nn.utils.clip_grad_norm_(self.parameters[param_group], max_norm)\n optimizer.step()\n\n def scheduler_step_all(self, step: int) -> None:\n \"\"\"Run step for all schedulers.\n\n Args:\n step: the current step\n \"\"\"\n for param_group_name, scheduler in self.schedulers.items():\n scheduler.step()\n # TODO(ethan): clean this up. why is there indexing into a list?\n lr = scheduler.get_last_lr()[0]\n writer.put_scalar(name=f\"learning_rate/{param_group_name}\", scalar=lr, step=step)\n\n def load_optimizers(self, loaded_state: Dict[str, Any]) -> None:\n \"\"\"Helper to load the optimizer state from previous checkpoint\n\n Args:\n loaded_state: the state from the previous checkpoint\n \"\"\"\n for k, v in loaded_state.items():\n self.optimizers[k].load_state_dict(v)\n\n def load_schedulers(self, loaded_state: Dict[str, Any]) -> None:\n \"\"\"Helper to load the schedulers state from previous checkpoint\n\n Args:\n loaded_state: the state from the previous checkpoint\n \"\"\"\n for k, v in loaded_state.items():\n self.schedulers[k].load_state_dict(v)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/engine/schedulers.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Scheduler Classes\"\"\"\n\nfrom dataclasses import dataclass, field\nfrom typing import Any, Optional, Type\n\nimport numpy as np\nfrom torch.optim import Optimizer, lr_scheduler\n\nfrom nerfstudio.configs.base_config import InstantiateConfig\n\n\n@dataclass\nclass SchedulerConfig(InstantiateConfig):\n \"\"\"Basic scheduler config with self-defined exponential decay schedule\"\"\"\n\n _target: Type = field(default_factory=lambda: ExponentialDecaySchedule)\n lr_final: float = 0.000005\n max_steps: int = 1000000\n\n # TODO: somehow make this more generic. i dont like the idea of overriding the setup function\n # but also not sure how to go about passing things into predefined torch objects.\n def setup(self, optimizer=None, lr_init=None, **kwargs) -> Any:\n \"\"\"Returns the instantiated object using the config.\"\"\"\n return self._target(optimizer, lr_init, self.lr_final, self.max_steps)\n\n\nclass ExponentialDecaySchedule(lr_scheduler.LambdaLR):\n \"\"\"Exponential learning rate decay function.\n See https://github.com/google-research/google-research/blob/\n fd2cea8cdd86b3ed2c640cbe5561707639e682f3/jaxnerf/nerf/utils.py#L360\n for details.\n\n Args:\n optimizer: The optimizer to update.\n lr_init: The initial learning rate.\n lr_final: The final learning rate.\n max_steps: The maximum number of steps.\n lr_delay_steps: The number of steps to delay the learning rate.\n lr_delay_mult: The multiplier for the learning rate after the delay.\n \"\"\"\n\n config: SchedulerConfig\n\n def __init__(self, optimizer, lr_init, lr_final, max_steps, lr_delay_steps=0, lr_delay_mult=1.0) -> None:\n def func(step):\n if lr_delay_steps > 0:\n delay_rate = lr_delay_mult + (1 - lr_delay_mult) * np.sin(\n 0.5 * np.pi * np.clip(step / lr_delay_steps, 0, 1)\n )\n else:\n delay_rate = 1.0\n t = np.clip(step / max_steps, 0, 1)\n log_lerp = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t)\n multiplier = (\n log_lerp / lr_init\n ) # divided by lr_init because the multiplier is with the initial learning rate\n return delay_rate * multiplier\n\n super().__init__(optimizer, lr_lambda=func)\n\n\nclass DelayerScheduler(lr_scheduler.LambdaLR):\n \"\"\"Starts with a flat lr schedule until it reaches N epochs then applies a given scheduler\"\"\"\n\n def __init__(\n self,\n optimizer: Optimizer,\n lr_init, # pylint: disable=unused-argument\n lr_final, # pylint: disable=unused-argument\n max_steps, # pylint: disable=unused-argument\n delay_epochs: int = 500,\n after_scheduler: Optional[lr_scheduler.LambdaLR] = None,\n ) -> None:\n def func(step):\n if step > delay_epochs:\n if after_scheduler is not None:\n multiplier = after_scheduler.lr_lambdas[0](step - delay_epochs) # type: ignore\n return multiplier\n return 1.0\n return 0.0\n\n super().__init__(optimizer, lr_lambda=func)\n\n\nclass DelayedExponentialScheduler(DelayerScheduler):\n \"\"\"Delayer Scheduler with an Exponential Scheduler initialized afterwards.\"\"\"\n\n def __init__(\n self,\n optimizer: Optimizer,\n lr_init,\n lr_final,\n max_steps,\n delay_epochs: int = 200,\n ):\n after_scheduler = ExponentialDecaySchedule(\n optimizer,\n lr_init,\n lr_final,\n max_steps,\n )\n super().__init__(optimizer, lr_init, lr_final, max_steps, delay_epochs, after_scheduler=after_scheduler)\n\n\n@dataclass\nclass MultiStepSchedulerConfig(InstantiateConfig):\n \"\"\"Basic scheduler config with self-defined exponential decay schedule\"\"\"\n\n _target: Type = field(default_factory=lambda: lr_scheduler.MultiStepLR)\n max_steps: int = 1000000\n\n def setup(self, optimizer=None, lr_init=None, **kwargs) -> Any:\n \"\"\"Returns the instantiated object using the config.\"\"\"\n return self._target(\n optimizer,\n milestones=[self.max_steps // 2, self.max_steps * 3 // 4, self.max_steps * 9 // 10],\n gamma=0.33,\n )\n\n\n@dataclass\nclass ExponentialSchedulerConfig(InstantiateConfig):\n \"\"\"Basic scheduler config with self-defined exponential decay schedule\"\"\"\n\n _target: Type = field(default_factory=lambda: lr_scheduler.ExponentialLR)\n decay_rate: float = 0.1\n max_steps: int = 1000000\n\n def setup(self, optimizer=None, lr_init=None, **kwargs) -> Any:\n \"\"\"Returns the instantiated object using the config.\"\"\"\n return self._target(\n optimizer,\n self.decay_rate ** (1.0 / self.max_steps),\n )\n\n\n@dataclass\nclass NeuSSchedulerConfig(InstantiateConfig):\n \"\"\"Basic scheduler config with self-defined exponential decay schedule\"\"\"\n\n _target: Type = field(default_factory=lambda: NeuSScheduler)\n warm_up_end: int = 5000\n learning_rate_alpha: float = 0.05\n max_steps: int = 300000\n\n def setup(self, optimizer=None, lr_init=None, **kwargs) -> Any:\n \"\"\"Returns the instantiated object using the config.\"\"\"\n return self._target(\n optimizer,\n self.warm_up_end,\n self.learning_rate_alpha,\n self.max_steps,\n )\n\n\nclass NeuSScheduler(lr_scheduler.LambdaLR):\n \"\"\"Starts with a flat lr schedule until it reaches N epochs then applies a given scheduler\"\"\"\n\n def __init__(self, optimizer, warm_up_end, learning_rate_alpha, max_steps) -> None:\n def func(step):\n if step < warm_up_end:\n learning_factor = step / warm_up_end\n else:\n alpha = learning_rate_alpha\n progress = (step - warm_up_end) / (max_steps - warm_up_end)\n learning_factor = (np.cos(np.pi * progress) + 1.0) * 0.5 * (1 - alpha) + alpha\n return learning_factor\n\n super().__init__(optimizer, lr_lambda=func)\n\n\n@dataclass\nclass AutoReconSchedulerConfig(InstantiateConfig):\n \"\"\"Scheduler tuned for AutoRecon w/ progressive hash encoding training\"\"\"\n\n _target: Type = field(default_factory=lambda: AutoReconScheduler)\n warm_up_end: int = 250\n # learning_rate_alpha: float = 0.05\n # max_steps: int = 300000\n\n def setup(self, optimizer=None, lr_init=None, **kwargs) -> Any:\n \"\"\"Returns the instantiated object using the config.\"\"\"\n return self._target(\n optimizer,\n self.warm_up_end,\n )\n\n\nclass AutoReconScheduler(lr_scheduler.LambdaLR):\n \"\"\"warm up then decay\"\"\"\n\n def __init__(self, optimizer, warm_up_end) -> None:\n def func(step):\n if step < warm_up_end:\n learning_factor = step / warm_up_end\n else:\n learning_factor = 1.0\n return learning_factor\n\n super().__init__(optimizer, lr_lambda=func)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/engine/trainer.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nCode to train model.\n\"\"\"\nfrom __future__ import annotations\n\nimport dataclasses\nimport functools\nimport os\nimport time\nfrom typing import Dict, List, Tuple\n\nimport torch\nfrom rich.console import Console\nfrom torch.cuda.amp.grad_scaler import GradScaler\nfrom typing_extensions import Literal\n\nfrom nerfstudio.configs import base_config as cfg\nfrom nerfstudio.engine.callbacks import (\n TrainingCallback,\n TrainingCallbackAttributes,\n TrainingCallbackLocation,\n)\nfrom nerfstudio.engine.optimizers import Optimizers, setup_optimizers\nfrom nerfstudio.pipelines.base_pipeline import VanillaPipeline\nfrom nerfstudio.utils import profiler, writer\nfrom nerfstudio.utils.decorators import (\n check_eval_enabled,\n check_main_thread,\n check_viewer_enabled,\n)\nfrom nerfstudio.utils.misc import step_check\nfrom nerfstudio.utils.writer import EventName, TimeWriter\nfrom nerfstudio.viewer.server import viewer_utils\n\nCONSOLE = Console(width=120)\n\n\nclass Trainer:\n \"\"\"Trainer class\n\n Args:\n config: The configuration object.\n local_rank: Local rank of the process.\n world_size: World size of the process.\n\n Attributes:\n config: The configuration object.\n local_rank: Local rank of the process.\n world_size: World size of the process.\n device: The device to run the training on.\n pipeline: The pipeline object.\n optimizers: The optimizers object.\n callbacks: The callbacks object.\n \"\"\"\n\n pipeline: VanillaPipeline\n optimizers: Optimizers\n callbacks: List[TrainingCallback]\n\n def __init__(self, config: cfg.Config, local_rank: int = 0, world_size: int = 1):\n self.config = config\n self.local_rank = local_rank\n self.world_size = world_size\n self.device = \"cpu\" if world_size == 0 else f\"cuda:{local_rank}\"\n self.mixed_precision = self.config.trainer.mixed_precision\n if self.device == \"cpu\":\n self.mixed_precision = False\n CONSOLE.print(\"Mixed precision is disabled for CPU training.\")\n self._start_step = 0\n # optimizers\n self.grad_scaler = GradScaler(enabled=self.mixed_precision)\n\n self.base_dir = config.get_base_dir()\n # directory to save checkpoints\n self.checkpoint_dir = config.get_checkpoint_dir()\n CONSOLE.log(f\"Saving checkpoints to: {self.checkpoint_dir}\")\n # set up viewer if enabled\n viewer_log_path = self.base_dir / config.viewer.relative_log_filename\n self.viewer_state, banner_messages = None, None\n if self.config.is_viewer_enabled() and local_rank == 0:\n self.viewer_state, banner_messages = viewer_utils.setup_viewer(config.viewer, log_filename=viewer_log_path)\n self._check_viewer_warnings()\n # set up writers/profilers if enabled\n writer_log_path = self.base_dir / config.logging.relative_log_dir\n writer.setup_event_writer(config, log_dir=writer_log_path)\n writer.setup_local_writer(\n config.logging, max_iter=config.trainer.max_num_iterations, banner_messages=banner_messages\n )\n writer.put_config(name=\"config\", config_dict=dataclasses.asdict(config), step=0)\n profiler.setup_profiler(config.logging)\n\n def setup(self, test_mode: Literal[\"test\", \"val\", \"inference\"] = \"val\"):\n \"\"\"Setup the Trainer by calling other setup functions.\n\n Args:\n test_mode:\n 'val': loads train/val datasets into memory\n 'test': loads train/test datset into memory\n 'inference': does not load any dataset into memory\n \"\"\"\n self.pipeline = self.config.pipeline.setup(\n device=self.device, test_mode=test_mode, world_size=self.world_size, local_rank=self.local_rank\n )\n self.optimizers = setup_optimizers(self.config, self.pipeline.get_param_groups())\n\n self._load_checkpoint()\n\n self.callbacks = self.pipeline.get_training_callbacks(\n TrainingCallbackAttributes(\n optimizers=self.optimizers, # type: ignore\n grad_scaler=self.grad_scaler, # type: ignore\n pipeline=self.pipeline, # type: ignore\n config=self.config.trainer, # type: ignore\n )\n )\n\n def train(self) -> None:\n \"\"\"Train the model.\"\"\"\n assert self.pipeline.datamanager.train_dataset is not None, \"Missing DatsetInputs\"\n self._init_viewer_state()\n with TimeWriter(writer, EventName.TOTAL_TRAIN_TIME):\n num_iterations = self.config.trainer.max_num_iterations\n step = 0\n for step in range(self._start_step, self._start_step + num_iterations):\n with TimeWriter(writer, EventName.ITER_TRAIN_TIME, step=step) as train_t:\n\n self.pipeline.train()\n\n # training callbacks before the training iteration\n for callback in self.callbacks:\n callback.run_callback_at_location(\n step, location=TrainingCallbackLocation.BEFORE_TRAIN_ITERATION\n )\n\n # time the forward pass\n loss, loss_dict, metrics_dict = self.train_iteration(step)\n\n # training callbacks after the training iteration\n for callback in self.callbacks:\n callback.run_callback_at_location(step, location=TrainingCallbackLocation.AFTER_TRAIN_ITERATION)\n\n # Skip the first two steps to avoid skewed timings that break the viewer rendering speed estimate.\n if step > 1:\n writer.put_time(\n name=EventName.TRAIN_RAYS_PER_SEC,\n duration=self.config.pipeline.datamanager.train_num_rays_per_batch / train_t.duration,\n step=step,\n avg_over_steps=True,\n )\n\n self._update_viewer_state(step)\n\n # a batch of train rays\n if step_check(step, self.config.logging.steps_per_log, run_at_zero=True):\n writer.put_scalar(name=\"Train Loss\", scalar=loss, step=step)\n writer.put_dict(name=\"Train Loss Dict\", scalar_dict=loss_dict, step=step)\n writer.put_dict(name=\"Train Metrics Dict\", scalar_dict=metrics_dict, step=step)\n\n self.eval_iteration(step)\n\n if step_check(step, self.config.trainer.steps_per_save):\n self.save_checkpoint(step)\n\n writer.write_out_storage()\n # save checkpoint at the end of training\n self.save_checkpoint(step)\n\n CONSOLE.rule()\n CONSOLE.print(\"[bold green]:tada: :tada: :tada: Training Finished :tada: :tada: :tada:\", justify=\"center\")\n if not self.config.viewer.quit_on_train_completion:\n CONSOLE.print(\"Use ctrl+c to quit\", justify=\"center\")\n self._always_render(step)\n\n @check_main_thread\n def _always_render(self, step):\n if self.config.is_viewer_enabled():\n while True:\n self.viewer_state.vis[\"renderingState/isTraining\"].write(False)\n self._update_viewer_state(step)\n\n @check_main_thread\n def _check_viewer_warnings(self) -> None:\n \"\"\"Helper to print out any warnings regarding the way the viewer/loggers are enabled\"\"\"\n if self.config.is_viewer_enabled():\n string = (\n \"[NOTE] Not running eval iterations since only viewer is enabled.\"\n \" Use [yellow]--vis wandb[/yellow] or [yellow]--vis tensorboard[/yellow] to run with eval instead.\"\n )\n CONSOLE.print(f\"{string}\")\n\n @check_viewer_enabled\n def _init_viewer_state(self) -> None:\n \"\"\"Initializes viewer scene with given train dataset\"\"\"\n assert self.viewer_state and self.pipeline.datamanager.train_dataset\n self.viewer_state.init_scene(\n dataset=self.pipeline.datamanager.train_dataset,\n start_train=self.config.viewer.start_train,\n )\n if not self.config.viewer.start_train:\n self._always_render(self._start_step)\n\n @check_viewer_enabled\n def _update_viewer_state(self, step: int):\n \"\"\"Updates the viewer state by rendering out scene with current pipeline\n Returns the time taken to render scene.\n\n Args:\n step: current train step\n \"\"\"\n assert self.viewer_state is not None\n with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as _:\n num_rays_per_batch = self.config.pipeline.datamanager.train_num_rays_per_batch\n try:\n self.viewer_state.update_scene(self, step, self.pipeline.model, num_rays_per_batch)\n except RuntimeError:\n time.sleep(0.03) # sleep to allow buffer to reset\n assert self.viewer_state.vis is not None\n self.viewer_state.vis[\"renderingState/log_errors\"].write(\n \"Error: GPU out of memory. Reduce resolution to prevent viewer from crashing.\"\n )\n\n @check_viewer_enabled\n def _update_viewer_rays_per_sec(self, train_t: TimeWriter, vis_t: TimeWriter, step: int):\n \"\"\"Performs update on rays/sec calclation for training\n\n Args:\n train_t: timer object carrying time to execute total training iteration\n vis_t: timer object carrying time to execute visualization step\n step: current step\n \"\"\"\n train_num_rays_per_batch = self.config.pipeline.datamanager.train_num_rays_per_batch\n writer.put_time(\n name=EventName.TRAIN_RAYS_PER_SEC,\n duration=train_num_rays_per_batch / (train_t.duration - vis_t.duration),\n step=step,\n avg_over_steps=True,\n )\n\n def _load_checkpoint(self) -> None:\n \"\"\"Helper function to load pipeline and optimizer from prespecified checkpoint\"\"\"\n load_dir = self.config.trainer.load_dir\n if load_dir is not None:\n load_step = self.config.trainer.load_step\n if load_step is None:\n print(\"Loading latest checkpoint from load_dir\")\n # NOTE: this is specific to the checkpoint name format\n load_step = sorted(int(x[x.find(\"-\") + 1: x.find(\".\")]) for x in os.listdir(load_dir))[-1]\n load_path = load_dir / f\"step-{load_step:09d}.ckpt\"\n assert load_path.exists(), f\"Checkpoint {load_path} does not exist\"\n loaded_state = torch.load(load_path, map_location=\"cpu\")\n \n # load pipeline ckpt\n self.pipeline.load_pipeline(loaded_state[\"pipeline\"], strict=self.config.trainer.load_pipeline_ckpt_strict)\n if not self.config.trainer.load_pipeline_ckpt_strict:\n CONSOLE.print(\"[yellow] Pipeline checkpoint loaded with strict=False\")\n \n # load ckpt for step, optimizer, scheduler, grad_scaler\n if not self.config.trainer.load_pipeline_ckpt_only:\n self._start_step = loaded_state[\"step\"] + 1\n self.optimizers.load_optimizers(loaded_state[\"optimizers\"])\n if \"schedulers\" in loaded_state and self.config.trainer.load_scheduler:\n self.optimizers.load_schedulers(loaded_state[\"schedulers\"])\n self.grad_scaler.load_state_dict(loaded_state[\"scalers\"])\n else:\n CONSOLE.print(\"[yellow] Only loading pipeline checkpoint, skipping optimizer and gradient scalar\")\n \n CONSOLE.print(f\"done loading checkpoint from {load_path}\")\n else:\n CONSOLE.print(\"No checkpoints to load, training from scratch\")\n\n @check_main_thread\n def save_checkpoint(self, step: int) -> None:\n \"\"\"Save the model and optimizers\n\n Args:\n step: number of steps in training for given checkpoint\n \"\"\"\n # possibly make the checkpoint directory\n if not self.checkpoint_dir.exists():\n self.checkpoint_dir.mkdir(parents=True, exist_ok=True)\n # save the checkpoint\n ckpt_path = self.checkpoint_dir / f\"step-{step:09d}.ckpt\"\n torch.save(\n {\n \"step\": step,\n \"pipeline\": self.pipeline.module.state_dict() # type: ignore\n if hasattr(self.pipeline, \"module\")\n else self.pipeline.state_dict(),\n \"optimizers\": {k: v.state_dict() for (k, v) in self.optimizers.optimizers.items()},\n \"schedulers\": {k: v.state_dict() for (k, v) in self.optimizers.schedulers.items()},\n \"scalers\": self.grad_scaler.state_dict(),\n },\n ckpt_path,\n )\n # possibly delete old checkpoints\n if self.config.trainer.save_only_latest_checkpoint:\n # delete everything else in the checkpoint folder\n for f in self.checkpoint_dir.glob(\"*\"):\n if f != ckpt_path:\n f.unlink()\n\n @profiler.time_function\n def train_iteration(self, step: int) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:\n \"\"\"Run one iteration with a batch of inputs. Returns dictionary of model losses.\n\n Args:\n step: Current training step.\n \"\"\"\n self.optimizers.zero_grad_all()\n cpu_or_cuda_str = self.device.split(\":\")[0]\n with torch.autocast(device_type=cpu_or_cuda_str, enabled=self.mixed_precision):\n _, loss_dict, metrics_dict = self.pipeline.get_train_loss_dict(step=step)\n loss = functools.reduce(torch.add, loss_dict.values())\n self.grad_scaler.scale(loss).backward() # type: ignore\n self.optimizers.optimizer_scaler_step_all(self.grad_scaler)\n self.grad_scaler.update()\n self.optimizers.scheduler_step_all(step)\n\n # Merging loss and metrics dict into a single output.\n return loss, loss_dict, metrics_dict\n\n @check_eval_enabled # only run this routine when using wandb / tensorboard\n @profiler.time_function\n def eval_iteration(self, step):\n \"\"\"Run one iteration with different batch/image/all image evaluations depending on step size.\n\n Args:\n step: Current training step.\n \"\"\"\n # a batch of eval rays\n if step_check(step, self.config.trainer.steps_per_eval_batch):\n _, eval_loss_dict, eval_metrics_dict = self.pipeline.get_eval_loss_dict(step=step)\n eval_loss = functools.reduce(torch.add, eval_loss_dict.values())\n writer.put_scalar(name=\"Eval Loss\", scalar=eval_loss, step=step)\n writer.put_dict(name=\"Eval Loss Dict\", scalar_dict=eval_loss_dict, step=step)\n writer.put_dict(name=\"Eval Metrics Dict\", scalar_dict=eval_metrics_dict, step=step)\n\n # one eval image\n if step_check(step, self.config.trainer.steps_per_eval_image):\n with TimeWriter(writer, EventName.TEST_RAYS_PER_SEC, write=False) as test_t:\n metrics_dict, images_dict = self.pipeline.get_eval_image_metrics_and_images(step=step)\n writer.put_time(\n name=EventName.TEST_RAYS_PER_SEC,\n duration=metrics_dict[\"num_rays\"] / test_t.duration,\n step=step,\n avg_over_steps=True,\n )\n writer.put_dict(name=\"Eval Images Metrics\", scalar_dict=metrics_dict, step=step)\n group = \"Eval Images\"\n for image_name, image in images_dict.items():\n writer.put_image(name=group + \"/\" + image_name, image=image, step=step)\n\n # all eval images\n if step_check(step, self.config.trainer.steps_per_eval_all_images):\n metrics_dict = self.pipeline.get_average_eval_image_metrics(step=step)\n writer.put_dict(name=\"Eval Images Metrics Dict (all images)\", scalar_dict=metrics_dict, step=step)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/exporter/__init__.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/exporter/exporter_utils.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nExport utils such as structs, point cloud generation, and rendering code.\n\"\"\"\n\n# pylint: disable=no-member\n\nfrom __future__ import annotations\n\nimport sys\nfrom dataclasses import dataclass\nfrom typing import List, Optional, Tuple\n\nimport numpy as np\nimport open3d as o3d\nimport pymeshlab\nimport torch\nfrom rich.console import Console\nfrom rich.progress import (\n BarColumn,\n Progress,\n TaskProgressColumn,\n TextColumn,\n TimeRemainingColumn,\n)\nfrom torchtyping import TensorType\n\nfrom nerfstudio.cameras.cameras import Cameras\nfrom nerfstudio.configs.base_config import Config # pylint: disable=unused-import\nfrom nerfstudio.pipelines.base_pipeline import Pipeline\nfrom nerfstudio.utils.rich_utils import ItersPerSecColumn\n\nCONSOLE = Console(width=120)\n\n\n@dataclass\nclass Mesh:\n \"\"\"Class for a mesh.\"\"\"\n\n vertices: TensorType[\"num_verts\", 3]\n \"\"\"Vertices of the mesh.\"\"\"\n faces: TensorType[\"num_faces\", 3]\n \"\"\"Faces of the mesh.\"\"\"\n normals: TensorType[\"num_verts\", 3]\n \"\"\"Normals of the mesh vertices.\"\"\"\n face_normals: TensorType[\"num_faces\", 3] = None\n \"\"\"Normals of the mesh faces.\"\"\"\n colors: Optional[TensorType[\"num_verts\", 3]] = None\n \"\"\"Colors of the mesh.\"\"\"\n\n\ndef get_mesh_from_pymeshlab_mesh(mesh: pymeshlab.Mesh) -> Mesh:\n \"\"\"Get a Mesh from a pymeshlab mesh.\n See https://pymeshlab.readthedocs.io/en/0.1.5/classes/mesh.html for details.\n \"\"\"\n return Mesh(\n vertices=torch.from_numpy(mesh.vertex_matrix()).float(),\n faces=torch.from_numpy(mesh.face_matrix()).long(),\n normals=torch.from_numpy(np.copy(mesh.vertex_normal_matrix())).float(),\n face_normals=torch.from_numpy(np.copy(mesh.face_normal_matrix())).float(),\n colors=torch.from_numpy(mesh.vertex_color_matrix()).float(),\n )\n\n\ndef get_mesh_from_filename(filename: str, target_num_faces: Optional[int] = None) -> Mesh:\n \"\"\"Get a Mesh from a filename.\"\"\"\n ms = pymeshlab.MeshSet()\n ms.load_new_mesh(filename)\n if target_num_faces is not None:\n CONSOLE.print(\"Running meshing decimation with quadric edge collapse\")\n ms.meshing_decimation_quadric_edge_collapse(targetfacenum=target_num_faces)\n mesh = ms.current_mesh()\n return get_mesh_from_pymeshlab_mesh(mesh)\n\n\ndef generate_point_cloud(\n pipeline: Pipeline,\n num_points: int = 1000000,\n remove_outliers: bool = True,\n estimate_normals: bool = False,\n rgb_output_name: str = \"rgb\",\n depth_output_name: str = \"depth\",\n normal_output_name: Optional[str] = None,\n use_bounding_box: bool = True,\n bounding_box_min: Tuple[float, float, float] = (-1.0, -1.0, -1.0),\n bounding_box_max: Tuple[float, float, float] = (1.0, 1.0, 1.0),\n std_ratio: float = 10.0,\n) -> o3d.geometry.PointCloud:\n \"\"\"Generate a point cloud from a nerf.\n\n Args:\n pipeline: Pipeline to evaluate with.\n num_points: Number of points to generate. May result in less if outlier removal is used.\n remove_outliers: Whether to remove outliers.\n estimate_normals: Whether to estimate normals.\n rgb_output_name: Name of the RGB output.\n depth_output_name: Name of the depth output.\n normal_output_name: Name of the normal output.\n use_bounding_box: Whether to use a bounding box to sample points.\n bounding_box_min: Minimum of the bounding box.\n bounding_box_max: Maximum of the bounding box.\n std_ratio: Threshold based on STD of the average distances across the point cloud to remove outliers.\n\n Returns:\n Point cloud.\n \"\"\"\n\n # pylint: disable=too-many-statements\n\n progress = Progress(\n TextColumn(\":cloud: Computing Point Cloud :cloud:\"),\n BarColumn(),\n TaskProgressColumn(show_speed=True),\n TimeRemainingColumn(elapsed_when_finished=True, compact=True),\n )\n points = []\n rgbs = []\n normals = []\n with progress as progress_bar:\n task = progress_bar.add_task(\"Generating Point Cloud\", total=num_points)\n while not progress_bar.finished:\n torch.cuda.empty_cache()\n with torch.no_grad():\n ray_bundle, _ = pipeline.datamanager.next_train(0)\n outputs = pipeline.model(ray_bundle)\n if rgb_output_name not in outputs:\n CONSOLE.rule(\"Error\", style=\"red\")\n CONSOLE.print(f\"Could not find {rgb_output_name} in the model outputs\", justify=\"center\")\n CONSOLE.print(f\"Please set --rgb_output_name to one of: {outputs.keys()}\", justify=\"center\")\n sys.exit(1)\n if depth_output_name not in outputs:\n CONSOLE.rule(\"Error\", style=\"red\")\n CONSOLE.print(f\"Could not find {depth_output_name} in the model outputs\", justify=\"center\")\n CONSOLE.print(f\"Please set --depth_output_name to one of: {outputs.keys()}\", justify=\"center\")\n sys.exit(1)\n rgb = outputs[rgb_output_name]\n depth = outputs[depth_output_name]\n if normal_output_name is not None:\n if normal_output_name not in outputs:\n CONSOLE.rule(\"Error\", style=\"red\")\n CONSOLE.print(f\"Could not find {normal_output_name} in the model outputs\", justify=\"center\")\n CONSOLE.print(f\"Please set --normal_output_name to one of: {outputs.keys()}\", justify=\"center\")\n sys.exit(1)\n normal = outputs[normal_output_name]\n point = ray_bundle.origins + ray_bundle.directions * depth\n\n if use_bounding_box:\n comp_l = torch.tensor(bounding_box_min, device=point.device)\n comp_m = torch.tensor(bounding_box_max, device=point.device)\n assert torch.all(\n comp_l < comp_m\n ), f\"Bounding box min {bounding_box_min} must be smaller than max {bounding_box_max}\"\n mask = torch.all(torch.concat([point > comp_l, point < comp_m], dim=-1), dim=-1)\n point = point[mask]\n rgb = rgb[mask]\n if normal_output_name is not None:\n normal = normal[mask]\n\n points.append(point)\n rgbs.append(rgb)\n if normal_output_name is not None:\n normals.append(normal)\n progress.advance(task, point.shape[0])\n points = torch.cat(points, dim=0)\n rgbs = torch.cat(rgbs, dim=0)\n\n pcd = o3d.geometry.PointCloud()\n pcd.points = o3d.utility.Vector3dVector(points.float().cpu().numpy())\n pcd.colors = o3d.utility.Vector3dVector(rgbs.float().cpu().numpy())\n\n ind = None\n if remove_outliers:\n CONSOLE.print(\"Cleaning Point Cloud\")\n pcd, ind = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=std_ratio)\n print(\"\\033[A\\033[A\")\n CONSOLE.print(\"[bold green]:white_check_mark: Cleaning Point Cloud\")\n\n # either estimate_normals or normal_output_name, not both\n if estimate_normals:\n if normal_output_name is not None:\n CONSOLE.rule(\"Error\", style=\"red\")\n CONSOLE.print(\"Cannot estimate normals and use normal_output_name at the same time\", justify=\"center\")\n sys.exit(1)\n CONSOLE.print(\"Estimating Point Cloud Normals\")\n pcd.estimate_normals()\n print(\"\\033[A\\033[A\")\n CONSOLE.print(\"[bold green]:white_check_mark: Estimating Point Cloud Normals\")\n elif normal_output_name is not None:\n normals = torch.cat(normals, dim=0)\n if ind is not None:\n # mask out normals for points that were removed with remove_outliers\n normals = normals[ind]\n pcd.normals = o3d.utility.Vector3dVector(normals.float().cpu().numpy())\n\n return pcd\n\n\ndef render_trajectory(\n pipeline: Pipeline,\n cameras: Cameras,\n rgb_output_name: str,\n depth_output_name: str,\n rendered_resolution_scaling_factor: float = 1.0,\n disable_distortion: bool = False,\n) -> Tuple[List[np.ndarray], List[np.ndarray]]:\n \"\"\"Helper function to create a video of a trajectory.\n\n Args:\n pipeline: Pipeline to evaluate with.\n cameras: Cameras to render.\n rgb_output_name: Name of the RGB output.\n depth_output_name: Name of the depth output.\n rendered_resolution_scaling_factor: Scaling factor to apply to the camera image resolution.\n disable_distortion: Whether to disable distortion.\n\n Returns:\n List of rgb images, list of depth images.\n \"\"\"\n images = []\n depths = []\n cameras.rescale_output_resolution(rendered_resolution_scaling_factor)\n\n progress = Progress(\n TextColumn(\":cloud: Computing rgb and depth images :cloud:\"),\n BarColumn(),\n TaskProgressColumn(show_speed=True),\n ItersPerSecColumn(suffix=\"fps\"),\n TimeRemainingColumn(elapsed_when_finished=True, compact=True),\n )\n with progress:\n for camera_idx in progress.track(range(cameras.size), description=\"\"):\n camera_ray_bundle = cameras.generate_rays(\n camera_indices=camera_idx, disable_distortion=disable_distortion\n ).to(pipeline.device)\n with torch.no_grad():\n outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle)\n if rgb_output_name not in outputs:\n CONSOLE.rule(\"Error\", style=\"red\")\n CONSOLE.print(f\"Could not find {rgb_output_name} in the model outputs\", justify=\"center\")\n CONSOLE.print(f\"Please set --rgb_output_name to one of: {outputs.keys()}\", justify=\"center\")\n sys.exit(1)\n if depth_output_name not in outputs:\n CONSOLE.rule(\"Error\", style=\"red\")\n CONSOLE.print(f\"Could not find {depth_output_name} in the model outputs\", justify=\"center\")\n CONSOLE.print(f\"Please set --depth_output_name to one of: {outputs.keys()}\", justify=\"center\")\n sys.exit(1)\n images.append(outputs[rgb_output_name].cpu().numpy())\n depths.append(outputs[depth_output_name].cpu().numpy())\n return images, depths\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/exporter/mesh_culling_utils.py",
"content": "\nfrom __future__ import annotations\n\nfrom pathlib import Path\nfrom typing import Optional, Tuple, Type\nfrom typing_extensions import Literal\nfrom torchtyping import TensorType\nfrom dataclasses import dataclass, field\nfrom rich.console import Console\n\nimport torch\nimport trimesh\nimport pymeshlab\nimport numpy as np\nfrom torch.nn import functional as F\n\nfrom nerfstudio.cameras.rays import RayBundle\nfrom nerfstudio.exporter.exporter_utils import Mesh\nfrom nerfstudio.pipelines.base_pipeline import Pipeline\nfrom nerfstudio.model_components.ray_samplers import Sampler, UniformSampler, DummySampler\nfrom nerfstudio.utils.rich_utils import get_progress\nfrom nerfstudio.configs.base_config import InstantiateConfig\n\nCONSOLE = Console(width=120)\n\n\n@dataclass\nclass MeshCullerConfig(InstantiateConfig):\n \n _target: Type = field(default_factory=lambda: MeshCuller)\n \n raylen_method: Literal[\"edge\", \"none\"] = \"edge\"\n facewise_raylen: bool = False\n raylen_multiplier: float = 10.0\n \"\"\"e.g., raylen = face_edge_len * raylen_multiplier\"\"\"\n use_train_ray_sampler: bool = True\n ray_sampler: Literal[\"uniform\", \"dummy\"] = \"uniform\"\n \"\"\"type of the new ray_sampler for mesh culling\"\"\"\n num_ray_samples: Optional[int] = 10 # TODO: need to specify the RaySampler as well\n \"\"\"number of ray samples\"\"\"\n segmentation_binarization_thr: float = 0.5\n mesh_close_holes: bool = True\n \"\"\"close holes on the mesh caused by incorrect mesh culling (bad segmentation query results)\"\"\"\n max_hole_size: int = 1000\n \"\"\"maximum hole size to be closed in terms of number of edges composing the hole boundary\"\"\"\n simplify_mesh: bool = False\n \"\"\"simplify mesh with decimation_quadric_edge_collapse\"\"\"\n\n\nclass MeshCuller:\n config: MeshCullerConfig\n \n def __init__(self, config: MeshCullerConfig):\n self.config = config\n self.ray_sampler = self._get_ray_sampler()\n \n @torch.no_grad()\n def main(self, mesh: Mesh, pipeline: Pipeline, output_path: Path):\n \"\"\"Cull mesh by rendering the segmentation of each face.\"\"\"\n device = pipeline.device\n\n vertices = mesh.vertices.to(device) # (nv, 3)\n faces = mesh.faces.to(device) # (nf, 3)\n face_normals = mesh.face_normals.to(device) # (nf, 3)\n ray_bundle = self._build_ray_bundle(vertices, faces, face_normals)\n \n outputs = pipeline.model.get_outputs_for_mesh_culling(ray_bundle, ray_sampler=self.ray_sampler)\n \n # TODO: support only cull faces near to the ground plane\n \n # TODO: if a faces is not visible from any viewpoint (or only visible from a very few viewpoints)\n # we could just keep it, since its feature is not learned by feature field.\n \n fg_seg_prob = outputs[\"fg_seg_prob\"]\n fg_faces_mask = (fg_seg_prob > self.config.segmentation_binarization_thr)[..., 0]\n CONSOLE.print(f\"{len(fg_faces_mask) - fg_faces_mask.sum()}/{len(fg_faces_mask)} faces get culled!\")\n faces = faces[fg_faces_mask] # (nf', 3)\n face_normals = face_normals[fg_faces_mask] # (nf', 3)\n \n culled_mesh = trimesh.Trimesh(vertices=vertices.cpu().numpy(),\n faces=faces.cpu().numpy(),\n face_normals=face_normals.cpu().numpy())\n culled_mesh.export(output_path)\n CONSOLE.print(f'Culled mesh saved: {output_path}')\n \n if self.config.mesh_close_holes:\n hole_closed_path = output_path.parent / (output_path.stem + \"_hole-closed\" + output_path.suffix)\n ms = pymeshlab.MeshSet()\n ms.load_new_mesh(str(output_path))\n CONSOLE.print('Running mesh hole closing...')\n ms.meshing_close_holes(maxholesize=self.config.max_hole_size)\n ms.save_current_mesh(str(hole_closed_path), save_face_color=False)\n CONSOLE.print(f'Hole-closed mesh saved: {hole_closed_path}')\n output_path = hole_closed_path\n \n if self.config.simplify_mesh:\n simplified_path = output_path.parent / (output_path.stem + \"_simplified\" + output_path.suffix)\n ms = pymeshlab.MeshSet()\n ms.load_new_mesh(str(output_path))\n\n CONSOLE.print('Running remashing...')\n ms.meshing_decimation_quadric_edge_collapse(targetfacenum=2000000)\n ms.save_current_mesh(str(simplified_path), save_face_color=False)\n CONSOLE.print(f'Simplified mesh saved: {simplified_path}')\n output_path = simplified_path\n return output_path\n\n def _get_ray_sampler(self):\n if self.config.use_train_ray_sampler:\n return None\n \n ray_sampler_type = self.config.ray_sampler\n num_samples = self.config.num_ray_samples\n if ray_sampler_type == \"uniform\":\n _log_str = f\"UniformSampler(num_samples={num_samples})\"\n sampler = UniformSampler(num_samples=num_samples)\n elif ray_sampler_type == \"dummy\":\n _log_str = \"DummySampler\"\n sampler = DummySampler()\n else:\n raise ValueError(f\"Unknown RaySampler type: {ray_sampler_type}\")\n CONSOLE.print(f\"Replace the original sampler with {_log_str}.\")\n return sampler\n \n def _build_ray_bundle(self, vertices, faces, face_normals) -> RayBundle:\n face_vertices = vertices[faces] # (nf, 3, 3)\n raylen = self._get_ray_length(vertices, faces) # (nf, )\n origins = face_vertices.mean(dim=1) # (nf, 3)\n directions = F.normalize(face_normals, dim=-1)\n if self.config.ray_sampler != \"dummy\":\n origins = origins - 0.5 * raylen * directions\n ray_bundle = RayBundle(\n origins=origins,\n directions=directions,\n pixel_area=torch.ones_like(origins[..., 0:1]),\n camera_indices=torch.zeros_like(origins[..., 0:1]),\n nears=torch.zeros_like(origins[..., 0:1]),\n fars=torch.ones_like(origins[..., 0:1]) * raylen\n )\n return ray_bundle\n \n def _get_ray_length(self, vertices, faces):\n raylen_method = self.config.raylen_method\n facewise_raylen = self.config.facewise_raylen\n raylen_multiplier = self.config.raylen_multiplier\n \n if raylen_method == \"edge\":\n face_vertices = vertices[faces]\n # compute the length of the rays we want to render\n # we make a reasonable approximation by using the mean length of one edge per face\n raylen = raylen_multiplier * torch.norm(face_vertices[:, 1, :] - face_vertices[:, 0, :], dim=-1) # (nf, )\n if not facewise_raylen:\n raylen = torch.full_like(raylen, float(torch.mean(raylen)))\n else:\n raise ValueError(f\"Ray length method {raylen_method} not supported.\")\n return raylen[..., None]\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/exporter/texture_utils.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nTexture utils.\n\"\"\"\n\n# pylint: disable=no-member\n\nfrom __future__ import annotations\n\nimport math\nfrom pathlib import Path\nfrom typing import Optional, Tuple\n\nimport mediapy as media\nimport numpy as np\nimport torch\nimport xatlas\nfrom rich.console import Console\nfrom torchtyping import TensorType\nfrom typing_extensions import Literal\n\nfrom nerfstudio.cameras.rays import RayBundle\nfrom nerfstudio.exporter.exporter_utils import Mesh\nfrom nerfstudio.pipelines.base_pipeline import Pipeline\nfrom nerfstudio.utils.rich_utils import get_progress\n\nCONSOLE = Console(width=120)\n\n\ndef get_parallelogram_area(\n p: TensorType[\"bs\":..., 2], v0: TensorType[\"bs\":..., 2], v1: TensorType[\"bs\":..., 2]\n) -> TensorType[\"bs\":...]:\n \"\"\"Given three 2D points, return the area defined by the parallelogram. I.e., 2x the triangle area.\n\n Args:\n p: The origin of the parallelogram.\n v0: The first vector of the parallelogram.\n v1: The second vector of the parallelogram.\n\n Returns:\n The area of the parallelogram.\n \"\"\"\n return (p[..., 0] - v0[..., 0]) * (v1[..., 1] - v0[..., 1]) - (p[..., 1] - v0[..., 1]) * (v1[..., 0] - v0[..., 0])\n\n\ndef get_texture_image(num_pixels_w, num_pixels_h, device):\n \"\"\"Get a texture image.\"\"\"\n px_w = 1.0 / num_pixels_w\n px_h = 1.0 / num_pixels_h\n uv_indices = torch.stack(\n torch.meshgrid(\n torch.arange(num_pixels_w, device=device), torch.arange(num_pixels_h, device=device), indexing=\"xy\"\n ),\n dim=-1,\n )\n linspace_h = torch.linspace(px_h / 2, 1 - px_h / 2, num_pixels_h, device=device)\n linspace_w = torch.linspace(px_w / 2, 1 - px_w / 2, num_pixels_w, device=device)\n uv_coords = torch.stack(\n torch.meshgrid(linspace_w, linspace_h, indexing=\"xy\"), dim=-1\n ) # (num_pixels_h, num_pixels_w, 2)\n\n return uv_coords, uv_indices\n\n\ndef unwrap_mesh_per_uv_triangle(\n vertices: TensorType[\"num_verts\", 3],\n faces: TensorType[\"num_faces\", 3],\n vertex_normals: TensorType[\"num_verts\", 3],\n px_per_uv_triangle: int,\n):\n \"\"\"Unwrap a mesh to a UV texture. This is done by making a grid of rectangles in the UV texture map\n and then having two triangles per rectangle. Then the texture image is rasterized and uses barycentric\n interpolation to get the origins and directions, per pixel, that are needed to render the NeRF with.\n\n Args:\n vertices: The vertices of the mesh.\n faces: The faces of the mesh.\n vertex_normals: The vertex normals of the mesh.\n px_per_uv_triangle: The number of pixels per UV triangle.\n \"\"\"\n\n # pylint: disable=too-many-statements\n\n assert len(vertices) == len(vertex_normals), \"Number of vertices and vertex normals must be equal\"\n device = vertices.device\n\n # calculate the number of rectangles needed\n triangle_padding = 3\n num_squares = math.ceil(len(faces) / 2)\n squares_per_side_w = math.ceil(math.sqrt(num_squares))\n squares_per_side_h = math.ceil(num_squares / squares_per_side_w)\n px_per_square_w = px_per_uv_triangle + triangle_padding\n px_per_square_h = px_per_uv_triangle\n num_pixels_w = squares_per_side_w * px_per_square_w\n num_pixels_h = squares_per_side_h * px_per_square_h\n\n # Construct what one square would look like\n # The height is equal to px_per_uv_triangle pixels.\n # The width is equal to px_per_uv_triangle + 3 pixels.\n # v0---------------v1------------------------v2\n # --Triangle 1---------------------------------\n # -----------------3px gap---------------------\n # --------------------------------Triangle 2---\n # v2-----------------------v1----------------v0\n\n lr_w = (px_per_uv_triangle + triangle_padding) / num_pixels_w\n lr_h = (px_per_uv_triangle) / num_pixels_h\n lr = torch.tensor([lr_w, lr_h], device=device)\n px_w = 1.0 / num_pixels_w\n px_h = 1.0 / num_pixels_h\n px = torch.tensor([px_w, px_h], device=device)\n edge_len_w = px_per_uv_triangle / num_pixels_w\n edge_len_h = px_per_uv_triangle / num_pixels_h\n scalar = (px_per_uv_triangle - 1) / px_per_uv_triangle\n # uv coords (upper left and lower right)\n uv_coords_upper_left = torch.tensor([[0, 0], [edge_len_w, 0], [0, edge_len_h]], device=device)\n # scale for bilinear interpolation reasons\n uv_coords_upper_left = uv_coords_upper_left * scalar + px / 2\n lower_right = [lr_w, lr_h]\n uv_coords_lower_right = torch.tensor(\n [\n lower_right, # lower right\n [3 * px_w, lr_h], # lower left\n [lr_w, 0], # upper right\n ],\n device=device,\n )\n # scale for bilinear interpolation reasons\n uv_coords_lower_right = (\n (uv_coords_lower_right - torch.tensor(lower_right, device=device)) * scalar\n + torch.tensor(lower_right, device=device)\n - px / 2\n )\n\n # Tile this pattern across the entire texture\n uv_coords_square = torch.stack([uv_coords_upper_left, uv_coords_lower_right], dim=0) # (2, 3, 2)\n uv_coords_square = uv_coords_square.reshape(1, 1, 6, 2) # (6, 2)\n square_offsets = (\n torch.stack(\n torch.meshgrid(\n torch.arange(squares_per_side_w, device=device),\n torch.arange(squares_per_side_h, device=device),\n indexing=\"xy\",\n ),\n dim=-1,\n )\n * lr\n )\n uv_coords_square = uv_coords_square + square_offsets.view(\n squares_per_side_h, squares_per_side_w, 1, 2\n ) # (num_squares_h, num_squares_w, 6, 2)\n texture_coordinates = uv_coords_square.view(-1, 3, 2)[: len(faces)] # (num_faces, 3, 2)\n\n # Now find the triangle indices for every pixel and the barycentric coordinates\n # which can be used to interpolate the XYZ and normal values to then query with NeRF\n uv_coords, uv_indices = get_texture_image(num_pixels_w, num_pixels_h, device)\n\n u_index = torch.div(uv_indices[..., 0], px_per_square_w, rounding_mode=\"floor\")\n v_index = torch.div(uv_indices[..., 1], px_per_square_h, rounding_mode=\"floor\")\n square_index = v_index * squares_per_side_w + u_index\n u_offset = uv_indices[..., 0] % px_per_square_w\n v_offset = uv_indices[..., 1] % px_per_square_h\n lower_right = (u_offset + v_offset) >= (px_per_square_w - 2)\n triangle_index = square_index * 2 + lower_right\n triangle_index = torch.clamp(triangle_index, min=0, max=len(faces) - 1)\n\n nearby_uv_coords = texture_coordinates[triangle_index] # (num_pixels_h, num_pixels_w, 3, 2)\n nearby_vertices = vertices[faces[triangle_index]] # (num_pixels_h, num_pixels_w, 3, 3)\n nearby_normals = vertex_normals[faces[triangle_index]] # (num_pixels_h, num_pixels_w, 3, 3)\n\n # compute barycentric coordinates\n v0 = nearby_uv_coords[..., 0, :] # (num_pixels, num_pixels, 2)\n v1 = nearby_uv_coords[..., 1, :] # (num_pixels, num_pixels, 2)\n v2 = nearby_uv_coords[..., 2, :] # (num_pixels, num_pixels, 2)\n p = uv_coords # (num_pixels, num_pixels, 2)\n area = get_parallelogram_area(v2, v0, v1) # 2x face area.\n w0 = get_parallelogram_area(p, v1, v2) / area\n w1 = get_parallelogram_area(p, v2, v0) / area\n w2 = get_parallelogram_area(p, v0, v1) / area\n\n origins = (\n nearby_vertices[..., 0, :] * w0[..., None]\n + nearby_vertices[..., 1, :] * w1[..., None]\n + nearby_vertices[..., 2, :] * w2[..., None]\n ).float()\n directions = -(\n nearby_normals[..., 0, :] * w0[..., None]\n + nearby_normals[..., 1, :] * w1[..., None]\n + nearby_normals[..., 2, :] * w2[..., None]\n ).float()\n # normalize the direction vector to make it a unit vector\n directions = torch.nn.functional.normalize(directions, dim=-1)\n\n return texture_coordinates, origins, directions\n\n\ndef unwrap_mesh_with_xatlas(\n vertices: TensorType[\"num_verts\", 3],\n faces: TensorType[\"num_faces\", 3, torch.long],\n vertex_normals: TensorType[\"num_verts\", 3],\n num_pixels_per_side=1024,\n num_faces_per_barycentric_chunk=10,\n) -> Tuple[\n TensorType[\"num_faces\", 3, 2],\n TensorType[\"num_pixels\", \"num_pixels\", 3],\n TensorType[\"num_pixels\", \"num_pixels\", \"num_pixels\"],\n]:\n \"\"\"Unwrap a mesh using xatlas. We use xatlas to unwrap the mesh with UV coordinates.\n Then we rasterize the mesh with a square pattern. We interpolate the XYZ and normal\n values for every pixel in the texture image. We return the texture coordinates, the\n origins, and the directions for every pixel.\n\n Args:\n vertices: Tensor of mesh vertices.\n faces: Tensor of mesh faces.\n vertex_normals: Tensor of mesh vertex normals.\n num_pixels_per_side: Number of pixels per side of the texture image. We use a square.\n num_faces_per_barycentric_chunk: Number of faces to use for barycentric chunk computation.\n\n Returns:\n texture_coordinates: Tensor of texture coordinates for every face.\n origins: Tensor of origins for every pixel.\n directions: Tensor of directions for every pixel.\n \"\"\"\n\n # pylint: disable=unused-variable\n # pylint: disable=too-many-statements\n\n device = vertices.device\n\n # unwrap the mesh\n vertices_np = vertices.cpu().numpy()\n faces_np = faces.cpu().numpy()\n vertex_normals_np = vertex_normals.cpu().cpu().numpy()\n vmapping, indices, uvs = xatlas.parametrize( # pylint: disable=c-extension-no-member\n vertices_np, faces_np, vertex_normals_np\n )\n\n # vertices texture coordinates\n vertices_tc = torch.from_numpy(uvs.astype(np.float32)).to(device)\n\n # render uv maps\n vertices_tc = vertices_tc * 2.0 - 1.0 # uvs to range [-1, 1]\n vertices_tc = torch.cat(\n (vertices_tc, torch.zeros_like(vertices_tc[..., :1]), torch.ones_like(vertices_tc[..., :1])), dim=-1\n ) # [num_verts, 4]\n\n texture_coordinates = torch.from_numpy(uvs[indices]).to(device) # (num_faces, 3, 2)\n\n # Now find the triangle indices for every pixel and the barycentric coordinates\n # which can be used to interpolate the XYZ and normal values to then query with NeRF\n uv_coords, _ = get_texture_image(num_pixels_per_side, num_pixels_per_side, device)\n uv_coords_shape = uv_coords.shape\n p = uv_coords.reshape(1, -1, 2) # (1, N, 2)\n num_vertices = p.shape[1]\n num_faces = texture_coordinates.shape[0]\n triangle_distances = torch.ones_like(p[..., 0]) * torch.finfo(torch.float32).max # (1, N)\n triangle_indices = torch.zeros_like(p[..., 0]).long() # (1, N)\n triangle_w0 = torch.zeros_like(p[..., 0]) # (1, N)\n triangle_w1 = torch.zeros_like(p[..., 0]) # (1, N)\n triangle_w2 = torch.zeros_like(p[..., 0]) # (1, N)\n arange_list = torch.arange(num_vertices, device=device)\n progress = get_progress(\"Chunking faces for rasterization\")\n with progress:\n for i in progress.track(range(num_faces // num_faces_per_barycentric_chunk)):\n s = i * num_faces_per_barycentric_chunk\n e = min((i + 1) * num_faces_per_barycentric_chunk, num_faces)\n v0 = texture_coordinates[s:e, 0:1, :] # (F, 1, 2)\n v1 = texture_coordinates[s:e, 1:2, :] # (F, 1, 2)\n v2 = texture_coordinates[s:e, 2:3, :] # (F, 1, 2)\n # NOTE: could try clockwise vs counter clockwise\n area = get_parallelogram_area(v2, v0, v1) # 2x face area.\n w0 = get_parallelogram_area(p, v1, v2) / area # (num_faces_per_barycentric_chunk, N)\n w1 = get_parallelogram_area(p, v2, v0) / area\n w2 = get_parallelogram_area(p, v0, v1) / area\n # get distance from center of triangle\n dist_to_center = torch.abs(w0) + torch.abs(w1) + torch.abs(w2)\n d_values, d_indices = torch.min(dist_to_center, dim=0, keepdim=True)\n d_indices_with_offset = d_indices + s # add offset\n condition = d_values < triangle_distances\n triangle_distances = torch.where(condition, d_values, triangle_distances)\n triangle_indices = torch.where(condition, d_indices_with_offset, triangle_indices)\n w0_selected = w0[d_indices[0], arange_list].unsqueeze(0) # (1, N)\n w1_selected = w1[d_indices[0], arange_list].unsqueeze(0) # (1, N)\n w2_selected = w2[d_indices[0], arange_list].unsqueeze(0) # (1, N)\n triangle_w0 = torch.where(condition, w0_selected, triangle_w0)\n triangle_w1 = torch.where(condition, w1_selected, triangle_w1)\n triangle_w2 = torch.where(condition, w2_selected, triangle_w2)\n\n nearby_vertices = vertices[faces[triangle_indices[0]]] # (N, 3, 3)\n nearby_normals = vertex_normals[faces[triangle_indices[0]]] # (N, 3, 3)\n\n origins = (\n nearby_vertices[..., 0, :] * triangle_w0[0, :, None]\n + nearby_vertices[..., 1, :] * triangle_w1[0, :, None]\n + nearby_vertices[..., 2, :] * triangle_w2[0, :, None]\n ).float()\n directions = -(\n nearby_normals[..., 0, :] * triangle_w0[0, :, None]\n + nearby_normals[..., 1, :] * triangle_w1[0, :, None]\n + nearby_normals[..., 2, :] * triangle_w2[0, :, None]\n ).float()\n\n origins = origins.reshape(uv_coords_shape[0], uv_coords_shape[1], 3)\n directions = directions.reshape(uv_coords_shape[0], uv_coords_shape[1], 3)\n\n # normalize the direction vector to make it a unit vector\n directions = torch.nn.functional.normalize(directions, dim=-1)\n\n return texture_coordinates, origins, directions\n\n\ndef export_textured_mesh(\n mesh: Mesh,\n pipeline: Pipeline,\n output_dir: Path,\n px_per_uv_triangle: Optional[int] = None,\n unwrap_method: Literal[\"xatlas\", \"custom\"] = \"xatlas\",\n raylen_method: Literal[\"edge\", \"none\"] = \"edge\",\n num_pixels_per_side=1024,\n):\n \"\"\"Textures a mesh using the radiance field from the Pipeline.\n The mesh is written to an OBJ file in the output directory,\n along with the corresponding material and texture files.\n Operations will occur on the same device as the Pipeline.\n\n Args:\n mesh: The mesh to texture.\n pipeline: The pipeline to use for texturing.\n output_dir: The directory to write the textured mesh to.\n px_per_uv_triangle: The number of pixels per side of UV triangle.\n unwrap_method: The method to use for unwrapping the mesh.\n offset_method: The method to use for computing the ray length to render.\n num_pixels_per_side: The number of pixels per side of the texture image.\n \"\"\"\n\n # pylint: disable=too-many-statements\n\n device = pipeline.device\n\n vertices = mesh.vertices.to(device)\n faces = mesh.faces.to(device)\n vertex_normals = mesh.normals.to(device)\n\n summary_log = []\n summary_log.append(f\"Unwrapped mesh using {unwrap_method} method.\")\n summary_log.append(f\"Mesh has {len(vertices)} vertices and {len(faces)} faces.\")\n\n if unwrap_method == \"xatlas\":\n CONSOLE.print(\"Unwrapping mesh with xatlas method... this may take a while.\")\n texture_coordinates, origins, directions = unwrap_mesh_with_xatlas(\n vertices, faces, vertex_normals, num_pixels_per_side=num_pixels_per_side\n )\n print(\"\\033[A\\033[A\")\n CONSOLE.print(\"[bold green]:white_check_mark: Unwrapped mesh with xatlas method\")\n elif unwrap_method == \"custom\":\n CONSOLE.print(\"Unwrapping mesh with custom method...\")\n texture_coordinates, origins, directions = unwrap_mesh_per_uv_triangle(\n vertices, faces, vertex_normals, px_per_uv_triangle\n )\n print(\"\\033[A\\033[A\")\n CONSOLE.print(\"[bold green]:white_check_mark: Unwrapped mesh with custom method\")\n else:\n raise ValueError(f\"Unwrap method {unwrap_method} not supported.\")\n\n if raylen_method == \"edge\":\n face_vertices = vertices[faces]\n # compute the length of the rays we want to render\n # we make a reasonable approximation by using the mean length of one edge per face\n raylen = 2.0 * torch.mean(torch.norm(face_vertices[:, 1, :] - face_vertices[:, 0, :], dim=-1)).float()\n elif raylen_method == \"none\":\n raylen = 0.0\n else:\n raise ValueError(f\"Ray length method {raylen_method} not supported.\")\n\n summary_log.append(f\"Length of rendered rays to compute texture values: {raylen}\")\n\n origins = origins - 0.5 * raylen * directions\n pixel_area = torch.ones_like(origins[..., 0:1])\n camera_indices = torch.zeros_like(origins[..., 0:1])\n nears = torch.zeros_like(origins[..., 0:1])\n fars = torch.ones_like(origins[..., 0:1]) * raylen\n camera_ray_bundle = RayBundle(\n origins=origins,\n directions=directions,\n pixel_area=pixel_area,\n camera_indices=camera_indices,\n nears=nears,\n fars=fars,\n )\n\n CONSOLE.print(\"Creating texture image by rendering with NeRF...\")\n with torch.no_grad():\n outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle)\n\n # save the texture image\n texture_image = outputs[\"rgb\"].cpu().numpy()\n media.write_image(str(output_dir / \"material_0.png\"), texture_image)\n\n CONSOLE.print(\"Writing relevant OBJ information to files...\")\n # create the .mtl file\n lines_mtl = [\n \"# Generated with nerfstudio\",\n \"newmtl material_0\",\n \"Ka 1.000 1.000 1.000\",\n \"Kd 1.000 1.000 1.000\",\n \"Ks 0.000 0.000 0.000\",\n \"d 1.0\",\n \"illum 2\",\n \"Ns 1.00000000\",\n \"map_Kd material_0.png\",\n ]\n lines_mtl = [line + \"\\n\" for line in lines_mtl]\n file_mtl = open(output_dir / \"material_0.mtl\", \"w\", encoding=\"utf-8\") # pylint: disable=consider-using-with\n file_mtl.writelines(lines_mtl)\n file_mtl.close()\n\n # create the .obj file\n lines_obj = [\"# Generated with nerfstudio\", \"mtllib material_0.mtl\", \"usemtl material_0\"]\n lines_obj = [line + \"\\n\" for line in lines_obj]\n file_obj = open(output_dir / \"mesh.obj\", \"w\", encoding=\"utf-8\") # pylint: disable=consider-using-with\n file_obj.writelines(lines_obj)\n\n # write the geometric vertices\n vertices = vertices.cpu().numpy()\n progress = get_progress(\"Writing vertices to file\", suffix=\"lines-per-sec\")\n with progress:\n for i in progress.track(range(len(vertices))):\n vertex = vertices[i]\n line = f\"v {vertex[0]} {vertex[1]} {vertex[2]}\\n\"\n file_obj.write(line)\n\n # write the texture coordinates\n texture_coordinates = texture_coordinates.cpu().numpy()\n with progress:\n progress = get_progress(\"Writing texture coordinates to file\", suffix=\"lines-per-sec\")\n for i in progress.track(range(len(faces))):\n for uv in texture_coordinates[i]:\n line = f\"vt {uv[0]} {1.0 - uv[1]}\\n\"\n file_obj.write(line)\n\n # write the vertex normals\n vertex_normals = vertex_normals.cpu().numpy()\n progress = get_progress(\"Writing vertex normals to file\", suffix=\"lines-per-sec\")\n with progress:\n for i in progress.track(range(len(vertex_normals))):\n normal = vertex_normals[i]\n line = f\"vn {normal[0]} {normal[1]} {normal[2]}\\n\"\n file_obj.write(line)\n\n # write the faces\n faces = faces.cpu().numpy()\n progress = get_progress(\"Writing faces to file\", suffix=\"lines-per-sec\")\n with progress:\n for i in progress.track(range(len(faces))):\n face = faces[i]\n v1 = face[0] + 1\n v2 = face[1] + 1\n v3 = face[2] + 1\n vt1 = i * 3 + 1\n vt2 = i * 3 + 2\n vt3 = i * 3 + 3\n vn1 = v1\n vn2 = v2\n vn3 = v3\n line = f\"f {v1}/{vt1}/{vn1} {v2}/{vt2}/{vn2} {v3}/{vt3}/{vn3}\\n\"\n file_obj.write(line)\n\n file_obj.close()\n\n summary_log.append(f\"OBJ file saved to {output_dir / 'mesh.obj'}\")\n summary_log.append(f\"MTL file saved to {output_dir / 'material_0.mtl'}\")\n summary_log.append(\n f\"Texture image saved to {output_dir / 'material_0.png'} \"\n f\"with resolution {texture_image.shape[1]}x{texture_image.shape[0]} (WxH)\"\n )\n\n CONSOLE.rule(\"[bold green]:tada: :tada: :tada: All DONE :tada: :tada: :tada:\")\n for summary in summary_log:\n CONSOLE.print(summary, justify=\"center\")\n CONSOLE.rule()\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/exporter/tsdf_utils.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nTSDF utils.\n\"\"\"\n\n# pylint: disable=no-member\n\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass, field\nfrom pathlib import Path\nfrom typing import List, Optional, Tuple, Union\n\nimport numpy as np\nimport pymeshlab\nimport torch\nimport torch.nn.functional as F\nfrom rich.console import Console\nfrom skimage import measure\nfrom torchtyping import TensorType\n\nfrom nerfstudio.exporter.exporter_utils import Mesh, render_trajectory\nfrom nerfstudio.pipelines.base_pipeline import Pipeline\n\nCONSOLE = Console(width=120)\n\n\n@dataclass\nclass TSDF:\n \"\"\"\n Class for creating TSDFs.\n \"\"\"\n\n voxel_coords: TensorType[3, \"xdim\", \"ydim\", \"zdim\"]\n \"\"\"Coordinates of each voxel in the TSDF.\"\"\"\n values: TensorType[\"xdim\", \"ydim\", \"zdim\"]\n \"\"\"TSDF values for each voxel.\"\"\"\n weights: TensorType[\"xdim\", \"ydim\", \"zdim\"]\n \"\"\"TSDF weights for each voxel.\"\"\"\n colors: TensorType[\"xdim\", \"ydim\", \"zdim\", 3]\n \"\"\"TSDF colors for each voxel.\"\"\"\n voxel_size: TensorType[3]\n \"\"\"Size of each voxel in the TSDF. [x, y, z] size.\"\"\"\n origin: TensorType[3]\n \"\"\"Origin of the TSDF [xmin, ymin, zmin].\"\"\"\n truncation_margin: float = 5.0\n \"\"\"Margin for truncation.\"\"\"\n\n def to(self, device: str):\n \"\"\"Move the tensors to the specified device.\n\n Args:\n device: The device to move the tensors to. E.g., \"cuda:0\" or \"cpu\".\n \"\"\"\n self.voxel_coords = self.voxel_coords.to(device)\n self.values = self.values.to(device)\n self.weights = self.weights.to(device)\n self.colors = self.colors.to(device)\n self.voxel_size = self.voxel_size.to(device)\n self.origin = self.origin.to(device)\n return self\n\n @property\n def device(self):\n \"\"\"Returns the device that voxel_coords is on.\"\"\"\n return self.voxel_coords.device\n\n @property\n def truncation(self):\n \"\"\"Returns the truncation distance.\"\"\"\n # TODO: clean this up\n truncation = self.voxel_size[0] * self.truncation_margin\n return truncation\n\n @staticmethod\n def from_aabb(aabb: TensorType[2, 3], volume_dims: TensorType[3]):\n \"\"\"Returns an instance of TSDF from an axis-aligned bounding box and volume dimensions.\n\n Args:\n aabb: The axis-aligned bounding box with shape [[xmin, ymin, zmin], [xmax, ymax, zmax]].\n volume_dims: The volume dimensions with shape [xdim, ydim, zdim].\n \"\"\"\n\n origin = aabb[0]\n voxel_size = (aabb[1] - aabb[0]) / volume_dims\n\n # create the voxel coordinates\n xdim = torch.arange(volume_dims[0])\n ydim = torch.arange(volume_dims[1])\n zdim = torch.arange(volume_dims[2])\n grid = torch.stack(torch.meshgrid([xdim, ydim, zdim], indexing=\"ij\"), dim=0)\n voxel_coords = origin.view(3, 1, 1, 1) + grid * voxel_size.view(3, 1, 1, 1)\n\n # initialize the values and weights\n values = -torch.ones(volume_dims.tolist())\n weights = torch.zeros(volume_dims.tolist())\n colors = torch.zeros(volume_dims.tolist() + [3])\n\n # TODO: move to device\n\n return TSDF(voxel_coords, values, weights, colors, voxel_size, origin)\n\n def get_mesh(self) -> Mesh:\n \"\"\"Extracts a mesh using marching cubes.\"\"\"\n\n device = self.values.device\n\n # run marching cubes on CPU\n tsdf_values_np = self.values.clamp(-1, 1).cpu().numpy()\n vertices, faces, normals, _ = measure.marching_cubes(tsdf_values_np, level=0, allow_degenerate=False)\n\n vertices_indices = np.round(vertices).astype(int)\n colors = self.colors[vertices_indices[:, 0], vertices_indices[:, 1], vertices_indices[:, 2]]\n\n # move back to original device\n vertices = torch.from_numpy(vertices.copy()).to(device)\n faces = torch.from_numpy(faces.copy()).to(device)\n normals = torch.from_numpy(normals.copy()).to(device)\n\n # move vertices back to world space\n vertices = self.origin.view(1, 3) + vertices * self.voxel_size.view(1, 3)\n\n return Mesh(vertices=vertices, faces=faces, normals=normals, colors=colors)\n\n @classmethod\n def export_mesh(cls, mesh: Mesh, filename: str):\n \"\"\"Exports the mesh to a file.\n We use pymeshlab to export the mesh as a PLY file.\n\n Args:\n mesh: The mesh to export.\n filename: The filename to export the mesh to.\n \"\"\"\n vertex_matrix = mesh.vertices.cpu().numpy().astype(\"float64\")\n face_matrix = mesh.faces.cpu().numpy().astype(\"int32\")\n v_normals_matrix = mesh.normals.cpu().numpy().astype(\"float64\")\n v_color_matrix = mesh.colors.cpu().numpy().astype(\"float64\")\n # colors need an alpha channel\n v_color_matrix = np.concatenate([v_color_matrix, np.ones((v_color_matrix.shape[0], 1))], axis=-1)\n\n # create a new Mesh\n m = pymeshlab.Mesh(\n vertex_matrix=vertex_matrix,\n face_matrix=face_matrix,\n v_normals_matrix=v_normals_matrix,\n v_color_matrix=v_color_matrix,\n )\n # create a new MeshSet\n ms = pymeshlab.MeshSet()\n # add the mesh to the MeshSet\n ms.add_mesh(m, \"mesh\")\n # save the current mesh\n ms.save_current_mesh(filename)\n\n def integrate_tsdf(\n self,\n c2w: TensorType[\"batch\", 4, 4],\n K: TensorType[\"batch\", 3, 3],\n depth_images: TensorType[\"batch\", 1, \"height\", \"width\"],\n color_images: Optional[TensorType[\"batch\", 3, \"height\", \"width\"]] = None,\n mask_images: Optional[TensorType[\"batch\", 1, \"height\", \"width\"]] = None,\n ):\n \"\"\"Integrates a batch of depth images into the TSDF.\n\n Args:\n c2w: The camera extrinsics.\n K: The camera intrinsics.\n depth_images: The depth images to integrate.\n color_images: The color images to integrate.\n mask_images: The mask images to integrate.\n \"\"\"\n\n if mask_images is not None:\n raise NotImplementedError(\"Mask images are not supported yet.\")\n\n batch_size = c2w.shape[0]\n shape = self.voxel_coords.shape[1:]\n\n # Project voxel_coords into image space...\n\n image_size = torch.tensor(\n [depth_images.shape[-1], depth_images.shape[-2]], device=self.device\n ) # [width, height]\n\n # make voxel_coords homogeneous\n voxel_world_coords = self.voxel_coords.view(3, -1)\n voxel_world_coords = torch.cat(\n [voxel_world_coords, torch.ones(1, voxel_world_coords.shape[1], device=self.device)], dim=0\n )\n voxel_world_coords = voxel_world_coords.unsqueeze(0) # [1, 4, N]\n voxel_world_coords = voxel_world_coords.expand(batch_size, *voxel_world_coords.shape[1:]) # [batch, 4, N]\n\n voxel_cam_coords = torch.bmm(torch.inverse(c2w), voxel_world_coords) # [batch, 4, N]\n\n # flip the z axis\n voxel_cam_coords[:, 2, :] = -voxel_cam_coords[:, 2, :]\n # flip the y axis\n voxel_cam_coords[:, 1, :] = -voxel_cam_coords[:, 1, :]\n\n # we need the distance of the point to the camera, not the z coordinate\n voxel_depth = torch.sqrt(torch.sum(voxel_cam_coords[:, :3, :] ** 2, dim=-2, keepdim=True)) # [batch, 1, N]\n\n voxel_cam_coords_z = voxel_cam_coords[:, 2:3, :]\n voxel_cam_points = torch.bmm(K, voxel_cam_coords[:, 0:3, :] / voxel_cam_coords_z) # [batch, 3, N]\n voxel_pixel_coords = voxel_cam_points[:, :2, :] # [batch, 2, N]\n\n # Sample the depth images with grid sample...\n\n grid = voxel_pixel_coords.permute(0, 2, 1) # [batch, N, 2]\n # normalize grid to [-1, 1]\n grid = 2.0 * grid / image_size.view(1, 1, 2) - 1.0 # [batch, N, 2]\n grid = grid[:, None] # [batch, 1, N, 2]\n # depth\n sampled_depth = F.grid_sample(\n input=depth_images, grid=grid, mode=\"nearest\", padding_mode=\"zeros\", align_corners=False\n ) # [batch, N, 1]\n sampled_depth = sampled_depth.squeeze(2) # [batch, 1, N]\n # colors\n if color_images is not None:\n sampled_colors = F.grid_sample(\n input=color_images, grid=grid, mode=\"nearest\", padding_mode=\"zeros\", align_corners=False\n ) # [batch, N, 3]\n sampled_colors = sampled_colors.squeeze(2) # [batch, 3, N]\n\n dist = sampled_depth - voxel_depth # [batch, 1, N]\n tsdf_values = torch.clamp(dist / self.truncation, min=-1.0, max=1.0) # [batch, 1, N]\n valid_points = (voxel_depth > 0) & (sampled_depth > 0) & (dist > -self.truncation) # [batch, 1, N]\n\n # Sequentially update the TSDF...\n\n for i in range(batch_size):\n\n valid_points_i = valid_points[i]\n valid_points_i_shape = valid_points_i.view(*shape) # [xdim, ydim, zdim]\n\n # the old values\n old_tsdf_values_i = self.values[valid_points_i_shape]\n old_weights_i = self.weights[valid_points_i_shape]\n\n # the new values\n # TODO: let the new weight be configurable\n new_tsdf_values_i = tsdf_values[i][valid_points_i]\n new_weights_i = 1.0\n\n total_weights = old_weights_i + new_weights_i\n\n self.values[valid_points_i_shape] = (\n old_tsdf_values_i * old_weights_i + new_tsdf_values_i * new_weights_i\n ) / total_weights\n self.weights[valid_points_i_shape] = torch.clamp(total_weights, max=1.0)\n\n if color_images is not None:\n old_colors_i = self.colors[valid_points_i_shape] # [M, 3]\n new_colors_i = sampled_colors[i][:, valid_points_i.squeeze(0)].permute(1, 0) # [M, 3]\n self.colors[valid_points_i_shape] = (\n old_colors_i * old_weights_i[:, None] + new_colors_i * new_weights_i\n ) / total_weights[:, None]\n\n\ndef export_tsdf_mesh(\n pipeline: Pipeline,\n output_dir: Path,\n downscale_factor: int = 2,\n depth_output_name: str = \"depth\",\n rgb_output_name: str = \"rgb\",\n resolution: Union[int, List[int]] = field(default_factory=lambda: [256, 256, 256]),\n batch_size: int = 10,\n use_bounding_box: bool = True,\n bounding_box_min: Tuple[float, float, float] = (-1.0, -1.0, -1.0),\n bounding_box_max: Tuple[float, float, float] = (1.0, 1.0, 1.0),\n):\n \"\"\"Export a TSDF mesh from a pipeline.\n\n Args:\n pipeline: The pipeline to export the mesh from.\n output_dir: The directory to save the mesh to.\n downscale_factor: Downscale factor for the images.\n depth_output_name: Name of the depth output.\n rgb_output_name: Name of the RGB output.\n resolution: Resolution of the TSDF volume or [x, y, z] resolutions individually.\n batch_size: How many depth images to integrate per batch.\n use_bounding_box: Whether to use a bounding box for the TSDF volume.\n bounding_box_min: Minimum coordinates of the bounding box.\n bounding_box_max: Maximum coordinates of the bounding box.\n \"\"\"\n\n device = pipeline.device\n\n dataparser_outputs = pipeline.datamanager.train_dataset._dataparser_outputs # pylint: disable=protected-access\n\n # initialize the TSDF volume\n if not use_bounding_box:\n aabb = dataparser_outputs.scene_box.aabb\n else:\n aabb = torch.tensor([bounding_box_min, bounding_box_max])\n if isinstance(resolution, int):\n volume_dims = torch.tensor([resolution] * 3)\n elif isinstance(resolution, List):\n volume_dims = torch.tensor(resolution)\n else:\n raise ValueError(\"Resolution must be an int or a list.\")\n tsdf = TSDF.from_aabb(aabb, volume_dims=volume_dims)\n # move TSDF to device\n tsdf.to(device)\n\n cameras = dataparser_outputs.cameras\n # we turn off distortion when populating the TSDF\n color_images, depth_images = render_trajectory(\n pipeline,\n cameras,\n rgb_output_name=rgb_output_name,\n depth_output_name=depth_output_name,\n rendered_resolution_scaling_factor=1.0 / downscale_factor,\n disable_distortion=True,\n )\n\n # camera extrinsics and intrinsics\n c2w: TensorType[\"N\", 3, 4] = cameras.camera_to_worlds.to(device)\n # make c2w homogeneous\n c2w = torch.cat([c2w, torch.zeros(c2w.shape[0], 1, 4, device=device)], dim=1)\n c2w[:, 3, 3] = 1\n K: TensorType[\"N\", 3, 3] = cameras.get_intrinsics_matrices().to(device)\n color_images = torch.tensor(np.array(color_images), device=device).permute(0, 3, 1, 2) # shape (N, 3, H, W)\n depth_images = torch.tensor(np.array(depth_images), device=device).permute(0, 3, 1, 2) # shape (N, 1, H, W)\n\n CONSOLE.print(\"Integrating the TSDF\")\n for i in range(0, len(c2w), batch_size):\n tsdf.integrate_tsdf(\n c2w[i : i + batch_size],\n K[i : i + batch_size],\n depth_images[i : i + batch_size],\n color_images=color_images[i : i + batch_size],\n )\n\n CONSOLE.print(\"Computing Mesh\")\n mesh = tsdf.get_mesh()\n CONSOLE.print(\"Saving TSDF Mesh\")\n tsdf.export_mesh(mesh, filename=str(output_dir / \"tsdf_mesh.ply\"))\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/field_components/__init__.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"init field modules\"\"\"\nfrom .base_field_component import FieldComponent\nfrom .encodings import Encoding, ScalingAndOffset\nfrom .mlp import MLP\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/field_components/activations.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nSpecial activation functions.\n\"\"\"\n\nimport torch\nfrom torch.autograd import Function\nfrom torch.cuda.amp import custom_bwd, custom_fwd\n\n\nclass _TruncExp(Function): # pylint: disable=abstract-method\n # Implementation from torch-ngp:\n # https://github.com/ashawkey/torch-ngp/blob/93b08a0d4ec1cc6e69d85df7f0acdfb99603b628/activation.py\n @staticmethod\n @custom_fwd(cast_inputs=torch.float32)\n def forward(ctx, x): # pylint: disable=arguments-differ\n ctx.save_for_backward(x)\n return torch.exp(x)\n\n @staticmethod\n @custom_bwd\n def backward(ctx, g): # pylint: disable=arguments-differ\n x = ctx.saved_tensors[0]\n return g * torch.exp(x.clamp(-15, 15))\n\n\ntrunc_exp = _TruncExp.apply\n\"\"\"Same as torch.exp, but with the backward pass clipped to prevent vanishing/exploding\ngradients.\"\"\"\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/field_components/base_field_component.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nThe field module baseclass.\n\"\"\"\nfrom abc import abstractmethod\nfrom typing import Optional\n\nfrom torch import nn\nfrom torchtyping import TensorType\n\n\nclass FieldComponent(nn.Module):\n \"\"\"Field modules that can be combined to store and compute the fields.\n\n Args:\n in_dim: Input dimension to module.\n out_dim: Ouput dimension to module.\n \"\"\"\n\n def __init__(self, in_dim: Optional[int] = None, out_dim: Optional[int] = None) -> None:\n super().__init__()\n self.in_dim = in_dim\n self.out_dim = out_dim\n\n def build_nn_modules(self) -> None:\n \"\"\"Function instantiates any torch.nn members within the module.\n If none exist, do nothing.\"\"\"\n\n def set_in_dim(self, in_dim: int) -> None:\n \"\"\"Sets input dimension of encoding\n\n Args:\n in_dim: input dimension\n \"\"\"\n if in_dim <= 0:\n raise ValueError(\"Input dimension should be greater than zero\")\n self.in_dim = in_dim\n\n def get_out_dim(self) -> int:\n \"\"\"Calculates output dimension of encoding.\"\"\"\n if self.out_dim is None:\n raise ValueError(\"Output dimension has not been set\")\n return self.out_dim\n\n @abstractmethod\n def forward(self, in_tensor: TensorType[\"bs\":..., \"input_dim\"]) -> TensorType[\"bs\":..., \"output_dim\"]:\n \"\"\"\n Returns processed tensor\n\n Args:\n in_tensor: Input tensor to process\n \"\"\"\n raise NotImplementedError\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/field_components/embedding.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nCode for embeddings.\n\"\"\"\n\n\nimport torch\nfrom torchtyping import TensorType\n\nfrom nerfstudio.field_components.base_field_component import FieldComponent\n\n\nclass Embedding(FieldComponent):\n \"\"\"Index into embeddings.\n # TODO: add different types of initializations\n\n Args:\n in_dim: Number of embeddings\n out_dim: Dimension of the embedding vectors\n \"\"\"\n\n def __init__(self, in_dim: int, out_dim: int) -> None:\n super().__init__()\n self.in_dim = in_dim\n self.out_dim = out_dim\n self.build_nn_modules()\n\n def build_nn_modules(self) -> None:\n self.embedding = torch.nn.Embedding(self.in_dim, self.out_dim)\n\n def mean(self, dim=0):\n \"\"\"Return the mean of the embedding weights along a dim.\"\"\"\n return self.embedding.weight.mean(dim)\n\n def forward(self, in_tensor: TensorType[..., \"input_dim\"]) -> TensorType[..., \"output_dim\"]:\n \"\"\"Call forward\n\n Args:\n in_tensor: input tensor to process\n \"\"\"\n return self.embedding(in_tensor)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/field_components/encodings.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nEncoding functions\n\"\"\"\n\nfrom abc import abstractmethod\nfrom typing import Optional\n\nimport numpy as np\nimport torch\nimport torch.nn.functional as F\nfrom torch import nn\nfrom torchtyping import TensorType\nfrom typing_extensions import Literal\n\nfrom nerfstudio.field_components.base_field_component import FieldComponent\nfrom nerfstudio.utils.math import components_from_spherical_harmonics, expected_sin\nfrom nerfstudio.utils.printing import print_tcnn_speed_warning\n\ntry:\n import tinycudann as tcnn\n\n TCNN_EXISTS = True\nexcept ImportError:\n TCNN_EXISTS = False\n\n\nclass Encoding(FieldComponent):\n \"\"\"Encode an input tensor. Intended to be subclassed\n\n Args:\n in_dim: Input dimension of tensor\n \"\"\"\n\n def __init__(self, in_dim: int) -> None:\n if in_dim <= 0:\n raise ValueError(\"Input dimension should be greater than zero\")\n super().__init__(in_dim=in_dim)\n\n @abstractmethod\n def forward(self, in_tensor: TensorType[\"bs\":..., \"input_dim\"]) -> TensorType[\"bs\":..., \"output_dim\"]:\n \"\"\"Call forward and returns and processed tensor\n\n Args:\n in_tensor: the input tensor to process\n \"\"\"\n raise NotImplementedError\n\n\nclass Identity(Encoding):\n \"\"\"Identity encoding (Does not modify input)\"\"\"\n\n def get_out_dim(self) -> int:\n if self.in_dim is None:\n raise ValueError(\"Input dimension has not been set\")\n return self.in_dim\n\n def forward(self, in_tensor: TensorType[\"bs\":..., \"input_dim\"]) -> TensorType[\"bs\":..., \"output_dim\"]:\n return in_tensor\n\n\nclass ScalingAndOffset(Encoding):\n \"\"\"Simple scaling and offet to input\n\n Args:\n in_dim: Input dimension of tensor\n scaling: Scaling applied to tensor.\n offset: Offset applied to tensor.\n \"\"\"\n\n def __init__(self, in_dim: int, scaling: float = 1.0, offset: float = 0.0) -> None:\n super().__init__(in_dim)\n\n self.scaling = scaling\n self.offset = offset\n\n def get_out_dim(self) -> int:\n if self.in_dim is None:\n raise ValueError(\"Input dimension has not been set\")\n return self.in_dim\n\n def forward(self, in_tensor: TensorType[\"bs\":..., \"input_dim\"]) -> TensorType[\"bs\":..., \"output_dim\"]:\n return self.scaling * in_tensor + self.offset\n\n\nclass NeRFEncoding(Encoding):\n \"\"\"Multi-scale sinusoidal encodings. Support ``integrated positional encodings`` if covariances are provided.\n Each axis is encoded with frequencies ranging from 2^min_freq_exp to 2^max_freq_exp.\n\n Args:\n in_dim: Input dimension of tensor\n num_frequencies: Number of encoded frequencies per axis\n min_freq_exp: Minimum frequency exponent\n max_freq_exp: Maximum frequency exponent\n include_input: Append the input coordinate to the encoding\n \"\"\"\n\n def __init__(\n self, in_dim: int, num_frequencies: int, min_freq_exp: float, max_freq_exp: float, include_input: bool = False\n ) -> None:\n super().__init__(in_dim)\n\n self.num_frequencies = num_frequencies\n self.min_freq = min_freq_exp\n self.max_freq = max_freq_exp\n self.include_input = include_input\n\n def get_out_dim(self) -> int:\n if self.in_dim is None:\n raise ValueError(\"Input dimension has not been set\")\n out_dim = self.in_dim * self.num_frequencies * 2\n if self.include_input:\n out_dim += self.in_dim\n return out_dim\n\n def forward(\n self,\n in_tensor: TensorType[\"bs\":..., \"input_dim\"],\n covs: Optional[TensorType[\"bs\":..., \"input_dim\", \"input_dim\"]] = None,\n ) -> TensorType[\"bs\":..., \"output_dim\"]:\n \"\"\"Calculates NeRF encoding. If covariances are provided the encodings will be integrated as proposed\n in mip-NeRF.\n\n Args:\n in_tensor: For best performance, the input tensor should be between 0 and 1.\n covs: Covariances of input points.\n Returns:\n Output values will be between -1 and 1\n \"\"\"\n in_tensor = 2 * torch.pi * in_tensor # scale to [0, 2pi]\n freqs = 2 ** torch.linspace(self.min_freq, self.max_freq, self.num_frequencies).to(in_tensor.device)\n scaled_inputs = in_tensor[..., None] * freqs # [..., \"input_dim\", \"num_scales\"]\n scaled_inputs = scaled_inputs.view(*scaled_inputs.shape[:-2], -1) # [..., \"input_dim\" * \"num_scales\"]\n\n if covs is None:\n encoded_inputs = torch.sin(torch.cat([scaled_inputs, scaled_inputs + torch.pi / 2.0], dim=-1))\n else:\n input_var = torch.diagonal(covs, dim1=-2, dim2=-1)[..., :, None] * freqs[None, :] ** 2\n input_var = input_var.reshape((*input_var.shape[:-2], -1))\n encoded_inputs = expected_sin(\n torch.cat([scaled_inputs, scaled_inputs + torch.pi / 2.0], dim=-1), torch.cat(2 * [input_var], dim=-1)\n )\n\n if self.include_input:\n encoded_inputs = torch.cat([encoded_inputs, in_tensor], dim=-1)\n return encoded_inputs\n\n\nclass RFFEncoding(Encoding):\n \"\"\"Random Fourier Feature encoding. Supports integrated encodings.\n\n Args:\n in_dim: Input dimension of tensor\n num_frequencies: Number of encoding frequencies\n scale: Std of Gaussian to sample frequencies. Must be greater than zero\n include_input: Append the input coordinate to the encoding\n \"\"\"\n\n def __init__(self, in_dim: int, num_frequencies: int, scale: float, include_input: bool = False) -> None:\n super().__init__(in_dim)\n\n self.num_frequencies = num_frequencies\n if not scale > 0:\n raise ValueError(\"RFF encoding scale should be greater than zero\")\n self.scale = scale\n if self.in_dim is None:\n raise ValueError(\"Input dimension has not been set\")\n b_matrix = torch.normal(mean=0, std=self.scale, size=(self.in_dim, self.num_frequencies))\n self.register_buffer(name=\"b_matrix\", tensor=b_matrix)\n self.include_input = include_input\n\n def get_out_dim(self) -> int:\n return self.num_frequencies * 2\n\n def forward(\n self,\n in_tensor: TensorType[\"bs\":..., \"input_dim\"],\n covs: Optional[TensorType[\"bs\":..., \"input_dim\", \"input_dim\"]] = None,\n ) -> TensorType[\"bs\":..., \"output_dim\"]:\n \"\"\"Calculates RFF encoding. If covariances are provided the encodings will be integrated as proposed\n in mip-NeRF.\n\n Args:\n in_tensor: For best performance, the input tensor should be between 0 and 1.\n covs: Covariances of input points.\n\n Returns:\n Output values will be between -1 and 1\n \"\"\"\n in_tensor = 2 * torch.pi * in_tensor # scale to [0, 2pi]\n scaled_inputs = in_tensor @ self.b_matrix # [..., \"num_frequencies\"]\n\n if covs is None:\n encoded_inputs = torch.sin(torch.cat([scaled_inputs, scaled_inputs + torch.pi / 2.0], dim=-1))\n else:\n input_var = torch.sum((covs @ self.b_matrix) * self.b_matrix, -2)\n encoded_inputs = expected_sin(\n torch.cat([scaled_inputs, scaled_inputs + torch.pi / 2.0], dim=-1), torch.cat(2 * [input_var], dim=-1)\n )\n\n if self.include_input:\n encoded_inputs = torch.cat([encoded_inputs, in_tensor], dim=-1)\n\n return encoded_inputs\n\n\nclass HashEncoding(Encoding):\n \"\"\"Hash encoding\n\n Args:\n num_levels: Number of feature grids.\n min_res: Resolution of smallest feature grid.\n max_res: Resolution of largest feature grid.\n log2_hashmap_size: Size of hash map is 2^log2_hashmap_size.\n features_per_level: Number of features per level.\n hash_init_scale: Value to initialize hash grid.\n implementation: Implementation of hash encoding. Fallback to torch if tcnn not available.\n interpolation: Interpolation override for tcnn hashgrid. Not supported for torch unless linear.\n \"\"\"\n\n def __init__(\n self,\n num_levels: int = 16,\n min_res: int = 16,\n max_res: int = 1024,\n log2_hashmap_size: int = 19,\n features_per_level: int = 2,\n hash_init_scale: float = 0.001,\n implementation: Literal[\"tcnn\", \"torch\"] = \"tcnn\",\n interpolation: Optional[Literal[\"Nearest\", \"Linear\", \"Smoothstep\"]] = None,\n ) -> None:\n\n super().__init__(in_dim=3)\n self.num_levels = num_levels\n self.features_per_level = features_per_level\n self.log2_hashmap_size = log2_hashmap_size\n self.hash_table_size = 2**log2_hashmap_size\n\n levels = torch.arange(num_levels)\n growth_factor = np.exp((np.log(max_res) - np.log(min_res)) / (num_levels - 1))\n self.scalings = torch.floor(min_res * growth_factor**levels)\n\n self.hash_offset = levels * self.hash_table_size\n self.hash_table = torch.rand(size=(self.hash_table_size * num_levels, features_per_level)) * 2 - 1\n self.hash_table *= hash_init_scale\n self.hash_table = nn.Parameter(self.hash_table)\n\n self.tcnn_encoding = None\n if not TCNN_EXISTS and implementation == \"tcnn\":\n print_tcnn_speed_warning(\"HashEncoding\")\n elif implementation == \"tcnn\":\n encoding_config = {\n \"otype\": \"HashGrid\",\n \"n_levels\": self.num_levels,\n \"n_features_per_level\": self.features_per_level,\n \"log2_hashmap_size\": self.log2_hashmap_size,\n \"base_resolution\": min_res,\n \"per_level_scale\": growth_factor,\n }\n if interpolation is not None:\n encoding_config[\"interpolation\"] = interpolation\n\n self.tcnn_encoding = tcnn.Encoding(\n n_input_dims=3,\n encoding_config=encoding_config,\n )\n\n if not TCNN_EXISTS or self.tcnn_encoding is None:\n assert (\n interpolation is None or interpolation == \"Linear\"\n ), f\"interpolation '{interpolation}' is not supported for torch encoding backend\"\n\n def get_out_dim(self) -> int:\n return self.num_levels * self.features_per_level\n\n def hash_fn(self, in_tensor: TensorType[\"bs\":..., \"num_levels\", 3]) -> TensorType[\"bs\":..., \"num_levels\"]:\n \"\"\"Returns hash tensor using method described in Instant-NGP\n\n Args:\n in_tensor: Tensor to be hashed\n \"\"\"\n\n # min_val = torch.min(in_tensor)\n # max_val = torch.max(in_tensor)\n # assert min_val >= 0.0\n # assert max_val <= 1.0\n\n in_tensor = in_tensor * torch.tensor([1, 2654435761, 805459861]).to(in_tensor.device)\n x = torch.bitwise_xor(in_tensor[..., 0], in_tensor[..., 1])\n x = torch.bitwise_xor(x, in_tensor[..., 2])\n x %= self.hash_table_size\n x += self.hash_offset.to(x.device)\n return x\n\n def pytorch_fwd(self, in_tensor: TensorType[\"bs\":..., \"input_dim\"]) -> TensorType[\"bs\":..., \"output_dim\"]:\n \"\"\"Forward pass using pytorch. Significantly slower than TCNN implementation.\"\"\"\n\n assert in_tensor.shape[-1] == 3\n in_tensor = in_tensor[..., None, :] # [..., 1, 3]\n scaled = in_tensor * self.scalings.view(-1, 1).to(in_tensor.device) # [..., L, 3]\n scaled_c = torch.ceil(scaled).type(torch.int32)\n scaled_f = torch.floor(scaled).type(torch.int32)\n\n offset = scaled - scaled_f\n\n hashed_0 = self.hash_fn(scaled_c) # [..., num_levels]\n hashed_1 = self.hash_fn(torch.cat([scaled_c[..., 0:1], scaled_f[..., 1:2], scaled_c[..., 2:3]], dim=-1))\n hashed_2 = self.hash_fn(torch.cat([scaled_f[..., 0:1], scaled_f[..., 1:2], scaled_c[..., 2:3]], dim=-1))\n hashed_3 = self.hash_fn(torch.cat([scaled_f[..., 0:1], scaled_c[..., 1:2], scaled_c[..., 2:3]], dim=-1))\n hashed_4 = self.hash_fn(torch.cat([scaled_c[..., 0:1], scaled_c[..., 1:2], scaled_f[..., 2:3]], dim=-1))\n hashed_5 = self.hash_fn(torch.cat([scaled_c[..., 0:1], scaled_f[..., 1:2], scaled_f[..., 2:3]], dim=-1))\n hashed_6 = self.hash_fn(scaled_f)\n hashed_7 = self.hash_fn(torch.cat([scaled_f[..., 0:1], scaled_c[..., 1:2], scaled_f[..., 2:3]], dim=-1))\n\n f_0 = self.hash_table[hashed_0] # [..., num_levels, features_per_level]\n f_1 = self.hash_table[hashed_1]\n f_2 = self.hash_table[hashed_2]\n f_3 = self.hash_table[hashed_3]\n f_4 = self.hash_table[hashed_4]\n f_5 = self.hash_table[hashed_5]\n f_6 = self.hash_table[hashed_6]\n f_7 = self.hash_table[hashed_7]\n\n f_03 = f_0 * offset[..., 0:1] + f_3 * (1 - offset[..., 0:1])\n f_12 = f_1 * offset[..., 0:1] + f_2 * (1 - offset[..., 0:1])\n f_56 = f_5 * offset[..., 0:1] + f_6 * (1 - offset[..., 0:1])\n f_47 = f_4 * offset[..., 0:1] + f_7 * (1 - offset[..., 0:1])\n\n f0312 = f_03 * offset[..., 1:2] + f_12 * (1 - offset[..., 1:2])\n f4756 = f_47 * offset[..., 1:2] + f_56 * (1 - offset[..., 1:2])\n\n encoded_value = f0312 * offset[..., 2:3] + f4756 * (\n 1 - offset[..., 2:3]\n ) # [..., num_levels, features_per_level]\n\n return torch.flatten(encoded_value, start_dim=-2, end_dim=-1) # [..., num_levels * features_per_level]\n\n def forward(self, in_tensor: TensorType[\"bs\":..., \"input_dim\"]) -> TensorType[\"bs\":..., \"output_dim\"]:\n if TCNN_EXISTS and self.tcnn_encoding is not None:\n return self.tcnn_encoding(in_tensor)\n return self.pytorch_fwd(in_tensor)\n\n\nclass TensorCPEncoding(Encoding):\n \"\"\"Learned CANDECOMP/PARFAC (CP) decomposition encoding used in TensoRF\n\n Args:\n resolution: Resolution of grid.\n num_components: Number of components per dimension.\n init_scale: Initialization scale.\n \"\"\"\n\n def __init__(self, resolution: int = 256, num_components: int = 24, init_scale: float = 0.1) -> None:\n super().__init__(in_dim=3)\n\n self.resolution = resolution\n self.num_components = num_components\n\n # TODO Learning rates should be different for these\n self.line_coef = nn.Parameter(init_scale * torch.randn((3, num_components, resolution, 1)))\n\n def get_out_dim(self) -> int:\n return self.num_components\n\n def forward(self, in_tensor: TensorType[\"bs\":..., \"input_dim\"]) -> TensorType[\"bs\":..., \"output_dim\"]:\n line_coord = torch.stack([in_tensor[..., 2], in_tensor[..., 1], in_tensor[..., 0]]) # [3, ...]\n line_coord = torch.stack([torch.zeros_like(line_coord), line_coord], dim=-1) # [3, ...., 2]\n\n # Stop gradients from going to sampler\n line_coord = line_coord.view(3, -1, 1, 2).detach()\n\n line_features = F.grid_sample(self.line_coef, line_coord, align_corners=True) # [3, Components, -1, 1]\n\n features = torch.prod(line_features, dim=0)\n features = torch.moveaxis(features.view(self.num_components, *in_tensor.shape[:-1]), 0, -1)\n\n return features # [..., Components]\n\n @torch.no_grad()\n def upsample_grid(self, resolution: int) -> None:\n \"\"\"Upsamples underyling feature grid\n\n Args:\n resolution: Target resolution.\n \"\"\"\n\n self.line_coef.data = F.interpolate(\n self.line_coef.data, size=(resolution, 1), mode=\"bilinear\", align_corners=True\n )\n\n self.resolution = resolution\n\n\nclass TensorVMEncoding(Encoding):\n \"\"\"Learned vector-matrix encoding proposed by TensoRF\n\n Args:\n resolution: Resolution of grid.\n num_components: Number of components per dimension.\n init_scale: Initialization scale.\n \"\"\"\n\n plane_coef: TensorType[3, \"num_components\", \"resolution\", \"resolution\"]\n line_coef: TensorType[3, \"num_components\", \"resolution\", 1]\n\n def __init__(\n self,\n resolution: int = 128,\n num_components: int = 24,\n init_scale: float = 0.1,\n smoothstep: bool = False,\n ) -> None:\n super().__init__(in_dim=3)\n\n self.resolution = resolution\n self.num_components = num_components\n self.smoothstep = smoothstep\n\n self.plane_coef = nn.Parameter(init_scale * torch.randn((3 * resolution * resolution, num_components)))\n self.line_coef = nn.Parameter(init_scale * torch.randn((3 * resolution, num_components)))\n\n self.n_output_dims = self.get_out_dim()\n\n def get_out_dim(self) -> int:\n return self.num_components * 3\n\n def index_fn(self, x: torch.Tensor, y: torch.Tensor, width: int, height: int):\n y.clamp_max_(height - 1)\n x.clamp_max_(width - 1)\n\n if y.max() >= height or x.max() >= width:\n breakpoint()\n\n index = y * width + x\n feature_offset = width * height * torch.arange(3)\n index += feature_offset.to(x.device)[:, None, None]\n\n return index.long()\n\n def grid_sample_2d(self, feature, coord, type=\"plane\"):\n if type == \"plane\":\n height, width = self.resolution, self.resolution\n else:\n height, width = self.resolution, 1\n\n scaled = coord * torch.tensor([width, height]).to(coord.device)[None, None]\n scaled_c = torch.ceil(scaled).type(torch.int32)\n scaled_f = torch.floor(scaled).type(torch.int32)\n\n offset = scaled - scaled_f\n\n # smooth version of offset\n if self.smoothstep:\n offset = offset * offset * (3.0 - 2.0 * offset)\n\n offset = offset[..., None, :]\n\n index_0 = self.index_fn(scaled_c[..., 0:1], scaled_c[..., 1:2], height, width) # [..., num_levels]\n index_2 = self.index_fn(scaled_f[..., 0:1], scaled_f[..., 1:2], height, width)\n if type == \"plane\":\n index_1 = self.index_fn(scaled_c[..., 0:1], scaled_f[..., 1:2], height, width)\n index_3 = self.index_fn(scaled_f[..., 0:1], scaled_c[..., 1:2], height, width)\n\n # breakpoint()\n if type == \"plane\":\n f_0 = feature[index_0] # [..., num_levels, features_per_level]\n f_1 = feature[index_1]\n f_2 = feature[index_2]\n f_3 = feature[index_3]\n\n f_03 = f_0 * offset[..., 0:1] + f_3 * (1 - offset[..., 0:1])\n f_12 = f_1 * offset[..., 0:1] + f_2 * (1 - offset[..., 0:1])\n\n f0312 = f_03 * offset[..., 1:2] + f_12 * (1 - offset[..., 1:2])\n\n return f0312\n else:\n f_0 = feature[index_0] # [..., num_levels, features_per_level]\n f_2 = feature[index_2]\n f_02 = f_0 * offset[..., 0:1] + f_2 * (1 - offset[..., 0:1])\n return f_02\n\n def forward(self, in_tensor: TensorType[\"bs\":..., \"input_dim\"]) -> TensorType[\"bs\":..., \"output_dim\"]:\n \"\"\"Compute encoding for each position in in_positions\n\n Args:\n in_tensor: position inside bounds in range [-1,1],\n\n Returns: Encoded position\n \"\"\"\n plane_coord = torch.stack([in_tensor[..., [0, 1]], in_tensor[..., [0, 2]], in_tensor[..., [1, 2]]]) # [3,...,2]\n line_coord = torch.stack([in_tensor[..., 2], in_tensor[..., 1], in_tensor[..., 0]]) # [3, ...]\n line_coord = torch.stack([line_coord, torch.zeros_like(line_coord)], dim=-1) # [3, ...., 2]\n\n # Stop gradients from going to sampler\n # plane_coord = plane_coord.view(3, -1, 1, 2).detach()\n # line_coord = line_coord.view(3, -1, 1, 2).detach()\n\n # plane_features = F.grid_sample(self.plane_coef, plane_coord, align_corners=True) # [3, Components, -1, 1]\n # line_features = F.grid_sample(self.line_coef, line_coord, align_corners=True) # [3, Components, -1, 1]\n\n # diff grid_sample\n\n plane_features = self.grid_sample_2d(self.plane_coef, plane_coord, type=\"plane\") # [3, -1, 1, Components]\n # line_features = self.grid_sample_2d(self.line_coef, line_coord, type=\"line\") # [3, -1, 1, Components]\n\n # features = plane_features * line_features # [3, -1, 1, components]\n features = plane_features\n features = torch.moveaxis(features, 0, 1).reshape(-1, 3 * self.num_components)\n\n # features = torch.moveaxis(features.view(3 * self.num_components, *in_tensor.shape[:-1]), 0, -1)\n\n return features # [..., 3 * Components]\n\n @torch.no_grad()\n def upsample_grid(self, resolution: int) -> None:\n \"\"\"Upsamples underyling feature grid\n\n Args:\n resolution: Target resolution.\n \"\"\"\n plane_coef = F.interpolate(\n self.plane_coef.data, size=(resolution, resolution), mode=\"bilinear\", align_corners=True\n )\n line_coef = F.interpolate(self.line_coef.data, size=(resolution, 1), mode=\"bilinear\", align_corners=True)\n\n # TODO(ethan): are these torch.nn.Parameters needed?\n self.plane_coef, self.line_coef = torch.nn.Parameter(plane_coef), torch.nn.Parameter(line_coef)\n self.resolution = resolution\n\n\nclass SHEncoding(Encoding):\n \"\"\"Spherical harmonic encoding\n\n Args:\n levels: Number of spherical hamonic levels to encode.\n \"\"\"\n\n def __init__(self, levels: int = 4) -> None:\n super().__init__(in_dim=3)\n\n if levels <= 0 or levels > 4:\n raise ValueError(f\"Spherical harmonic encoding only suports 1 to 4 levels, requested {levels}\")\n\n self.levels = levels\n\n def get_out_dim(self) -> int:\n return self.levels**2\n\n @torch.no_grad()\n def forward(self, in_tensor: TensorType[\"bs\":..., \"input_dim\"]) -> TensorType[\"bs\":..., \"output_dim\"]:\n return components_from_spherical_harmonics(levels=self.levels, directions=in_tensor)\n\n\nclass PeriodicVolumeEncoding(Encoding):\n \"\"\"Periodic Volume encoding\n\n Args:\n num_levels: Number of feature grids.\n min_res: Resolution of smallest feature grid.\n max_res: Resolution of largest feature grid.\n log2_hashmap_size: Size of hash map is 2^log2_hashmap_size.\n features_per_level: Number of features per level.\n hash_init_scale: Value to initialize hash grid.\n implementation: Implementation of hash encoding. Fallback to torch if tcnn not available.\n \"\"\"\n\n def __init__(\n self,\n num_levels: int = 16,\n min_res: int = 16,\n max_res: int = 1024,\n log2_hashmap_size: int = 19,\n features_per_level: int = 2,\n hash_init_scale: float = 0.001,\n smoothstep: bool = False,\n ) -> None:\n\n super().__init__(in_dim=3)\n self.num_levels = num_levels\n self.features_per_level = features_per_level\n self.log2_hashmap_size = log2_hashmap_size\n assert log2_hashmap_size % 3 == 0\n self.hash_table_size = 2**log2_hashmap_size\n self.n_output_dims = num_levels * features_per_level\n self.smoothstep = smoothstep\n\n levels = torch.arange(num_levels)\n growth_factor = np.exp((np.log(max_res) - np.log(min_res)) / (num_levels - 1))\n self.scalings = torch.floor(min_res * growth_factor**levels)\n\n self.periodic_volume_resolution = 2 ** (log2_hashmap_size // 3)\n # self.periodic_resolution = torch.minimum(torch.floor(self.scalings), periodic_volume_resolution)\n\n self.hash_offset = levels * self.hash_table_size\n self.hash_table = torch.rand(size=(self.hash_table_size * num_levels, features_per_level)) * 2 - 1\n self.hash_table *= hash_init_scale\n self.hash_table = nn.Parameter(self.hash_table)\n\n # TODO weight loss by level?\n self.per_level_weights = 1.0\n\n def get_out_dim(self) -> int:\n return self.num_levels * self.features_per_level\n\n def hash_fn(self, in_tensor: TensorType[\"bs\":..., \"num_levels\", 3]) -> TensorType[\"bs\":..., \"num_levels\"]:\n \"\"\"Returns hash tensor using method described in Instant-NGP\n\n Args:\n in_tensor: Tensor to be hashed\n \"\"\"\n\n # round to make it perioidic\n x = in_tensor\n x %= self.periodic_volume_resolution\n # xyz to index\n x = (\n x[..., 0] * (self.periodic_volume_resolution**2)\n + x[..., 1] * (self.periodic_volume_resolution)\n + x[..., 2]\n )\n # offset by feature levels\n x += self.hash_offset.to(x.device)\n\n return x.long()\n\n def pytorch_fwd(self, in_tensor: TensorType[\"bs\":..., \"input_dim\"]) -> TensorType[\"bs\":..., \"output_dim\"]:\n \"\"\"Forward pass using pytorch. Significantly slower than TCNN implementation.\"\"\"\n\n assert in_tensor.shape[-1] == 3\n in_tensor = in_tensor[..., None, :] # [..., 1, 3]\n scaled = in_tensor * self.scalings.view(-1, 1).to(in_tensor.device) # [..., L, 3]\n scaled_c = torch.ceil(scaled).type(torch.int32)\n scaled_f = torch.floor(scaled).type(torch.int32)\n\n offset = scaled - scaled_f\n\n if self.smoothstep:\n offset = offset * offset * (3.0 - 2.0 * offset)\n\n hashed_0 = self.hash_fn(scaled_c) # [..., num_levels]\n hashed_1 = self.hash_fn(torch.cat([scaled_c[..., 0:1], scaled_f[..., 1:2], scaled_c[..., 2:3]], dim=-1))\n hashed_2 = self.hash_fn(torch.cat([scaled_f[..., 0:1], scaled_f[..., 1:2], scaled_c[..., 2:3]], dim=-1))\n hashed_3 = self.hash_fn(torch.cat([scaled_f[..., 0:1], scaled_c[..., 1:2], scaled_c[..., 2:3]], dim=-1))\n hashed_4 = self.hash_fn(torch.cat([scaled_c[..., 0:1], scaled_c[..., 1:2], scaled_f[..., 2:3]], dim=-1))\n hashed_5 = self.hash_fn(torch.cat([scaled_c[..., 0:1], scaled_f[..., 1:2], scaled_f[..., 2:3]], dim=-1))\n hashed_6 = self.hash_fn(scaled_f)\n hashed_7 = self.hash_fn(torch.cat([scaled_f[..., 0:1], scaled_c[..., 1:2], scaled_f[..., 2:3]], dim=-1))\n\n f_0 = self.hash_table[hashed_0] # [..., num_levels, features_per_level]\n f_1 = self.hash_table[hashed_1]\n f_2 = self.hash_table[hashed_2]\n f_3 = self.hash_table[hashed_3]\n f_4 = self.hash_table[hashed_4]\n f_5 = self.hash_table[hashed_5]\n f_6 = self.hash_table[hashed_6]\n f_7 = self.hash_table[hashed_7]\n\n f_03 = f_0 * offset[..., 0:1] + f_3 * (1 - offset[..., 0:1])\n f_12 = f_1 * offset[..., 0:1] + f_2 * (1 - offset[..., 0:1])\n f_56 = f_5 * offset[..., 0:1] + f_6 * (1 - offset[..., 0:1])\n f_47 = f_4 * offset[..., 0:1] + f_7 * (1 - offset[..., 0:1])\n\n f0312 = f_03 * offset[..., 1:2] + f_12 * (1 - offset[..., 1:2])\n f4756 = f_47 * offset[..., 1:2] + f_56 * (1 - offset[..., 1:2])\n\n encoded_value = f0312 * offset[..., 2:3] + f4756 * (\n 1 - offset[..., 2:3]\n ) # [..., num_levels, features_per_level]\n\n return torch.flatten(encoded_value, start_dim=-2, end_dim=-1) # [..., num_levels * features_per_level]\n\n def forward(self, in_tensor: TensorType[\"bs\":..., \"input_dim\"]) -> TensorType[\"bs\":..., \"output_dim\"]:\n return self.pytorch_fwd(in_tensor)\n\n def get_total_variation_loss(self):\n \"\"\"Compute the total variation loss for the feature volume.\"\"\"\n feature_volume = self.hash_table.reshape(\n self.num_levels,\n self.periodic_volume_resolution,\n self.periodic_volume_resolution,\n self.periodic_volume_resolution,\n self.features_per_level,\n )\n diffx = feature_volume[:, 1:, :, :, :] - feature_volume[:, :-1, :, :, :]\n diffy = feature_volume[:, :, 1:, :, :] - feature_volume[:, :, :-1, :, :]\n diffz = feature_volume[:, :, :, 1:, :] - feature_volume[:, :, :, :-1, :]\n\n # TODO how to sum here or should we use mask?\n resx = diffx.abs().mean(dim=(1, 2, 3, 4))\n resy = diffy.abs().mean(dim=(1, 2, 3, 4))\n resz = diffz.abs().mean(dim=(1, 2, 3, 4))\n\n return ((resx + resy + resz) * self.per_level_weights).mean()\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/field_components/field_heads.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nCollection of render heads\n\"\"\"\nfrom enum import Enum\nfrom typing import Callable, Optional, Union\n\nimport torch\nfrom torch import nn\nfrom torchtyping import TensorType\n\nfrom nerfstudio.field_components.base_field_component import FieldComponent\n\n\nclass FieldHeadNames(Enum):\n \"\"\"Possible field outputs\"\"\"\n\n RGB = \"rgb\"\n SH = \"sh\"\n DENSITY = \"density\"\n NORMALS = \"normals\"\n PRED_NORMALS = \"pred_normals\"\n UNCERTAINTY = \"uncertainty\"\n TRANSIENT_RGB = \"transient_rgb\"\n TRANSIENT_DENSITY = \"transient_density\"\n SEMANTICS = \"semantics\"\n NORMAL = \"normal\"\n SDF = \"sdf\"\n ALPHA = \"alpha\"\n GRADIENT = \"gradient\"\n OCCUPANCY = \"occupancy\"\n FEATURE = \"feature\"\n FG_SEG = \"fg_segmentation\"\n\n\nclass FieldHead(FieldComponent):\n \"\"\"Base field output\n\n Args:\n out_dim: output dimension for renderer\n field_head_name: Field type\n in_dim: input dimension. If not defined in constructor, it must be set later.\n activation: output head activation\n \"\"\"\n\n def __init__(\n self,\n out_dim: int,\n field_head_name: FieldHeadNames,\n in_dim: Optional[int] = None,\n activation: Optional[Union[nn.Module, Callable]] = None,\n ) -> None:\n\n super().__init__()\n self.out_dim = out_dim\n self.activation = activation\n self.field_head_name = field_head_name\n self.net = None\n if in_dim is not None:\n self.in_dim = in_dim\n self._construct_net()\n\n def set_in_dim(self, in_dim: int) -> None:\n \"\"\"Set input dimension of Field Head\"\"\"\n self.in_dim = in_dim\n self._construct_net()\n\n def _construct_net(self):\n self.net = nn.Linear(self.in_dim, self.out_dim)\n\n def forward(self, in_tensor: TensorType[\"bs\":..., \"in_dim\"]) -> TensorType[\"bs\":..., \"out_dim\"]:\n \"\"\"Process network output for renderer\n\n Args:\n in_tensor: Network input\n\n Returns:\n Render head output\n \"\"\"\n if not self.net:\n raise SystemError(\"in_dim not set. Must be provided to construtor, or set_in_dim() should be called.\")\n out_tensor = self.net(in_tensor)\n if self.activation:\n out_tensor = self.activation(out_tensor)\n return out_tensor\n\n\nclass DensityFieldHead(FieldHead):\n \"\"\"Density output\n\n Args:\n in_dim: input dimension. If not defined in constructor, it must be set later.\n activation: output head activation\n \"\"\"\n\n def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Softplus()) -> None:\n super().__init__(in_dim=in_dim, out_dim=1, field_head_name=FieldHeadNames.DENSITY, activation=activation)\n\n\nclass RGBFieldHead(FieldHead):\n \"\"\"RGB output\n\n Args:\n in_dim: input dimension. If not defined in constructor, it must be set later.\n activation: output head activation\n \"\"\"\n\n def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Sigmoid()) -> None:\n super().__init__(in_dim=in_dim, out_dim=3, field_head_name=FieldHeadNames.RGB, activation=activation)\n\n\nclass SHFieldHead(FieldHead):\n \"\"\"Spherical harmonics output\n\n Args:\n in_dim: input dimension. If not defined in constructor, it must be set later.\n levels: Number of spherical harmonics layers.\n channels: Number of channels. Defaults to 3 (ie RGB).\n activation: Output activation.\n \"\"\"\n\n def __init__(\n self, in_dim: Optional[int] = None, levels: int = 3, channels: int = 3, activation: Optional[nn.Module] = None\n ) -> None:\n\n out_dim = channels * levels**2\n super().__init__(in_dim=in_dim, out_dim=out_dim, field_head_name=FieldHeadNames.SH, activation=activation)\n\n\nclass UncertaintyFieldHead(FieldHead):\n \"\"\"Uncertainty output\n\n Args:\n in_dim: input dimension. If not defined in constructor, it must be set later.\n activation: output head activation\n \"\"\"\n\n def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Softplus()) -> None:\n super().__init__(in_dim=in_dim, out_dim=1, field_head_name=FieldHeadNames.UNCERTAINTY, activation=activation)\n\n\nclass TransientRGBFieldHead(FieldHead):\n \"\"\"Transient RGB output\n\n Args:\n in_dim: input dimension. If not defined in constructor, it must be set later.\n activation: output head activation\n \"\"\"\n\n def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Sigmoid()) -> None:\n super().__init__(in_dim=in_dim, out_dim=3, field_head_name=FieldHeadNames.TRANSIENT_RGB, activation=activation)\n\n\nclass TransientDensityFieldHead(FieldHead):\n \"\"\"Transient density output\n\n Args:\n in_dim: input dimension. If not defined in constructor, it must be set later.\n activation: output head activation\n \"\"\"\n\n def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Softplus()) -> None:\n super().__init__(\n in_dim=in_dim, out_dim=1, field_head_name=FieldHeadNames.TRANSIENT_DENSITY, activation=activation\n )\n\n\nclass SemanticFieldHead(FieldHead):\n \"\"\"Semantic output\n\n Args:\n num_classes: Number of semantic classes\n in_dim: input dimension. If not defined in constructor, it must be set later.\n activation: output head activation\n \"\"\"\n\n def __init__(self, num_classes: int, in_dim: Optional[int] = None) -> None:\n super().__init__(in_dim=in_dim, out_dim=num_classes, field_head_name=FieldHeadNames.SEMANTICS, activation=None)\n\n\nclass PredNormalsFieldHead(FieldHead):\n \"\"\"Predicted normals output.\n\n Args:\n in_dim: input dimension. If not defined in constructor, it must be set later.\n activation: output head activation\n \"\"\"\n\n def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Tanh()) -> None:\n super().__init__(in_dim=in_dim, out_dim=3, field_head_name=FieldHeadNames.PRED_NORMALS, activation=activation)\n\n def forward(self, in_tensor: TensorType[\"bs\":..., \"in_dim\"]) -> TensorType[\"bs\":..., \"out_dim\"]:\n \"\"\"Needed to normalize the output into valid normals.\"\"\"\n out_tensor = super().forward(in_tensor)\n out_tensor = torch.nn.functional.normalize(out_tensor, dim=-1)\n return out_tensor\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/field_components/mlp.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nMulti Layer Perceptron\n\"\"\"\nfrom typing import Optional, Set, Tuple\n\nimport torch\nfrom torch import nn\nfrom torchtyping import TensorType\n\nfrom nerfstudio.field_components.base_field_component import FieldComponent\n\n\nclass MLP(FieldComponent):\n \"\"\"Multilayer perceptron\n\n Args:\n in_dim: Input layer dimension\n num_layers: Number of network layers\n layer_width: Width of each MLP layer\n out_dim: Ouput layer dimension. Uses layer_width if None.\n activation: intermediate layer activation function.\n out_activation: output activation function.\n \"\"\"\n\n def __init__(\n self,\n in_dim: int,\n num_layers: int,\n layer_width: int,\n out_dim: Optional[int] = None,\n skip_connections: Optional[Tuple[int]] = None,\n activation: Optional[nn.Module] = nn.ReLU(),\n out_activation: Optional[nn.Module] = None,\n ) -> None:\n\n super().__init__()\n self.in_dim = in_dim\n assert self.in_dim > 0\n self.out_dim = out_dim if out_dim is not None else layer_width\n self.num_layers = num_layers\n self.layer_width = layer_width\n self.skip_connections = skip_connections\n self._skip_connections: Set[int] = set(skip_connections) if skip_connections else set()\n self.activation = activation\n self.out_activation = out_activation\n self.net = None\n self.build_nn_modules()\n\n def build_nn_modules(self) -> None:\n \"\"\"Initialize multi-layer perceptron.\"\"\"\n layers = []\n if self.num_layers == 1:\n layers.append(nn.Linear(self.in_dim, self.out_dim))\n else:\n for i in range(self.num_layers - 1):\n if i == 0:\n assert i not in self._skip_connections, \"Skip connection at layer 0 doesn't make sense.\"\n layers.append(nn.Linear(self.in_dim, self.layer_width))\n elif i in self._skip_connections:\n layers.append(nn.Linear(self.layer_width + self.in_dim, self.layer_width))\n else:\n layers.append(nn.Linear(self.layer_width, self.layer_width))\n layers.append(nn.Linear(self.layer_width, self.out_dim))\n self.layers = nn.ModuleList(layers)\n\n def forward(self, in_tensor: TensorType[\"bs\":..., \"in_dim\"]) -> TensorType[\"bs\":..., \"out_dim\"]:\n \"\"\"Process input with a multilayer perceptron.\n\n Args:\n in_tensor: Network input\n\n Returns:\n MLP network output\n \"\"\"\n x = in_tensor\n for i, layer in enumerate(self.layers):\n # as checked in `build_nn_modules`, 0 should not be in `_skip_connections`\n if i in self._skip_connections:\n x = torch.cat([in_tensor, x], -1)\n x = layer(x)\n if self.activation is not None and i < len(self.layers) - 1:\n x = self.activation(x)\n if self.out_activation is not None:\n x = self.out_activation(x)\n return x\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/field_components/progressive_encoding.py",
"content": "\nfrom rich.console import Console\n\nimport torch\nimport tinycudann as tcnn\nfrom torch import nn\n\nCONSOLE = Console(width=120)\n\n\nclass ProgressiveEncoding(nn.Module):\n \"\"\"Progressive multi-resolution encoding proposed in NeuS2.\"\"\"\n \n def __init__(\n self,\n encoding: tcnn.Encoding,\n n_levels: int = 16, # number of resolution levels\n n_levels_init: int = 1, # number of resolution levels to use in the beginning\n n_features_per_level: int = 2,\n n_scheduler_steps: int = 10000 # total number of steps of the scheduler\n ) -> None:\n super().__init__()\n self.encoding = encoding\n self.n_levels = n_levels\n if n_scheduler_steps == -1: # enable progressive encoding but set all weights to 1\n self.n_levels_init = self.n_levels # useful for testing / using pretrained models\n else:\n self.n_levels_init = n_levels_init\n self.n_features_per_level = n_features_per_level\n self.n_output_dims = self.encoding.n_output_dims\n assert self.encoding.n_output_dims == n_levels * n_features_per_level\n \n self.n_scheduler_steps = n_scheduler_steps\n self.step_size = n_scheduler_steps // n_levels\n self.cur_iter = 0 # TODO: make ProgressiveEncoding stateless and pass in cur_iter as an argument\n self.cur_level: int = -1\n # self.register_buffer(\"weights\", self._get_init_weights(), persistent=False) # persistent=True / nn.Parameter\n self.weights = nn.Parameter(self._get_init_weights(), requires_grad=False)\n CONSOLE.print(f'[bold green]Using progressive hash encoding training (n_iters={n_scheduler_steps})')\n \n def _get_init_weights(self):\n weights = torch.zeros(self.n_levels, self.n_features_per_level, dtype=torch.float32)\n weights[:self.n_levels_init, :] = 1.0\n return weights\n \n def forward(self, positions):\n if self.n_scheduler_steps == -1:\n return self.encoding(positions)\n \n cur_level = min(int(self.cur_iter // self.step_size) + self.n_levels_init - 1, self.n_levels - 1)\n if self.cur_level != cur_level:\n self.cur_level = cur_level\n self.weights[cur_level, :] = 1.0\n encoding = self.encoding(positions) # (..., c)\n encoding = encoding * self.weights.view(-1)\n \n if self.training:\n self.cur_iter += 1\n return encoding\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/field_components/spatial_distortions.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Space distortions.\"\"\"\n\nfrom typing import Optional, Union\n\nimport torch\nfrom functorch import jacrev, vmap\nfrom torch import nn\nfrom torchtyping import TensorType\n\nfrom nerfstudio.utils.math import Gaussians\n\n\nclass SpatialDistortion(nn.Module):\n \"\"\"Apply spatial distortions\"\"\"\n\n def forward(\n self, positions: Union[TensorType[\"bs\":..., 3], Gaussians]\n ) -> Union[TensorType[\"bs\":..., 3], Gaussians]:\n \"\"\"\n Args:\n positions: Sample to distort\n\n Returns:\n Union: distorted sample\n \"\"\"\n\n\nclass SceneContraction(SpatialDistortion):\n \"\"\"Contract unbounded space using the contraction was proposed in MipNeRF-360.\n We use the following contraction equation:\n\n .. math::\n\n f(x) = \\\\begin{cases}\n x & ||x|| \\\\leq 1 \\\\\\\\\n (2 - \\\\frac{1}{||x||})(\\\\frac{x}{||x||}) & ||x|| > 1\n \\\\end{cases}\n\n If the order is not specified, we use the Frobenius norm, this will contract the space to a sphere of\n radius 1. If the order is L_inf (order=float(\"inf\")), we will contract the space to a cube of side length 2.\n If using voxel based encodings such as the Hash encoder, we recommend using the L_inf norm.\n\n Args:\n order: Order of the norm. Default to the Frobenius norm. Must be set to None for Gaussians.\n scale_factor: Scale positions by this factor before applying contraction, default to 1.0.\n This is useful when the [-1, 1] region (no contraction) is too small and only covers a small portion of the scene.\n \"\"\"\n\n def __init__(\n self,\n order: Optional[Union[float, int]] = None,\n scale_factor: float = 1.0\n ) -> None:\n super().__init__()\n self.order = order\n self.scale_factor = scale_factor\n\n def forward(self, positions):\n def contract(x):\n x = x * self.scale_factor\n \n mag = torch.linalg.norm(x, ord=self.order, dim=-1)\n mask = mag >= 1\n x_new = x.clone()\n x_new[mask] = (2 - (1 / mag[mask][..., None])) * (x[mask] / mag[mask][..., None])\n\n return x_new\n\n if isinstance(positions, Gaussians):\n means = contract(positions.mean.clone())\n\n if self.scale_factor != 1.0:\n raise NotImplementedError('Scaling of gaussian covariances is not implemented.')\n \n # FIXME: covariances are not handled for now!\n contract = lambda x: (2 - (1 / torch.linalg.norm(x, ord=self.order, dim=-1, keepdim=True))) * (\n x / torch.linalg.norm(x, ord=self.order, dim=-1, keepdim=True)\n )\n jc_means = vmap(jacrev(contract))(positions.mean.view(-1, positions.mean.shape[-1]))\n jc_means = jc_means.view(list(positions.mean.shape) + [positions.mean.shape[-1]])\n\n # Only update covariances on positions outside the unit sphere\n mag = positions.mean.norm(dim=-1)\n mask = mag >= 1\n cov = positions.cov.clone()\n cov[mask] = jc_means[mask] @ positions.cov[mask] @ torch.transpose(jc_means[mask], -2, -1)\n\n return Gaussians(mean=means, cov=cov)\n\n return contract(positions)\n\n\nclass ForegroundAwareSceneContraction(SpatialDistortion):\n \"\"\"Contract unbounded space using a modified version of the contraction proposed in MipNeRF-360.\n \n 1. Keep the \"||x|| <= 1 / x in aabb\" region unchanged, i.e., to [-1, 1] / [aabb[0], aabb[1]], which should be modeled by the foreground field.\n 2. Linearly contract the 1 < ||x|| <= alpha / x not in aabb but x <= alpha region to [[-1, -1/alpha], [1/alpha, 1]] (alpha>=1), which should be modeled with a separate background field.\n 3. Non-linearly contract the ||x|| > alpha region to [[-2, -1], [1, 2]], which should be modeled with the background field.\n \n # FIXME: covariances are not handled for now!\n \"\"\"\n def __init__(\n self,\n order: Optional[Union[float, int]] = None,\n aabb: Optional[torch.Tensor] = None, # (2, 3)\n alpha: float = 5.0,\n scale_factor: float = 1.0\n ) -> None:\n super().__init__()\n self.order = order\n self.alpha = alpha\n self.scale_factor = scale_factor\n self.aabb = nn.Parameter(aabb, requires_grad=False) if aabb is not None else None\n if aabb is not None:\n assert aabb.abs().max() <= 1.0, f\"Invalid aabb: aabb={aabb}\"\n assert order == float(\"inf\"), f\"Must use inf norm with aabb: order={order}\"\n assert self.scale_factor == 1.0\n assert alpha >= 1.0, f\"Invalid hparam setting: alpha={alpha}\"\n\n def forward(self, positions):\n def _contract(x):\n mag = torch.linalg.norm(x, ord=self.order, dim=-1)\n mask = mag >= 1\n x_new = x.clone()\n x_new[mask] = (2 - (1 / mag[mask][..., None])) * (x[mask] / mag[mask][..., None])\n\n return x_new\n\n def contract(x):\n x = x * self.scale_factor\n if self.aabb is None:\n mag = torch.linalg.norm(x, ord=self.order, dim=-1)\n mask = mag >= 1\n else:\n in_aabb_mask = ((x >= self.aabb[0]) & (x <= self.aabb[1])).all(-1) # whether including = ?\n mask = ~in_aabb_mask\n x_new = x.clone()\n x_new[mask] = x[mask] / self.alpha\n return _contract(x_new)\n\n if isinstance(positions, Gaussians):\n means = contract(positions.mean.clone())\n # FIXME: covariances are not handled for now!\n if self.scale_factor != 1.0:\n raise NotImplementedError('Scaling of gaussian covariances is not implemented.')\n contract = lambda x: (2 - (1 / torch.linalg.norm(x, ord=self.order, dim=-1, keepdim=True))) * (\n x / torch.linalg.norm(x, ord=self.order, dim=-1, keepdim=True)\n )\n jc_means = vmap(jacrev(contract))(positions.mean.view(-1, positions.mean.shape[-1]))\n jc_means = jc_means.view(list(positions.mean.shape) + [positions.mean.shape[-1]])\n\n # Only update covariances on positions outside the unit sphere\n mag = positions.mean.norm(dim=-1)\n mask = mag >= 1\n cov = positions.cov.clone()\n cov[mask] = jc_means[mask] @ positions.cov[mask] @ torch.transpose(jc_means[mask], -2, -1)\n\n return Gaussians(mean=means, cov=cov)\n\n return contract(positions)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/field_components/temporal_distortions.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Space distortions which occur as a function of time.\"\"\"\n\nfrom enum import Enum\nfrom typing import Any, Dict, Optional, Tuple\n\nimport torch\nfrom torch import nn\nfrom torchtyping import TensorType\n\nfrom nerfstudio.field_components.encodings import Encoding, NeRFEncoding\nfrom nerfstudio.field_components.mlp import MLP\n\n\nclass TemporalDistortion(nn.Module):\n \"\"\"Apply spatial distortions as a function of time\"\"\"\n\n def forward(self, positions: TensorType[\"bs\":..., 3], times: Optional[TensorType[1]]) -> TensorType[\"bs\":..., 3]:\n \"\"\"\n Args:\n positions: Samples to translate as a function of time\n times: times for each sample\n\n Returns:\n Translated positions.\n \"\"\"\n\n\nclass TemporalDistortionKind(Enum):\n \"\"\"Possible temporal distortion names\"\"\"\n\n DNERF = \"dnerf\"\n\n def to_temporal_distortion(self, config: Dict[str, Any]) -> TemporalDistortion:\n \"\"\"Converts this kind to a temporal distortion\"\"\"\n if self == TemporalDistortionKind.DNERF:\n return DNeRFDistortion(**config)\n raise NotImplementedError(f\"Unknown temporal distortion kind {self}\")\n\n\nclass DNeRFDistortion(TemporalDistortion):\n \"\"\"Optimizable temporal deformation using an MLP.\n Args:\n position_encoding: An encoding for the XYZ of distortion\n temporal_encoding: An encoding for the time of distortion\n mlp_num_layers: Number of layers in distortion MLP\n mlp_layer_width: Size of hidden layer for the MLP\n skip_connections: Number of layers for skip connections in the MLP\n \"\"\"\n\n def __init__(\n self,\n position_encoding: Encoding = NeRFEncoding(\n in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0, include_input=True\n ),\n temporal_encoding: Encoding = NeRFEncoding(\n in_dim=1, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0, include_input=True\n ),\n mlp_num_layers: int = 4,\n mlp_layer_width: int = 256,\n skip_connections: Tuple[int] = (4,),\n ) -> None:\n super().__init__()\n self.position_encoding = position_encoding\n self.temporal_encoding = temporal_encoding\n self.mlp_deform = MLP(\n in_dim=self.position_encoding.get_out_dim() + self.temporal_encoding.get_out_dim(),\n out_dim=3,\n num_layers=mlp_num_layers,\n layer_width=mlp_layer_width,\n skip_connections=skip_connections,\n )\n\n def forward(self, positions, times=None):\n if times is None:\n return None\n p = self.position_encoding(positions)\n t = self.temporal_encoding(times)\n return self.mlp_deform(torch.cat([p, t], dim=-1))\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/fields/__init__.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/fields/base_field.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nBase class for the graphs.\n\"\"\"\n\nfrom abc import abstractmethod\nfrom typing import Dict, Optional, Tuple, Type\nfrom dataclasses import dataclass, field\n\nimport torch\nfrom torch import nn\nfrom torchtyping import TensorType\n\nfrom nerfstudio.configs.base_config import InstantiateConfig\nfrom nerfstudio.cameras.rays import Frustums, RaySamples\nfrom nerfstudio.field_components.field_heads import FieldHeadNames\n\n# Field related configs\n@dataclass\nclass FieldConfig(InstantiateConfig):\n \"\"\"Configuration for model instantiation\"\"\"\n\n _target: Type = field(default_factory=lambda: Field)\n \"\"\"target class to instantiate\"\"\"\n\n\nclass Field(nn.Module):\n \"\"\"Base class for fields.\"\"\"\n\n def __init__(self) -> None:\n super().__init__()\n self._sample_locations = None\n self._density_before_activation = None\n\n def density_fn(self, positions: TensorType[\"bs\":..., 3]) -> TensorType[\"bs\":..., 1]:\n \"\"\"Returns only the density. Used primarily with the density grid.\n\n Args:\n positions: the origin of the samples/frustums\n \"\"\"\n # Need to figure out a better way to describe positions with a ray.\n ray_samples = RaySamples(\n frustums=Frustums(\n origins=positions,\n directions=torch.ones_like(positions),\n starts=torch.zeros_like(positions[..., :1]),\n ends=torch.zeros_like(positions[..., :1]),\n pixel_area=torch.ones_like(positions[..., :1]),\n )\n )\n density, _ = self.get_density(ray_samples)\n return density\n\n @abstractmethod\n def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType[..., 1], TensorType[..., \"num_features\"]]:\n \"\"\"Computes and returns the densities. Returns a tensor of densities and a tensor of features.\n\n Args:\n ray_samples: Samples locations to compute density.\n \"\"\"\n\n def get_normals(self) -> TensorType[..., 3]:\n \"\"\"Computes and returns a tensor of normals.\n\n Args:\n density: Tensor of densities.\n \"\"\"\n assert self._sample_locations is not None, \"Sample locations must be set before calling get_normals.\"\n assert self._density_before_activation is not None, \"Density must be set before calling get_normals.\"\n assert (\n self._sample_locations.shape[:-1] == self._density_before_activation.shape[:-1]\n ), \"Sample locations and density must have the same shape besides the last dimension.\"\n\n self._density_before_activation.backward(\n gradient=torch.ones_like(self._density_before_activation), inputs=self._sample_locations, retain_graph=True\n )\n normals = -torch.nn.functional.normalize(self._sample_locations.grad, dim=-1)\n return normals\n\n @abstractmethod\n def get_outputs(\n self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None\n ) -> Dict[FieldHeadNames, TensorType]:\n \"\"\"Computes and returns the colors. Returns output field values.\n\n Args:\n ray_samples: Samples locations to compute outputs.\n density_embedding: Density embeddings to condition on.\n \"\"\"\n\n def forward(self, ray_samples: RaySamples, compute_normals: bool = False):\n \"\"\"Evaluates the field at points along the ray.\n\n Args:\n ray_samples: Samples to evaluate field on.\n \"\"\"\n if compute_normals:\n with torch.enable_grad():\n density, density_embedding = self.get_density(ray_samples)\n else:\n density, density_embedding = self.get_density(ray_samples)\n\n field_outputs = self.get_outputs(ray_samples, density_embedding=density_embedding)\n field_outputs[FieldHeadNames.DENSITY] = density # type: ignore\n\n if compute_normals:\n with torch.enable_grad():\n normals = self.get_normals()\n field_outputs[FieldHeadNames.NORMALS] = normals # type: ignore\n return field_outputs\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/fields/density_fields.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nProposal network field.\n\"\"\"\n\n\nfrom typing import Optional\nfrom typing_extensions import Literal\n\nimport numpy as np\nimport torch\nfrom torch.nn.parameter import Parameter\nfrom torchtyping import TensorType\n\nfrom nerfstudio.cameras.rays import RaySamples\nfrom nerfstudio.data.scene_box import SceneBox\nfrom nerfstudio.field_components.activations import trunc_exp\nfrom nerfstudio.field_components.spatial_distortions import SpatialDistortion\nfrom nerfstudio.fields.base_field import Field\n\ntry:\n import tinycudann as tcnn\nexcept ImportError:\n # tinycudann module doesn't exist\n pass\n\n\nclass HashMLPDensityField(Field):\n \"\"\"A lightweight density field module.\n\n Args:\n aabb: parameters of scene aabb bounds\n num_layers: number of hidden layers\n hidden_dim: dimension of hidden layers\n spatial_distortion: spatial distortion module\n use_linear: whether to skip the MLP and use a single linear layer instead\n spatial_normalization_region: if 'full', normalize coorinates from [-2, 2] to [0, 1];\n if 'fg', normalize coordinates from [-1, 1] to [0, 1]\n \"\"\"\n\n def __init__(\n self,\n aabb,\n num_layers: int = 2,\n hidden_dim: int = 64,\n spatial_distortion: Optional[SpatialDistortion] = None,\n use_linear=False,\n num_levels=8,\n max_res=1024,\n base_res=16,\n log2_hashmap_size=18,\n features_per_level=2,\n spatial_normalization_region: Literal['full', 'fg'] = 'full'\n ) -> None:\n super().__init__()\n self.aabb = Parameter(aabb, requires_grad=False)\n self.spatial_distortion = spatial_distortion\n self.use_linear = use_linear\n self.spatial_normalization_region = spatial_normalization_region\n growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1))\n\n config = {\n \"encoding\": {\n \"otype\": \"HashGrid\",\n \"n_levels\": num_levels,\n \"n_features_per_level\": features_per_level,\n \"log2_hashmap_size\": log2_hashmap_size,\n \"base_resolution\": base_res,\n \"per_level_scale\": growth_factor,\n },\n \"network\": {\n \"otype\": \"FullyFusedMLP\",\n \"activation\": \"ReLU\",\n \"output_activation\": \"None\",\n \"n_neurons\": hidden_dim,\n \"n_hidden_layers\": num_layers - 1,\n },\n }\n\n if not self.use_linear:\n self.mlp_base = tcnn.NetworkWithInputEncoding(\n n_input_dims=3,\n n_output_dims=1,\n encoding_config=config[\"encoding\"],\n network_config=config[\"network\"],\n )\n else:\n self.encoding = tcnn.Encoding(n_input_dims=3, encoding_config=config[\"encoding\"])\n self.linear = torch.nn.Linear(self.encoding.n_output_dims, 1)\n\n def _get_normalized_positions(self, positions):\n if self.spatial_normalization_region == 'full':\n return (positions + 2.0) / 4.0\n elif self.spatial_normalization_region == 'fg':\n return (positions + 1.0) / 2.0\n else:\n raise ValueError(f\"Unknown spatial normalization region: {self.spatial_normalization_region}\")\n \n def get_density(self, ray_samples: RaySamples):\n if self.spatial_distortion is not None:\n positions = self.spatial_distortion(ray_samples.frustums.get_positions())\n positions = self._get_normalized_positions(positions)\n else:\n positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)\n positions_flat = positions.view(-1, 3)\n if not self.use_linear:\n density_before_activation = (\n self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1).to(positions)\n )\n else:\n x = self.encoding(positions_flat).to(positions)\n density_before_activation = self.linear(x).view(*ray_samples.frustums.shape, -1)\n\n # Rectifying the density with an exponential is much more stable than a ReLU or\n # softplus, because it enables high post-activation (float32) density outputs\n # from smaller internal (float16) parameters.\n density = trunc_exp(density_before_activation)\n \n return density, None\n\n def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None):\n return {}\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/fields/feature_field.py",
"content": "from dataclasses import dataclass, field\nfrom typing import Dict, Optional, Tuple, Type\n\nfrom rich.console import Console\nfrom torchtyping import TensorType\nfrom typing_extensions import Literal\n\ntry:\n import faiss\nexcept ImportError as err:\n pass\nimport numpy as np\nimport torch\nfrom torch.nn import Parameter\nfrom torch.nn import functional as F\n\nfrom nerfstudio.cameras.rays import RaySamples\nfrom nerfstudio.field_components.field_heads import FieldHeadNames\nfrom nerfstudio.field_components.spatial_distortions import SpatialDistortion\nfrom nerfstudio.fields.base_field import Field, FieldConfig\nfrom nerfstudio.utils.pointclouds import BasicPointClouds\n\ntry:\n import tinycudann as tcnn\nexcept ImportError:\n # tinycudann module doesn't exist\n pass\n\nCONSOLE = Console(width=128)\n\n\n@dataclass\nclass FeatureFieldConfig(FieldConfig):\n \"\"\"Feature Field Config\"\"\"\n\n _target: Type = field(default_factory=lambda: FeatureField)\n\n nc_position: int = 3\n \"\"\"Number of channels for positional inputs\"\"\"\n nc_feature: int = 384\n nc_hidden: int = 128\n n_hidden_layers: int = 2\n use_viewdirs: bool = False\n use_grid_feature: bool = True\n \"\"\"Whether to use multi-resolution feature grids (only support hash encoding for now)\"\"\"\n hash_grid_max_res: int = 128\n \"\"\"Maximum resolution of the feature grid (use a smaller default value since ViT features is of lower frequency)\"\"\"\n hash_grid_log2_hashmap_size: int = 19\n # hash_grid_precision: Literal[\"float32\", \"float16\"] = \"float32\"\n use_position_encoding: bool = True\n \"\"\"Whether to include positional encodings beyond the grid features (for better smoothness)\"\"\"\n use_fused_encoding_mlp: bool = True\n \"\"\"Whether to use TCNN fused encoding & mlp\"\"\"\n\n\nclass FeatureField(Field):\n \"\"\"A indivisual feature field without feature sharing with the main field.\"\"\"\n\n config: FeatureFieldConfig\n\n def __init__(self, config: FeatureFieldConfig, spatial_distortion: Optional[SpatialDistortion] = None) -> None:\n super().__init__()\n self.config = config\n\n self.spatial_distortion = spatial_distortion\n self.use_grid_feature = self.config.use_grid_feature\n self.use_viewdirs = self.config.use_viewdirs\n\n # feature field hash grid\n if self.config.use_viewdirs:\n raise NotImplementedError\n\n # if self.config.use_position_encoding:\n # raise NotImplementedError # tcnn nested positional encoding not working\n\n if self.config.use_grid_feature:\n self.build_ngp()\n else:\n self.build()\n\n def build_ngp(self):\n L, F, N_min = 16, 2, 16\n log2_T = self.config.hash_grid_log2_hashmap_size\n N_max = self.config.hash_grid_max_res\n b = np.exp(np.log(N_max / N_min) / (L - 1)) # feature is of lower-frequency -> N_max=128\n nc_pos = self.config.nc_position\n\n hash_encoding_config = {\n \"otype\": \"Grid\",\n \"type\": \"Hash\",\n \"n_levels\": L,\n \"n_features_per_level\": F,\n \"log2_hashmap_size\": log2_T,\n \"base_resolution\": N_min,\n \"per_level_scale\": b,\n \"n_dims_to_encode\": nc_pos,\n \"interpolation\": \"Linear\", # TODO: support \"Smoothstep\"\n }\n\n if self.config.use_fused_encoding_mlp: # tcnn encoding + mlp\n encoding_config = {\n \"otype\": \"Composite\",\n # Hash for the first nc_pos dims; Frequency for the later nc_pos dims\n \"nested\": [\n hash_encoding_config,\n {\n \"otype\": \"Frequency\",\n \"n_frequencies\": 8,\n \"n_dims_to_encode\": nc_pos,\n },\n ],\n }\n\n # when tinycudann is built w/ float16 support, it will use float16 automatically\n self.feature_encoder = tcnn.NetworkWithInputEncoding(\n n_input_dims=nc_pos * 2\n if self.config.use_position_encoding\n else nc_pos, # 2 * nc_pos for smooth encoding\n n_output_dims=self.config.nc_feature,\n encoding_config=encoding_config,\n network_config={\n \"otype\": \"FullyFusedMLP\",\n \"activation\": \"ReLU\",\n \"output_activation\": \"None\",\n \"n_neurons\": self.config.nc_hidden,\n \"n_hidden_layers\": self.config.n_hidden_layers,\n },\n # dtype=torch.float32 if self.config.hash_grid_precision == \"float32\" else torch.float16\n )\n else: # tcnn encoding + torch MLP\n raise NotImplementedError\n # self.pos_encoder = tcnn.Encoding(n_input_dims=nc_pos,\n # encoding_config=hash_encoding_config)\n # pos_dim = self.pos_encoder.n_output_dims\n\n # self.mlp = nn.Sequential(\n # fc_block(pos_dim, self.config.nc_hidden),\n # *[fc_block(self.config.nc_hidden, self.config.nc_hidden) for _ in range(self.config.n_hidden_layers)],\n # nn.Linear(self.config.nc_hidden, self.config.nc_feature) # TODO: support output activation\n # )\n\n def build(self):\n raise NotImplementedError\n\n def query_features(\n self, positions: TensorType[\"n_rays\", \"d_pos\"], viewdirs: Optional[TensorType[\"n_rays\", \"d_dir\"]] = None\n ) -> TensorType[\"n_rays\", \"d_feature\"]:\n \"\"\"Qeury the feature field\"\"\"\n x = self._get_input_positions(positions)\n\n if self.config.use_position_encoding:\n # FIXME: raise runtime_error after loading checkpoints to other fields.\n features = self.feature_encoder(torch.cat([x, x], dim=-1))\n else:\n features = self.feature_encoder(x)\n return features\n\n def get_outputs(self, ray_samples: RaySamples) -> Dict[str, torch.Tensor]:\n outputs = {}\n\n inputs = ray_samples.frustums.get_start_positions()\n inputs = inputs.view(-1, 3)\n\n directions = ray_samples.frustums.directions\n directions_flat = directions.reshape(-1, 3)\n features = self.query_features(inputs, directions_flat) # NOTE: feautures.dtype = torch.float16\n\n features = features.view(*ray_samples.frustums.directions.shape[:-1], -1)\n\n outputs = {FieldHeadNames.FEATURE: features}\n return outputs\n\n def forward(self, ray_samples: RaySamples):\n \"\"\"Evaluates the field at the given ray samples\"\"\"\n field_outputs = self.get_outputs(ray_samples)\n return field_outputs\n\n def _get_input_positions(self, positions: torch.Tensor):\n \"\"\"compute input positions (for positional encoding / feature encoding)\n from raw ray sample positions.\n\n NOTE: assume the whole scene is represented by a single feature field\n \"\"\"\n positions = self.spatial_distortion(positions) # [[-2, -1], [-1, 1], [1, 2]]\n positions = (positions + 2.0) / 4.0 # Full normalization: [-2, 2] -> [0, 1]\n # FIXME: handle positions out of [0, 1] after spatial_distortion & normalization\n return positions\n\n\n@dataclass\nclass FeatureSegFieldConfig(FieldConfig):\n \"\"\"FeatureSegField Config\"\"\"\n\n _target: Type = field(default_factory=lambda: FeatureSegField)\n\n knn: int = 10 # TOOD: seg fg & bg knn separately?\n \"\"\"number of reference points to retrieve for each point to be segmented\"\"\"\n distance_weighted_average: bool = False\n \"\"\"compute distance-weighted average of cosine similairties of the knn retrieved points\"\"\"\n in_fg_aabb_only: bool = True\n \"\"\"if a point is not in the fg aabb, it is directly considered as background\"\"\"\n contrast_fg_bg: bool = True\n \"\"\"whether to derive segmentation field by contrast between foreground and background points\"\"\"\n separate_bg_plane: bool = False\n \"\"\"when computing segmentation probs, treat background and plane as two different classes\"\"\"\n\n\nclass FeatureSegField(Field):\n \"\"\"calculate segmentation field from feature field\"\"\"\n\n config: FeatureSegFieldConfig\n\n def __init__(self, config: FeatureFieldConfig, ptcd_data: BasicPointClouds, aabb: TensorType[2, 3]) -> None:\n super().__init__()\n self.config = config\n self.ptcd_data = ptcd_data\n self.aabb = Parameter(aabb, requires_grad=False)\n self.ref_feats_fg = ptcd_data.fg_feats\n self.ref_feats_bg = torch.cat([ptcd_data.bg_feats, ptcd_data.plane_feats], 0)\n self.ref_pts_fg = ptcd_data.fg_pts\n self.ref_pts_bg = torch.cat([ptcd_data.bg_pts, ptcd_data.plane_pts], 0)\n self.fg_nn_index = ptcd_data.fg_nn_index\n self.bg_nn_index = ptcd_data.bg_nn_index\n\n if self.config.separate_bg_plane:\n self.plane_nn_index = ptcd_data.plane_nn_index\n self.ref_feats_bg, self.ref_pts_bg = ptcd_data.bg_feats, ptcd_data.bg_pts\n self.ref_feats_plane, self.ref_pts_plane = ptcd_data.plane_feats, ptcd_data.plane_pts\n\n self._register_all_tensors_as_buffers()\n\n @torch.no_grad()\n def compute_seg_prob(\n self,\n points: Optional[TensorType[\"num_samples\", 3]], # unnormalized, raw positions\n features: TensorType[\"num_samples\", \"nc_features\"],\n ptcd_data: Optional[BasicPointClouds] = None,\n ) -> TensorType[\"num_samples\", 1]:\n if ptcd_data is not None: # Override self.ptcd_data\n raise NotImplementedError\n\n features = features.detach()\n if points is not None:\n points = points.detach()\n if self.config.in_fg_aabb_only:\n assert points is not None\n in_fg_aabb_mask = ((points >= self.aabb[0]) & (points <= self.aabb[1])).all(-1) # (n, )\n # TODO: for pts not in fg aabb, omit the seg_probs computation directly.\n\n feats_l2_normalized = F.normalize(features, p=2, dim=-1)\n\n knn_inds_fg, knn_feats_fg, knn_pts_fg = self.search_knn_fg(feats_l2_normalized)\n knn_cos_sims_avg_fg = self.compute_cos_sim(points, feats_l2_normalized, knn_feats_fg, knn_pts_fg)\n\n if not self.config.contrast_fg_bg:\n seg_probs = knn_cos_sims_avg_fg * 0.5 + 0.5\n else:\n knn_inds_bg, knn_feats_bg, knn_pts_bg = self.search_knn_bg(features)\n knn_cos_sims_avg_bg = self.compute_cos_sim(points, feats_l2_normalized, knn_feats_bg, knn_pts_bg)\n\n # contrastive seg_probs\n if self.config.separate_bg_plane:\n knn_inds_plane, knn_feats_plane, knn_pts_plane = self.search_knn_plane(features)\n knn_cos_sims_avg_plane = self.compute_cos_sim(\n points, feats_l2_normalized, knn_feats_plane, knn_pts_plane\n ) # (n, 1)\n knn_cos_sims_avg_bg = torch.maximum(knn_cos_sims_avg_bg, knn_cos_sims_avg_plane)\n\n # TODO: tune temperature and offset cos_sims before applying softmax (since most cos_sims > -0.1)\n seg_probs = torch.softmax(torch.cat([knn_cos_sims_avg_fg, knn_cos_sims_avg_bg], -1), dim=-1)\n seg_probs = seg_probs[..., [0]]\n\n if self.config.in_fg_aabb_only:\n seg_probs[~in_fg_aabb_mask] = 0.0\n return seg_probs\n\n def search_knn_fg(self, feats: TensorType[\"n\", \"c\"]):\n return self._search_knn(feats, self.ref_feats_fg, self.ref_pts_fg, self.fg_nn_index)\n\n def search_knn_bg(self, feats: TensorType[\"n\", \"c\"]):\n return self._search_knn(feats, self.ref_feats_bg, self.ref_pts_bg, self.bg_nn_index)\n\n def search_knn_plane(self, feats: TensorType[\"n\", \"c\"]):\n return self._search_knn(feats, self.ref_feats_plane, self.ref_pts_plane, self.plane_nn_index)\n\n def _search_knn(\n self,\n feats: TensorType[\"n\", \"c\"],\n ref_feats: TensorType[\"m\", \"c\"],\n ref_pts: TensorType[\"m\", 3],\n nn_index: faiss.Index,\n ) -> Tuple[TensorType[\"n\", \"k\"], TensorType[\"n\", \"k\", \"c\"], TensorType[\"n\", \"k\", 3]]:\n feats = feats.to(torch.float32)\n knn_dists, knn_inds = nn_index.search(feats, self.config.knn) # (n, k)\n knn_feats = torch.index_select(ref_feats, 0, knn_inds.reshape(-1)).reshape(*knn_inds.shape, -1) # (n, k, c)\n knn_pts = torch.index_select(ref_pts, 0, knn_inds.reshape(-1)).reshape(*knn_inds.shape, -1) # (n, k, c)\n return knn_inds, knn_feats, knn_pts\n\n def compute_cos_sim(\n self,\n points: Optional[TensorType[\"n\", 3]],\n feats: TensorType[\"n\", \"c\"],\n knn_feats: TensorType[\"n\", \"k\", \"c\"],\n knn_pts: TensorType[\"n\", \"k\", 3],\n ) -> TensorType[\"n\", 1]:\n knn_cos_sims = F.cosine_similarity(feats[:, None], knn_feats, dim=-1) # (n, k)\n if not self.config.distance_weighted_average:\n return knn_cos_sims.mean(-1, keepdim=True)\n\n knn_pts_dists = torch.linalg.norm(points[:, None] - knn_pts, ord=2, dim=-1)\n knn_pts_weights = torch.exp(-knn_pts_dists) # TODO: tune temperature\n knn_pts_weights = knn_pts_weights / knn_pts_weights.sum(-1, keepdim=True) # normalized weights\n # use unnormalized weights would down weight bg_cos_sims too much\n knn_weighed_cos_sims = (knn_cos_sims * knn_pts_weights).sum(-1, keepdim=True)\n return knn_weighed_cos_sims\n\n def _register_all_tensors_as_buffers(self):\n \"\"\"register all Tensor attributes as buffers to automatically move them to the target device\"\"\"\n _attr_names = list(vars(self).keys())\n for _name in _attr_names:\n _attr = getattr(self, _name)\n if not isinstance(_attr, torch.Tensor):\n continue\n\n delattr(self, _name)\n self.register_buffer(_name, _attr, persistent=False)\n\n @torch.no_grad()\n def forward(\n self,\n ray_samples: Optional[RaySamples],\n features: TensorType[\"bs\":..., \"num_samples\", \"nc_features\"],\n ptcd_data: Optional[BasicPointClouds] = None,\n ) -> Dict[str, TensorType[\"bs\":..., 1]]:\n \"\"\"Calculate segmentation along the ray.\"\"\"\n prefix_dims, nc_feat = features.shape[:-1], features.shape[-1]\n features = features.view(-1, nc_feat)\n\n raw_positions = ray_samples.frustums.get_start_positions().view(-1, 3) if ray_samples is not None else None\n\n seg_probs = self.compute_seg_prob(raw_positions, features, ptcd_data=ptcd_data).reshape(*prefix_dims, 1)\n\n outputs = {FieldHeadNames.FG_SEG: seg_probs}\n return outputs\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/fields/instant_ngp_field.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nInstant-NGP field implementations using tiny-cuda-nn, torch, ....\n\"\"\"\n\n\nfrom typing import Optional\n\nimport torch\nfrom nerfacc import ContractionType, contract\nfrom torch.nn.parameter import Parameter\nfrom torchtyping import TensorType\n\nfrom nerfstudio.cameras.rays import RaySamples\nfrom nerfstudio.data.scene_box import SceneBox\nfrom nerfstudio.field_components.activations import trunc_exp\nfrom nerfstudio.field_components.embedding import Embedding\nfrom nerfstudio.field_components.field_heads import FieldHeadNames\nfrom nerfstudio.fields.base_field import Field\n\ntry:\n import tinycudann as tcnn\nexcept ImportError:\n # tinycudann module doesn't exist\n pass\n\n\ndef get_normalized_directions(directions: TensorType[\"bs\":..., 3]):\n \"\"\"SH encoding must be in the range [0, 1]\n\n Args:\n directions: batch of directions\n \"\"\"\n return (directions + 1.0) / 2.0\n\n\nclass TCNNInstantNGPField(Field):\n \"\"\"TCNN implementation of the Instant-NGP field.\n\n Args:\n aabb: parameters of scene aabb bounds\n num_layers: number of hidden layers\n hidden_dim: dimension of hidden layers\n geo_feat_dim: output geo feat dimensions\n num_layers_color: number of hidden layers for color network\n hidden_dim_color: dimension of hidden layers for color network\n use_appearance_embedding: whether to use appearance embedding\n num_images: number of images, requried if use_appearance_embedding is True\n appearance_embedding_dim: dimension of appearance embedding\n contraction_type: type of contraction\n \"\"\"\n\n def __init__(\n self,\n aabb,\n num_layers: int = 2,\n hidden_dim: int = 64,\n geo_feat_dim: int = 15,\n num_layers_color: int = 3,\n hidden_dim_color: int = 64,\n use_appearance_embedding: bool = False,\n num_images: Optional[int] = None,\n appearance_embedding_dim: int = 32,\n contraction_type: ContractionType = ContractionType.UN_BOUNDED_SPHERE,\n ) -> None:\n super().__init__()\n\n self.aabb = Parameter(aabb, requires_grad=False)\n self.geo_feat_dim = geo_feat_dim\n self.contraction_type = contraction_type\n\n self.use_appearance_embedding = use_appearance_embedding\n if use_appearance_embedding:\n assert num_images is not None\n self.appearance_embedding_dim = appearance_embedding_dim\n self.appearance_embedding = Embedding(num_images, appearance_embedding_dim)\n\n # TODO: set this properly based on the aabb\n per_level_scale = 1.4472692012786865\n\n self.direction_encoding = tcnn.Encoding(\n n_input_dims=3,\n encoding_config={\n \"otype\": \"SphericalHarmonics\",\n \"degree\": 4,\n },\n )\n\n self.mlp_base = tcnn.NetworkWithInputEncoding(\n n_input_dims=3,\n n_output_dims=1 + self.geo_feat_dim,\n encoding_config={\n \"otype\": \"HashGrid\",\n \"n_levels\": 16,\n \"n_features_per_level\": 2,\n \"log2_hashmap_size\": 19,\n \"base_resolution\": 16,\n \"per_level_scale\": per_level_scale,\n },\n network_config={\n \"otype\": \"FullyFusedMLP\",\n \"activation\": \"ReLU\",\n \"output_activation\": \"None\",\n \"n_neurons\": hidden_dim,\n \"n_hidden_layers\": num_layers - 1,\n },\n )\n\n in_dim = self.direction_encoding.n_output_dims + self.geo_feat_dim\n if self.use_appearance_embedding:\n in_dim += self.appearance_embedding_dim\n self.mlp_head = tcnn.Network(\n n_input_dims=in_dim,\n n_output_dims=3,\n network_config={\n \"otype\": \"FullyFusedMLP\",\n \"activation\": \"ReLU\",\n \"output_activation\": \"Sigmoid\",\n \"n_neurons\": hidden_dim_color,\n \"n_hidden_layers\": num_layers_color - 1,\n },\n )\n\n def get_density(self, ray_samples: RaySamples):\n positions = ray_samples.frustums.get_positions()\n positions_flat = positions.view(-1, 3)\n positions_flat = contract(x=positions_flat, roi=self.aabb, type=self.contraction_type)\n\n h = self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1)\n density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1)\n\n # Rectifying the density with an exponential is much more stable than a ReLU or\n # softplus, because it enables high post-activation (float32) density outputs\n # from smaller internal (float16) parameters.\n density = trunc_exp(density_before_activation.to(positions))\n return density, base_mlp_out\n\n def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None):\n directions = get_normalized_directions(ray_samples.frustums.directions)\n directions_flat = directions.view(-1, 3)\n\n d = self.direction_encoding(directions_flat)\n if density_embedding is None:\n positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)\n h = torch.cat([d, positions.view(-1, 3)], dim=-1)\n else:\n h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1)\n\n if self.use_appearance_embedding:\n if ray_samples.camera_indices is None:\n raise AttributeError(\"Camera indices are not provided.\")\n camera_indices = ray_samples.camera_indices.squeeze()\n if self.training:\n embedded_appearance = self.appearance_embedding(camera_indices)\n else:\n embedded_appearance = torch.zeros(\n (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device\n )\n h = torch.cat([h, embedded_appearance.view(-1, self.appearance_embedding_dim)], dim=-1)\n\n rgb = self.mlp_head(h).view(*ray_samples.frustums.directions.shape[:-1], -1).to(directions)\n return {FieldHeadNames.RGB: rgb}\n\n def get_opacity(self, positions: TensorType[\"bs\":..., 3], step_size) -> TensorType[\"bs\":..., 1]:\n \"\"\"Returns the opacity for a position. Used primarily by the occupancy grid.\n\n Args:\n positions: the positions to evaluate the opacity at.\n step_size: the step size to use for the opacity evaluation.\n \"\"\"\n density = self.density_fn(positions)\n ## TODO: We should scale step size based on the distortion. Currently it uses too much memory.\n # aabb_min, aabb_max = self.aabb[0], self.aabb[1]\n # if self.contraction_type is not ContractionType.AABB:\n # x = (positions - aabb_min) / (aabb_max - aabb_min)\n # x = x * 2 - 1 # aabb is at [-1, 1]\n # mag = x.norm(dim=-1, keepdim=True)\n # mask = mag.squeeze(-1) > 1\n\n # dev = (2 * mag - 1) / mag**2 + 2 * x**2 * (1 / mag**3 - (2 * mag - 1) / mag**4)\n # dev[~mask] = 1.0\n # dev = torch.clamp(dev, min=1e-6)\n # step_size = step_size / dev.norm(dim=-1, keepdim=True)\n # else:\n # step_size = step_size * (aabb_max - aabb_min)\n\n opacity = density * step_size\n return opacity\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/fields/nerfacto_field.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nField for compound nerf model, adds scene contraction and image embeddings to instant ngp\n\"\"\"\n\n\nfrom typing import Dict, Optional, Tuple\nfrom typing_extensions import Literal\n\nimport numpy as np\nimport torch\nfrom torch import nn\nfrom torch.nn.parameter import Parameter\nfrom torchtyping import TensorType\n\nfrom nerfstudio.cameras.rays import RaySamples\nfrom nerfstudio.data.scene_box import SceneBox\nfrom nerfstudio.field_components.activations import trunc_exp\nfrom nerfstudio.field_components.embedding import Embedding\nfrom nerfstudio.field_components.encodings import Encoding, HashEncoding, SHEncoding\nfrom nerfstudio.field_components.field_heads import (\n DensityFieldHead,\n FieldHead,\n FieldHeadNames,\n PredNormalsFieldHead,\n RGBFieldHead,\n SemanticFieldHead,\n TransientDensityFieldHead,\n TransientRGBFieldHead,\n UncertaintyFieldHead,\n)\nfrom nerfstudio.field_components.mlp import MLP\nfrom nerfstudio.field_components.spatial_distortions import (\n SceneContraction,\n SpatialDistortion,\n)\nfrom nerfstudio.fields.base_field import Field\n\ntry:\n import tinycudann as tcnn\nexcept ImportError:\n # tinycudann module doesn't exist\n pass\n\n\ndef get_normalized_directions(directions: TensorType[\"bs\":..., 3]):\n \"\"\"SH encoding must be in the range [0, 1]\n\n Args:\n directions: batch of directions\n \"\"\"\n return (directions + 1.0) / 2.0\n\n\nclass TCNNNerfactoField(Field):\n \"\"\"Compound Field that uses TCNN\n\n Args:\n aabb: parameters of scene aabb bounds\n num_layers: number of hidden layers\n hidden_dim: dimension of hidden layers\n geo_feat_dim: output geo feat dimensions\n num_levels: number of levels of the hashmap for the base mlp\n max_res: maximum resolution of the hashmap for the base mlp\n log2_hashmap_size: size of the hashmap for the base mlp\n num_layers_color: number of hidden layers for color network\n hidden_dim_color: dimension of hidden layers for color network\n appearance_embedding_dim: dimension of appearance embedding\n transient_embedding_dim: dimension of transient embedding\n use_average_appearance_embedding: whether to use average appearance embedding or zeros for inference\n spatial_distortion: spatial distortion to apply to the scene\n spatial_normalization_region: specify the reigion get normalized to the [[0, 0, 0], [1, 1, 1]] feature grid\n \"\"\"\n\n def __init__(\n self,\n aabb,\n num_images: int,\n num_layers: int = 2,\n hidden_dim: int = 64,\n geo_feat_dim: int = 15,\n num_levels: int = 16,\n max_res: int = 1024,\n log2_hashmap_size: int = 19,\n num_layers_color: int = 3,\n num_layers_transient: int = 2,\n hidden_dim_color: int = 64,\n hidden_dim_transient: int = 64,\n appearance_embedding_dim: int = 32,\n transient_embedding_dim: int = 16,\n use_transient_embedding: bool = False,\n use_semantics: bool = False,\n num_semantic_classes: int = 100,\n use_pred_normals: bool = False,\n use_appearance_embedding: bool = True,\n use_average_appearance_embedding: bool = False,\n spatial_distortion: Optional[SpatialDistortion] = None,\n spatial_normalization_region: Literal[\"full\", \"fg\", \"aabb\"] = \"full\",\n ) -> None:\n super().__init__()\n\n self.aabb = Parameter(aabb, requires_grad=False)\n self.geo_feat_dim = geo_feat_dim\n\n self.spatial_distortion = spatial_distortion\n self.spatial_normalization_region = spatial_normalization_region\n self.num_images = num_images\n self.use_appearance_embedding = use_appearance_embedding\n self.appearance_embedding_dim = appearance_embedding_dim if use_appearance_embedding else 0\n self.embedding_appearance = Embedding(self.num_images, self.appearance_embedding_dim)\n self.use_average_appearance_embedding = use_average_appearance_embedding\n self.use_transient_embedding = use_transient_embedding\n self.use_semantics = use_semantics\n self.use_pred_normals = use_pred_normals\n\n base_res = 16\n features_per_level = 2\n growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1))\n\n self.direction_encoding = tcnn.Encoding(\n n_input_dims=3,\n encoding_config={\n \"otype\": \"SphericalHarmonics\",\n \"degree\": 4,\n },\n )\n\n self.position_encoding = tcnn.Encoding(\n n_input_dims=3,\n encoding_config={\"otype\": \"Frequency\", \"n_frequencies\": 2},\n )\n\n self.mlp_base = tcnn.NetworkWithInputEncoding(\n n_input_dims=3,\n n_output_dims=1 + self.geo_feat_dim,\n encoding_config={\n \"otype\": \"HashGrid\",\n \"n_levels\": num_levels,\n \"n_features_per_level\": features_per_level,\n \"log2_hashmap_size\": log2_hashmap_size,\n \"base_resolution\": base_res,\n \"per_level_scale\": growth_factor,\n },\n network_config={\n \"otype\": \"FullyFusedMLP\",\n \"activation\": \"ReLU\",\n \"output_activation\": \"None\",\n \"n_neurons\": hidden_dim,\n \"n_hidden_layers\": num_layers - 1,\n },\n )\n\n # transients\n if self.use_transient_embedding:\n self.transient_embedding_dim = transient_embedding_dim\n self.embedding_transient = Embedding(self.num_images, self.transient_embedding_dim)\n self.mlp_transient = tcnn.Network(\n n_input_dims=self.geo_feat_dim + self.transient_embedding_dim,\n n_output_dims=hidden_dim_transient,\n network_config={\n \"otype\": \"FullyFusedMLP\",\n \"activation\": \"ReLU\",\n \"output_activation\": \"None\",\n \"n_neurons\": hidden_dim_transient,\n \"n_hidden_layers\": num_layers_transient - 1,\n },\n )\n self.field_head_transient_uncertainty = UncertaintyFieldHead(in_dim=self.mlp_transient.n_output_dims)\n self.field_head_transient_rgb = TransientRGBFieldHead(in_dim=self.mlp_transient.n_output_dims)\n self.field_head_transient_density = TransientDensityFieldHead(in_dim=self.mlp_transient.n_output_dims)\n\n # semantics\n if self.use_semantics:\n self.mlp_semantics = tcnn.Network(\n n_input_dims=self.geo_feat_dim,\n n_output_dims=hidden_dim_transient,\n network_config={\n \"otype\": \"FullyFusedMLP\",\n \"activation\": \"ReLU\",\n \"output_activation\": \"None\",\n \"n_neurons\": 64,\n \"n_hidden_layers\": 1,\n },\n )\n self.field_head_semantics = SemanticFieldHead(\n in_dim=self.mlp_semantics.n_output_dims, num_classes=num_semantic_classes\n )\n\n # predicted normals\n if self.use_pred_normals:\n self.mlp_pred_normals = tcnn.Network(\n n_input_dims=self.geo_feat_dim + self.position_encoding.n_output_dims,\n n_output_dims=hidden_dim_transient,\n network_config={\n \"otype\": \"FullyFusedMLP\",\n \"activation\": \"ReLU\",\n \"output_activation\": \"None\",\n \"n_neurons\": 64,\n \"n_hidden_layers\": 2,\n },\n )\n self.field_head_pred_normals = PredNormalsFieldHead(in_dim=self.mlp_pred_normals.n_output_dims)\n\n self.mlp_head = tcnn.Network(\n n_input_dims=self.direction_encoding.n_output_dims + self.geo_feat_dim + self.appearance_embedding_dim,\n n_output_dims=3,\n network_config={\n \"otype\": \"FullyFusedMLP\",\n \"activation\": \"ReLU\",\n \"output_activation\": \"Sigmoid\",\n \"n_neurons\": hidden_dim_color,\n \"n_hidden_layers\": num_layers_color - 1,\n },\n )\n\n def get_density(self, ray_samples: RaySamples):\n \"\"\"Computes and returns the densities.\"\"\"\n if self.spatial_distortion is not None:\n positions = ray_samples.frustums.get_positions()\n positions = self.spatial_distortion(positions)\n positions = self._get_normalized_positions(positions)\n else:\n positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)\n self._sample_locations = positions\n if not self._sample_locations.requires_grad:\n self._sample_locations.requires_grad = True\n positions_flat = positions.view(-1, 3)\n h = self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1)\n density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1)\n self._density_before_activation = density_before_activation\n\n # Rectifying the density with an exponential is much more stable than a ReLU or\n # softplus, because it enables high post-activation (float32) density outputs\n # from smaller internal (float16) parameters.\n density = trunc_exp(density_before_activation.to(positions))\n return density, base_mlp_out\n\n def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None):\n assert density_embedding is not None\n outputs = {}\n if ray_samples.camera_indices is None:\n raise AttributeError(\"Camera indices are not provided.\")\n camera_indices = ray_samples.camera_indices.squeeze()\n directions = get_normalized_directions(ray_samples.frustums.directions)\n directions_flat = directions.view(-1, 3)\n d = self.direction_encoding(directions_flat)\n\n outputs_shape = ray_samples.frustums.directions.shape[:-1]\n\n # per-image appearance embedding\n embedded_appearance = self._get_appeanreance_embedding(camera_indices, directions)\n\n # transients\n if self.use_transient_embedding and self.training:\n embedded_transient = self.embedding_transient(camera_indices)\n transient_input = torch.cat(\n [\n density_embedding.view(-1, self.geo_feat_dim),\n embedded_transient.view(-1, self.transient_embedding_dim),\n ],\n dim=-1,\n )\n x = self.mlp_transient(transient_input).view(*outputs_shape, -1).to(directions)\n outputs[FieldHeadNames.UNCERTAINTY] = self.field_head_transient_uncertainty(x)\n outputs[FieldHeadNames.TRANSIENT_RGB] = self.field_head_transient_rgb(x)\n outputs[FieldHeadNames.TRANSIENT_DENSITY] = self.field_head_transient_density(x)\n\n # semantics\n if self.use_semantics:\n density_embedding_copy = density_embedding.clone().detach()\n semantics_input = torch.cat(\n [\n density_embedding_copy.view(-1, self.geo_feat_dim),\n ],\n dim=-1,\n )\n x = self.mlp_semantics(semantics_input).view(*outputs_shape, -1).to(directions)\n outputs[FieldHeadNames.SEMANTICS] = self.field_head_semantics(x)\n\n # predicted normals\n if self.use_pred_normals:\n positions = ray_samples.frustums.get_positions()\n\n positions_flat = self.position_encoding(positions.view(-1, 3))\n pred_normals_inp = torch.cat([positions_flat, density_embedding.view(-1, self.geo_feat_dim)], dim=-1)\n\n x = self.mlp_pred_normals(pred_normals_inp).view(*outputs_shape, -1).to(directions)\n outputs[FieldHeadNames.PRED_NORMALS] = self.field_head_pred_normals(x)\n\n features = [d, density_embedding.view(-1, self.geo_feat_dim)]\n if self.appearance_embedding_dim > 0:\n features.append(embedded_appearance.view(-1, self.appearance_embedding_dim))\n \n h = torch.cat(features, dim=-1)\n rgb = self.mlp_head(h).view(*outputs_shape, -1).to(directions)\n outputs.update({FieldHeadNames.RGB: rgb})\n\n return outputs\n\n def _get_appeanreance_embedding(self, camera_indices: TensorType, directions: TensorType) -> TensorType:\n if self.training:\n if self.use_appearance_embedding:\n embedded_appearance = self.embedding_appearance(camera_indices)\n else: # empty embedding\n embedded_appearance = torch.zeros(\n (*directions.shape[:-1], 0), device=directions.device)\n else: # evaluation\n if self.use_appearance_embedding:\n if self.use_average_appearance_embedding: # average embedding\n embedded_appearance = torch.ones(\n (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device\n ) * self.embedding_appearance.mean(dim=0)\n else: # zero embedding\n embedded_appearance = torch.zeros(\n (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device\n )\n else: # empty embedding\n embedded_appearance = torch.zeros(\n (*directions.shape[:-1], 0), device=directions.device)\n return embedded_appearance\n\n def _get_normalized_positions(self, positions):\n if self.spatial_normalization_region == 'full':\n return (positions + 2.0) / 4.0\n elif self.spatial_normalization_region == 'fg':\n return (positions + 1.0) / 2.0\n elif self.config.spatial_normalization_region == 'aabb':\n return (positions - self.aabb[0]) / (self.aabb[1] - self.aabb[0])\n else:\n raise ValueError(f\"Unknown spatial normalization region: {self.spatial_normalization_region}\")\n\n\nclass TorchNerfactoField(Field):\n \"\"\"\n PyTorch implementation of the compound field.\n \"\"\"\n\n def __init__(\n self,\n aabb,\n num_images: int,\n position_encoding: Encoding = HashEncoding(),\n direction_encoding: Encoding = SHEncoding(),\n base_mlp_num_layers: int = 3,\n base_mlp_layer_width: int = 64,\n head_mlp_num_layers: int = 2,\n head_mlp_layer_width: int = 32,\n appearance_embedding_dim: int = 40,\n skip_connections: Tuple = (4,),\n field_heads: Tuple[FieldHead] = (RGBFieldHead(),),\n spatial_distortion: SpatialDistortion = SceneContraction(),\n ) -> None:\n super().__init__()\n self.aabb = Parameter(aabb, requires_grad=False)\n self.spatial_distortion = spatial_distortion\n self.num_images = num_images\n # FIXME: align w/ TcnnNerfactoField: use_appearance_embedding\n self.appearance_embedding_dim = appearance_embedding_dim\n self.embedding_appearance = Embedding(self.num_images, self.appearance_embedding_dim)\n\n self.position_encoding = position_encoding\n self.direction_encoding = direction_encoding\n\n self.mlp_base = MLP(\n in_dim=self.position_encoding.get_out_dim(),\n num_layers=base_mlp_num_layers,\n layer_width=base_mlp_layer_width,\n skip_connections=skip_connections,\n out_activation=nn.ReLU(),\n )\n\n self.mlp_head = MLP(\n in_dim=self.mlp_base.get_out_dim() + self.direction_encoding.get_out_dim() + self.appearance_embedding_dim,\n num_layers=head_mlp_num_layers,\n layer_width=head_mlp_layer_width,\n out_activation=nn.ReLU(),\n )\n\n self.field_output_density = DensityFieldHead(in_dim=self.mlp_base.get_out_dim())\n self.field_heads = nn.ModuleList(field_heads)\n for field_head in self.field_heads:\n field_head.set_in_dim(self.mlp_head.get_out_dim()) # type: ignore\n\n def get_density(self, ray_samples: RaySamples):\n if self.spatial_distortion is not None:\n positions = ray_samples.frustums.get_positions()\n positions = self.spatial_distortion(positions)\n else:\n positions = ray_samples.frustums.get_positions()\n encoded_xyz = self.position_encoding(positions)\n base_mlp_out = self.mlp_base(encoded_xyz)\n density = self.field_output_density(base_mlp_out)\n return density, base_mlp_out\n\n def get_outputs(\n self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None\n ) -> Dict[FieldHeadNames, TensorType]:\n\n outputs_shape = ray_samples.frustums.directions.shape[:-1]\n\n if ray_samples.camera_indices is None:\n raise AttributeError(\"Camera indices are not provided.\")\n camera_indices = ray_samples.camera_indices.squeeze()\n if self.training:\n embedded_appearance = self.embedding_appearance(camera_indices)\n else:\n embedded_appearance = torch.zeros(\n (*outputs_shape, self.appearance_embedding_dim),\n device=ray_samples.frustums.directions.device,\n )\n\n outputs = {}\n for field_head in self.field_heads:\n encoded_dir = self.direction_encoding(ray_samples.frustums.directions)\n \n features = [encoded_dir, density_embedding] # type: ignore\n if self.appearance_embedding_dim > 0:\n features.append(embedded_appearance.view(-1, self.appearance_embedding_dim))\n \n # FIXME: mlp_out is computed repeatedly for each field head?\n mlp_out = self.mlp_head(\n torch.cat(features, dim=-1)\n )\n outputs[field_head.field_head_name] = field_head(mlp_out)\n return outputs\n\n\nfield_implementation_to_class = {\"tcnn\": TCNNNerfactoField, \"torch\": TorchNerfactoField}\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/fields/nerfw_field.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Fields for nerf-w\"\"\"\n\nfrom typing import Dict, Optional, Tuple\n\nimport torch\nfrom torch import nn\nfrom torchtyping import TensorType\n\nfrom nerfstudio.cameras.rays import RaySamples\nfrom nerfstudio.field_components.embedding import Embedding\nfrom nerfstudio.field_components.encodings import Encoding, Identity\nfrom nerfstudio.field_components.field_heads import (\n DensityFieldHead,\n FieldHeadNames,\n RGBFieldHead,\n TransientDensityFieldHead,\n TransientRGBFieldHead,\n UncertaintyFieldHead,\n)\nfrom nerfstudio.field_components.mlp import MLP\nfrom nerfstudio.fields.base_field import Field\n\n\nclass VanillaNerfWField(Field):\n \"\"\"The NeRF-W field which has appearance and transient conditioning.\n\n Args:\n num_images: How many images exist in the dataset.\n position_encoding: Position encoder.\n direction_encoding: Direction encoder.\n base_mlp_num_layers: Number of layers for base MLP.\n base_mlp_layer_width: Width of base MLP layers.\n head_mlp_num_layers: Number of layer for ourput head MLP.\n head_mlp_layer_width: Width of output head MLP layers.\n appearance_embedding_dim:: Dimension of appearance embedding.\n transient_embedding_dim:: Dimension of transient embedding.\n skip_connections: Where to add skip connection in base MLP.\n \"\"\"\n\n def __init__(\n self,\n num_images: int,\n position_encoding: Encoding = Identity(in_dim=3),\n direction_encoding: Encoding = Identity(in_dim=3),\n base_mlp_num_layers: int = 8,\n base_mlp_layer_width: int = 256,\n head_mlp_num_layers: int = 2,\n head_mlp_layer_width: int = 128,\n appearance_embedding_dim: int = 48,\n transient_embedding_dim: int = 16,\n skip_connections: Tuple[int] = (4,),\n ) -> None:\n super().__init__()\n self.num_images = num_images\n self.position_encoding = position_encoding\n self.direction_encoding = direction_encoding\n self.base_mlp_num_layers = base_mlp_num_layers\n self.base_mlp_layer_width = base_mlp_layer_width\n self.head_mlp_num_layers = head_mlp_num_layers\n self.head_mlp_layer_width = head_mlp_layer_width\n self.appearance_embedding_dim = appearance_embedding_dim\n self.transient_embedding_dim = transient_embedding_dim\n\n self.embedding_appearance = Embedding(self.num_images, self.appearance_embedding_dim)\n self.embedding_transient = Embedding(self.num_images, self.transient_embedding_dim)\n\n self.mlp_base = MLP(\n in_dim=self.position_encoding.get_out_dim(),\n num_layers=base_mlp_num_layers,\n layer_width=base_mlp_layer_width,\n skip_connections=skip_connections,\n out_activation=nn.ReLU(),\n )\n self.mlp_transient = MLP(\n in_dim=self.mlp_base.get_out_dim() + self.embedding_transient.get_out_dim(),\n out_dim=base_mlp_layer_width // 2,\n num_layers=4,\n layer_width=base_mlp_layer_width,\n activation=nn.ReLU(),\n out_activation=nn.ReLU(),\n )\n self.mlp_head = MLP(\n in_dim=self.mlp_base.get_out_dim()\n + self.direction_encoding.get_out_dim()\n + self.embedding_appearance.get_out_dim(),\n num_layers=head_mlp_num_layers,\n layer_width=head_mlp_layer_width,\n out_activation=nn.ReLU(),\n )\n\n self.field_head_density = DensityFieldHead(in_dim=self.mlp_base.get_out_dim())\n self.field_head_rgb = RGBFieldHead(in_dim=self.mlp_head.get_out_dim())\n\n self.field_head_transient_uncertainty = UncertaintyFieldHead(in_dim=self.mlp_transient.get_out_dim())\n self.field_head_transient_rgb = TransientRGBFieldHead(in_dim=self.mlp_transient.get_out_dim())\n self.field_head_transient_density = TransientDensityFieldHead(in_dim=self.mlp_transient.get_out_dim())\n\n def get_density(self, ray_samples: RaySamples):\n \"\"\"Computes and returns the densities.\"\"\"\n encoded_xyz = self.position_encoding(ray_samples.frustums.get_positions())\n encoded_xyz = self.position_encoding(ray_samples.frustums.get_positions())\n base_mlp_out = self.mlp_base(encoded_xyz)\n density = self.field_head_density(base_mlp_out)\n return density, base_mlp_out\n\n def get_outputs(\n self, ray_samples: RaySamples, density_embedding: Optional[TensorType[..., \"embedding_size\"]] = None\n ) -> Dict[FieldHeadNames, TensorType]:\n \"\"\"Returns the outputs of the NeRF-W field.\n\n Args:\n ray_samples (RaySamples): Ray samples.\n density_embedding (TensorType[..., \"embedding_size\"], optional): Density embedding.\n\n Returns:\n Dict[FieldHeadNames, TensorType]: Outputs of the NeRF-W field.\n \"\"\"\n outputs = {}\n encoded_dir = self.direction_encoding(ray_samples.frustums.directions)\n if ray_samples.camera_indices is None:\n raise AttributeError(\"Camera indices are not provided.\")\n camera_indices = ray_samples.camera_indices.squeeze().to(ray_samples.frustums.origins.device)\n embedded_appearance = self.embedding_appearance(camera_indices)\n mlp_in = torch.cat([density_embedding, encoded_dir, embedded_appearance], dim=-1) # type: ignore\n mlp_head_out = self.mlp_head(mlp_in)\n outputs[self.field_head_rgb.field_head_name] = self.field_head_rgb(mlp_head_out) # static rgb\n embedded_transient = self.embedding_transient(camera_indices)\n transient_mlp_in = torch.cat([density_embedding, embedded_transient], dim=-1) # type: ignore\n transient_mlp_out = self.mlp_transient(transient_mlp_in)\n outputs[self.field_head_transient_uncertainty.field_head_name] = self.field_head_transient_uncertainty(\n transient_mlp_out\n ) # uncertainty\n outputs[self.field_head_transient_rgb.field_head_name] = self.field_head_transient_rgb(\n transient_mlp_out\n ) # transient rgb\n outputs[self.field_head_transient_density.field_head_name] = self.field_head_transient_density(\n transient_mlp_out\n ) # transient density\n return outputs\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/fields/sdf_field.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nField for compound nerf model, adds scene contraction and image embeddings to instant ngp\n\"\"\"\n\nfrom dataclasses import dataclass, field\nfrom typing import Optional, Type, Union\nfrom typing_extensions import Literal\n\nimport numpy as np\nimport torch\nimport torch.nn.functional as F\nfrom torch import nn\nfrom torch.nn.parameter import Parameter\nfrom torchtyping import TensorType\nfrom rich.console import Console\n\nfrom nerfstudio.cameras.rays import RaySamples\nfrom nerfstudio.field_components.embedding import Embedding\nfrom nerfstudio.field_components.encodings import (\n Encoding,\n NeRFEncoding,\n PeriodicVolumeEncoding,\n TensorVMEncoding,\n SHEncoding\n)\nfrom nerfstudio.field_components.field_heads import FieldHeadNames\nfrom nerfstudio.field_components.spatial_distortions import SpatialDistortion\nfrom nerfstudio.field_components.progressive_encoding import ProgressiveEncoding\nfrom nerfstudio.fields.base_field import Field, FieldConfig\n\ntry:\n import tinycudann as tcnn\nexcept ImportError:\n # tinycudann module doesn't exist\n pass\n\nCONSOLE = Console(width=128)\n\n\nclass LaplaceDensity(nn.Module): # alpha * Laplace(loc=0, scale=beta).cdf(-sdf)\n \"\"\"Laplace density from VolSDF\"\"\"\n\n def __init__(self, init_val, beta_min=0.0001):\n super().__init__()\n self.register_parameter(\"beta_min\", nn.Parameter(beta_min * torch.ones(1), requires_grad=False))\n self.register_parameter(\"beta\", nn.Parameter(init_val * torch.ones(1), requires_grad=True))\n\n def forward(\n self, sdf: TensorType[\"bs\":...], beta: Union[TensorType[\"bs\":...], None] = None\n ) -> TensorType[\"bs\":...]:\n \"\"\"convert sdf value to density value with beta, if beta is missing, then use learable beta\"\"\"\n\n if beta is None:\n beta = self.get_beta()\n\n alpha = 1.0 / beta\n return alpha * (0.5 + 0.5 * sdf.sign() * torch.expm1(-sdf.abs() / beta))\n\n def get_beta(self):\n \"\"\"return current beta value\"\"\"\n beta = self.beta.abs() + self.beta_min\n return beta\n\n\nclass SigmoidDensity(nn.Module): # alpha * Laplace(loc=0, scale=beta).cdf(-sdf)\n \"\"\"Sigmoid density from VolSDF\"\"\"\n\n def __init__(self, init_val, beta_min=0.0001):\n super().__init__()\n self.register_parameter(\"beta_min\", nn.Parameter(beta_min * torch.ones(1), requires_grad=False))\n self.register_parameter(\"beta\", nn.Parameter(init_val * torch.ones(1), requires_grad=True))\n\n def forward(\n self, sdf: TensorType[\"bs\":...], beta: Union[TensorType[\"bs\":...], None] = None\n ) -> TensorType[\"bs\":...]:\n \"\"\"convert sdf value to density value with beta, if beta is missing, then use learable beta\"\"\"\n\n if beta is None:\n beta = self.get_beta()\n\n alpha = 1.0 / beta\n\n # negtive sdf will have large density\n return alpha * torch.sigmoid(-sdf * alpha)\n\n def get_beta(self):\n \"\"\"return current beta value\"\"\"\n beta = self.beta.abs() + self.beta_min\n return beta\n\n\nclass SingleVarianceNetwork(nn.Module):\n \"\"\"Variance network in NeuS\n\n Args:\n nn (_type_): init value in NeuS variance network. smaller value leads to larger varaince and\n smoother initial surface (density field).\n Returns:\n inverse std value of logistic distribution (\"s\" in NeuS's S-density)\n \"\"\"\n\n def __init__(self, init_val):\n super(SingleVarianceNetwork, self).__init__()\n self.register_parameter(\"variance\", nn.Parameter(init_val * torch.ones(1), requires_grad=True))\n\n def forward(self, x):\n \"\"\"Returns current variance value\"\"\"\n return torch.ones([len(x), 1], device=x.device) * torch.exp(self.variance * 10.0)\n\n def get_variance(self):\n \"\"\"return current variance value\"\"\"\n return torch.exp(self.variance * 10.0).clip(1e-6, 1e6)\n\n\n@dataclass\nclass SDFFieldConfig(FieldConfig):\n \"\"\"Nerfacto Model Config\"\"\"\n\n _target: Type = field(default_factory=lambda: SDFField)\n num_layers: int = 8\n \"\"\"Number of layers for geometric network\"\"\"\n hidden_dim: int = 256\n \"\"\"Number of hidden dimension of geometric network\"\"\"\n geo_feat_dim: int = 256\n \"\"\"Dimension of geometric feature\"\"\"\n num_layers_color: int = 4\n \"\"\"Number of layers for color network\"\"\"\n hidden_dim_color: int = 256\n \"\"\"Number of hidden dimension of color network\"\"\"\n appearance_embedding_dim: int = 32\n \"\"\"Dimension of appearance embedding\"\"\"\n use_appearance_embedding: bool = False\n \"\"\"Dimension of appearance embedding\"\"\"\n bias: float = 0.8\n \"\"\"sphere size of geometric initializaion\"\"\"\n geometric_init: bool = True\n \"\"\"Whether to use geometric initialization\"\"\"\n dynamic_geometric_init: bool = False\n \"\"\"whether to infer the best geometric init radius from the given aabb\"\"\"\n inside_outside: bool = True\n \"\"\"whether to revert signed distance value, set to True for indoor scene\"\"\"\n weight_norm: bool = True\n \"\"\"Whether to use weight norm for linear laer\"\"\"\n use_normalized_raw_coords: bool = False\n \"\"\"if True, use normalized positions as raw coordinates of the network input, for geo-init to work properly\"\"\"\n use_grid_feature: bool = False\n \"\"\"Whether to use multi-resolution feature grids\"\"\"\n # TODO: build a saperate config for hash grid\n hash_grid_precision: Literal[\"float16\", \"float32\"] = \"float16\"\n \"\"\"The floating point precision of the hash grid features (hash only)\n when using float16, the grid features are stored in float16, but would be transfored to flaot32 for inference\n \"\"\"\n hash_grid_num_levels: int = 16\n hash_grid_max_res: int = 2048\n hash_grid_log2_hashmap_size: int = 19\n hash_grid_smoothstep: bool = True\n hash_grid_progressive_training: bool = False\n hash_grid_progressive_training_iters: int = 10000\n hash_grid_progressive_training_n_levels_init: int = 1\n use_position_encoding: bool = False\n \"\"\"Whether to use positional encoding (even when using feature grids)\"\"\"\n pos_enc_n_freqs: int = 6\n pos_enc_min_freq: int = 0\n pos_enc_max_freq: int = 5\n divide_factor: float = 1.0\n \"\"\"Normalization factor for multi-resolution grids\"\"\"\n beta_init: float = 0.1\n \"\"\"Init learnable beta value for transformation of sdf to density\"\"\"\n encoding_type: Literal[\"hash\", \"periodic\", \"tensorf_vm\"] = \"hash\"\n \"\"\"Encoding type for feature grids\"\"\"\n direction_encoding_type: Literal[\"nerf\", \"sh\"] = \"nerf\"\n color_network_include_sdf: bool = False\n \"\"\"Whether to include sdf value in inputs of the color network\"\"\"\n \"\"\"Encoding type for direction\"\"\"\n spatial_normalization_region: Literal[\"full\", \"fg\", \"aabb\"] = \"fg\"\n \"\"\"specify the reigion get normalized to the [[0, 0, 0], [1, 1, 1]] feature grid,\n if set to \"fg\", normalize [-1, 1] to [0, 1]; if set to \"full\", normalize [-2, 2] to [0, 1]\"\"\"\n compute_grad_in_contraction_space: bool = False\n \"\"\"compute gradients in the contraction space instead of the euclidean space.\n the gradients are used as inputs to the color network and computing eikonal loss.\n \"\"\"\n curvature_perturb_step: float = 1e-3 # TODO: tune\n \"\"\"perturbation step along a random direction on the tangent plane when computing curvature loss,\n assuming the fg object is enclosed with a unit cube.\n \"\"\"\n\n\nclass SDFField(Field):\n \"\"\"_summary_\n\n Args:\n Field (_type_): _description_\n \"\"\"\n\n config: SDFFieldConfig\n\n def __init__(\n self,\n config: SDFFieldConfig,\n aabb,\n num_images: int,\n use_average_appearance_embedding: bool = False,\n spatial_distortion: Optional[SpatialDistortion] = None,\n compute_curvature_loss: bool = False,\n ) -> None:\n super().__init__()\n self.config = config\n aabb_enlarge = 0.01 # 0.01 is not enough for some instances, some query points might be outside the aabb a little (due to numerical issues?),\n # and the reflection padding would be used, causing incorrect sdf values near to the border\n self.aabb = Parameter(torch.stack([aabb[0] - aabb_enlarge, aabb[1] + aabb_enlarge], 0),\n requires_grad=False) # enlarge the aabb a little\n self.spatial_distortion = spatial_distortion\n self.num_images = num_images\n self.compute_curvature_loss = compute_curvature_loss\n \n self._build_geo_encodings()\n self._build_color_encodings(use_average_appearance_embedding)\n self._build_geo_networks()\n self._build_color_networks()\n \n # laplace function for transform sdf to density from VolSDF\n self.laplace_density = LaplaceDensity(init_val=self.config.beta_init)\n # self.laplace_density = SigmoidDensity(init_val=self.config.beta_init)\n\n # TODO use different name for beta_init for config\n # deviation_network to compute alpha from sdf from NeuS\n self.deviation_network = SingleVarianceNetwork(init_val=self.config.beta_init)\n\n self.softplus = nn.Softplus(beta=100)\n self.relu = nn.ReLU()\n self.sigmoid = torch.nn.Sigmoid()\n\n self._cos_anneal_ratio = 1.0\n\n def _build_geo_encodings(self) -> None:\n self.divide_factor = self.config.divide_factor\n self.use_grid_feature = self.config.use_grid_feature\n \n num_levels = self.config.hash_grid_num_levels # default: 16\n max_res = self.config.hash_grid_max_res # default: 2048\n base_res = 16\n log2_hashmap_size = self.config.hash_grid_log2_hashmap_size # default: 19\n features_per_level = 2\n use_hash = True\n smoothstep = self.config.hash_grid_smoothstep\n growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1))\n \n if self.config.encoding_type == \"hash\":\n # feature encoding\n self.encoding = tcnn.Encoding(\n n_input_dims=3,\n encoding_config={\n \"otype\": \"HashGrid\" if use_hash else \"DenseGrid\",\n \"n_levels\": num_levels,\n \"n_features_per_level\": features_per_level,\n \"log2_hashmap_size\": log2_hashmap_size,\n \"base_resolution\": base_res,\n \"per_level_scale\": growth_factor,\n \"interpolation\": \"Smoothstep\" if smoothstep else \"Linear\", # why Smoothstep? for 2nd derivatives?\n },\n dtype=torch.float32 if self.config.hash_grid_precision == \"float32\" else None, # cannot use torch.float16 here\n )\n if self.config.hash_grid_progressive_training:\n self.encoding = ProgressiveEncoding(\n self.encoding,\n n_levels=num_levels,\n n_levels_init=self.config.hash_grid_progressive_training_n_levels_init,\n n_features_per_level=features_per_level,\n n_scheduler_steps=self.config.hash_grid_progressive_training_iters,\n )\n elif self.config.encoding_type == \"periodic\":\n print(\"using periodic encoding\")\n self.encoding = PeriodicVolumeEncoding(\n num_levels=num_levels,\n min_res=base_res,\n max_res=max_res,\n log2_hashmap_size=18, # 64 ** 3 = 2^18\n features_per_level=features_per_level,\n smoothstep=smoothstep,\n )\n elif self.config.encoding_type == \"tensorf_vm\":\n print(\"using tensor vm\")\n self.encoding = TensorVMEncoding(128, 24, smoothstep=smoothstep)\n\n # we concat inputs position ourselves\n if self.config.use_position_encoding:\n self.position_encoding = NeRFEncoding(\n in_dim=3, num_frequencies=self.config.pos_enc_n_freqs,\n min_freq_exp=self.config.pos_enc_min_freq,\n max_freq_exp=self.config.pos_enc_max_freq, include_input=False\n )\n \n def _build_geo_networks(self):\n dims = [self.config.hidden_dim for _ in range(self.config.num_layers)]\n in_dim = self._get_input_dims()\n dims = [in_dim] + dims + [1 + self.config.geo_feat_dim]\n self.num_layers = len(dims)\n self.skip_in = [4] # TODO: make configurable\n # build MLP with geometric initialization\n for l in range(0, self.num_layers - 1):\n if l + 1 in self.skip_in:\n out_dim = dims[l + 1] - dims[0]\n else:\n out_dim = dims[l + 1]\n\n lin = nn.Linear(dims[l], out_dim)\n\n if self.config.geometric_init:\n geo_init_radius = self._get_geometric_init_radius()\n if l == self.num_layers - 2:\n if not self.config.inside_outside:\n torch.nn.init.normal_(lin.weight, mean=np.sqrt(np.pi) / np.sqrt(dims[l]), std=0.0001)\n torch.nn.init.constant_(lin.bias, -geo_init_radius)\n else:\n torch.nn.init.normal_(lin.weight, mean=-np.sqrt(np.pi) / np.sqrt(dims[l]), std=0.0001)\n torch.nn.init.constant_(lin.bias, geo_init_radius)\n elif l == 0:\n torch.nn.init.constant_(lin.bias, 0.0)\n torch.nn.init.constant_(lin.weight[:, 3:], 0.0)\n torch.nn.init.normal_(lin.weight[:, :3], 0.0, np.sqrt(2) / np.sqrt(out_dim))\n elif l in self.skip_in:\n torch.nn.init.constant_(lin.bias, 0.0)\n torch.nn.init.normal_(lin.weight, 0.0, np.sqrt(2) / np.sqrt(out_dim))\n torch.nn.init.constant_(lin.weight[:, -(dims[0] - 3):], 0.0)\n else:\n torch.nn.init.constant_(lin.bias, 0.0)\n torch.nn.init.normal_(lin.weight, 0.0, np.sqrt(2) / np.sqrt(out_dim))\n\n if self.config.weight_norm:\n lin = nn.utils.weight_norm(lin)\n # print(\"=======\", lin.weight.shape)\n setattr(self, \"glin\" + str(l), lin)\n \n def _build_color_encodings(self, use_average_appearance_embedding):\n # per-frame appearance embedding\n self.appearance_embedding_dim = self.config.appearance_embedding_dim \\\n if self.config.use_appearance_embedding else 0\n self.embedding_appearance = Embedding(self.num_images, self.appearance_embedding_dim)\n self.use_average_appearance_embedding = use_average_appearance_embedding\n \n # directional embedding\n self.direction_encoding = self._get_direction_encoding()\n \n def _build_color_networks(self):\n # call self._build_color_encodings() first!\n dims = [self.config.hidden_dim_color for _ in range(self.config.num_layers_color)]\n # point, view_direction, normal, feature, embedding\n in_dim = (\n 3\n + self.direction_encoding.get_out_dim()\n + 3\n + self.config.geo_feat_dim\n + self.embedding_appearance.get_out_dim()\n )\n if self.config.color_network_include_sdf:\n in_dim = in_dim + 1\n dims = [in_dim] + dims + [3]\n self.num_layers_color = len(dims)\n\n for l in range(0, self.num_layers_color - 1):\n out_dim = dims[l + 1]\n lin = nn.Linear(dims[l], out_dim)\n\n if self.config.weight_norm:\n lin = nn.utils.weight_norm(lin)\n # print(\"=======\", lin.weight.shape)\n setattr(self, \"clin\" + str(l), lin)\n \n def set_cos_anneal_ratio(self, anneal: float) -> None:\n \"\"\"Set the anneal value for the proposal network.\"\"\"\n self._cos_anneal_ratio = anneal\n\n def forward_geonetwork(self, inputs):\n \"\"\"forward the geonetwork\"\"\"\n positions = self._get_input_positions(inputs) # in range(0, 1)\n if self.use_grid_feature:\n feature = self.encoding(positions)\n\n coord_inputs = positions if self.config.use_normalized_raw_coords else inputs\n combined_inputs = [coord_inputs, ] # must contain raw coordinates for geo-init to work properly\n if self.config.use_position_encoding:\n combined_inputs.append(self.position_encoding(positions))\n if self.use_grid_feature:\n combined_inputs.append(feature)\n combined_inputs = torch.cat(combined_inputs, dim=-1).to(torch.float32)\n x = combined_inputs\n \n for l in range(0, self.num_layers - 1):\n lin = getattr(self, \"glin\" + str(l))\n\n if l in self.skip_in:\n x = torch.cat([x, combined_inputs], 1) / np.sqrt(2)\n\n x = lin(x)\n\n if l < self.num_layers - 2:\n x = self.softplus(x)\n return x\n\n def get_sdf(self, ray_samples: RaySamples):\n \"\"\"predict the sdf value for ray samples\"\"\"\n positions = ray_samples.frustums.get_start_positions()\n positions_flat = positions.view(-1, 3)\n h = self.forward_geonetwork(positions_flat).view(*ray_samples.frustums.shape, -1)\n sdf, _ = torch.split(h, [1, self.config.geo_feat_dim], dim=-1)\n return sdf\n\n def get_gradient(self, ray_samples: RaySamples):\n positions = ray_samples.frustums.get_start_positions()\n positions_flat = positions.view(-1, 3)\n return self.gradient(positions_flat).view(*positions.shape)\n \n def gradient(self, x):\n \"\"\"compute the gradient of the ray\"\"\"\n x.requires_grad_(True)\n y = self.forward_geonetwork(x)[:, :1]\n d_output = torch.ones_like(y, requires_grad=False, device=y.device)\n gradients = torch.autograd.grad(\n outputs=y, inputs=x, grad_outputs=d_output, create_graph=True, retain_graph=True, only_inputs=True\n )[0]\n return gradients\n\n def get_density(self, ray_samples: RaySamples):\n \"\"\"Computes and returns the densities.\"\"\"\n positions = ray_samples.frustums.get_start_positions()\n positions_flat = positions.view(-1, 3)\n h = self.forward_geonetwork(positions_flat).view(*ray_samples.frustums.shape, -1)\n sdf, geo_feature = torch.split(h, [1, self.config.geo_feat_dim], dim=-1)\n density = self.laplace_density(sdf)\n return density, geo_feature\n\n def get_alpha(self, ray_samples: RaySamples, sdf=None, gradients=None):\n \"\"\"compute alpha from sdf as in NeuS\"\"\"\n if sdf is None or gradients is None:\n inputs = ray_samples.frustums.get_start_positions()\n inputs.requires_grad_(True)\n with torch.enable_grad():\n h = self.forward_geonetwork(inputs)\n sdf, _ = torch.split(h, [1, self.config.geo_feat_dim], dim=-1)\n d_output = torch.ones_like(sdf, requires_grad=False, device=sdf.device)\n gradients = torch.autograd.grad(\n outputs=sdf,\n inputs=inputs,\n grad_outputs=d_output,\n create_graph=True,\n retain_graph=True,\n only_inputs=True,\n )[0]\n\n # TODO: support manual scaling of inv_s\n inv_s = self.deviation_network.get_variance() # Single parameter\n\n true_cos = (ray_samples.frustums.directions * gradients).sum(-1, keepdim=True)\n\n # anneal as NeuS\n cos_anneal_ratio = self._cos_anneal_ratio\n\n # \"cos_anneal_ratio\" grows from 0 to 1 in the beginning training iterations. The anneal strategy below makes\n # the cos value \"not dead\" at the beginning training iterations, for better convergence.\n iter_cos = - (\n F.relu(-true_cos * 0.5 + 0.5) * (1.0 - cos_anneal_ratio) + F.relu(-true_cos) * cos_anneal_ratio\n ) # always non-positive\n\n # Estimate signed distances at section points\n estimated_next_sdf = sdf + iter_cos * ray_samples.deltas * 0.5\n estimated_prev_sdf = sdf - iter_cos * ray_samples.deltas * 0.5\n\n prev_cdf = torch.sigmoid(estimated_prev_sdf * inv_s)\n next_cdf = torch.sigmoid(estimated_next_sdf * inv_s)\n\n p = prev_cdf - next_cdf\n c = prev_cdf\n\n alpha = ((p + 1e-5) / (c + 1e-5)).clip(0.0, 1.0)\n\n # HF-NeuS\n # # sigma\n # cdf = torch.sigmoid(sdf * inv_s)\n # e = inv_s * (1 - cdf) * (-iter_cos) * ray_samples.deltas\n # alpha = (1 - torch.exp(-e)).clip(0.0, 1.0)\n\n return alpha\n\n def get_occupancy(self, sdf):\n \"\"\"compute occupancy as in UniSurf\"\"\"\n occupancy = self.sigmoid(-10.0 * sdf)\n return occupancy\n\n def get_colors(self, points, directions, normals, geo_features, camera_indices):\n \"\"\"compute colors\"\"\"\n d = self.direction_encoding(directions)\n\n # per-image appearance embedding\n embedded_appearance = self._get_appeanreance_embedding(camera_indices, directions)\n\n features = [points, d, normals, geo_features]\n if self.appearance_embedding_dim > 0:\n features.append(embedded_appearance.view(-1, self.appearance_embedding_dim))\n \n h = torch.cat(features, dim=-1)\n\n for l in range(0, self.num_layers_color - 1):\n lin = getattr(self, \"clin\" + str(l))\n\n h = lin(h)\n\n if l < self.num_layers_color - 2:\n h = self.relu(h)\n\n rgb = self.sigmoid(h)\n\n # unisurf\n # rgb = self.tanh(h) * 0.5 + 0.5\n\n return rgb\n\n def get_outputs(self, ray_samples: RaySamples, return_alphas=False, return_occupancy=False):\n \"\"\"compute output of ray samples\"\"\"\n if ray_samples.camera_indices is None:\n raise AttributeError(\"Camera indices are not provided.\")\n\n _prefix_dims = ray_samples.frustums.directions.shape[:-1]\n outputs = {}\n\n camera_indices = ray_samples.camera_indices.squeeze()\n\n inputs = ray_samples.frustums.get_start_positions()\n inputs = inputs.view(-1, 3)\n\n directions = ray_samples.frustums.directions\n directions_flat = directions.reshape(-1, 3)\n\n inputs.requires_grad_(True)\n with torch.enable_grad():\n sdf_geo_feature = self.forward_geonetwork(inputs)\n sdf, geo_feature = torch.split(sdf_geo_feature, [1, self.config.geo_feat_dim], dim=-1)\n d_output = torch.ones_like(sdf, requires_grad=False, device=sdf.device)\n gradients = torch.autograd.grad(\n outputs=sdf, inputs=inputs, grad_outputs=d_output,\n create_graph=True, retain_graph=True, only_inputs=True\n )[0]\n \n rgb = self.get_colors(inputs, directions_flat, gradients,\n sdf_geo_feature if self.config.color_network_include_sdf else geo_feature,\n camera_indices)\n\n density = self.laplace_density(sdf) # density of VolSDF\n\n rgb = rgb.view(*_prefix_dims, -1)\n sdf = sdf.view(*_prefix_dims, -1)\n density = density.view(*_prefix_dims, -1)\n gradients = gradients.view(*_prefix_dims, -1)\n normals = F.normalize(gradients, p=2, dim=-1)\n\n outputs.update(\n {\n FieldHeadNames.RGB: rgb,\n FieldHeadNames.DENSITY: density,\n FieldHeadNames.SDF: sdf,\n FieldHeadNames.NORMAL: normals,\n FieldHeadNames.GRADIENT: gradients,\n }\n )\n\n if return_alphas: # alpha of NeuS\n # TODO use mid point sdf for NeuS\n alphas = self.get_alpha(ray_samples, sdf, gradients)\n outputs.update({FieldHeadNames.ALPHA: alphas})\n\n if return_occupancy:\n occupancy = self.get_occupancy(sdf)\n outputs.update({FieldHeadNames.OCCUPANCY: occupancy})\n \n if self.compute_curvature_loss:\n curvature_losses = self._compute_curvature_loss(inputs, normals.view(-1, 3))\n outputs.update({\"curvature_losses\": curvature_losses.view(*_prefix_dims, -1)})\n\n return outputs\n\n def forward(self, ray_samples: RaySamples, return_alphas=False, return_occupancy=False):\n \"\"\"Evaluates the field at points along the ray.\n\n Args:\n ray_samples: Samples to evaluate field on.\n \"\"\"\n field_outputs = self.get_outputs(\n ray_samples, return_alphas=return_alphas, return_occupancy=return_occupancy\n )\n return field_outputs\n\n def _get_appeanreance_embedding(self, camera_indices: TensorType, directions: TensorType) -> TensorType:\n if self.training:\n if self.config.use_appearance_embedding:\n embedded_appearance = self.embedding_appearance(camera_indices)\n else: # empty embedding\n # NOTE: break previously trained models\n embedded_appearance = torch.zeros(\n (*directions.shape[:-1], 0), device=directions.device)\n else: # evaluation\n if self.config.use_appearance_embedding:\n if self.use_average_appearance_embedding: # average embedding\n embedded_appearance = torch.ones(\n (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device\n ) * self.embedding_appearance.mean(dim=0)\n else: # zero embedding\n embedded_appearance = torch.zeros(\n (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device\n )\n else: # empty embedding\n embedded_appearance = torch.zeros(\n (*directions.shape[:-1], 0), device=directions.device)\n return embedded_appearance\n \n def _get_direction_encoding(self):\n if self.config.direction_encoding_type == \"nerf\":\n return NeRFEncoding(\n in_dim=3, num_frequencies=4, min_freq_exp=0.0, max_freq_exp=3.0, include_input=True\n )\n elif self.config.direction_encoding_type == \"sh\":\n # TODO: use tcnn.encoding?\n return SHEncoding(levels=4) # nc=16\n else:\n raise ValueError(f\"Unknown direction encoding: {self.config.direction_encoding_type}\")\n \n def _get_input_dims(self):\n n_input_dims = 3 # must contain raw coordinates\n if self.use_grid_feature:\n n_input_dims += self.encoding.n_output_dims\n if self.config.use_position_encoding:\n n_input_dims += self.position_encoding.get_out_dim()\n assert n_input_dims > 0\n return n_input_dims\n \n def _get_geometric_init_radius(self):\n \"\"\"compute the geometric initialization radius\"\"\"\n if not self.config.dynamic_geometric_init:\n CONSOLE.print(f\"Using pre-defined geometric init radius: {self.config.bias:.3f}\")\n return self.config.bias\n \n aabb_center = (self.aabb[0, :] + self.aabb[1, :]) / 2.0\n assert (aabb_center == 0.0).all(), \"The dynamic geometric initialization radius is only valid for the origin-centered AABB.\"\n aabb_side_lens = self.aabb[1, :] - self.aabb[0, :]\n radius = float(aabb_side_lens.min()) / 2.0\n CONSOLE.print(f\"Using dynamic geometric init radius: {radius:.3f}\")\n return radius\n\n def _get_input_positions(self, inputs: torch.Tensor):\n \"\"\"compute input positions (for positional encoding / feature encoding)\n from raw ray sample positions.\n \"\"\"\n # positions = inputs / self.divide_factor\n positions = self.spatial_distortion(inputs) # [[-2, -1], [-1, 1], [1, 2]]\n positions = self._get_normalized_positions(positions)\n # FIXME: handle positions out of [0, 1] after spatial_distortion & normalization\n return positions\n \n def _get_normalized_positions(self, positions):\n if self.config.spatial_normalization_region == 'full':\n return (positions + 2.0) / 4.0\n elif self.config.spatial_normalization_region == 'fg':\n return (positions + 1.0) / 2.0\n elif self.config.spatial_normalization_region == 'aabb':\n return (positions - self.aabb[0]) / (self.aabb[1] - self.aabb[0])\n else:\n raise ValueError(f\"Unknown spatial normalization region: {self.config.spatial_normalization_region}\")\n\n def _compute_curvature_loss(self, inputs: TensorType[\"bs\", 3], normals: TensorType[\"bs\", 3]):\n \"\"\"compute point-wise approximation of curvature loss used in HashSDF\"\"\"\n # sample locally around inputs\n perturb_step = self.config.curvature_perturb_step\n rand_dirs = F.normalize(torch.rand_like(inputs) - 0.5, p=2, dim=-1)\n tangent_dirs = F.normalize(torch.linalg.cross(normals, rand_dirs), p=2, dim=-1)\n new_inputs = (inputs + perturb_step * tangent_dirs) # .detach()\n \n # compute new normals\n new_inputs.requires_grad_(True)\n with torch.enable_grad():\n sdf_geo_feature = self.forward_geonetwork(new_inputs)\n sdf, geo_feature = torch.split(sdf_geo_feature, [1, self.config.geo_feat_dim], dim=-1)\n d_output = torch.ones_like(sdf, requires_grad=False, device=sdf.device)\n gradients = torch.autograd.grad(\n outputs=sdf, inputs=new_inputs, grad_outputs=d_output,\n create_graph=True, retain_graph=True, only_inputs=True\n )[0]\n new_normals = F.normalize(gradients, p=2, dim=-1)\n \n losses = ((normals * new_normals).sum(dim=-1, keepdim=True) - 1) ** 2\n return losses[..., None]\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/fields/semantic_nerf_field.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nSemantic NeRF field implementation.\n\"\"\"\nfrom typing import Dict, Optional, Tuple\n\nimport torch\nfrom torch import nn\nfrom torchtyping import TensorType\n\nfrom nerfstudio.cameras.rays import RaySamples\nfrom nerfstudio.field_components.encodings import Encoding, Identity\nfrom nerfstudio.field_components.field_heads import (\n DensityFieldHead,\n FieldHeadNames,\n RGBFieldHead,\n SemanticFieldHead,\n)\nfrom nerfstudio.field_components.mlp import MLP\nfrom nerfstudio.fields.base_field import Field\n\n\nclass SemanticNerfField(Field):\n \"\"\"Semantic-NeRF field\n\n Args:\n num_semantic_classes: Number of distinct semantic classes.\n position_encoding: Position encoder.\n direction_encoding: Direction encoder.\n base_mlp_num_layers: Number of layers for base MLP.\n base_mlp_layer_width: Width of base MLP layers.\n head_mlp_num_layers: Number of layer for ourput head MLP.\n head_mlp_layer_width: Width of output head MLP layers.\n skip_connections: Where to add skip connection in base MLP.\n \"\"\"\n\n def __init__(\n self,\n num_semantic_classes: int,\n position_encoding: Encoding = Identity(in_dim=3),\n direction_encoding: Encoding = Identity(in_dim=3),\n base_mlp_num_layers: int = 8,\n base_mlp_layer_width: int = 256,\n head_mlp_num_layers: int = 2,\n head_mlp_layer_width: int = 128,\n skip_connections: Tuple[int] = (4,),\n ) -> None:\n super().__init__()\n self.num_semantic_classes = num_semantic_classes\n self.position_encoding = position_encoding\n self.direction_encoding = direction_encoding\n self.mlp_base = MLP(\n in_dim=self.position_encoding.get_out_dim(),\n num_layers=base_mlp_num_layers,\n layer_width=base_mlp_layer_width,\n skip_connections=skip_connections,\n out_activation=nn.ReLU(),\n )\n self.mlp_head = MLP(\n in_dim=self.mlp_base.get_out_dim() + self.direction_encoding.get_out_dim(),\n num_layers=head_mlp_num_layers,\n layer_width=head_mlp_layer_width,\n out_activation=nn.ReLU(),\n )\n self.mlp_semantic = MLP(\n in_dim=self.mlp_head.get_out_dim(),\n layer_width=self.mlp_head.layer_width // 2,\n num_layers=1,\n activation=nn.ReLU(),\n out_activation=nn.ReLU(),\n )\n self.field_head_density = DensityFieldHead(in_dim=self.mlp_base.get_out_dim())\n self.field_head_rgb = RGBFieldHead(in_dim=self.mlp_head.get_out_dim())\n self.field_head_semantic = SemanticFieldHead(\n in_dim=self.mlp_semantic.get_out_dim(), num_classes=self.num_semantic_classes\n )\n\n def get_density(self, ray_samples: RaySamples):\n encoded_xyz = self.position_encoding(ray_samples.frustums.get_positions())\n base_mlp_out = self.mlp_base(encoded_xyz)\n density = self.field_head_density(base_mlp_out)\n return density, base_mlp_out\n\n def get_outputs(\n self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None\n ) -> Dict[FieldHeadNames, TensorType]:\n encoded_dir = self.direction_encoding(ray_samples.frustums.directions)\n mlp_out = self.mlp_head(torch.cat([encoded_dir, density_embedding], dim=-1)) # type: ignore\n outputs = {}\n # rgb\n outputs[self.field_head_rgb.field_head_name] = self.field_head_rgb(mlp_out)\n # semantic\n mlp_out_sem = self.mlp_semantic(mlp_out)\n outputs[self.field_head_semantic.field_head_name] = self.field_head_semantic(mlp_out_sem)\n return outputs\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/fields/tensorf_field.py",
"content": "# Copyright 2022 The Plenoptix Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"TensoRF Field\"\"\"\n\n\nfrom typing import Optional\n\nimport torch\nfrom torch import nn\nfrom torch.nn.parameter import Parameter\nfrom torchtyping import TensorType\n\nfrom nerfstudio.cameras.rays import RaySamples\nfrom nerfstudio.data.scene_box import SceneBox\nfrom nerfstudio.field_components.encodings import Encoding, Identity, SHEncoding\nfrom nerfstudio.field_components.field_heads import FieldHeadNames, RGBFieldHead\nfrom nerfstudio.field_components.mlp import MLP\nfrom nerfstudio.fields.base_field import Field\n\n\nclass TensoRFField(Field):\n \"\"\"TensoRF Field\"\"\"\n\n def __init__(\n self,\n aabb,\n # the aabb bounding box of the dataset\n feature_encoding: Encoding = Identity(in_dim=3),\n # the encoding method used for appearance encoding outputs\n direction_encoding: Encoding = Identity(in_dim=3),\n # the encoding method used for ray direction\n density_encoding: Encoding = Identity(in_dim=3),\n # the tensor encoding method used for scene density\n color_encoding: Encoding = Identity(in_dim=3),\n # the tensor encoding method used for scene color\n appearance_dim: int = 27,\n # the number of dimensions for the appearance embedding\n head_mlp_num_layers: int = 2,\n # number of layers for the MLP\n head_mlp_layer_width: int = 128,\n # layer width for the MLP\n use_sh: bool = False,\n # whether to use spherical harmonics as the feature decoding function\n sh_levels: int = 2,\n # number of levels to use for spherical harmonics\n ) -> None:\n super().__init__()\n self.aabb = Parameter(aabb, requires_grad=False)\n self.feature_encoding = feature_encoding\n self.direction_encoding = direction_encoding\n self.density_encoding = density_encoding\n self.color_encoding = color_encoding\n\n self.mlp_head = MLP(\n in_dim=appearance_dim + 3 + self.direction_encoding.get_out_dim() + self.feature_encoding.get_out_dim(),\n num_layers=head_mlp_num_layers,\n layer_width=head_mlp_layer_width,\n activation=nn.ReLU(),\n out_activation=nn.ReLU(),\n )\n\n self.use_sh = use_sh\n\n if self.use_sh:\n self.sh = SHEncoding(sh_levels)\n self.B = nn.Linear(\n in_features=self.color_encoding.get_out_dim(), out_features=3 * self.sh.get_out_dim(), bias=False\n )\n else:\n self.B = nn.Linear(in_features=self.color_encoding.get_out_dim(), out_features=appearance_dim, bias=False)\n\n self.field_output_rgb = RGBFieldHead(in_dim=self.mlp_head.get_out_dim(), activation=nn.Sigmoid())\n\n def get_density(self, ray_samples: RaySamples):\n positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)\n positions = positions * 2 - 1\n density = self.density_encoding(positions)\n density_enc = torch.sum(density, dim=-1)[:, :, None]\n relu = torch.nn.ReLU()\n density_enc = relu(density_enc)\n return density_enc\n\n def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None) -> TensorType:\n d = ray_samples.frustums.directions\n positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)\n positions = positions * 2 - 1\n rgb_features = self.color_encoding(positions)\n rgb_features = self.B(rgb_features)\n\n d_encoded = self.direction_encoding(d)\n rgb_features_encoded = self.feature_encoding(rgb_features)\n\n if self.use_sh:\n sh_mult = self.sh(d)[:, :, None]\n rgb_sh = rgb_features.view(sh_mult.shape[0], sh_mult.shape[1], 3, sh_mult.shape[-1])\n rgb = torch.relu(torch.sum(sh_mult * rgb_sh, dim=-1) + 0.5)\n else:\n out = self.mlp_head(torch.cat([rgb_features, d, rgb_features_encoded, d_encoded], dim=-1)) # type: ignore\n rgb = self.field_output_rgb(out)\n\n return rgb\n\n def forward(\n self,\n ray_samples: RaySamples,\n compute_normals: bool = False,\n mask: Optional[TensorType] = None,\n bg_color: Optional[TensorType] = None,\n ):\n if compute_normals is True:\n raise ValueError(\"Surface normals are not currently supported with TensoRF\")\n if mask is not None and bg_color is not None:\n base_density = torch.zeros(ray_samples.shape)[:, :, None].to(mask.device)\n base_rgb = bg_color.repeat(ray_samples[:, :, None].shape)\n if mask.any():\n input_rays = ray_samples[mask, :]\n density = self.get_density(input_rays)\n rgb = self.get_outputs(input_rays, None)\n\n base_density[mask] = density\n base_rgb[mask] = rgb\n\n base_density.requires_grad_()\n base_rgb.requires_grad_()\n\n density = base_density\n rgb = base_rgb\n else:\n density = self.get_density(ray_samples)\n rgb = self.get_outputs(ray_samples, None)\n\n return {FieldHeadNames.DENSITY: density, FieldHeadNames.RGB: rgb}\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/fields/vanilla_nerf_field.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Classic NeRF field\"\"\"\n\n\nfrom typing import Dict, Optional, Tuple\nfrom typing_extensions import Literal\n\nimport torch\nfrom torch import nn\nfrom torchtyping import TensorType\n\nfrom nerfstudio.cameras.rays import RaySamples\nfrom nerfstudio.field_components.encodings import Encoding, Identity\nfrom nerfstudio.field_components.field_heads import (\n DensityFieldHead,\n FieldHead,\n FieldHeadNames,\n RGBFieldHead,\n)\nfrom nerfstudio.field_components.mlp import MLP\nfrom nerfstudio.field_components.spatial_distortions import SpatialDistortion\nfrom nerfstudio.fields.base_field import Field\n\n\nclass NeRFField(Field):\n \"\"\"NeRF Field\n\n Args:\n position_encoding: Position encoder.\n direction_encoding: Direction encoder.\n base_mlp_num_layers: Number of layers for base MLP.\n base_mlp_layer_width: Width of base MLP layers.\n head_mlp_num_layers: Number of layer for ourput head MLP.\n head_mlp_layer_width: Width of output head MLP layers.\n skip_connections: Where to add skip connection in base MLP.\n use_integrated_encoding: Used integrated samples as encoding input.\n spatial_distortion: Spatial distortion.\n spatial_normalization_region: normalization after spatial distortion.\n \"\"\"\n\n def __init__(\n self,\n position_encoding: Encoding = Identity(in_dim=3),\n direction_encoding: Encoding = Identity(in_dim=3),\n base_mlp_num_layers: int = 8,\n base_mlp_layer_width: int = 256,\n head_mlp_num_layers: int = 2,\n head_mlp_layer_width: int = 128,\n skip_connections: Tuple[int] = (4,),\n field_heads: Tuple[FieldHead] = (RGBFieldHead(),),\n use_integrated_encoding: bool = False,\n spatial_distortion: Optional[SpatialDistortion] = None,\n spatial_normalization_region: Literal[\"full\", \"fg\", \"none\"] = \"none\",\n ) -> None:\n super().__init__()\n self.position_encoding = position_encoding\n self.direction_encoding = direction_encoding\n self.use_integrated_encoding = use_integrated_encoding\n self.spatial_distortion = spatial_distortion\n self.spatial_normalization_region = spatial_normalization_region\n if (use_integrated_encoding or spatial_distortion is None) and spatial_normalization_region != \"none\":\n raise NotImplementedError\n\n self.mlp_base = MLP(\n in_dim=self.position_encoding.get_out_dim(),\n num_layers=base_mlp_num_layers,\n layer_width=base_mlp_layer_width,\n skip_connections=skip_connections,\n out_activation=nn.ReLU(),\n )\n\n self.mlp_head = MLP(\n in_dim=self.mlp_base.get_out_dim() + self.direction_encoding.get_out_dim(),\n num_layers=head_mlp_num_layers,\n layer_width=head_mlp_layer_width,\n out_activation=nn.ReLU(),\n )\n\n self.field_output_density = DensityFieldHead(in_dim=self.mlp_base.get_out_dim())\n self.field_heads = nn.ModuleList(field_heads)\n for field_head in self.field_heads:\n field_head.set_in_dim(self.mlp_head.get_out_dim()) # type: ignore\n\n def get_density(self, ray_samples: RaySamples):\n if self.use_integrated_encoding:\n gaussian_samples = ray_samples.frustums.get_gaussian_blob()\n if self.spatial_distortion is not None:\n gaussian_samples = self.spatial_distortion(gaussian_samples)\n encoded_xyz = self.position_encoding(gaussian_samples.mean, covs=gaussian_samples.cov)\n else:\n positions = ray_samples.frustums.get_positions()\n if self.spatial_distortion is not None:\n positions = self.spatial_distortion(positions)\n positions = self._get_normalized_positions(positions)\n encoded_xyz = self.position_encoding(positions)\n base_mlp_out = self.mlp_base(encoded_xyz)\n density = self.field_output_density(base_mlp_out)\n return density, base_mlp_out\n\n def get_outputs(\n self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None\n ) -> Dict[FieldHeadNames, TensorType]:\n outputs = {}\n for field_head in self.field_heads:\n encoded_dir = self.direction_encoding(ray_samples.frustums.directions)\n mlp_out = self.mlp_head(torch.cat([encoded_dir, density_embedding], dim=-1)) # type: ignore\n outputs[field_head.field_head_name] = field_head(mlp_out)\n return outputs\n\n def _get_normalized_positions(self, positions):\n if self.spatial_normalization_region == 'full':\n return (positions + 2.0) / 4.0\n elif self.spatial_normalization_region == 'fg':\n return (positions + 1.0) / 2.0\n elif self.config.spatial_normalization_region == 'none':\n return positions\n else:\n raise ValueError(f\"Unknown spatial normalization region: {self.spatial_normalization_region}\")\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/model_components/__init__.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/model_components/losses.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nCollection of Losses.\n\"\"\"\n\nimport torch\nimport torch.nn.functional as F\nfrom torch import nn\nfrom torchtyping import TensorType\n\nfrom nerfstudio.cameras.rays import RaySamples\nfrom nerfstudio.field_components.field_heads import FieldHeadNames\n\nL1Loss = nn.L1Loss\nMSELoss = nn.MSELoss\n\nLOSSES = {\"L1\": L1Loss, \"MSE\": MSELoss}\n\nEPS = 1.0e-7\n\n\ndef outer(\n t0_starts: TensorType[..., \"num_samples_0\"],\n t0_ends: TensorType[..., \"num_samples_0\"],\n t1_starts: TensorType[..., \"num_samples_1\"],\n t1_ends: TensorType[..., \"num_samples_1\"],\n y1: TensorType[..., \"num_samples_1\"],\n) -> TensorType[..., \"num_samples_0\"]:\n \"\"\"Faster version of\n\n https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/helper.py#L117\n https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/stepfun.py#L64\n\n Args:\n t0_starts: start of the interval edges\n t0_ends: end of the interval edges\n t1_starts: start of the interval edges\n t1_ends: end of the interval edges\n y1: weights\n \"\"\"\n cy1 = torch.cat([torch.zeros_like(y1[..., :1]), torch.cumsum(y1, dim=-1)], dim=-1)\n\n idx_lo = torch.searchsorted(t1_starts.contiguous(), t0_starts.contiguous(), side=\"right\") - 1\n idx_lo = torch.clamp(idx_lo, min=0, max=y1.shape[-1] - 1)\n idx_hi = torch.searchsorted(t1_ends.contiguous(), t0_ends.contiguous(), side=\"right\")\n idx_hi = torch.clamp(idx_hi, min=0, max=y1.shape[-1] - 1)\n cy1_lo = torch.take_along_dim(cy1[..., :-1], idx_lo, dim=-1)\n cy1_hi = torch.take_along_dim(cy1[..., 1:], idx_hi, dim=-1)\n y0_outer = cy1_hi - cy1_lo\n\n return y0_outer\n\n\ndef lossfun_outer(\n t: TensorType[..., \"num_samples+1\"],\n w: TensorType[..., \"num_samples\"],\n t_env: TensorType[..., \"num_samples+1\"],\n w_env: TensorType[..., \"num_samples\"],\n):\n \"\"\"\n https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/helper.py#L136\n https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/stepfun.py#L80\n\n Args:\n t: interval edges\n w: weights\n t_env: interval edges of the upper bound enveloping historgram\n w_env: weights that should upper bound the inner (t,w) histogram\n \"\"\"\n w_outer = outer(t[..., :-1], t[..., 1:], t_env[..., :-1], t_env[..., 1:], w_env)\n return torch.clip(w - w_outer, min=0) ** 2 / (w + EPS)\n\n\ndef ray_samples_to_sdist(ray_samples):\n \"\"\"Convert ray samples to s space\"\"\"\n starts = ray_samples.spacing_starts\n ends = ray_samples.spacing_ends\n sdist = torch.cat([starts[..., 0], ends[..., -1:, 0]], dim=-1) # (num_rays, num_samples + 1)\n return sdist\n\n\ndef interlevel_loss(weights_list, ray_samples_list):\n \"\"\"Calculates the proposal loss in the MipNeRF-360 paper.\n\n https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/model.py#L515\n https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/train_utils.py#L133\n \"\"\"\n c = ray_samples_to_sdist(ray_samples_list[-1]).detach()\n w = weights_list[-1][..., 0].detach()\n loss_interlevel = 0.0\n for ray_samples, weights in zip(ray_samples_list[:-1], weights_list[:-1]):\n sdist = ray_samples_to_sdist(ray_samples)\n cp = sdist # (num_rays, num_samples + 1)\n wp = weights[..., 0] # (num_rays, num_samples)\n loss_interlevel += torch.mean(lossfun_outer(c, w, cp, wp))\n return loss_interlevel\n\n\n# Verified\ndef lossfun_distortion(t, w):\n \"\"\"\n https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/helper.py#L142\n https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/stepfun.py#L266\n \"\"\"\n ut = (t[..., 1:] + t[..., :-1]) / 2\n dut = torch.abs(ut[..., :, None] - ut[..., None, :])\n loss_inter = torch.sum(w * torch.sum(w[..., None, :] * dut, dim=-1), dim=-1)\n\n loss_intra = torch.sum(w**2 * (t[..., 1:] - t[..., :-1]), dim=-1) / 3\n\n return loss_inter + loss_intra\n\n\ndef distortion_loss(weights_list, ray_samples_list):\n \"\"\"From mipnerf360\"\"\"\n c = ray_samples_to_sdist(ray_samples_list[-1])\n w = weights_list[-1][..., 0]\n loss = torch.mean(lossfun_distortion(c, w))\n return loss\n\n\ndef nerfstudio_distortion_loss(\n ray_samples: RaySamples,\n densities: TensorType[\"bs\":..., \"num_samples\", 1] = None,\n weights: TensorType[\"bs\":..., \"num_samples\", 1] = None,\n) -> TensorType[\"bs\":..., 1]:\n \"\"\"Ray based distortion loss proposed in MipNeRF-360. Returns distortion Loss.\n\n .. math::\n\n \\\\mathcal{L}(\\\\mathbf{s}, \\\\mathbf{w}) =\\\\iint\\\\limits_{-\\\\infty}^{\\\\,\\\\,\\\\,\\\\infty}\n \\\\mathbf{w}_\\\\mathbf{s}(u)\\\\mathbf{w}_\\\\mathbf{s}(v)|u - v|\\\\,d_{u}\\\\,d_{v}\n\n where :math:`\\\\mathbf{w}_\\\\mathbf{s}(u)=\\\\sum_i w_i \\\\mathbb{1}_{[\\\\mathbf{s}_i, \\\\mathbf{s}_{i+1})}(u)`\n is the weight at location :math:`u` between bin locations :math:`s_i` and :math:`s_{i+1}`.\n\n Args:\n ray_samples: Ray samples to compute loss over\n densities: Predicted sample densities\n weights: Predicted weights from densities and sample locations\n \"\"\"\n if torch.is_tensor(densities):\n assert not torch.is_tensor(weights), \"Cannot use both densities and weights\"\n # Compute the weight at each sample location\n weights = ray_samples.get_weights(densities)\n if torch.is_tensor(weights):\n assert not torch.is_tensor(densities), \"Cannot use both densities and weights\"\n\n starts = ray_samples.spacing_starts\n ends = ray_samples.spacing_ends\n\n assert starts is not None and ends is not None, \"Ray samples must have spacing starts and ends\"\n midpoints = (starts + ends) / 2.0 # (..., num_samples, 1)\n\n loss = (\n weights * weights[..., None, :, 0] * torch.abs(midpoints - midpoints[..., None, :, 0])\n ) # (..., num_samples, num_samples)\n loss = torch.sum(loss, dim=(-1, -2))[..., None] # (..., num_samples)\n loss = loss + 1 / 3.0 * torch.sum(weights**2 * (ends - starts), dim=-2)\n\n return loss\n\n\ndef orientation_loss(\n weights: TensorType[\"bs\":..., \"num_samples\", 1],\n normals: TensorType[\"bs\":..., \"num_samples\", 3],\n viewdirs: TensorType[\"bs\":..., 3],\n):\n \"\"\"Orientation loss proposed in Ref-NeRF.\n Loss that encourages that all visible normals are facing towards the camera.\n \"\"\"\n w = weights\n n = normals\n v = viewdirs\n n_dot_v = (n * v[..., None, :]).sum(axis=-1)\n return (w[..., 0] * torch.fmin(torch.zeros_like(n_dot_v), n_dot_v) ** 2).sum(dim=-1)\n\n\ndef pred_normal_loss(\n weights: TensorType[\"bs\":..., \"num_samples\", 1],\n normals: TensorType[\"bs\":..., \"num_samples\", 3],\n pred_normals: TensorType[\"bs\":..., \"num_samples\", 3],\n):\n \"\"\"Loss between normals calculated from density and normals from prediction network.\"\"\"\n return (weights[..., 0] * (1.0 - torch.sum(normals * pred_normals, dim=-1))).sum(dim=-1)\n\n\ndef monosdf_normal_loss(normal_pred: torch.Tensor, normal_gt: torch.Tensor):\n \"\"\"normal consistency loss as monosdf\n\n Args:\n normal_pred (torch.Tensor): volume rendered normal\n normal_gt (torch.Tensor): monocular normal\n \"\"\"\n normal_gt = torch.nn.functional.normalize(normal_gt, p=2, dim=-1)\n normal_pred = torch.nn.functional.normalize(normal_pred, p=2, dim=-1)\n l1 = torch.abs(normal_pred - normal_gt).sum(dim=-1).mean()\n cos = (1.0 - torch.sum(normal_pred * normal_gt, dim=-1)).mean()\n return l1 + cos\n\n\n# copy from MiDaS\ndef compute_scale_and_shift(prediction, target, mask):\n # system matrix: A = [[a_00, a_01], [a_10, a_11]]\n a_00 = torch.sum(mask * prediction * prediction, (1, 2))\n a_01 = torch.sum(mask * prediction, (1, 2))\n a_11 = torch.sum(mask, (1, 2))\n\n # right hand side: b = [b_0, b_1]\n b_0 = torch.sum(mask * prediction * target, (1, 2))\n b_1 = torch.sum(mask * target, (1, 2))\n\n # solution: x = A^-1 . b = [[a_11, -a_01], [-a_10, a_00]] / (a_00 * a_11 - a_01 * a_10) . b\n x_0 = torch.zeros_like(b_0)\n x_1 = torch.zeros_like(b_1)\n\n det = a_00 * a_11 - a_01 * a_01\n valid = det.nonzero()\n\n x_0[valid] = (a_11[valid] * b_0[valid] - a_01[valid] * b_1[valid]) / det[valid]\n x_1[valid] = (-a_01[valid] * b_0[valid] + a_00[valid] * b_1[valid]) / det[valid]\n\n return x_0, x_1\n\n\ndef reduction_batch_based(image_loss, M):\n # average of all valid pixels of the batch\n\n # avoid division by 0 (if sum(M) = sum(sum(mask)) = 0: sum(image_loss) = 0)\n divisor = torch.sum(M)\n\n if divisor == 0:\n return 0\n else:\n return torch.sum(image_loss) / divisor\n\n\ndef reduction_image_based(image_loss, M):\n # mean of average of valid pixels of an image\n\n # avoid division by 0 (if M = sum(mask) = 0: image_loss = 0)\n valid = M.nonzero()\n\n image_loss[valid] = image_loss[valid] / M[valid]\n\n return torch.mean(image_loss)\n\n\ndef mse_loss(prediction, target, mask, reduction=reduction_batch_based):\n\n M = torch.sum(mask, (1, 2))\n res = prediction - target\n image_loss = torch.sum(mask * res * res, (1, 2))\n\n return reduction(image_loss, 2 * M)\n\n\ndef gradient_loss(prediction, target, mask, reduction=reduction_batch_based):\n\n M = torch.sum(mask, (1, 2))\n\n diff = prediction - target\n diff = torch.mul(mask, diff)\n\n grad_x = torch.abs(diff[:, :, 1:] - diff[:, :, :-1])\n mask_x = torch.mul(mask[:, :, 1:], mask[:, :, :-1])\n grad_x = torch.mul(mask_x, grad_x)\n\n grad_y = torch.abs(diff[:, 1:, :] - diff[:, :-1, :])\n mask_y = torch.mul(mask[:, 1:, :], mask[:, :-1, :])\n grad_y = torch.mul(mask_y, grad_y)\n\n image_loss = torch.sum(grad_x, (1, 2)) + torch.sum(grad_y, (1, 2))\n\n return reduction(image_loss, M)\n\n\nclass MiDaSMSELoss(nn.Module):\n def __init__(self, reduction=\"batch-based\"):\n super().__init__()\n\n if reduction == \"batch-based\":\n self.__reduction = reduction_batch_based\n else:\n self.__reduction = reduction_image_based\n\n def forward(self, prediction, target, mask):\n return mse_loss(prediction, target, mask, reduction=self.__reduction)\n\n\nclass GradientLoss(nn.Module):\n def __init__(self, scales=4, reduction=\"batch-based\"):\n super().__init__()\n\n if reduction == \"batch-based\":\n self.__reduction = reduction_batch_based\n else:\n self.__reduction = reduction_image_based\n\n self.__scales = scales\n\n def forward(self, prediction, target, mask):\n total = 0\n\n for scale in range(self.__scales):\n step = pow(2, scale)\n\n total += gradient_loss(\n prediction[:, ::step, ::step],\n target[:, ::step, ::step],\n mask[:, ::step, ::step],\n reduction=self.__reduction,\n )\n\n return total\n\n\nclass ScaleAndShiftInvariantLoss(nn.Module):\n def __init__(self, alpha=0.5, scales=4, reduction=\"batch-based\"):\n super().__init__()\n\n self.__data_loss = MiDaSMSELoss(reduction=reduction)\n self.__regularization_loss = GradientLoss(scales=scales, reduction=reduction)\n self.__alpha = alpha\n\n self.__prediction_ssi = None\n\n def forward(self, prediction, target, mask):\n\n scale, shift = compute_scale_and_shift(prediction, target, mask)\n self.__prediction_ssi = scale.view(-1, 1, 1) * prediction + shift.view(-1, 1, 1)\n\n total = self.__data_loss(self.__prediction_ssi, target, mask)\n if self.__alpha > 0:\n total += self.__alpha * self.__regularization_loss(self.__prediction_ssi, target, mask)\n\n return total\n\n def __get_prediction_ssi(self):\n return self.__prediction_ssi\n\n prediction_ssi = property(__get_prediction_ssi)\n\n\n# end copy\n\n\n# copy from https://github.com/svip-lab/Indoor-SfMLearner/blob/0d682b7ce292484e5e3e2161fc9fc07e2f5ca8d1/layers.py#L218\nclass SSIM(nn.Module):\n \"\"\"Layer to compute the SSIM loss between a pair of images\"\"\"\n\n def __init__(self, patch_size):\n super(SSIM, self).__init__()\n self.mu_x_pool = nn.AvgPool2d(patch_size, 1)\n self.mu_y_pool = nn.AvgPool2d(patch_size, 1)\n self.sig_x_pool = nn.AvgPool2d(patch_size, 1)\n self.sig_y_pool = nn.AvgPool2d(patch_size, 1)\n self.sig_xy_pool = nn.AvgPool2d(patch_size, 1)\n\n self.refl = nn.ReflectionPad2d(patch_size // 2)\n\n self.C1 = 0.01**2\n self.C2 = 0.03**2\n\n def forward(self, x, y):\n x = self.refl(x)\n y = self.refl(y)\n\n mu_x = self.mu_x_pool(x)\n mu_y = self.mu_y_pool(y)\n\n sigma_x = self.sig_x_pool(x**2) - mu_x**2\n sigma_y = self.sig_y_pool(y**2) - mu_y**2\n sigma_xy = self.sig_xy_pool(x * y) - mu_x * mu_y\n\n SSIM_n = (2 * mu_x * mu_y + self.C1) * (2 * sigma_xy + self.C2)\n SSIM_d = (mu_x**2 + mu_y**2 + self.C1) * (sigma_x + sigma_y + self.C2)\n\n return torch.clamp((1 - SSIM_n / SSIM_d) / 2, 0, 1)\n\n\n# TODO test different losses\nclass NCC(nn.Module):\n \"\"\"Layer to compute the normalization cross correlation (NCC) of patches\"\"\"\n\n def __init__(self, patch_size: int = 11, min_patch_variance: float = 0.01):\n super(NCC, self).__init__()\n self.patch_size = patch_size\n self.min_patch_variance = min_patch_variance\n\n def forward(self, x, y):\n # TODO if we use gray image we should do it right after loading the image to save computations\n # to gray image\n x = torch.mean(x, dim=1)\n y = torch.mean(y, dim=1)\n\n x_mean = torch.mean(x, dim=(1, 2), keepdim=True)\n y_mean = torch.mean(y, dim=(1, 2), keepdim=True)\n\n x_normalized = x - x_mean\n y_normalized = y - y_mean\n\n norm = torch.sum(x_normalized * y_normalized, dim=(1, 2))\n var = torch.square(x_normalized).sum(dim=(1, 2)) * torch.square(y_normalized).sum(dim=(1, 2))\n denom = torch.sqrt(var + 1e-6)\n\n ncc = norm / (denom + 1e-6)\n\n # ignore pathces with low variances\n not_valid = (torch.square(x_normalized).sum(dim=(1, 2)) < self.min_patch_variance) | (\n torch.square(y_normalized).sum(dim=(1, 2)) < self.min_patch_variance\n )\n ncc[not_valid] = 1.0\n\n score = 1 - ncc.clip(-1.0, 1.0) # 0->2: smaller, better\n return score[:, None, None, None]\n\n\nclass MultiViewLoss(nn.Module):\n \"\"\"compute multi-view consistency loss\"\"\"\n\n def __init__(self, patch_size: int = 11, topk: int = 4, min_patch_variance: float = 0.01):\n super(MultiViewLoss, self).__init__()\n self.patch_size = patch_size\n self.topk = topk\n self.min_patch_variance = min_patch_variance\n # TODO make metric configurable\n # self.ssim = SSIM(patch_size=patch_size)\n # self.ncc = NCC(patch_size=patch_size)\n self.ssim = NCC(patch_size=patch_size, min_patch_variance=min_patch_variance)\n\n self.iter = 0\n\n def forward(self, patches: torch.Tensor, valid: torch.Tensor):\n \"\"\"take the mim\n\n Args:\n patches (torch.Tensor): _description_\n valid (torch.Tensor): _description_\n\n Returns:\n _type_: _description_\n \"\"\"\n num_imgs, num_rays, _, num_channels = patches.shape\n\n if num_rays <= 0:\n return torch.tensor(0.0).to(patches.device)\n\n ref_patches = (\n patches[:1, ...]\n .reshape(1, num_rays, self.patch_size, self.patch_size, num_channels)\n .expand(num_imgs - 1, num_rays, self.patch_size, self.patch_size, num_channels)\n .reshape(-1, self.patch_size, self.patch_size, num_channels)\n .permute(0, 3, 1, 2)\n ) # [N_src*N_rays, 3, patch_size, patch_size]\n src_patches = (\n patches[1:, ...]\n .reshape(num_imgs - 1, num_rays, self.patch_size, self.patch_size, num_channels)\n .reshape(-1, self.patch_size, self.patch_size, num_channels)\n .permute(0, 3, 1, 2)\n ) # [N_src*N_rays, 3, patch_size, patch_size]\n\n # apply same reshape to the valid mask\n src_patches_valid = (\n valid[1:, ...]\n .reshape(num_imgs - 1, num_rays, self.patch_size, self.patch_size, 1)\n .reshape(-1, self.patch_size, self.patch_size, 1)\n .permute(0, 3, 1, 2)\n ) # [N_src*N_rays, 1, patch_size, patch_size]\n\n ssim = self.ssim(ref_patches.detach(), src_patches)\n ssim = torch.mean(ssim, dim=(1, 2, 3))\n ssim = ssim.reshape(num_imgs - 1, num_rays)\n\n # ignore invalid patch by setting ssim error to very large value\n ssim_valid = (\n src_patches_valid.reshape(-1, self.patch_size * self.patch_size).all(dim=-1).reshape(num_imgs - 1, num_rays)\n )\n # we should mask the error after we select the topk value, otherwise we might select far way patches that happens to be inside the image\n # ssim[torch.logical_not(ssim_valid)] = 1.1 # max ssim_error is 1\n\n min_ssim, idx = torch.topk(ssim, k=self.topk, largest=False, dim=0, sorted=True)\n\n min_ssim_valid = ssim_valid[idx, torch.arange(num_rays)[None].expand_as(idx)]\n # TODO how to set this value for better visualization\n min_ssim[torch.logical_not(min_ssim_valid)] = 0.0 # max ssim_error is 1\n\n if False:\n\n # visualization of topK error computations\n\n import cv2\n import numpy as np\n\n vis_patch_num = num_rays\n K = min(100, vis_patch_num)\n\n image = (\n patches[:, :vis_patch_num, :, :]\n .reshape(-1, vis_patch_num, self.patch_size, self.patch_size, 3)\n .permute(1, 2, 0, 3, 4)\n .reshape(vis_patch_num * self.patch_size, -1, 3)\n )\n\n src_patches_reshaped = src_patches.reshape(\n num_imgs - 1, num_rays, 3, self.patch_size, self.patch_size\n ).permute(1, 0, 3, 4, 2)\n idx = idx.permute(1, 0)\n\n selected_patch = (\n src_patches_reshaped[torch.arange(num_rays)[:, None].expand(idx.shape), idx]\n .permute(0, 2, 1, 3, 4)\n .reshape(num_rays, self.patch_size, self.topk * self.patch_size, 3)[:vis_patch_num]\n .reshape(-1, self.topk * self.patch_size, 3)\n )\n\n # apply same reshape to the valid mask\n src_patches_valid_reshaped = src_patches_valid.reshape(\n num_imgs - 1, num_rays, 1, self.patch_size, self.patch_size\n ).permute(1, 0, 3, 4, 2)\n\n selected_patch_valid = (\n src_patches_valid_reshaped[torch.arange(num_rays)[:, None].expand(idx.shape), idx]\n .permute(0, 2, 1, 3, 4)\n .reshape(num_rays, self.patch_size, self.topk * self.patch_size, 1)[:vis_patch_num]\n .reshape(-1, self.topk * self.patch_size, 1)\n )\n # valid to image\n selected_patch_valid = selected_patch_valid.expand_as(selected_patch).float()\n # breakpoint()\n\n image = torch.cat([selected_patch_valid, selected_patch, image], dim=1)\n # select top rays with highest errors\n\n image = image.reshape(num_rays, self.patch_size, -1, 3)\n\n _, idx2 = torch.topk(\n torch.sum(min_ssim, dim=0) / (min_ssim_valid.float().sum(dim=0) + 1e-6),\n k=K,\n largest=True,\n dim=0,\n sorted=True,\n )\n\n image = image[idx2].reshape(K * self.patch_size, -1, 3)\n\n cv2.imwrite(f\"vis/{self.iter}.png\", (image.detach().cpu().numpy() * 255).astype(np.uint8)[..., ::-1])\n self.iter += 1\n if self.iter == 9:\n breakpoint()\n\n return torch.sum(min_ssim) / (min_ssim_valid.float().sum() + 1e-6)\n\n\n# sensor depth loss, adapted from https://github.com/dazinovic/neural-rgbd-surface-reconstruction/blob/main/losses.py\nclass SensorDepthLoss(nn.Module):\n \"\"\"Sensor Depth loss\"\"\"\n\n def __init__(self, truncation: float):\n super(SensorDepthLoss, self).__init__()\n self.truncation = truncation # 0.05 * 0.3 5cm scaled\n\n def forward(self, batch, outputs):\n \"\"\"take the mim\n\n Args:\n batch (Dict): inputs\n outputs (Dict): outputs data from surface model\n\n Returns:\n l1_loss: l1 loss\n freespace_loss: free space loss\n sdf_loss: sdf loss\n \"\"\"\n depth_pred = outputs[\"depth\"]\n depth_gt = batch[\"sensor_depth\"].to(depth_pred.device)[..., None]\n valid_gt_mask = depth_gt > 0.0\n\n l1_loss = torch.sum(valid_gt_mask * torch.abs(depth_gt - depth_pred)) / (valid_gt_mask.sum() + 1e-6)\n\n # free space loss and sdf loss\n ray_samples = outputs[\"ray_samples\"]\n filed_outputs = outputs[\"field_outputs\"]\n pred_sdf = filed_outputs[FieldHeadNames.SDF][..., 0]\n directions_norm = outputs[\"directions_norm\"]\n\n z_vals = ray_samples.frustums.starts[..., 0] / directions_norm\n\n truncation = self.truncation\n front_mask = valid_gt_mask & (z_vals < (depth_gt - truncation))\n back_mask = valid_gt_mask & (z_vals > (depth_gt + truncation))\n sdf_mask = valid_gt_mask & (~front_mask) & (~back_mask)\n\n num_fs_samples = front_mask.sum()\n num_sdf_samples = sdf_mask.sum()\n num_samples = num_fs_samples + num_sdf_samples + 1e-6\n fs_weight = 1.0 - num_fs_samples / num_samples\n sdf_weight = 1.0 - num_sdf_samples / num_samples\n\n free_space_loss = torch.mean((F.relu(truncation - pred_sdf) * front_mask) ** 2) * fs_weight\n\n sdf_loss = torch.mean(((z_vals + pred_sdf) - depth_gt) ** 2 * sdf_mask) * sdf_weight\n\n return l1_loss, free_space_loss, sdf_loss\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/model_components/patch_warping.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nImplementation of patch warping for multi-view consistency loss.\n\"\"\"\n\nimport numpy as np\nimport torch\nfrom torch import nn\n\nfrom nerfstudio.cameras.cameras import Cameras\nfrom nerfstudio.cameras.rays import RaySamples\n\n\ndef get_intersection_points(\n ray_samples: RaySamples, sdf: torch.Tensor, normal: torch.Tensor, in_image_mask: torch.Tensor\n):\n \"\"\"compute intersection points\n\n Args:\n ray_samples (RaySamples): _description_\n sdf (torch.Tensor): _description_\n normal (torch.Tensor): _description_\n in_image_mask (torch.Tensor): we only use the rays in the range of [half_patch:h-half_path, half_patch:w-half_path]\n Returns:\n _type_: _description_\n \"\"\"\n # TODO we should support different ways to compute intersections\n\n # Calculate if sign change occurred and concat 1 (no sign change) in\n # last dimension\n n_rays, n_samples = ray_samples.shape\n starts = ray_samples.frustums.starts\n sign_matrix = torch.cat([torch.sign(sdf[:, :-1, 0] * sdf[:, 1:, 0]), torch.ones(n_rays, 1).to(sdf.device)], dim=-1)\n cost_matrix = sign_matrix * torch.arange(n_samples, 0, -1).float().to(sdf.device)\n\n # Get first sign change and mask for values where a.) a sign changed\n # occurred and b.) no a neg to pos sign change occurred (meaning from\n # inside surface to outside)\n values, indices = torch.min(cost_matrix, -1)\n mask_sign_change = values < 0\n mask_pos_to_neg = sdf[torch.arange(n_rays), indices, 0] > 0\n\n # Define mask where a valid depth value is found\n mask = mask_sign_change & mask_pos_to_neg & in_image_mask\n\n # Get depth values and function values for the interval\n d_low = starts[torch.arange(n_rays), indices, 0][mask]\n v_low = sdf[torch.arange(n_rays), indices, 0][mask]\n n_low = normal[torch.arange(n_rays), indices, :][mask]\n\n indices = torch.clamp(indices + 1, max=n_samples - 1)\n d_high = starts[torch.arange(n_rays), indices, 0][mask]\n v_high = sdf[torch.arange(n_rays), indices, 0][mask]\n n_high = normal[torch.arange(n_rays), indices, :][mask]\n\n # linear-interpolations or run secant method to refine depth\n z = (v_low * d_high - v_high * d_low) / (v_low - v_high)\n\n # make this simpler\n origins = ray_samples.frustums.origins[torch.arange(n_rays), indices, :][mask]\n directions = ray_samples.frustums.directions[torch.arange(n_rays), indices, :][mask]\n\n intersection_points = origins + directions * z[..., None]\n\n # interpolate normal for simplicity so we don't need to call the model again\n points_normal = (v_low[..., None] * n_high - v_high[..., None] * n_low) / (v_low[..., None] - v_high[..., None])\n\n points_normal = torch.nn.functional.normalize(points_normal, dim=-1, p=2)\n\n # filter normals that are perpendicular to view directions\n valid = (points_normal * directions).sum(dim=-1).abs() > 0.1\n intersection_points = intersection_points[valid]\n points_normal = points_normal[valid]\n new_mask = mask.clone()\n new_mask[mask] &= valid\n\n return intersection_points, points_normal, new_mask\n\n\ndef get_homography(\n intersection_points: torch.Tensor, normal: torch.Tensor, cameras: Cameras, valid_angle_thres: float = 0.3\n):\n \"\"\"get homography\n\n Args:\n intersection_points (torch.Tensor): _description_\n normal (torch.Tensor): _description_\n cameras (Cameras): _description_\n \"\"\"\n device = intersection_points.device\n\n # construct homography\n c2w = cameras.camera_to_worlds.to(device)\n K = cameras.get_intrinsics_matrices().to(device)\n K_inv = torch.linalg.inv(K)\n\n # convert camera to opencv format\n c2w[:, :3, 1:3] *= -1\n\n w2c_r = c2w[:, :3, :3].transpose(1, 2)\n w2c_t = -w2c_r @ c2w[:, :3, 3:]\n w2c = torch.cat([w2c_r, w2c_t], dim=-1)\n\n R_rel = w2c[:, :3, :3] @ c2w[:1, :3, :3] # [N, 3, 3]\n t_rel = w2c[:, :3, :3] @ c2w[:1, :3, 3:] + w2c[:1, :3, 3:] # [N, 3, 1]\n\n p_ref = w2c[0, :3, :3] @ intersection_points.transpose(1, 0) + w2c[0, :3, 3:] # [3, n_pts]\n n_ref = w2c[0, :3, :3] @ normal.transpose(1, 0) # [3, n_pts]\n\n d = torch.sum(n_ref * p_ref, dim=0, keepdims=True)\n # TODO make this clear\n H = R_rel[:, None, :, :] + t_rel[:, None, :, :] @ n_ref.transpose(1, 0)[None, :, None, :] / d[..., None, None]\n\n H = K[:, None] @ H @ K_inv[None, :1] # [n_cameras, n_pts, 3, 3]\n\n # compute valid mask for homograpy, we should filter normal that are prependicular to source viewing ray directions\n dir_src = torch.nn.functional.normalize(c2w[:, None, :, 3] - intersection_points[None], dim=-1)\n valid = (dir_src * normal[None]).sum(dim=-1) > valid_angle_thres\n\n # point should be in front of cameras\n p_src = w2c[:, :3, :3] @ intersection_points.transpose(1, 0) + w2c[:, :3, 3:] # [:, 3, n_pts]\n valid_2 = p_src[:, 2, :] > 0.01\n\n return H, valid & valid_2\n\n\nclass PatchWarping(nn.Module):\n \"\"\"Standard patch warping.\"\"\"\n\n def __init__(self, patch_size: int = 31, pixel_offset: float = 0.5, valid_angle_thres: float = 0.3):\n super().__init__()\n\n self.patch_size = patch_size\n half_size = patch_size // 2\n self.valid_angle_thres = valid_angle_thres\n\n # generate pattern\n patch_coords = torch.meshgrid(\n torch.arange(-half_size, half_size + 1), torch.arange(-half_size, half_size + 1), indexing=\"xy\"\n )\n\n patch_coords = torch.stack(patch_coords, dim=-1) + pixel_offset # stored as (y, x) coordinates\n self.patch_coords = torch.cat([patch_coords, torch.zeros_like(patch_coords[..., :1])], dim=-1)\n\n def forward(\n self,\n ray_samples: RaySamples,\n sdf: torch.Tensor,\n normal: torch.Tensor,\n cameras: Cameras,\n images: torch.Tensor,\n pix_indices: torch.Tensor,\n ):\n\n device = sdf.device\n\n cameras = cameras.to(device)\n\n # filter out the patches that are outside the boarder of image\n in_image_mask = (\n (pix_indices[:, 0] > self.patch_size // 2)\n & (pix_indices[:, 1] > self.patch_size // 2)\n & (pix_indices[:, 0] < (cameras.image_height[0] - self.patch_size // 2 - 1))\n & (pix_indices[:, 1] < (cameras.image_width[0] - self.patch_size // 2 - 1))\n )\n # [n_imgs, n_rays_valid, patch_h*patch_w]\n\n # find intersection points and normals\n intersection_points, normal, mask = get_intersection_points(ray_samples, sdf, normal, in_image_mask)\n\n # Attention: we construct homography with OPENCV coordinate system\n H, H_valid_mask = get_homography(intersection_points, normal, cameras, self.valid_angle_thres)\n\n # Attention uv is (y, x) and we should change to (x, y) for homography\n pix_indices = torch.flip(pix_indices, dims=[-1])[mask].float()\n pix_indices = torch.cat([pix_indices, torch.ones(pix_indices.shape[0], 1).to(device)], dim=-1) # [n_pts, 3]\n\n pix_indices = pix_indices[:, None, None, :] + self.patch_coords[None].to(device) # [n_pts, patch_h, patch_w, 3]\n pix_indices = pix_indices.permute(0, 3, 1, 2).reshape(\n 1, -1, 3, self.patch_size**2\n ) # [1, n_pts, 3, patch_h*patch_w]\n\n warped_indices = H @ pix_indices\n\n # patches after warping should have positive depth\n positive_depth_mask = warped_indices[:, :, 2, :] >= 0.2\n warped_indices[:, :, 2, :] *= positive_depth_mask\n warped_indices = warped_indices[:, :, :2, :] / (warped_indices[:, :, 2:, :] + 1e-6)\n\n pix_coords = warped_indices.permute(0, 1, 3, 2).contiguous() # [..., :2]\n pix_coords[..., 0:1] /= cameras.image_width[:, None, None].to(device) - 1\n pix_coords[..., 1:2] /= cameras.image_height[:, None, None].to(device) - 1\n pix_coords = (pix_coords - 0.5) * 2\n\n # valid\n valid = (\n (pix_coords[..., 0] > -1.0)\n & (pix_coords[..., 0] < 1.0)\n & (pix_coords[..., 1] > -1.0)\n & (pix_coords[..., 1] < 1.0)\n ) # [n_imgs, n_rays_valid, patch_h*patch_w]\n\n # combine valid with H\n valid = valid & H_valid_mask[..., None] & positive_depth_mask\n\n rgb = torch.nn.functional.grid_sample(\n images.permute(0, 3, 1, 2).to(sdf.device),\n pix_coords,\n mode=\"bilinear\",\n padding_mode=\"zeros\",\n align_corners=True,\n )\n\n rgb = rgb.permute(0, 2, 3, 1) # [n_imgs, n_rays_valid, patch_h*patch_w, 3]\n\n # save as visualization\n if False:\n import cv2\n\n vis_patch_num = 60\n image = (\n rgb[:, :vis_patch_num, :, :]\n .reshape(-1, vis_patch_num, self.patch_size, self.patch_size, 3)\n .permute(1, 2, 0, 3, 4)\n .reshape(vis_patch_num * self.patch_size, -1, 3)\n )\n\n cv2.imwrite(\"vis.png\", (image.detach().cpu().numpy() * 255).astype(np.uint8)[..., ::-1])\n return rgb, valid\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/model_components/ray_generators.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nRay generator.\n\"\"\"\nfrom torch import nn\nfrom torchtyping import TensorType\n\nfrom nerfstudio.cameras.camera_optimizers import CameraOptimizer\nfrom nerfstudio.cameras.cameras import Cameras\nfrom nerfstudio.cameras.rays import RayBundle\n\n\nclass RayGenerator(nn.Module):\n \"\"\"torch.nn Module for generating rays.\n This class is the interface between the scene's cameras/camera optimizer and the ray sampler.\n\n Args:\n cameras: Camera objects containing camera info.\n pose_optimizer: pose optimization module, for optimizing noisy camera intrisics/extrinsics.\n \"\"\"\n\n def __init__(self, cameras: Cameras, pose_optimizer: CameraOptimizer) -> None:\n super().__init__()\n self.cameras = cameras\n self.pose_optimizer = pose_optimizer\n self.image_coords = nn.Parameter(cameras.get_image_coords(), requires_grad=False)\n\n def forward(self, ray_indices: TensorType[\"num_rays\", 3]) -> RayBundle:\n \"\"\"Index into the cameras to generate the rays.\n\n Args:\n ray_indices: Contains camera, row, and col indicies for target rays.\n \"\"\"\n c = ray_indices[:, 0] # camera indices\n y = ray_indices[:, 1] # row indices\n x = ray_indices[:, 2] # col indices\n coords = self.image_coords[y, x]\n\n camera_opt_to_camera = self.pose_optimizer(c)\n\n ray_bundle = self.cameras.generate_rays(\n camera_indices=c.unsqueeze(-1),\n coords=coords,\n camera_opt_to_camera=camera_opt_to_camera,\n )\n return ray_bundle\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/model_components/ray_samplers.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nCollection of sampling strategies\n\"\"\"\n\nimport math\nfrom abc import abstractmethod\nfrom typing import Callable, List, Optional, Tuple, Union\nfrom torchtyping import TensorType\nfrom rich.console import Console\nfrom copy import deepcopy\n\nimport torch\nimport nerfacc\nfrom torch import nn\nfrom nerfacc import OccupancyGrid\n\nfrom nerfstudio.cameras.rays import Frustums, RayBundle, RaySamples\n\nCONSOLE = Console(width=120)\n\nNEGATIVE_MIN_DELTA_WARNING = \"\"\"[bold yellow]RaySamples.deltas.min() < 0, which might be caused by\n nears > fars / bad weights (e.g., all zeros)\n / numerical issues caused by nears being too closed with fars.\n This will lead to negative weights and NaN interlevel loss.\"\"\"\n\n\nclass Sampler(nn.Module):\n \"\"\"Generate Samples\n\n Args:\n num_samples: number of samples to take\n \"\"\"\n\n def __init__(\n self,\n num_samples: Optional[int] = None,\n ) -> None:\n super().__init__()\n self.num_samples = num_samples\n\n @abstractmethod\n def generate_ray_samples(self) -> RaySamples:\n \"\"\"Generate Ray Samples\"\"\"\n\n def forward(self, *args, **kwargs) -> RaySamples:\n \"\"\"Generate ray samples\"\"\"\n return self.generate_ray_samples(*args, **kwargs)\n\n\nclass SpacedSampler(Sampler):\n \"\"\"Sample points according to a function.\n\n Args:\n num_samples: Number of samples per ray\n spacing_fn: Function that dictates sample spacing (ie `lambda x : x` is uniform).\n spacing_fn_inv: The inverse of spacing_fn.\n train_stratified: Use stratified sampling during training. Defults to True\n single_jitter: Use a same random jitter for all samples along a ray. Defaults to False\n \"\"\"\n\n def __init__(\n self,\n spacing_fn: Callable,\n spacing_fn_inv: Callable,\n num_samples: Optional[int] = None,\n train_stratified=True,\n single_jitter=False,\n ) -> None:\n super().__init__(num_samples=num_samples)\n self.train_stratified = train_stratified\n self.single_jitter = single_jitter\n self.spacing_fn = spacing_fn\n self.spacing_fn_inv = spacing_fn_inv\n\n def generate_ray_samples(\n self,\n ray_bundle: Optional[RayBundle] = None,\n num_samples: Optional[int] = None,\n ) -> RaySamples:\n \"\"\"Generates position samples accoring to spacing function.\n\n Args:\n ray_bundle: Rays to generate samples for\n num_samples: Number of samples per ray\n\n Returns:\n Positions and deltas for samples along a ray\n \"\"\"\n assert ray_bundle is not None\n assert ray_bundle.nears is not None\n assert ray_bundle.fars is not None\n\n num_samples = num_samples or self.num_samples\n assert num_samples is not None\n num_rays = ray_bundle.origins.shape[0]\n\n bins = torch.linspace(0.0, 1.0, num_samples + 1).to(ray_bundle.origins.device)[None, ...] # [1, num_samples+1]\n\n # TODO More complicated than it needs to be.\n if self.train_stratified and self.training:\n if self.single_jitter:\n t_rand = torch.rand((num_rays, 1), dtype=bins.dtype, device=bins.device)\n else:\n t_rand = torch.rand((num_rays, num_samples + 1), dtype=bins.dtype, device=bins.device)\n bin_centers = (bins[..., 1:] + bins[..., :-1]) / 2.0\n bin_upper = torch.cat([bin_centers, bins[..., -1:]], -1)\n bin_lower = torch.cat([bins[..., :1], bin_centers], -1)\n bins = bin_lower + (bin_upper - bin_lower) * t_rand\n\n s_near, s_far = (self.spacing_fn(x) for x in (ray_bundle.nears.clone(), ray_bundle.fars.clone()))\n spacing_to_euclidean_fn = lambda x: self.spacing_fn_inv(x * s_far + (1 - x) * s_near)\n euclidean_bins = spacing_to_euclidean_fn(bins) # [num_rays, num_samples+1]\n\n # _deltas = euclidean_bins[..., 1:, None] - euclidean_bins[..., :-1, None]\n # if _deltas.min() < 0:\n # CONSOLE.log(NEGATIVE_MIN_DELTA_WARNING)\n # # if nears and fars are too close to each other (e.g., the ray is intersecting the edge region of an aabb),\n # # the spacing between ray samples might not be handled properly using float32.\n \n ray_samples = ray_bundle.get_ray_samples(\n bin_starts=euclidean_bins[..., :-1, None],\n bin_ends=euclidean_bins[..., 1:, None],\n spacing_starts=bins[..., :-1, None],\n spacing_ends=bins[..., 1:, None],\n spacing_to_euclidean_fn=spacing_to_euclidean_fn,\n )\n\n return ray_samples\n\n\nclass UniformSampler(SpacedSampler):\n \"\"\"Sample uniformly along a ray\n\n Args:\n num_samples: Number of samples per ray\n train_stratified: Use stratified sampling during training. Defults to True\n single_jitter: Use a same random jitter for all samples along a ray. Defaults to False\n \"\"\"\n\n def __init__(\n self,\n num_samples: Optional[int] = None,\n train_stratified=True,\n single_jitter=False,\n ) -> None:\n super().__init__(\n num_samples=num_samples,\n spacing_fn=lambda x: x,\n spacing_fn_inv=lambda x: x,\n train_stratified=train_stratified,\n single_jitter=single_jitter,\n )\n\n\nclass LinearDisparitySampler(SpacedSampler):\n \"\"\"Sample linearly in disparity along a ray\n\n Args:\n num_samples: Number of samples per ray\n train_stratified: Use stratified sampling during training. Defults to True\n single_jitter: Use a same random jitter for all samples along a ray. Defaults to False\n \"\"\"\n\n def __init__(\n self,\n num_samples: Optional[int] = None,\n train_stratified=True,\n single_jitter=False,\n ) -> None:\n super().__init__(\n num_samples=num_samples,\n spacing_fn=lambda x: 1 / x,\n spacing_fn_inv=lambda x: 1 / x,\n train_stratified=train_stratified,\n single_jitter=single_jitter,\n )\n\n\nclass SqrtSampler(SpacedSampler):\n \"\"\"Square root sampler along a ray\n\n Args:\n num_samples: Number of samples per ray\n train_stratified: Use stratified sampling during training. Defults to True\n \"\"\"\n\n def __init__(\n self,\n num_samples: Optional[int] = None,\n train_stratified=True,\n single_jitter=False,\n ) -> None:\n super().__init__(\n num_samples=num_samples,\n spacing_fn=torch.sqrt,\n spacing_fn_inv=lambda x: x**2,\n train_stratified=train_stratified,\n single_jitter=single_jitter,\n )\n\n\nclass LogSampler(SpacedSampler):\n \"\"\"Log sampler along a ray\n\n Args:\n num_samples: Number of samples per ray\n train_stratified: Use stratified sampling during training. Defults to True\n \"\"\"\n\n def __init__(\n self,\n num_samples: Optional[int] = None,\n train_stratified=True,\n single_jitter=False,\n ) -> None:\n super().__init__(\n num_samples=num_samples,\n spacing_fn=torch.log,\n spacing_fn_inv=torch.exp,\n train_stratified=train_stratified,\n single_jitter=single_jitter,\n )\n\n\nclass UniformLinDispPiecewiseSampler(SpacedSampler):\n \"\"\"Piecewise sampler along a ray that allocates the first half of the samples uniformly and the second half\n using linearly in disparity spacing.\n\n\n Args:\n num_samples: Number of samples per ray\n train_stratified: Use stratified sampling during training. Defults to True\n single_jitter: Use a same random jitter for all samples along a ray. Defaults to False\n \"\"\"\n\n def __init__(\n self,\n num_samples: Optional[int] = None,\n train_stratified=True,\n single_jitter=False,\n ) -> None:\n super().__init__(\n num_samples=num_samples,\n spacing_fn=lambda x: torch.where(x < 1, x / 2, 1 - 1 / (2 * x)),\n spacing_fn_inv=lambda x: torch.where(x < 0.5, 2 * x, 1 / (2 - 2 * x)),\n train_stratified=train_stratified,\n single_jitter=single_jitter,\n )\n\n\nclass PDFSampler(Sampler):\n \"\"\"Sample based on probability distribution\n\n Args:\n num_samples: Number of samples per ray\n train_stratified: Randomize location within each bin during training.\n single_jitter: Use a same random jitter for all samples along a ray. Defaults to False\n include_original: Add original samples to ray.\n histogram_padding: Amount to weights prior to computing PDF.\n \"\"\"\n\n def __init__(\n self,\n num_samples: Optional[int] = None,\n train_stratified: bool = True,\n single_jitter: bool = False,\n include_original: bool = True,\n histogram_padding: float = 0.01,\n ) -> None:\n super().__init__(num_samples=num_samples)\n self.train_stratified = train_stratified\n self.include_original = include_original\n self.histogram_padding = histogram_padding\n self.single_jitter = single_jitter\n\n def generate_ray_samples(\n self,\n ray_bundle: Optional[RayBundle] = None,\n ray_samples: Optional[RaySamples] = None,\n weights: TensorType[..., \"num_samples\", 1] = None,\n num_samples: Optional[int] = None,\n eps: float = 1e-5,\n ) -> RaySamples:\n \"\"\"Generates position samples given a distribution.\n\n Args:\n ray_bundle: Rays to generate samples for\n ray_samples: Existing ray samples\n weights: Weights for each bin\n num_samples: Number of samples per ray\n eps: Small value to prevent numerical issues.\n\n Returns:\n Positions and deltas for samples along a ray\n \"\"\"\n\n if ray_samples is None or ray_bundle is None:\n raise ValueError(\"ray_samples and ray_bundle must be provided\")\n\n num_samples = num_samples or self.num_samples\n assert num_samples is not None\n num_bins = num_samples + 1\n\n weights = weights[..., 0] + self.histogram_padding\n\n # Add small offset to rays with zero weight to prevent NaNs\n weights_sum = torch.sum(weights, dim=-1, keepdim=True)\n padding = torch.relu(eps - weights_sum)\n weights = weights + padding / weights.shape[-1]\n weights_sum += padding\n\n pdf = weights / weights_sum\n cdf = torch.min(torch.ones_like(pdf), torch.cumsum(pdf, dim=-1))\n cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], dim=-1)\n\n if self.train_stratified and self.training:\n # Stratified samples between 0 and 1\n u = torch.linspace(0.0, 1.0 - (1.0 / num_bins), steps=num_bins, device=cdf.device)\n u = u.expand(size=(*cdf.shape[:-1], num_bins))\n if self.single_jitter:\n rand = torch.rand((*cdf.shape[:-1], 1), device=cdf.device) / num_bins\n else:\n rand = torch.rand((*cdf.shape[:-1], num_samples + 1), device=cdf.device) / num_bins\n u = u + rand\n else:\n # Uniform samples between 0 and 1\n u = torch.linspace(0.0, 1.0 - (1.0 / num_bins), steps=num_bins, device=cdf.device)\n u = u + 1.0 / (2 * num_bins)\n u = u.expand(size=(*cdf.shape[:-1], num_bins))\n u = u.contiguous()\n\n assert (\n ray_samples.spacing_starts is not None and ray_samples.spacing_ends is not None\n ), \"ray_sample spacing_starts and spacing_ends must be provided\"\n assert ray_samples.spacing_to_euclidean_fn is not None, \"ray_samples.spacing_to_euclidean_fn must be provided\"\n existing_bins = torch.cat(\n [\n ray_samples.spacing_starts[..., 0],\n ray_samples.spacing_ends[..., -1:, 0],\n ],\n dim=-1,\n )\n\n inds = torch.searchsorted(cdf, u, side=\"right\")\n below = torch.clamp(inds - 1, 0, existing_bins.shape[-1] - 1)\n above = torch.clamp(inds, 0, existing_bins.shape[-1] - 1)\n cdf_g0 = torch.gather(cdf, -1, below)\n bins_g0 = torch.gather(existing_bins, -1, below)\n cdf_g1 = torch.gather(cdf, -1, above)\n bins_g1 = torch.gather(existing_bins, -1, above)\n\n t = torch.clip(torch.nan_to_num((u - cdf_g0) / (cdf_g1 - cdf_g0), 0), 0, 1)\n bins = bins_g0 + t * (bins_g1 - bins_g0)\n\n if self.include_original:\n bins, _ = torch.sort(torch.cat([existing_bins, bins], -1), -1)\n\n # Stop gradients\n bins = bins.detach()\n\n euclidean_bins = ray_samples.spacing_to_euclidean_fn(bins)\n # NOTE: numerical issues: deltas directly computed from bins are always positive, but deltas computed from euclidean_bins can be negative\n # spcaing_to_euclidean_fn = lambda x: x\n # _deltas = euclidean_bins[..., 1:, None] - euclidean_bins[..., :-1, None]\n # if _deltas.min() < 0:\n # if _deltas.min() < 0:\n # CONSOLE.log(NEGATIVE_MIN_DELTA_WARNING)\n # # if nears and fars are too close to each other (e.g., the ray is intersecting the edge region of an aabb),\n # # the spacing between ray samples might not be handled properly using float32.\n # __import__('ipdb').set_trace()\n\n ray_samples = ray_bundle.get_ray_samples(\n bin_starts=euclidean_bins[..., :-1, None],\n bin_ends=euclidean_bins[..., 1:, None],\n spacing_starts=bins[..., :-1, None],\n spacing_ends=bins[..., 1:, None],\n spacing_to_euclidean_fn=ray_samples.spacing_to_euclidean_fn,\n )\n\n return ray_samples\n\n\nclass VolumetricSampler(Sampler):\n \"\"\"Sampler inspired by the one proposed in the Instant-NGP paper.\n Generates samples along a ray by sampling the occupancy field.\n Optionally removes occluded samples if the density_fn is provided.\n\n Args:\n occupancy_grid: Occupancy grid to sample from.\n density_fn: Function that evaluates density at a given point.\n scene_aabb: Axis-aligned bounding box of the scene, should be set to None if the scene is unbounded.\n \"\"\"\n\n def __init__(\n self,\n occupancy_grid: Optional[OccupancyGrid] = None,\n density_fn: Optional[Callable[[TensorType[..., 3]], TensorType[..., 1]]] = None,\n scene_aabb: Optional[TensorType[2, 3]] = None,\n ) -> None:\n\n super().__init__()\n self.scene_aabb = scene_aabb\n self.density_fn = density_fn\n self.occupancy_grid = occupancy_grid\n if self.scene_aabb is not None:\n self.scene_aabb = self.scene_aabb.to(\"cuda\").flatten()\n print(self.scene_aabb)\n\n def get_sigma_fn(self, origins, directions) -> Optional[Callable]:\n \"\"\"Returns a function that returns the density of a point.\n\n Args:\n origins: Origins of rays\n directions: Directions of rays\n Returns:\n Function that returns the density of a point or None if a density function is not provided.\n \"\"\"\n\n if self.density_fn is None or not self.training:\n return None\n\n density_fn = self.density_fn\n\n def sigma_fn(t_starts, t_ends, ray_indices):\n t_origins = origins[ray_indices]\n t_dirs = directions[ray_indices]\n positions = t_origins + t_dirs * (t_starts + t_ends) / 2.0\n return density_fn(positions)\n\n return sigma_fn\n\n def generate_ray_samples(self) -> RaySamples:\n raise RuntimeError(\n \"The VolumetricSampler fuses sample generation and density check together. Please call forward() directly.\"\n )\n\n # pylint: disable=arguments-differ\n def forward(\n self,\n ray_bundle: RayBundle,\n render_step_size: float,\n near_plane: float = 0.0,\n far_plane: Optional[float] = None,\n cone_angle: float = 0.0,\n alpha_thre: float = 1e-2,\n ) -> Tuple[RaySamples, TensorType[\"total_samples\", 3], TensorType[\"total_samples\", 2]]:\n \"\"\"Generate ray samples in a bounding box.\n\n Args:\n ray_bundle: Rays to generate samples for\n render_step_size: Minimum step size to use for rendering\n near_plane: Near plane for raymarching\n far_plane: Far plane for raymarching\n cone_angle: Cone angle for raymarching, set to 0 for uniform marching.\n\n Returns:\n a tuple of (ray_samples, packed_info, ray_indices)\n The ray_samples are packed, only storing the valid samples.\n The packed_info contains all the information to recover packed samples into unpacked mode for rendering.\n The ray_indices contains the indices of the rays that each sample belongs to.\n \"\"\"\n\n rays_o = ray_bundle.origins.contiguous()\n rays_d = ray_bundle.directions.contiguous()\n if ray_bundle.camera_indices is not None:\n camera_indices = ray_bundle.camera_indices.contiguous()\n else:\n camera_indices = None\n\n packed_info, starts, ends = nerfacc.ray_marching(\n rays_o=rays_o,\n rays_d=rays_d,\n scene_aabb=self.scene_aabb,\n grid=self.occupancy_grid,\n sigma_fn=self.get_sigma_fn(rays_o, rays_d),\n render_step_size=render_step_size,\n near_plane=near_plane,\n far_plane=far_plane,\n stratified=self.training,\n cone_angle=cone_angle,\n alpha_thre=alpha_thre,\n )\n num_samples = starts.shape[0]\n if num_samples == 0:\n # create a single fake sample and update packed_info accordingly\n # this says the last ray in packed_info has 1 sample, which starts and ends at 1\n packed_info[-1, 1] = 1\n starts = torch.ones((1, 1), dtype=starts.dtype, device=rays_o.device)\n ends = torch.ones((1, 1), dtype=ends.dtype, device=rays_o.device)\n\n ray_indices = nerfacc.unpack_info(packed_info)\n origins = rays_o[ray_indices]\n dirs = rays_d[ray_indices]\n if camera_indices is not None:\n camera_indices = camera_indices[ray_indices]\n\n zeros = torch.zeros_like(origins[:, :1])\n ray_samples = RaySamples(\n frustums=Frustums(\n origins=origins,\n directions=dirs,\n starts=starts,\n ends=ends,\n pixel_area=zeros,\n ),\n camera_indices=camera_indices,\n )\n return ray_samples, packed_info, ray_indices\n\n\nclass ProposalNetworkSampler(Sampler):\n \"\"\"Sampler that uses a proposal network to generate samples.\"\"\"\n\n def __init__(\n self,\n num_proposal_samples_per_ray: Tuple[int] = (64,),\n num_nerf_samples_per_ray: int = 32,\n num_proposal_network_iterations: int = 2,\n use_uniform_sampler: bool = False,\n single_jitter: bool = False,\n update_sched: Callable = lambda x: 1,\n ) -> None:\n super().__init__()\n self.num_proposal_samples_per_ray = num_proposal_samples_per_ray\n self.num_nerf_samples_per_ray = num_nerf_samples_per_ray\n self.num_proposal_network_iterations = num_proposal_network_iterations\n self.update_sched = update_sched\n if self.num_proposal_network_iterations < 1:\n raise ValueError(\"num_proposal_network_iterations must be >= 1\")\n\n if use_uniform_sampler:\n self.initial_sampler = UniformSampler(single_jitter=single_jitter)\n else:\n self.initial_sampler = UniformLinDispPiecewiseSampler(single_jitter=single_jitter)\n\n self.pdf_sampler = PDFSampler(include_original=False, single_jitter=single_jitter)\n\n self._anneal = 1.0\n self._steps_since_update = 0\n self._step = 0\n\n def set_anneal(self, anneal: float) -> None:\n \"\"\"Set the anneal value for the proposal network.\"\"\"\n self._anneal = anneal\n\n def step_cb(self, step):\n \"\"\"Callback to register a training step has passed. This is used to keep track of the sampling schedule\"\"\"\n self._step = step\n self._steps_since_update += 1\n\n def generate_ray_samples(\n self,\n ray_bundle: Optional[RayBundle] = None,\n density_fns: Optional[List[Callable]] = None,\n ) -> Tuple[RaySamples, List, List]:\n assert ray_bundle is not None\n assert density_fns is not None\n\n weights_list = []\n ray_samples_list = []\n\n n = self.num_proposal_network_iterations\n weights = None\n ray_samples = None\n updated = self._steps_since_update > self.update_sched(self._step) or self._step < 10\n for i_level in range(n + 1):\n is_prop = i_level < n\n num_samples = self.num_proposal_samples_per_ray[i_level] if is_prop else self.num_nerf_samples_per_ray\n if i_level == 0:\n # Uniform sampling because we need to start with some samples\n ray_samples = self.initial_sampler(ray_bundle, num_samples=num_samples)\n else:\n # PDF sampling based on the last samples and their weights\n # Perform annealing to the weights. This will be a no-op if self._anneal is 1.0.\n assert weights is not None\n annealed_weights = torch.pow(weights, self._anneal)\n ray_samples = self.pdf_sampler(ray_bundle, ray_samples, annealed_weights, num_samples=num_samples)\n if is_prop:\n if updated:\n # always update on the first step or the inf check in grad scaling crashes\n density = density_fns[i_level](ray_samples.frustums.get_positions())\n else:\n with torch.no_grad():\n density = density_fns[i_level](ray_samples.frustums.get_positions())\n weights = ray_samples.get_weights(density)\n weights_list.append(weights) # (num_rays, num_samples)\n ray_samples_list.append(ray_samples)\n if updated:\n self._steps_since_update = 0\n\n assert ray_samples is not None\n return ray_samples, weights_list, ray_samples_list\n\n\nclass ErrorBoundedSampler(Sampler):\n \"\"\"VolSDF's error bounded sampler that uses a sdf network to generate samples.\"\"\"\n\n def __init__(\n self,\n num_samples: int = 64,\n num_samples_eval: int = 128,\n num_samples_extra: int = 32,\n eps: float = 0.1,\n beta_iters: int = 10,\n max_total_iters: int = 5,\n add_tiny: float = 1e-6,\n single_jitter: bool = False,\n ) -> None:\n super().__init__()\n self.num_samples = num_samples\n self.num_samples_eval = num_samples_eval\n self.num_samples_extra = num_samples_extra\n self.eps = eps\n self.beta_iters = beta_iters\n self.max_total_iters = max_total_iters\n self.add_tiny = add_tiny\n self.single_jitter = single_jitter\n\n # samplers\n self.uniform_sampler = UniformSampler(single_jitter=single_jitter)\n self.pdf_sampler = PDFSampler(\n include_original=False,\n single_jitter=single_jitter,\n histogram_padding=1e-5,\n )\n\n def generate_ray_samples(\n self,\n ray_bundle: Optional[RayBundle] = None,\n density_fn: Optional[Callable] = None,\n sdf_fn: Optional[Callable] = None,\n return_eikonal_points: bool = True,\n ) -> Union[Tuple[RaySamples, torch.Tensor], RaySamples]:\n assert ray_bundle is not None\n assert density_fn is not None\n assert sdf_fn is not None\n\n beta0 = density_fn.get_beta().detach()\n\n # Start with uniform sampling\n ray_samples = self.uniform_sampler(ray_bundle, num_samples=self.num_samples_eval)\n\n # Get maximum beta from the upper bound (Lemma 2)\n deltas = ray_samples.deltas.squeeze(-1)\n\n bound = (1.0 / (4.0 * torch.log(torch.tensor(self.eps + 1.0)))) * (deltas**2.0).sum(-1)\n beta = torch.sqrt(bound)\n\n total_iters, not_converge = 0, True\n sorted_index = None\n new_samples = ray_samples\n\n # Algorithm 1\n while not_converge and total_iters < self.max_total_iters:\n\n with torch.no_grad():\n new_sdf = sdf_fn(new_samples)\n\n # merge sdf predictions\n if sorted_index is not None:\n sdf_merge = torch.cat([sdf.squeeze(-1), new_sdf.squeeze(-1)], -1)\n sdf = torch.gather(sdf_merge, 1, sorted_index).unsqueeze(-1)\n else:\n sdf = new_sdf\n\n # Calculating the bound d* (Theorem 1)\n d_star = self.get_dstar(sdf, ray_samples)\n\n # Updating beta using line search\n beta = self.get_updated_beta(beta0, beta, density_fn, sdf, d_star, ray_samples)\n\n # Upsample more points\n density = density_fn(sdf.reshape(ray_samples.shape), beta=beta.unsqueeze(-1))\n\n weights, transmittance = ray_samples.get_weights_and_transmittance(density.unsqueeze(-1))\n\n # Check if we are done and this is the last sampling\n total_iters += 1\n not_converge = beta.max() > beta0\n\n if not_converge and total_iters < self.max_total_iters:\n # Sample more points proportional to the current error bound\n deltas = ray_samples.deltas.squeeze(-1)\n\n error_per_section = (\n torch.exp(-d_star / beta.unsqueeze(-1)) * (deltas**2.0) / (4 * beta.unsqueeze(-1) ** 2)\n )\n\n error_integral = torch.cumsum(error_per_section, dim=-1)\n weights = (torch.clamp(torch.exp(error_integral), max=1.0e6) - 1.0) * transmittance[..., 0]\n\n new_samples = self.pdf_sampler(\n ray_bundle, ray_samples, weights.unsqueeze(-1), num_samples=self.num_samples_eval\n )\n\n ray_samples, sorted_index = self.merge_ray_samples(ray_bundle, ray_samples, new_samples)\n\n else:\n # Sample the final sample set to be used in the volume rendering integral\n ray_samples = self.pdf_sampler(ray_bundle, ray_samples, weights, num_samples=self.num_samples)\n\n if return_eikonal_points:\n # sample some of the near surface points for eikonal loss\n sampled_points = ray_samples.frustums.get_positions().view(-1, 3)\n idx = torch.randint(sampled_points.shape[0], (ray_samples.shape[0] * 10,)).to(sampled_points.device)\n points = sampled_points[idx]\n\n # Add extra samples uniformly\n if self.num_samples_extra > 0:\n ray_samples_uniform = self.uniform_sampler(ray_bundle, num_samples=self.num_samples_extra)\n ray_samples, _ = self.merge_ray_samples(ray_bundle, ray_samples, ray_samples_uniform)\n\n if return_eikonal_points:\n return ray_samples, points\n\n return ray_samples\n\n def get_dstar(self, sdf, ray_samples: RaySamples):\n \"\"\"Calculating the bound d* (Theorem 1) from VolSDF\"\"\"\n d = sdf.reshape(ray_samples.shape)\n dists = ray_samples.deltas.squeeze(-1)\n a, b, c = dists[:, :-1], d[:, :-1].abs(), d[:, 1:].abs()\n first_cond = a.pow(2) + b.pow(2) <= c.pow(2)\n second_cond = a.pow(2) + c.pow(2) <= b.pow(2)\n d_star = torch.zeros(ray_samples.shape[0], ray_samples.shape[1] - 1).to(d.device)\n d_star[first_cond] = b[first_cond]\n d_star[second_cond] = c[second_cond]\n s = (a + b + c) / 2.0\n area_before_sqrt = s * (s - a) * (s - b) * (s - c)\n mask = ~first_cond & ~second_cond & (b + c - a > 0)\n d_star[mask] = (2.0 * torch.sqrt(area_before_sqrt[mask])) / (a[mask])\n d_star = (d[:, 1:].sign() * d[:, :-1].sign() == 1) * d_star # Fixing the sign\n\n # padding to make the same shape as ray_samples\n # d_star_left = torch.cat((d_star[:, :1], d_star), dim=-1)\n # d_star_right = torch.cat((d_star, d_star[:, -1:]), dim=-1)\n # d_star = torch.minimum(d_star_left, d_star_right)\n\n d_star = torch.cat((d_star, d_star[:, -1:]), dim=-1)\n return d_star\n\n def get_updated_beta(self, beta0, beta, density_fn, sdf, d_star, ray_samples: RaySamples):\n curr_error = self.get_error_bound(beta0, density_fn, sdf, d_star, ray_samples)\n beta[curr_error <= self.eps] = beta0\n beta_min, beta_max = beta0.repeat(ray_samples.shape[0]), beta\n for j in range(self.beta_iters):\n beta_mid = (beta_min + beta_max) / 2.0\n curr_error = self.get_error_bound(beta_mid.unsqueeze(-1), density_fn, sdf, d_star, ray_samples)\n beta_max[curr_error <= self.eps] = beta_mid[curr_error <= self.eps]\n beta_min[curr_error > self.eps] = beta_mid[curr_error > self.eps]\n beta = beta_max\n return beta\n\n def get_error_bound(self, beta, density_fn, sdf, d_star, ray_samples):\n \"\"\"Get error bound from VolSDF\"\"\"\n densities = density_fn(sdf.reshape(ray_samples.shape), beta=beta)\n\n deltas = ray_samples.deltas.squeeze(-1)\n delta_density = deltas * densities\n\n integral_estimation = torch.cumsum(delta_density[..., :-1], dim=-1)\n integral_estimation = torch.cat(\n [torch.zeros((*integral_estimation.shape[:1], 1), device=densities.device), integral_estimation], dim=-1\n )\n\n error_per_section = torch.exp(-d_star / beta) * (deltas**2.0) / (4 * beta**2)\n error_integral = torch.cumsum(error_per_section, dim=-1)\n bound_opacity = (torch.clamp(torch.exp(error_integral), max=1.0e6) - 1.0) * torch.exp(-integral_estimation)\n\n return bound_opacity.max(-1)[0]\n\n def merge_ray_samples(self, ray_bundle: RayBundle, ray_samples_1: RaySamples, ray_samples_2: RaySamples):\n \"\"\"Merge two set of ray samples and return sorted index which can be used to merge sdf values\n\n Args:\n ray_samples_1 : ray_samples to merge\n ray_samples_2 : ray_samples to merge\n \"\"\"\n\n starts_1 = ray_samples_1.spacing_starts[..., 0]\n starts_2 = ray_samples_2.spacing_starts[..., 0]\n\n ends = torch.maximum(ray_samples_1.spacing_ends[..., -1:, 0], ray_samples_2.spacing_ends[..., -1:, 0])\n\n bins, sorted_index = torch.sort(torch.cat([starts_1, starts_2], -1), -1)\n\n bins = torch.cat([bins, ends], dim=-1)\n\n # Stop gradients\n bins = bins.detach()\n\n euclidean_bins = ray_samples_1.spacing_to_euclidean_fn(bins)\n \n # _deltas = euclidean_bins[..., 1:, None] - euclidean_bins[..., :-1, None]\n # if _deltas.min() < 0:\n # CONSOLE.print(\"[bold yellow]Warning: RaySamples.deltas.min() < 0, which indicates nears > fars. This will lead to negative weights and NaN interlevel loss.\")\n # # NOTE: \"deltas.min() == 0\" might be caused by the nears and fars being too close to each other, which cannot be handled using float32.\n # __import__('ipdb').set_trace()\n\n ray_samples = ray_bundle.get_ray_samples(\n bin_starts=euclidean_bins[..., :-1, None],\n bin_ends=euclidean_bins[..., 1:, None],\n spacing_starts=bins[..., :-1, None],\n spacing_ends=bins[..., 1:, None],\n spacing_to_euclidean_fn=ray_samples_1.spacing_to_euclidean_fn,\n )\n\n return ray_samples, sorted_index\n\n\ndef save_points(path_save, pts, colors=None, normals=None, BRG2RGB=False):\n \"\"\"save points to point cloud using open3d\"\"\"\n assert len(pts) > 0\n if colors is not None:\n assert colors.shape[1] == 3\n assert pts.shape[1] == 3\n import numpy as np\n import open3d as o3d\n\n cloud = o3d.geometry.PointCloud()\n cloud.points = o3d.utility.Vector3dVector(pts)\n if colors is not None:\n # Open3D assumes the color values are of float type and in range [0, 1]\n if np.max(colors) > 1:\n colors = colors / np.max(colors)\n if BRG2RGB:\n colors = np.stack([colors[:, 2], colors[:, 1], colors[:, 0]], axis=-1)\n cloud.colors = o3d.utility.Vector3dVector(colors)\n if normals is not None:\n cloud.normals = o3d.utility.Vector3dVector(normals)\n\n o3d.io.write_point_cloud(path_save, cloud)\n\n\nclass NeuSSampler(Sampler):\n \"\"\"NeuS sampler that uses a sdf network to generate samples with fixed variance value in each iterations.\"\"\"\n\n def __init__(\n self,\n num_samples: int = 64,\n num_samples_importance: int = 64,\n num_samples_outside: int = 32,\n num_upsample_steps: int = 4,\n base_variance: float = 64,\n single_jitter: bool = True,\n ) -> None:\n super().__init__()\n self.num_samples = num_samples\n self.num_samples_importance = num_samples_importance\n self.num_samples_outside = num_samples_outside\n self.num_upsample_steps = num_upsample_steps\n self.base_variance = base_variance\n self.single_jitter = single_jitter\n\n # samplers\n self.uniform_sampler = UniformSampler(single_jitter=single_jitter)\n self.pdf_sampler = PDFSampler(\n include_original=False,\n single_jitter=single_jitter,\n histogram_padding=1e-5,\n )\n self.outside_sampler = LinearDisparitySampler()\n # TODO make it outside\n # for merge samples\n self.error_bounded_sampler = ErrorBoundedSampler()\n\n def generate_ray_samples(\n self,\n ray_bundle: Optional[RayBundle] = None,\n sdf_fn: Optional[Callable] = None,\n ray_samples: Optional[RaySamples] = None,\n ) -> Union[Tuple[RaySamples, torch.Tensor], RaySamples]:\n assert ray_bundle is not None\n assert sdf_fn is not None\n\n # Start with uniform sampling\n if ray_samples is None:\n ray_samples = self.uniform_sampler(ray_bundle, num_samples=self.num_samples)\n\n total_iters = 0\n sorted_index = None\n new_samples = ray_samples\n\n base_variance = self.base_variance\n\n while total_iters < self.num_upsample_steps:\n\n with torch.no_grad():\n new_sdf = sdf_fn(new_samples)\n \n # merge sdf predictions\n if sorted_index is not None:\n sdf_merge = torch.cat([sdf.squeeze(-1), new_sdf.squeeze(-1)], -1)\n sdf = torch.gather(sdf_merge, 1, sorted_index).unsqueeze(-1)\n else:\n sdf = new_sdf\n\n # compute with fix variances\n alphas = self.rendering_sdf_with_fixed_inv_s(\n ray_samples, sdf.reshape(ray_samples.shape), inv_s=base_variance * 2**total_iters\n )\n\n weights = ray_samples.get_weights_from_alphas(alphas[..., None])\n weights = torch.cat((weights, torch.zeros_like(weights[:, :1])), dim=1)\n\n new_samples = self.pdf_sampler(\n ray_bundle,\n ray_samples,\n weights,\n num_samples=self.num_samples_importance // self.num_upsample_steps,\n )\n\n ray_samples, sorted_index = self.error_bounded_sampler.merge_ray_samples(\n ray_bundle, ray_samples, new_samples\n )\n\n total_iters += 1\n\n # save_points(\"p.ply\", ray_samples.frustums.get_start_positions().detach().cpu().numpy().reshape(-1, 3))\n # exit(-1)\n # TODO\n # sample more points outside surface\n # if self.num_samples_outside > 0:\n # ray_samples_uniform = self.outside_sampler(ray_bundle, num_samples=self.num_samples_outside)\n # ray_samples, _ = self.merge_ray_samples(ray_bundle, ray_samples, ray_samples_uniform)\n\n return ray_samples\n\n def rendering_sdf_with_fixed_inv_s(self, ray_samples: RaySamples, sdf: torch.Tensor, inv_s):\n \"\"\"rendering given a fixed inv_s as NeuS\"\"\"\n batch_size = ray_samples.shape[0]\n prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]\n deltas = ray_samples.deltas[:, :-1, 0]\n mid_sdf = (prev_sdf + next_sdf) * 0.5\n cos_val = (next_sdf - prev_sdf) / (deltas + 1e-5)\n\n # ----------------------------------------------------------------------------------------------------------\n # Use min value of [ cos, prev_cos ]\n # Though it makes the sampling (not rendering) a little bit biased, this strategy can make the sampling more\n # robust when meeting situations like below:\n #\n # SDF\n # ^\n # |\\ -----x----...\n # | \\ /\n # | x x\n # |---\\----/-------------> 0 level\n # | \\ /\n # | \\/\n # |\n # ----------------------------------------------------------------------------------------------------------\n prev_cos_val = torch.cat([torch.zeros([batch_size, 1], device=sdf.device), cos_val[:, :-1]], dim=-1)\n cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)\n cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)\n cos_val = cos_val.clip(-1e3, 0.0)\n\n dist = deltas\n prev_esti_sdf = mid_sdf - cos_val * dist * 0.5\n next_esti_sdf = mid_sdf + cos_val * dist * 0.5\n prev_cdf = torch.sigmoid(prev_esti_sdf * inv_s)\n next_cdf = torch.sigmoid(next_esti_sdf * inv_s)\n alpha = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)\n\n return alpha\n\n\nclass UniSurfSampler(Sampler):\n \"\"\"NeuS sampler that uses a sdf network to generate samples with fixed variance value in each iterations.\"\"\"\n\n def __init__(\n self,\n num_samples_interval: int = 64,\n num_samples_outside: int = 32,\n num_samples_importance: int = 32,\n num_marching_steps: int = 256,\n num_secant_steps: int = 8,\n interval_start: float = 0.25,\n interval_end: float = 0.0125,\n interval_decay: float = 0.00005, # default value is 0.000015 and will reach the end value at 200K while 0.00005 will reach the end at 60K iter\n single_jitter: bool = False,\n ) -> None:\n super().__init__()\n self.num_samples_interval = num_samples_interval\n self.num_samples_outside = num_samples_outside\n self.num_samples_importance = num_samples_importance\n self.num_marching_steps = num_marching_steps\n self.num_secant_steps = num_secant_steps\n self.interval_start = interval_start\n self.interval_end = interval_end\n self.interval_decay = interval_decay\n self.single_jitter = single_jitter\n\n # samplers\n self.uniform_sampler = UniformSampler(single_jitter=single_jitter)\n self.outside_sampler = UniformSampler(single_jitter=single_jitter)\n self.pdf_sampler = PDFSampler(\n include_original=False,\n single_jitter=single_jitter,\n histogram_padding=1e-5,\n )\n # for merge samples\n self.error_bounded_sampler = ErrorBoundedSampler()\n\n # step counter\n self._step = 0\n self.delta = self.interval_start\n\n def step_cb(self, step):\n \"\"\"Callback to register a training step has passed. This is used to keep track of the sampling schedule\"\"\"\n self._step = step\n self.delta = max(self.interval_start * math.exp(-1 * self.interval_decay * self._step), self.interval_end)\n\n def generate_ray_samples(\n self,\n ray_bundle: Optional[RayBundle] = None,\n occupancy_fn: Optional[RayBundle] = None,\n sdf_fn: Optional[Callable] = None,\n return_surface_points: bool = False,\n ) -> Union[Tuple[RaySamples, torch.Tensor], RaySamples]:\n assert ray_bundle is not None\n assert sdf_fn is not None\n\n # Start with uniform sampling\n ray_samples = self.uniform_sampler(ray_bundle, num_samples=self.num_marching_steps)\n\n with torch.no_grad():\n sdf = sdf_fn(ray_samples)\n\n # importance sampling\n occupancy = occupancy_fn(sdf)\n weights = ray_samples.get_weights_from_alphas(occupancy)\n\n importance_samples = self.pdf_sampler(\n ray_bundle,\n ray_samples,\n weights,\n num_samples=self.num_samples_importance,\n )\n\n # samples uniformly with near and far\n ray_samples_uniform_outside = self.outside_sampler(ray_bundle, num_samples=self.num_samples_outside)\n\n # merge\n ray_samples_uniform_importance, _ = self.error_bounded_sampler.merge_ray_samples(\n ray_bundle, importance_samples, ray_samples_uniform_outside\n )\n\n # surface points\n # Calculate if sign change occurred and concat 1 (no sign change) in\n # last dimension\n n_rays, n_samples = ray_samples.shape\n starts = ray_samples.frustums.starts\n sign_matrix = torch.cat(\n [torch.sign(sdf[:, :-1, 0] * sdf[:, 1:, 0]), torch.ones(n_rays, 1).to(sdf.device)], dim=-1\n )\n cost_matrix = sign_matrix * torch.arange(n_samples, 0, -1).float().to(sdf.device)\n\n # Get first sign change and mask for values where a.) a sign changed\n # occurred and b.) no a neg to pos sign change occurred (meaning from\n # inside surface to outside)\n values, indices = torch.min(cost_matrix, -1)\n mask_sign_change = values < 0\n mask_pos_to_neg = sdf[torch.arange(n_rays), indices, 0] > 0\n\n # Define mask where a valid depth value is found\n mask = mask_sign_change & mask_pos_to_neg\n\n # Get depth values and function values for the interval\n d_low = starts[torch.arange(n_rays), indices, 0][mask]\n v_low = sdf[torch.arange(n_rays), indices, 0][mask]\n\n indices = torch.clamp(indices + 1, max=n_samples - 1)\n d_high = starts[torch.arange(n_rays), indices, 0][mask]\n v_high = sdf[torch.arange(n_rays), indices, 0][mask]\n\n # TODO secant method\n # linear-interpolations\n z = (v_low * d_high - v_high * d_low) / (v_low - v_high)\n\n # make this simpler\n origins = ray_samples.frustums.origins[torch.arange(n_rays), indices, :][mask]\n directions = ray_samples.frustums.directions[torch.arange(n_rays), indices, :][mask]\n surface_points = origins + directions * z[..., None]\n\n if surface_points.shape[0] <= 0:\n surface_points = torch.rand((1024, 3), device=sdf.device) - 0.5\n\n # modify near and far values according current schedule\n nears, fars = ray_bundle.nears.clone(), ray_bundle.fars.clone()\n dists = fars - nears\n\n ray_bundle.nears[mask] = z[:, None] - dists[mask] * self.delta\n ray_bundle.fars[mask] = z[:, None] + dists[mask] * self.delta\n # min max bound\n ray_bundle.nears = torch.maximum(ray_bundle.nears, nears)\n ray_bundle.fars = torch.minimum(ray_bundle.fars, fars)\n\n # samples uniformly with new surface interval\n ray_samples_interval = self.uniform_sampler(ray_bundle, num_samples=self.num_samples_interval)\n\n # change back to original values\n ray_bundle.nears = nears\n ray_bundle.fars = fars\n\n # merge sampled points\n ray_samples = self.merge_ray_samples_in_eculidean(\n ray_bundle, ray_samples_interval, ray_samples_uniform_importance\n )\n\n if return_surface_points:\n return ray_samples, surface_points\n\n return ray_samples\n\n def merge_ray_samples_in_eculidean(\n self, ray_bundle: RayBundle, ray_samples_1: RaySamples, ray_samples_2: RaySamples\n ):\n \"\"\"Merge two set of ray samples and return sorted index which can be used to merge sdf values\n\n Args:\n ray_samples_1 : ray_samples to merge\n ray_samples_2 : ray_samples to merge\n \"\"\"\n starts_1 = ray_samples_1.spacing_to_euclidean_fn(ray_samples_1.spacing_starts[..., 0])\n starts_2 = ray_samples_2.spacing_to_euclidean_fn(ray_samples_2.spacing_starts[..., 0])\n\n end_1 = ray_samples_1.spacing_to_euclidean_fn(ray_samples_1.spacing_ends[:, -1:, 0])\n end_2 = ray_samples_2.spacing_to_euclidean_fn(ray_samples_2.spacing_ends[:, -1:, 0])\n\n end = torch.maximum(end_1, end_2)\n\n euclidean_bins, _ = torch.sort(torch.cat([starts_1, starts_2], -1), -1)\n\n euclidean_bins = torch.cat([euclidean_bins, end], dim=-1)\n\n # Stop gradients\n euclidean_bins = euclidean_bins.detach()\n\n # TODO convert euclidean bins to spacing bins\n bins = euclidean_bins\n\n ray_samples = ray_bundle.get_ray_samples(\n bin_starts=euclidean_bins[..., :-1, None],\n bin_ends=euclidean_bins[..., 1:, None],\n spacing_starts=bins[..., :-1, None],\n spacing_ends=bins[..., 1:, None],\n spacing_to_euclidean_fn=ray_samples_1.spacing_to_euclidean_fn,\n )\n\n return ray_samples\n\n def secant_method(\n self,\n ray_bundle: Optional[RayBundle] = None,\n sdf_fn: Optional[Callable] = None,\n ):\n \"\"\"run secant method to refine surface points\"\"\"\n raise NotImplementedError\n\n\nclass NeuralReconWSampler(Sampler):\n \"\"\"Voxel surface guided sampler in NeuralReconW.\"\"\"\n\n def __init__(\n self,\n aabb,\n coarse_binary_grid,\n coarse_resolution: int = 32,\n fine_resolution: int = 512,\n num_samples: int = 8,\n num_samples_importance: int = 16,\n num_samples_boundary: int = 10,\n steps_per_grid_update: int = 5000,\n local_rank: int = 0,\n single_jitter: bool = False,\n ) -> None:\n super().__init__()\n self.aabb = aabb\n self.coarse_resolution = coarse_resolution\n self.fine_resolution = fine_resolution\n self.num_samples = num_samples\n self.num_samples_importance = num_samples_importance\n self.num_samples_boundary = num_samples_boundary\n self.single_jitter = single_jitter\n self.steps_per_grid_update = steps_per_grid_update\n self.local_rank = local_rank\n\n # TODO remvoe 2.0 and create cube in the initialization\n self.grid_size = self.coarse_resolution\n self.voxel_size = 2.0 / self.grid_size\n self.fine_grid_size = self.fine_resolution // self.grid_size\n\n # for voxel guided sampling\n self.uniform_sampler = UniformSampler(single_jitter=False)\n\n # for surface guided sampling\n self.neus_sampler = NeuSSampler(\n num_samples=num_samples,\n num_samples_importance=num_samples_importance,\n num_samples_outside=0,\n num_upsample_steps=2,\n base_variance=512,\n )\n\n # for merge samples\n self.unisurf_sampler = UniSurfSampler()\n\n self.grid = nerfacc.OccupancyGrid(aabb.reshape(-1), resolution=self.coarse_resolution)\n self._binary = coarse_binary_grid.reshape(\n self.coarse_resolution, self.coarse_resolution, self.coarse_resolution\n ).contiguous()\n self._binary_fine = None\n\n self.init_grid_coordinate()\n\n def init_grid_coordinate(self):\n # fine grid coordinates in each coarse voxel\n offset = torch.linspace(-1.0, 1.0, self.fine_grid_size * 2 + 1)[1::2]\n x, y, z = torch.meshgrid(offset, offset, offset, indexing=\"ij\")\n fine_offset_cube = torch.stack([x, y, z], dim=-1).reshape(-1, 3) * self.voxel_size * 0.5\n\n # coarse grid coordinates\n offset = torch.linspace(-1.0 + self.voxel_size / 2.0, 1.0 - self.voxel_size / 2.0, self.grid_size)\n x, y, z = torch.meshgrid(offset, offset, offset, indexing=\"ij\")\n coarse_offset_cube = torch.stack([x, y, z], dim=-1).reshape(-1, 3)\n\n self.register_buffer(\"coarse_offset_cube\", coarse_offset_cube)\n self.register_buffer(\"fine_offset_cube\", fine_offset_cube)\n\n @torch.no_grad()\n def update_binary_grid(self, step, sdf_fn=None):\n assert sdf_fn is not None\n # bootstrap should needs longer if using only one gpus\n if step >= self.steps_per_grid_update and step % self.steps_per_grid_update == 0:\n device = self.coarse_offset_cube.device\n\n mask = torch.zeros((self.grid_size**3, self.fine_grid_size**3), dtype=torch.bool, device=device)\n\n occupied_voxel = self.coarse_offset_cube[self._binary.reshape(-1)]\n fine_voxel = occupied_voxel[:, None] + self.fine_offset_cube[None, :]\n\n fine_voxel = fine_voxel.reshape(-1, 3)\n # save_points(\"fine_voxel.ply\", fine_voxel.cpu().numpy())\n\n def evaluate(points):\n z = []\n for _, pnts in enumerate(torch.split(points, 100000, dim=0)):\n z.append(sdf_fn(pnts))\n z = torch.cat(z, axis=0)\n return z\n\n sdf = evaluate(fine_voxel)\n sdf = sdf.reshape(occupied_voxel.shape[0], self.fine_grid_size**3)\n sdf_mask = sdf <= 0.0\n mask[self._binary.reshape(-1)] = sdf_mask\n\n self._binary_fine = (\n mask.reshape([self.grid_size] * 3 + [self.fine_grid_size] * 3)\n .permute(0, 3, 1, 4, 2, 5)\n .reshape(self.fine_resolution, self.fine_resolution, self.fine_resolution)\n .contiguous()\n )\n\n # offset = torch.linspace(-1.0, 1.0, self.fine_resolution * 2 + 1, device=device)[1::2]\n # x, y, z = torch.meshgrid(offset, offset, offset, indexing=\"ij\")\n # grid_coord = torch.stack([x, y, z], dim=-1).reshape(-1, 3)\n # save_points(\"fine_voxel_valid.ply\", grid_coord[self._binary_fine.reshape(-1)].cpu().numpy())\n\n def generate_ray_samples(\n self,\n ray_bundle: Optional[RayBundle] = None,\n sdf_fn: Optional[Callable] = None,\n ) -> Union[Tuple[RaySamples, torch.Tensor], RaySamples]:\n assert ray_bundle is not None\n assert sdf_fn is not None\n\n # near and far from occupancy grids\n packed_info, _, t_starts, t_ends = nerfacc.cuda.ray_marching(\n ray_bundle.origins.contiguous(),\n ray_bundle.directions.contiguous(),\n ray_bundle.nears[:, 0].contiguous(),\n ray_bundle.fars[:, 0].contiguous(),\n self.grid.roi_aabb.contiguous(),\n self._binary.to(ray_bundle.origins.device),\n self.grid.contraction_type.to_cpp_version(),\n 1e-2, # large value for coarse voxels\n 0.0,\n )\n\n tt_starts = nerfacc.unpack_data(packed_info, t_starts)\n tt_ends = nerfacc.unpack_data(packed_info, t_ends)\n\n hit_grid = (tt_starts > 0).any(dim=1)[:, 0]\n if hit_grid.float().sum() > 0:\n ray_bundle.nears[hit_grid] = tt_starts[hit_grid][:, 0]\n ray_bundle.fars[hit_grid] = tt_ends[hit_grid].max(dim=1)[0]\n\n # sample uniformly with currently nears and far\n voxel_samples = self.uniform_sampler(ray_bundle, num_samples=self.num_samples_boundary)\n\n if self._binary_fine is not None:\n # near from finer occupancy grids\n packed_info, _, t_starts, t_ends = nerfacc.cuda.ray_marching(\n ray_bundle.origins.contiguous(),\n ray_bundle.directions.contiguous(),\n ray_bundle.nears[:, 0].contiguous(),\n ray_bundle.fars[:, 0].contiguous(),\n self.grid.roi_aabb.contiguous(),\n self._binary_fine,\n self.grid.contraction_type.to_cpp_version(),\n 1e-3, # small value for fine voxels\n 0.0,\n )\n\n tt_starts = nerfacc.unpack_data(packed_info, t_starts)\n\n # update with near and far\n hit_grid = (tt_starts > 0).any(dim=1)[:, 0]\n if hit_grid.float().sum() > 0:\n ray_bundle.nears[hit_grid] = tt_starts[hit_grid][:, 0] - 0.03\n ray_bundle.fars[hit_grid] = tt_starts[hit_grid][:, 0] + 0.03\n else:\n print(\"waring not intersection\")\n print(\"sampling around surfaces\")\n\n # surface guided sampling\n surface_samples = self.neus_sampler(ray_bundle, sdf_fn=sdf_fn)\n\n # merge samples\n ray_samples = self.unisurf_sampler.merge_ray_samples_in_eculidean(ray_bundle, surface_samples, voxel_samples)\n\n return ray_samples\n\n\nclass NeuSAccSampler(Sampler):\n \"\"\"Voxel surface guided sampler in NeuralReconW.\"\"\"\n\n def __init__(\n self,\n aabb,\n neus_sampler: NeuSSampler = None,\n resolution: int = 128,\n num_samples: int = 8,\n num_samples_importance: int = 16,\n num_samples_boundary: int = 10,\n steps_warpup: int = 2000,\n steps_per_grid_update: int = 1000,\n importance_sampling: bool = False,\n local_rank: int = 0,\n single_jitter: bool = False,\n ) -> None:\n super().__init__()\n self.aabb = aabb\n self.resolution = resolution\n self.num_samples = num_samples\n self.num_samples_importance = num_samples_importance\n self.num_samples_boundary = num_samples_boundary\n self.single_jitter = single_jitter\n self.importance_sampling = importance_sampling\n self.steps_warpup = steps_warpup\n self.steps_per_grid_update = steps_per_grid_update\n self.local_rank = local_rank\n self.step_size = 0.01 / 5.0\n self.alpha_thres = 0.001\n\n # only supports cubic bbox for now\n assert aabb[0, 0] == aabb[0, 1] and aabb[0, 0] == aabb[0, 2]\n assert aabb[1, 0] == aabb[1, 1] and aabb[1, 0] == aabb[1, 2]\n self.grid_size = self.resolution\n self.voxel_size = (aabb[1, 0] - aabb[0, 0]) / self.grid_size\n\n # nesu_sampler at the begining of training\n self.neus_sampler = neus_sampler\n\n self.grid = nerfacc.OccupancyGrid(aabb.reshape(-1), resolution=self.resolution)\n self.register_buffer(\"_binary\", torch.ones((self.grid_size, self.grid_size, self.grid_size), dtype=torch.bool))\n self.register_buffer(\"_update_counter\", torch.zeros(1, dtype=torch.int32))\n\n self.init_grid_coordinate()\n\n def init_grid_coordinate(self):\n # coarse grid coordinates\n aabb = self.aabb\n offset_x = torch.linspace(\n aabb[0, 0] + self.voxel_size / 2.0, aabb[1, 0] - self.voxel_size / 2.0, self.grid_size\n )\n offset_y = torch.linspace(\n aabb[0, 1] + self.voxel_size / 2.0, aabb[1, 1] - self.voxel_size / 2.0, self.grid_size\n )\n offset_z = torch.linspace(\n aabb[0, 2] + self.voxel_size / 2.0, aabb[1, 2] - self.voxel_size / 2.0, self.grid_size\n )\n x, y, z = torch.meshgrid(offset_x, offset_y, offset_z, indexing=\"ij\")\n cube_coordinate = torch.stack([x, y, z], dim=-1).reshape(-1, 3)\n\n self.register_buffer(\"cube_coordinate\", cube_coordinate)\n\n def update_step_size(self, step, inv_s=None):\n assert inv_s is not None\n inv_s = inv_s().item()\n self.step_size = 14.0 / inv_s / 16\n\n @torch.no_grad()\n def update_binary_grid(self, step, sdf_fn=None, inv_s=None):\n assert sdf_fn is not None\n assert inv_s is not None\n\n if step >= self.steps_warpup and step % self.steps_per_grid_update == 0:\n\n mask = self._binary.reshape(-1)\n # TODO voxels can't be recovered once it is pruned\n occupied_voxel = self.cube_coordinate[mask.reshape(-1)]\n\n # save_points(\"occupied_voxel.ply\", occupied_voxel.cpu().numpy())\n\n def evaluate(points):\n z = []\n for _, pnts in enumerate(torch.split(points, 100000, dim=0)):\n z.append(sdf_fn(pnts))\n z = torch.cat(z, axis=0)\n return z\n\n sdf = evaluate(occupied_voxel)\n\n # sdf_mask = sdf.abs() <= 0.04\n\n # use maximum bound for sdf value\n bound = self.voxel_size * (3**0.5) / 2.0\n sdf = sdf.abs()\n sdf = torch.maximum(sdf - bound, torch.zeros_like(sdf))\n\n estimated_next_sdf = sdf - self.step_size * 0.5\n estimated_prev_sdf = sdf + self.step_size * 0.5\n inv_s = inv_s()\n prev_cdf = torch.sigmoid(estimated_prev_sdf * inv_s)\n next_cdf = torch.sigmoid(estimated_next_sdf * inv_s)\n\n p = prev_cdf - next_cdf\n c = prev_cdf\n\n alpha = ((p + 1e-5) / (c + 1e-5)).clip(0.0, 1.0)\n sdf_mask = alpha > self.alpha_thres\n\n mask[mask.reshape(-1).clone()] = sdf_mask\n\n self._binary = mask.reshape([self.grid_size] * 3).contiguous()\n\n # save_points(\"voxel_valid.ply\", self.cube_coordinate[self._binary.reshape(-1)].cpu().numpy())\n\n # TODO do we need dilation\n # F.max_pool3d(M.float(), kernel_size=3, padding=1, stride=1).bool()\n # save_points(\"voxel_valid_dilated.ply\", self.cube_coordinate[self._binary.reshape(-1)].cpu().numpy())\n self._update_counter += 1\n\n def create_ray_samples_from_ray_indices(self, ray_bundle: RayBundle, ray_indices, t_starts, t_ends):\n rays_o = ray_bundle.origins[ray_indices]\n rays_d = ray_bundle.directions[ray_indices]\n camera_indices = ray_bundle.camera_indices[ray_indices]\n deltas = t_ends - t_starts\n\n frustums = Frustums(\n origins=rays_o, # [..., 1, 3]\n directions=rays_d, # [..., 1, 3]\n starts=t_starts, # [..., num_samples, 1]\n ends=t_ends, # [..., num_samples, 1]\n pixel_area=torch.ones_like(t_starts), # [..., 1, 1]\n )\n\n ray_samples = RaySamples(\n frustums=frustums,\n camera_indices=camera_indices, # [..., 1, 1]\n deltas=deltas, # [..., num_samples, 1]\n )\n return ray_samples\n\n @torch.no_grad()\n def generate_ray_samples(\n self,\n ray_bundle: Optional[RayBundle] = None,\n sdf_fn: Optional[Callable] = None,\n alpha_fn: Optional[Callable] = None,\n ) -> Union[Tuple[RaySamples, torch.Tensor], RaySamples]:\n assert ray_bundle is not None\n assert sdf_fn is not None\n\n # sampler with original neus Sampler?\n if self._update_counter.item() <= 0:\n return self.neus_sampler(ray_bundle, sdf_fn=sdf_fn)\n\n assert alpha_fn is not None\n\n # sampling from occupancy grids\n packed_info, ray_indices, t_starts, t_ends = nerfacc.cuda.ray_marching(\n ray_bundle.origins.contiguous(),\n ray_bundle.directions.contiguous(),\n ray_bundle.nears[:, 0].contiguous(),\n ray_bundle.fars[:, 0].contiguous(),\n self.grid.roi_aabb.contiguous(),\n self._binary,\n self.grid.contraction_type.to_cpp_version(),\n self.step_size, # TODO stepsize based on inv_s value?\n 0.0,\n )\n\n # create ray_samples with the intersection\n ray_indices = ray_indices.long()\n ray_samples = self.create_ray_samples_from_ray_indices(ray_bundle, ray_indices, t_starts, t_ends)\n\n if self.importance_sampling and ray_samples.shape[0] > 0:\n # save_points(\"first.ply\", ray_samples.frustums.get_start_positions().cpu().numpy().reshape(-1, 3))\n\n alphas = alpha_fn(ray_samples)\n weights = nerfacc.render_weight_from_alpha(alphas, ray_indices=ray_indices, n_rays=ray_bundle.shape[0])\n\n # TODO make it configurable\n # re sample\n packed_info, t_starts, t_ends = nerfacc.ray_resampling(packed_info, t_starts, t_ends, weights[:, 0], 16)\n ray_indices = nerfacc.unpack_info(packed_info, t_starts.shape[0])\n ray_samples = self.create_ray_samples_from_ray_indices(ray_bundle, ray_indices, t_starts, t_ends)\n\n # save_points(\"second.ply\", ray_samples.frustums.get_start_positions().cpu().numpy().reshape(-1, 3))\n\n return ray_samples, ray_indices\n\n\nclass RayTracingSampler(Sampler):\n \"\"\"Find sample points with ray tracing.\n This sampler only finds the intersection points, assuming the SDFField is already trained,\n thus is not suitable for training.\n Handling of rays with non-convergent sphere-tracing results (e.g., in IDR) is ommitted.\n \n TODO: accelarate ray marching (lots of marching steps are too small)\n - performing 64 st steps for each ray is too slow\n - cuda implementation\n - when the convergence criterion is met, additionally check if the volrend weights at the intersection point is above a threshold, if not, keep marching forward\n NOTE:\n - the sdf learned with VolRend might not be a valid sdf everywhere, especially near to the border of the fg bbox\n - a ray might immediately converge near to the border of the fg bbox\n - if the step_sizes are too consevative, some rays hitting surfaces (e.g., the ground plane) might not converge\n \"\"\"\n def __init__(\n self,\n sdf_threshold: float = 1.0e-4, # TODO: tune the termination threshold\n mid_points_criteria: bool = True,\n mid_points_multiplier: float = 5.0, # multiply sdf_threshold with this scaler if using the center of two consecutive marching points as convergence criterion\n first_step_size: float = 1.0,\n conservative_step_size: float = 0.5, # use < 1.0 for conservative marching, since the sdf_field learnt with VolRend might not be a valid sdf everywhere\n line_search_step: float = 0.5, # IDR: backward stepping of overly-marched steps\n line_step_iters: int = 3, # IDR: backward stepping of overly-marched steps\n sphere_tracing_iters: int = 64, # TODO: tune maximum number of st iters\n conservative_iters: int = 32, # perform ray marching with convervative step size at the starting iters\n aabb: Optional[TensorType[2, 3]] = None,\n start_aabb_ratio: Optional[float] = 0.0075 # TODO: tune start marching from the ray-aabb-intersection-point + aabb_side_len * start_aabb_ratio to avoid bad sdf values near to the aabb boundary\n ) -> None:\n super().__init__()\n self.sdf_threshold = sdf_threshold\n self.mid_points_criteria = mid_points_criteria\n self.mid_pts_sdf_threshold = sdf_threshold * mid_points_multiplier\n self.first_step_size = first_step_size\n self.conservative_step_size = conservative_step_size\n self.line_search_step = line_search_step\n self.line_step_iters = line_step_iters\n self.sphere_tracing_iters = sphere_tracing_iters\n self.conservative_iters = conservative_iters\n self.aabb = nn.Parameter(aabb, requires_grad=False) if aabb is not None else None\n self.start_aabb_ratio = start_aabb_ratio\n \n if aabb is not None:\n # assert (aabb[0].abs() == aabb[1].abs()).all(), \"only support symmetric aabb for now\"\n extra_start_offset = (aabb[1] - aabb[0]) * start_aabb_ratio\n new_aabb = torch.stack([aabb[0] + extra_start_offset, aabb[1] - extra_start_offset], dim=0)\n self.extra_start_offset = nn.Parameter(extra_start_offset, requires_grad=False)\n self.new_aabb = nn.Parameter(new_aabb, requires_grad=False)\n self.axes = nn.Parameter(torch.eye(3, dtype=aabb.dtype), requires_grad=False)\n \n @torch.no_grad()\n def generate_ray_samples(\n self,\n ray_bundle: RayBundle,\n sdf_fn: Callable # (..., 3) -> (..., 1)\n ) -> Tuple[RaySamples, TensorType[\"bs\": ..., 1, 1]]:\n ray_bundle = self.adjust_aabb_intersection(ray_bundle)\n \n # sphere tracing from near to far\n curr_pts, acc_dist, non_convergent_masks = self.sphere_tracing(ray_bundle, sdf_fn) # (..., c)\n # compute non-convergent masks (assign a arbitrary intersection point, use zeros for now)\n # donnot update the original no_intersect_mask since volrend depends on it\n \n # build RaySamples with a single point per ray (should we sample around the intersection point?)\n shaped_ray_bundle = ray_bundle[..., None]\n frustums = Frustums(\n origins=shaped_ray_bundle.origins, # [..., 1, 3]\n directions=shaped_ray_bundle.directions, # [..., 1, 3]\n starts=acc_dist[..., None, [0]], # [..., 1, 1]\n ends=acc_dist[..., None, [0]], # [..., 1, 1]\n pixel_area=shaped_ray_bundle.pixel_area, # [..., 1, 1]\n )\n ray_samples = RaySamples(\n frustums=frustums,\n camera_indices=shaped_ray_bundle.camera_indices,\n metadata=shaped_ray_bundle.metadata,\n )\n return ray_samples, non_convergent_masks[..., None]\n \n def adjust_aabb_intersection(self, ray_bundle: RayBundle) -> RayBundle:\n \"\"\"#adjust the 1st ray-aabb intersection to avoid bad sdf around boundary\"\"\"\n if self.aabb is None:\n return ray_bundle\n \n aabb = self.aabb # (2, 3) assume symmetric\n assert (aabb[0].abs() == aabb[1].abs()).all(), \"only support symmetric aabb for now\"\n offset = self.extra_start_offset # (3, )\n ray_bundle = deepcopy(ray_bundle)\n ray_o, ray_d = ray_bundle.origins, ray_bundle.directions\n near, far = ray_bundle.nears, ray_bundle.fars\n intersect_mask = ~ray_bundle.metadata[\"no_intersect_mask\"][..., 0]\n ray_o, ray_d, near, far = map(lambda x: x[intersect_mask], [ray_o, ray_d, near, far])\n \n # compute offset along ray\n p = torch.addcmul(ray_o, ray_d, near)\n \n intersect_axis_idx = (p.abs() - aabb[1]).abs().argmin(dim=-1) # (n, )\n intersect_axis = self.axes[intersect_axis_idx] # (n, 3)\n _cos = ((ray_d * intersect_axis).sum(-1) / ray_d.norm(dim=-1)).abs() # (n, )\n ray_dist_offset = offset[intersect_axis_idx] / (_cos + 1e-8) # (n, )\n \n # some rays may not intersect with the new aabb\n new_p = torch.addcmul(ray_o, ray_d, near + ray_dist_offset[..., None])\n # _new_non_intersect = (new_p.abs() - self.new_aabb[1] > 1e-4).any(-1) # this causes some rays hitting the intersection of two plane to be mistakenly discarded\n _new_non_intersect = (new_p.abs() - aabb[1] > 1e-4).any(-1)\n non_intersect_mask = ~intersect_mask.clone()\n non_intersect_mask[intersect_mask] = _new_non_intersect\n \n # update ray_bundle\n ray_bundle.nears[intersect_mask] += ray_dist_offset[..., None]\n ray_bundle.metadata[\"no_intersect_mask\"] = non_intersect_mask[..., None]\n return ray_bundle\n \n def sphere_tracing(self, ray_bundle: RayBundle, sdf_fn: Callable):\n \"\"\"Perform iterative sphere tracing.\n ref: kaolin-wisp\n \n Returns:\n intersection_points: (..., 3)\n intersection_ray_dists: (..., 3)\n non_convergent_masks: (..., 1)\n \"\"\"\n prefix_dims = ray_bundle.origins.shape[:-1]\n rays_o, rays_d = ray_bundle.origins.view(-1, 3), ray_bundle.directions.view(-1, 3) # (n, 3)\n nears, fars = ray_bundle.nears.view(-1, 1), ray_bundle.fars.view(-1, 1) # (n, 1)\n zeros = torch.zeros_like(nears)\n _dtype = torch.float32 # sdf_fn might return float16\n \n # initialization\n if ray_bundle.metadata is not None and \"no_intersect_mask\" in ray_bundle.metadata:\n no_intersect_masks = ray_bundle.metadata[\"no_intersect_mask\"].view(-1)\n else: # there might not be \"no_intersect_mask\" (e.g., during mesh extraction)\n no_intersect_masks = torch.zeros_like(nears, dtype=torch.bool).view(-1)\n mask = (~no_intersect_masks).clone()[..., None] # (n, 1) non-convergent rays: only trace rays intersecting with the fg box\n x = torch.addcmul(rays_o, rays_d, nears) # (n, 3) current positions\n hit = torch.zeros_like(nears, dtype=torch.bool) # mask of rays hitting the surface\n acc_dist = torch.where(mask, nears, fars) # fars: maximum acc_dist\n if not mask.any(): # e.g., all ray samples are from the background\n return x, acc_dist, ~mask\n \n dist = zeros.clone() # current marching distance\n # the first step size should be especially conservative, since sdf near to the ray-aabb intersection might be invalid and causes a overshoot.\n dist[mask[..., 0]] = sdf_fn(x[mask[..., 0]]).to(_dtype) * self.first_step_size\n dist_prev = dist.clone()\n # __import__('pudb').set_trace() # 查看是否有初始点的sdf很接近0\n # iterative ray marching\n for i in range(self.sphere_tracing_iters):\n step_size = self.conservative_step_size if i < self.conservative_iters else 1.0\n \n # update non-convergent masks\n acc_dist += dist\n x = torch.where(mask, torch.addcmul(rays_o, rays_d, acc_dist), x) # (n, 3)\n hit = torch.where(mask, dist.abs() <= self.sdf_threshold, hit) # criteria-1\n if self.mid_points_criteria:\n hit |= torch.where(mask, ((dist + dist_prev) / 2).abs() <= self.mid_pts_sdf_threshold, hit) # criteria-2\n # TODO: for rays satisfying criteria-2, compute the accurate intersection pt w/ secant method or LinInterp\n mask = torch.where(mask, acc_dist < fars, mask) # criteria-3\n mask &= ~hit\n if not mask.any():\n break\n dist_prev = torch.where(mask, dist, dist_prev)\n \n # marching forward\n dist[~mask] = 0 # TODO: handle criteria-2 rays\n dist[mask[..., 0]] = sdf_fn(x[mask[..., 0]]).to(_dtype) * step_size\n \n # mask |= no_intersect_masks[..., None]\n non_hit_mask = ~hit | no_intersect_masks[..., None]\n \n # reshape with prefix_dims\n x = x.reshape(*prefix_dims, 3)\n acc_dist = acc_dist.reshape(*prefix_dims, 1)\n non_hit_mask = non_hit_mask.reshape(*prefix_dims, 1)\n return x, acc_dist, non_hit_mask\n \n def _sphere_tracing_idr(self, ray_bundle: RayBundle, sdf_fn: Callable):\n \"\"\"\n non_convergent_masks definition:\n 1. sphere-tracing is not convergent / intersection points falling beyond the far plane\n 2. 1 & forward tracing dist < backward tracing dist (IDR)\n \n Returns:\n intersection_points: (..., 3)\n intersection_ray_dists: (..., 3)\n non_convergent_masks: (..., 1)\n \"\"\"\n # TODO: share torch.zeros_like()\n prefix_dims = ray_bundle.origins.shape[:-1]\n origins, directions = ray_bundle.origins.view(-1, 3), ray_bundle.directions.view(-1, 3) # (n, 3)\n nears, fars = ray_bundle.nears.view(-1, 1), ray_bundle.fars.view(-1, 1) # (n, )\n \n forward_pts = origins + directions * nears # (n, 3)\n if ray_bundle.metadata is not None and \"no_intersect_mask\" in ray_bundle.metadata:\n no_intersect_masks = ray_bundle.metadata[\"no_intersect_mask\"].view(-1)\n else: # there might not be \"no_intersect_mask\" (e.g., during mesh extraction)\n no_intersect_masks = torch.zeros_like(nears, dtype=torch.bool).view(-1)\n non_convergent_masks = (~no_intersect_masks).clone() # (..., ) only trace rays intersecting with the fg box\n # initialize start positions with the 1st ray-aabb intersection\n curr_forward_pts = torch.where(non_convergent_masks[..., None],\n forward_pts, torch.zeros_like(forward_pts)) # (n, 3)\n acc_forward_dist = torch.where(non_convergent_masks[..., None],\n nears, torch.zeros_like(nears)) # (n, 1)\n if not non_convergent_masks.any():\n return curr_forward_pts, acc_forward_dist, non_convergent_masks\n \n # iterative sphere tracing\n i_sphere_tracing = 0\n next_sdf_forward = torch.zeros_like(acc_forward_dist)\n _sdf_dtype = next_sdf_forward.dtype\n next_sdf_forward[non_convergent_masks] = sdf_fn(curr_forward_pts[non_convergent_masks]).to(_sdf_dtype)\n curr_sdf_forward = next_sdf_forward.clone()\n while True:\n # TODO: when using the mid-pt as criteria, compute the intersection pt w/ LinInterp / Secant method\n curr_sdf_forward = torch.where(non_convergent_masks[..., None],\n next_sdf_forward, torch.zeros_like(next_sdf_forward))\n \n # update masks\n # non_convergent_masks = non_convergent_masks & (curr_sdf_forward > self.sdf_threshold)[..., 0]\n non_convergent_masks = non_convergent_masks & (curr_sdf_forward.abs() > self.sdf_threshold)[..., 0]\n \n # early stop\n if non_convergent_masks.sum() == 0 or i_sphere_tracing == self.sphere_tracing_iters:\n break\n i_sphere_tracing += 1\n \n # perform ray marching\n acc_forward_dist = acc_forward_dist + curr_sdf_forward # FIXME: the last curr_sdf_forward is not correct (manually set to zero)\n curr_forward_pts = origins + directions * acc_forward_dist\n \n # correct overshooting marching steps (pts falling inside the surface)\n next_sdf_forward = torch.zeros_like(acc_forward_dist)\n next_sdf_forward[non_convergent_masks] = sdf_fn(curr_forward_pts[non_convergent_masks]).to(_sdf_dtype)\n uncorrected_forward = (next_sdf_forward < 0)[..., 0]\n i_correction = 0\n # step backwards with increasing step size until the sdf is positive\n while uncorrected_forward.sum() > 0 and i_correction < self.line_step_iters:\n # step backward\n backward_step = ((1 - self.line_search_step) / (2**i_correction)) * curr_sdf_forward[uncorrected_forward]\n acc_forward_dist[uncorrected_forward] -= backward_step\n curr_forward_pts[uncorrected_forward] = (origins + directions * acc_forward_dist)[uncorrected_forward]\n # compute new sdf\n next_sdf_forward[uncorrected_forward] = sdf_fn(curr_forward_pts[uncorrected_forward]).to(_sdf_dtype)\n # update mask\n uncorrected_forward = (next_sdf_forward < 0)[..., 0]\n i_correction += 1\n \n # non_convergent_masks = non_convergent_masks & (acc_forward_dist < acc_backward_dist)[..., 0] # (n, 1)\n non_convergent_masks = non_convergent_masks | no_intersect_masks\n non_convergent_masks[(acc_forward_dist > fars)[..., 0]] = True\n \n # reshape with prefix_dims\n curr_forward_pts = curr_forward_pts.reshape(*prefix_dims, 3)\n acc_forward_dist = acc_forward_dist.reshape(*prefix_dims, 1)\n non_convergent_masks = non_convergent_masks.reshape(*prefix_dims, )\n return curr_forward_pts, acc_forward_dist, non_convergent_masks\n\n\nclass DummySampler(Sampler):\n \"\"\"A dummy sampler which just use the origins of ray_bundle as the ray samples \"\"\"\n def __init__(self):\n super().__init__()\n \n def generate_ray_samples(self, ray_bundle: RayBundle) -> RaySamples:\n shaped_ray_bundle = ray_bundle[..., None]\n origins = shaped_ray_bundle.origins # [..., 1, 3]\n zeros = torch.zeros_like(origins)\n \n # set zero starts and zero origins -> frustums.get_start_positions() = origins\n frustums = Frustums(\n origins=origins, # [..., 1, 3]\n directions=zeros, # [..., 1, 3]\n starts=zeros[..., [0]], # [..., 1, 1]\n ends=zeros[..., [0]],\n pixel_area=torch.ones_like(origins[..., [0]])\n )\n ray_samples = RaySamples(\n frustums=frustums,\n camera_indices=zeros[..., [0]],\n )\n return ray_samples\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/model_components/renderers.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nCollection of renderers\n\nExample:\n\n.. code-block:: python\n\n field_outputs = field(ray_sampler)\n weights = ray_sampler.get_weights(field_outputs[FieldHeadNames.DENSITY])\n\n rgb_renderer = RGBRenderer()\n rgb = rgb_renderer(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)\n\n\"\"\"\nimport math\nfrom typing import Optional, Union, Type\nfrom dataclasses import dataclass, field\n\nimport nerfacc\nimport torch\nfrom torch import nn\nfrom torchtyping import TensorType\nfrom typing_extensions import Literal\n\nfrom torch.nn import functional as F\n\nfrom nerfstudio.configs.base_config import InstantiateConfig\nfrom nerfstudio.cameras.rays import RaySamples\nfrom nerfstudio.utils.math import components_from_spherical_harmonics\nfrom nerfstudio.utils.pointclouds import BasicPointClouds\n\n\nclass RGBRenderer(nn.Module):\n \"\"\"Standard volumetic rendering.\n\n Args:\n background_color: Background color as RGB. Uses random colors if None.\n \"\"\"\n def __init__(\n self,\n background_color: Union[Literal[\"random\", \"last_sample\"], TensorType[3]] = \"random\"\n ) -> None:\n super().__init__()\n self.background_color = background_color\n\n @classmethod\n def combine_rgb(\n cls,\n rgb: TensorType[\"bs\":..., \"num_samples\", 3],\n weights: TensorType[\"bs\":..., \"num_samples\", 1],\n background_color: Union[Literal[\"random\", \"last_sample\"], TensorType[3]] = \"random\",\n ray_indices: Optional[TensorType[\"num_samples\"]] = None,\n num_rays: Optional[int] = None,\n ) -> TensorType[\"bs\":..., 3]:\n \"\"\"Composite samples along ray and render color image\n\n Args:\n rgb: RGB for each sample\n weights: Weights for each sample\n background_color: Background color as RGB.\n ray_indices: Ray index for each sample, used when samples are packed.\n num_rays: Number of rays, used when samples are packed.\n\n Returns:\n Outputs rgb values.\n \"\"\"\n if ray_indices is not None and num_rays is not None:\n # Necessary for packed samples from volumetric ray sampler\n if background_color == \"last_sample\":\n raise NotImplementedError(\n \"Background color 'last_sample' not implemented for packed samples.\"\n )\n comp_rgb = nerfacc.accumulate_along_rays(weights, ray_indices, rgb, num_rays)\n accumulated_weight = nerfacc.accumulate_along_rays(weights, ray_indices, None, num_rays)\n else:\n comp_rgb = torch.sum(weights * rgb, dim=-2)\n accumulated_weight = torch.sum(weights, dim=-2)\n\n if background_color == \"last_sample\":\n background_color = rgb[..., -1, :]\n if background_color == \"random\":\n background_color = torch.rand_like(comp_rgb).to(rgb.device)\n\n assert isinstance(background_color, torch.Tensor)\n comp_rgb = comp_rgb + background_color.to(weights.device) * (1.0 - accumulated_weight)\n\n return comp_rgb\n\n def forward(\n self,\n rgb: TensorType[\"bs\":..., \"num_samples\", 3],\n weights: TensorType[\"bs\":..., \"num_samples\", 1],\n ray_indices: Optional[TensorType[\"num_samples\"]] = None,\n num_rays: Optional[int] = None,\n ) -> TensorType[\"bs\":..., 3]:\n \"\"\"Composite samples along ray and render color image\n\n Args:\n rgb: RGB for each sample\n weights: Weights for each sample\n ray_indices: Ray index for each sample, used when samples are packed.\n num_rays: Number of rays, used when samples are packed.\n\n Returns:\n Outputs of rgb values.\n \"\"\"\n\n rgb = self.combine_rgb(\n rgb,\n weights,\n background_color=self.background_color,\n ray_indices=ray_indices,\n num_rays=num_rays\n )\n if not self.training:\n torch.clamp_(rgb, min=0.0, max=1.0)\n return rgb\n\n\nclass SHRenderer(nn.Module):\n \"\"\"Render RGB value from spherical harmonics.\n\n Args:\n background_color: Background color as RGB. Uses random colors if None\n activation: Output activation.\n \"\"\"\n def __init__(\n self,\n background_color: Union[Literal[\"random\", \"last_sample\"], TensorType[3]] = \"random\",\n activation: Optional[nn.Module] = nn.Sigmoid(),\n ) -> None:\n super().__init__()\n self.background_color = background_color\n self.activation = activation\n\n def forward(\n self,\n sh: TensorType[..., \"num_samples\", \"coeffs\"],\n directions: TensorType[..., \"num_samples\", 3],\n weights: TensorType[..., \"num_samples\", 1],\n ) -> TensorType[..., 3]:\n \"\"\"Composite samples along ray and render color image\n\n Args:\n sh: Spherical hamonics coefficients for each sample\n directions: Sample direction\n weights: Weights for each sample\n\n Returns:\n Outputs of rgb values.\n \"\"\"\n\n sh = sh.view(*sh.shape[:-1], 3, sh.shape[-1] // 3)\n\n levels = int(math.sqrt(sh.shape[-1]))\n components = components_from_spherical_harmonics(levels=levels, directions=directions)\n\n rgb = sh * components[..., None, :] # [..., num_samples, 3, sh_components]\n rgb = torch.sum(sh, dim=-1) + 0.5 # [..., num_samples, 3]\n\n if self.activation is not None:\n self.activation(rgb)\n\n rgb = RGBRenderer.combine_rgb(rgb, weights, background_color=self.background_color)\n\n return rgb\n\n\nclass AccumulationRenderer(nn.Module):\n \"\"\"Accumulated value along a ray.\"\"\"\n @classmethod\n def forward(\n cls,\n weights: TensorType[\"bs\":..., \"num_samples\", 1],\n ray_indices: Optional[TensorType[\"num_samples\"]] = None,\n num_rays: Optional[int] = None,\n ) -> TensorType[\"bs\":..., 1]:\n \"\"\"Composite samples along ray and calculate accumulation.\n\n Args:\n weights: Weights for each sample\n ray_indices: Ray index for each sample, used when samples are packed.\n num_rays: Number of rays, used when samples are packed.\n\n Returns:\n Outputs of accumulated values.\n \"\"\"\n\n if ray_indices is not None and num_rays is not None:\n # Necessary for packed samples from volumetric ray sampler\n accumulation = nerfacc.accumulate_along_rays(weights, ray_indices, None, num_rays)\n else:\n accumulation = torch.sum(weights, dim=-2)\n return accumulation\n\n\nclass DepthRenderer(nn.Module):\n \"\"\"Calculate depth along ray.\n\n Depth Method:\n - median: Depth is set to the distance where the accumulated weight reaches 0.5.\n - expected: Expected depth along ray. Same procedure as rendering rgb, but with depth.\n\n Args:\n method: Depth calculation method.\n \"\"\"\n def __init__(self, method: Literal[\"median\", \"expected\"] = \"median\") -> None:\n super().__init__()\n self.method = method\n\n def forward(\n self,\n weights: TensorType[..., \"num_samples\", 1],\n ray_samples: RaySamples,\n ray_indices: Optional[TensorType[\"num_samples\"]] = None,\n num_rays: Optional[int] = None,\n ) -> TensorType[..., 1]:\n \"\"\"Composite samples along ray and calculate depths.\n\n Args:\n weights: Weights for each sample.\n ray_samples: Set of ray samples.\n ray_indices: Ray index for each sample, used when samples are packed.\n num_rays: Number of rays, used when samples are packed.\n\n Returns:\n Outputs of depth values.\n \"\"\"\n\n if self.method == \"median\":\n steps = (ray_samples.frustums.starts + ray_samples.frustums.ends) / 2\n\n if ray_indices is not None and num_rays is not None:\n raise NotImplementedError(\n \"Median depth calculation is not implemented for packed samples.\"\n )\n cumulative_weights = torch.cumsum(weights[..., 0], dim=-1) # [..., num_samples]\n split = torch.ones((*weights.shape[:-2], 1), device=weights.device) * 0.5 # [..., 1]\n median_index = torch.searchsorted(cumulative_weights, split, side=\"left\") # [..., 1]\n median_index = torch.clamp(median_index, 0, steps.shape[-2] - 1) # [..., 1]\n median_depth = torch.gather(steps[..., 0], dim=-1, index=median_index) # [..., 1]\n return median_depth\n if self.method == \"expected\":\n eps = 1e-10\n steps = (ray_samples.frustums.starts + ray_samples.frustums.ends) / 2\n\n if ray_indices is not None and num_rays is not None:\n # Necessary for packed samples from volumetric ray sampler\n depth = nerfacc.accumulate_along_rays(weights, ray_indices, steps, num_rays)\n accumulation = nerfacc.accumulate_along_rays(weights, ray_indices, None, num_rays)\n depth = depth / (accumulation + eps)\n else:\n depth = torch.sum(weights * steps, dim=-2) / (torch.sum(weights, -2) + eps)\n\n depth = torch.clip(depth, steps.min(), steps.max())\n\n return depth\n\n raise NotImplementedError(f\"Method {self.method} not implemented\")\n\n\nclass UncertaintyRenderer(nn.Module):\n \"\"\"Calculate uncertainty along the ray.\"\"\"\n @classmethod\n def forward(\n cls, betas: TensorType[\"bs\":..., \"num_samples\", 1], weights: TensorType[\"bs\":...,\n \"num_samples\", 1]\n ) -> TensorType[\"bs\":..., 1]:\n \"\"\"Calculate uncertainty along the ray.\n\n Args:\n betas: Uncertainty betas for each sample.\n weights: Weights of each sample.\n\n Returns:\n Rendering of uncertainty.\n \"\"\"\n uncertainty = torch.sum(weights * betas, dim=-2)\n return uncertainty\n\n\nclass SemanticRenderer(nn.Module):\n \"\"\"Calculate semantics along the ray.\"\"\"\n @classmethod\n def forward(\n cls,\n semantics: TensorType[\"bs\":..., \"num_samples\", \"num_classes\"],\n weights: TensorType[\"bs\":..., \"num_samples\", 1],\n ) -> TensorType[\"bs\":..., \"num_classes\"]:\n \"\"\"Calculate semantics along the ray.\"\"\"\n sem = torch.sum(weights * semantics, dim=-2)\n return sem\n\n\nclass NormalsRenderer(nn.Module):\n \"\"\"Calculate normals along the ray.\"\"\"\n @classmethod\n def forward(\n cls,\n normals: TensorType[\"bs\":..., \"num_samples\", 3],\n weights: TensorType[\"bs\":..., \"num_samples\", 1],\n ) -> TensorType[\"bs\":..., 3]:\n \"\"\"Calculate normals along the ray.\"\"\"\n n = torch.sum(weights * normals, dim=-2)\n return n\n\n\nclass FeatureRenderer(nn.Module):\n \"\"\"Calculate feature along the ray.\"\"\"\n @classmethod\n def forward(\n cls,\n features: TensorType[\"bs\":..., \"num_samples\", \"nc_features\"],\n weights: TensorType[\"bs\":..., \"num_samples\", 1],\n ) -> TensorType[\"bs\":..., \"nc_features\"]:\n \"\"\"Calculate features along the ray.\"\"\"\n feat = torch.sum(weights * features, dim=-2)\n return feat\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/model_components/scene_colliders.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nScene Colliders\n\"\"\"\n\nfrom __future__ import annotations\n\nimport torch\nfrom torch import nn\nfrom torchtyping import TensorType\n\nfrom nerfstudio.cameras.rays import RayBundle\nfrom nerfstudio.data.scene_box import SceneBox\nfrom nerfstudio.utils import profiler\n\nclass SceneCollider(nn.Module):\n \"\"\"Module for setting near and far values for rays.\"\"\"\n\n def __init__(self, **kwargs) -> None:\n self.kwargs = kwargs\n super().__init__()\n\n def set_nears_and_fars(self, ray_bundle) -> RayBundle:\n \"\"\"To be implemented.\"\"\"\n raise NotImplementedError\n\n def forward(self, ray_bundle: RayBundle) -> RayBundle:\n \"\"\"Sets the nears and fars if they are not set already.\"\"\"\n if ray_bundle.nears is not None and ray_bundle.fars is not None:\n return ray_bundle\n return self.set_nears_and_fars(ray_bundle)\n\n\nclass AABBBoxCollider(SceneCollider):\n \"\"\"Module for colliding rays with the scene box to compute near and far values.\n\n Args:\n scene_box: scene box to apply to dataset\n \"\"\"\n\n def __init__(self, scene_box: SceneBox, near_plane: float = 0.0, **kwargs) -> None:\n super().__init__(**kwargs)\n self.scene_box = scene_box\n self.near_plane = near_plane\n\n def _intersect_with_aabb(\n self, rays_o: TensorType[\"num_rays\", 3], rays_d: TensorType[\"num_rays\", 3], aabb: TensorType[2, 3]\n ):\n \"\"\"Returns collection of valid rays within a specified near/far bounding box along with a mask\n specifying which rays are valid\n\n Args:\n rays_o: (num_rays, 3) ray origins\n rays_d: (num_rays, 3) ray directions\n aabb: (2, 3) This is [min point (x,y,z), max point (x,y,z)]\n \"\"\"\n # avoid divide by zero\n dir_fraction = 1.0 / (rays_d + 1e-6)\n\n # x\n t1 = (aabb[0, 0] - rays_o[:, 0:1]) * dir_fraction[:, 0:1]\n t2 = (aabb[1, 0] - rays_o[:, 0:1]) * dir_fraction[:, 0:1]\n # y\n t3 = (aabb[0, 1] - rays_o[:, 1:2]) * dir_fraction[:, 1:2]\n t4 = (aabb[1, 1] - rays_o[:, 1:2]) * dir_fraction[:, 1:2]\n # z\n t5 = (aabb[0, 2] - rays_o[:, 2:3]) * dir_fraction[:, 2:3]\n t6 = (aabb[1, 2] - rays_o[:, 2:3]) * dir_fraction[:, 2:3]\n\n nears = torch.max(\n torch.cat([torch.minimum(t1, t2), torch.minimum(t3, t4), torch.minimum(t5, t6)], dim=1), dim=1\n ).values\n fars = torch.min(\n torch.cat([torch.maximum(t1, t2), torch.maximum(t3, t4), torch.maximum(t5, t6)], dim=1), dim=1\n ).values\n\n # clamp to near plane\n near_plane = self.near_plane if self.training else 0\n nears = torch.clamp(nears, min=near_plane)\n fars = torch.maximum(fars, nears + 1e-6) # NOTE: this would cause numerical issue when sample along the ray\n\n return nears, fars\n\n def set_nears_and_fars(self, ray_bundle: RayBundle) -> RayBundle:\n \"\"\"Intersects the rays with the scene box and updates the near and far values.\n Populates nears and fars fields and returns the ray_bundle.\n\n Args:\n ray_bundle: specified ray bundle to operate on\n \"\"\"\n aabb = self.scene_box.aabb\n nears, fars = self._intersect_with_aabb(ray_bundle.origins, ray_bundle.directions, aabb)\n ray_bundle.nears = nears[..., None]\n ray_bundle.fars = fars[..., None]\n return ray_bundle\n\n\nclass NearFarCollider(SceneCollider):\n \"\"\"Sets the nears and fars with fixed values.\n\n Args:\n near_plane: distance to near plane\n far_plane: distance to far plane\n \"\"\"\n\n def __init__(self, near_plane: float, far_plane: float, **kwargs) -> None:\n self.near_plane = near_plane\n self.far_plane = far_plane\n super().__init__(**kwargs)\n\n def set_nears_and_fars(self, ray_bundle: RayBundle) -> RayBundle:\n ones = torch.ones_like(ray_bundle.origins[..., 0:1])\n near_plane = self.near_plane if self.training else 0\n ray_bundle.nears = ones * near_plane\n ray_bundle.fars = ones * self.far_plane\n return ray_bundle\n\n\nclass SphereCollider(SceneCollider):\n \"\"\"Sets the nears and fars with ray-sphere intersection.\n \"\"\"\n\n def __init__(self, radius: float = 1., near_plane: float = 0.1, far_plane: float = 1000.0, **kwargs) -> None:\n super().__init__(**kwargs)\n self.radius = radius\n self.near_plane = near_plane\n self.far_plane = far_plane\n\n def forward(self, ray_bundle: RayBundle) -> RayBundle:\n rr = self.radius ** 2\n ray_cam_dot = (ray_bundle.directions * ray_bundle.origins).sum(dim=-1, keepdims=True) # negative\n no_intersect_mask = ray_cam_dot >= 0\n dd = ray_bundle.origins.norm(p=2, dim=-1, keepdim=True) ** 2 - ray_cam_dot ** 2\n no_intersect_mask |= (dd >= rr)\n half_len_squared = rr - dd\n \n intersections = torch.sqrt(half_len_squared) * dd.new_tensor([-1, 1]) - ray_cam_dot\n nears, fars = intersections[:, [0]], intersections[:, [1]]\n nears = torch.where(no_intersect_mask, torch.full_like(nears, self.near_plane), nears)\n fars = torch.where(no_intersect_mask, torch.full_like(fars, self.far_plane), fars)\n \n ray_bundle.nears = nears\n ray_bundle.fars = fars\n if isinstance(ray_bundle.metadata, dict):\n ray_bundle.metadata[\"no_intersect_mask\"] = no_intersect_mask\n else:\n ray_bundle.metadata = {\"no_intersect_mask\": no_intersect_mask}\n \n return ray_bundle\n\n\nclass AABBBoxNearFarCollider(SceneCollider):\n \"\"\"Use AABBBoxCollider for rays intersecting the scene box and NearFarCollider otherwise.\n \n NOTE: the distance nears and fars inferred from ray-aabb intersection can be near, which cannot be handled properly using float32.\n \"\"\"\n \n def __init__(self, scene_box: SceneBox, near_plane: float = 0.1, far_plane: float = 1000.0, **kwargs) -> None:\n super().__init__(**kwargs)\n self.scene_box = scene_box\n self.near_plane = near_plane\n self.far_plane = far_plane\n \n def _intersect_with_aabb(\n self, rays_o: TensorType[\"num_rays\", 3], rays_d: TensorType[\"num_rays\", 3], aabb: TensorType[2, 3]\n ):\n \"\"\"Compute nears and fars for all rays. \"\"\"\n # avoid divide by zero\n dir_fraction = 1.0 / (rays_d + 1e-6)\n\n # x\n t1 = (aabb[0, 0] - rays_o[:, 0:1]) * dir_fraction[:, 0:1]\n t2 = (aabb[1, 0] - rays_o[:, 0:1]) * dir_fraction[:, 0:1]\n # y\n t3 = (aabb[0, 1] - rays_o[:, 1:2]) * dir_fraction[:, 1:2]\n t4 = (aabb[1, 1] - rays_o[:, 1:2]) * dir_fraction[:, 1:2]\n # z\n t5 = (aabb[0, 2] - rays_o[:, 2:3]) * dir_fraction[:, 2:3]\n t6 = (aabb[1, 2] - rays_o[:, 2:3]) * dir_fraction[:, 2:3]\n \n tx_min, tx_max = torch.minimum(t1, t2), torch.maximum(t1, t2)\n ty_min, ty_max = torch.minimum(t3, t4), torch.maximum(t3, t4)\n tz_min, tz_max = torch.minimum(t5, t6), torch.maximum(t5, t6)\n \n # no_intersect_mask = torch.logical_or(\n # torch.logical_or(tx_min > ty_max, tx_min > tz_max),\n # torch.logical_or(ty_min > tx_max, tz_min > tx_max)\n # )[..., 0]\n \n nears = torch.max(torch.cat([tx_min, ty_min, tz_min], dim=1), dim=1).values\n fars = torch.min(torch.cat([tx_max, ty_max, tz_max], dim=1), dim=1).values\n no_intersect_mask = torch.logical_or(nears - fars >= -1e-3, fars <= 0) # use nears >= fars might cause the nears and fars being too close, which finally leads to negative weights due to numerical issues.\n bottom_intersect_mask = (~no_intersect_mask) & (fars == t4[..., 0]) # fars == t4 & t4 == ty_max\n nears = torch.where(no_intersect_mask, torch.full_like(nears, self.near_plane), nears) # type: ignore\n fars = torch.where(no_intersect_mask, torch.full_like(fars, self.far_plane), fars) # type: ignore\n return nears, fars, no_intersect_mask, bottom_intersect_mask\n \n def set_nears_and_fars(self, ray_bundle: RayBundle) -> RayBundle:\n \"\"\"Intersects the rays with the scene box and updates the near and far values.\n Populates nears and fars fields and returns the ray_bundle.\n\n Args:\n ray_bundle: specified ray bundle to operate on\n \"\"\"\n # TODO: use NearFarCollider if not intersection found\n aabb = self.scene_box.aabb\n nears, fars, no_intersect_mask, bottom_intersect_mask = self._intersect_with_aabb(\n ray_bundle.origins, ray_bundle.directions, aabb\n )\n ray_bundle.nears = nears[..., None]\n ray_bundle.fars = fars[..., None]\n if isinstance(ray_bundle.metadata, dict):\n ray_bundle.metadata[\"no_intersect_mask\"] = no_intersect_mask[..., None]\n ray_bundle.metadata[\"bottom_intersect_mask\"] = bottom_intersect_mask[..., None]\n else:\n ray_bundle.metadata = {\n \"no_intersect_mask\": no_intersect_mask[..., None],\n \"bottom_intersect_mask\": bottom_intersect_mask[..., None],\n }\n return ray_bundle\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/models/__init__.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/models/base_model.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nBase Model implementation which takes in RayBundles\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom abc import abstractmethod\nfrom collections import defaultdict\nfrom dataclasses import dataclass, field\nfrom typing import Any, Dict, List, Optional, Tuple, Type\nfrom tqdm import tqdm\nfrom rich.console import Console\n\nimport torch\nfrom torch import nn\nfrom torch.nn import Parameter\n\nfrom nerfstudio.cameras.rays import RayBundle\nfrom nerfstudio.configs.base_config import InstantiateConfig\nfrom nerfstudio.configs.config_utils import to_immutable_dict\nfrom nerfstudio.data.scene_box import SceneBox\nfrom nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes\nfrom nerfstudio.model_components.scene_colliders import NearFarCollider\n\nCONSOLE = Console(width=120)\n\n# Model related configs\n@dataclass\nclass ModelConfig(InstantiateConfig):\n \"\"\"Configuration for model instantiation\"\"\"\n\n _target: Type = field(default_factory=lambda: Model)\n \"\"\"target class to instantiate\"\"\"\n enable_collider: bool = True\n \"\"\"Whether to create a scene collider to filter rays.\"\"\"\n collider_params: Optional[Dict[str, float]] = to_immutable_dict({\"near_plane\": 2.0, \"far_plane\": 6.0})\n \"\"\"parameters to instantiate scene collider with\"\"\"\n loss_coefficients: Dict[str, float] = to_immutable_dict({\"rgb_loss_coarse\": 1.0, \"rgb_loss_fine\": 1.0})\n \"\"\"parameters to instantiate density field with\"\"\"\n eval_num_rays_per_chunk: int = 4096\n \"\"\"specifies number of rays per chunk during eval\"\"\"\n\n\nclass Model(nn.Module):\n \"\"\"Model class\n Where everything (Fields, Optimizers, Samplers, Visualization, etc) is linked together. This should be\n subclassed for custom NeRF model.\n\n Args:\n config: configuration for instantiating model\n scene_box: dataset scene box\n \"\"\"\n\n config: ModelConfig\n\n def __init__(\n self,\n config: ModelConfig,\n scene_box: SceneBox,\n num_train_data: int,\n world_size: int = 1,\n local_rank: int = 0,\n **kwargs,\n ) -> None:\n super().__init__()\n self.config = config\n self.scene_box = scene_box\n self.num_train_data = num_train_data\n self.kwargs = kwargs\n self.collider = None\n self.world_size = world_size\n self.local_rank = local_rank\n\n self.populate_modules() # populate the modules\n self.callbacks = None\n # to keep track of which device the nn.Module is on\n self.device_indicator_param = nn.Parameter(torch.empty(0))\n\n @property\n def device(self):\n \"\"\"Returns the device that the model is on.\"\"\"\n return self.device_indicator_param.device\n\n def get_training_callbacks( # pylint:disable=no-self-use\n self, training_callback_attributes: TrainingCallbackAttributes # pylint: disable=unused-argument\n ) -> List[TrainingCallback]:\n \"\"\"Returns a list of callbacks that run functions at the specified training iterations.\"\"\"\n return []\n\n def populate_modules(self):\n \"\"\"Set the necessary modules to get the network working.\"\"\"\n # default instantiates optional modules that are common among many networks\n # NOTE: call `super().populate_modules()` in subclasses\n\n if self.config.enable_collider:\n self.collider = NearFarCollider(\n near_plane=self.config.collider_params[\"near_plane\"], far_plane=self.config.collider_params[\"far_plane\"]\n )\n\n @abstractmethod\n def get_param_groups(self) -> Dict[str, List[Parameter]]:\n \"\"\"Obtain the parameter groups for the optimizers\n\n Returns:\n Mapping of different parameter groups\n \"\"\"\n\n @abstractmethod\n def get_outputs(self, ray_bundle: RayBundle) -> Dict[str, torch.Tensor]:\n \"\"\"Takes in a Ray Bundle and returns a dictionary of outputs.\n\n Args:\n ray_bundle: Input bundle of rays. This raybundle should have all the\n needed information to compute the outputs.\n\n Returns:\n Outputs of model. (ie. rendered colors)\n \"\"\"\n\n def forward(self, ray_bundle: RayBundle) -> Dict[str, torch.Tensor]:\n \"\"\"Run forward starting with a ray bundle. This outputs different things depending on the configuration\n of the model and whether or not the batch is provided (whether or not we are training basically)\n\n Args:\n ray_bundle: containing all the information needed to render that ray latents included\n \"\"\"\n\n if self.collider is not None:\n ray_bundle = self.collider(ray_bundle)\n\n return self.get_outputs(ray_bundle)\n\n def get_metrics_dict(self, outputs, batch) -> Dict[str, torch.Tensor]:\n \"\"\"Compute and returns metrics.\n\n Args:\n outputs: the output to compute loss dict to\n batch: ground truth batch corresponding to outputs\n \"\"\"\n # pylint: disable=unused-argument\n # pylint: disable=no-self-use\n return {}\n\n @abstractmethod\n def get_loss_dict(self, outputs, batch, metrics_dict=None, **kwargs) -> Dict[str, torch.Tensor]:\n \"\"\"Computes and returns the losses dict.\n\n Args:\n outputs: the output to compute loss dict to\n batch: ground truth batch corresponding to outputs\n metrics_dict: dictionary of metrics, some of which we can use for loss\n \"\"\"\n\n @torch.no_grad()\n def get_outputs_for_camera_ray_bundle(self, camera_ray_bundle: RayBundle) -> Dict[str, torch.Tensor]:\n \"\"\"Takes in camera parameters and computes the output of the model.\n\n Args:\n camera_ray_bundle: ray bundle to calculate outputs over\n \"\"\"\n num_rays_per_chunk = self.config.eval_num_rays_per_chunk\n image_height, image_width = camera_ray_bundle.origins.shape[:2]\n num_rays = len(camera_ray_bundle)\n outputs_lists = defaultdict(list)\n for i in tqdm(range(0, num_rays, num_rays_per_chunk), desc='Rendering ray chunks'):\n # TODO: ignore some keys / only store specific keys\n start_idx = i\n end_idx = i + num_rays_per_chunk\n ray_bundle = camera_ray_bundle.get_row_major_sliced_ray_bundle(start_idx, end_idx)\n outputs = self.forward(ray_bundle=ray_bundle)\n for output_name, output in outputs.items(): # type: ignore\n outputs_lists[output_name].append(output)\n outputs = {}\n for output_name, outputs_list in outputs_lists.items():\n if not torch.is_tensor(outputs_list[0]):\n # TODO: handle lists of tensors as well\n continue\n try:\n outputs[output_name] = torch.cat(outputs_list).view(image_height, image_width, -1) # type: ignore\n except RuntimeError as err:\n CONSOLE.print(f\"[bold red] Incompatible output shape: {output_name}: {torch.cat(outputs_list).shape}\")\n raise err\n return outputs\n\n @abstractmethod\n def get_image_metrics_and_images(\n self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]\n ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:\n \"\"\"Writes the test image outputs.\n TODO: This shouldn't return a loss\n\n Args:\n image_idx: Index of the image.\n step: Current step.\n batch: Batch of data.\n outputs: Outputs of the model.\n\n Returns:\n A dictionary of metrics.\n \"\"\"\n\n def load_model(self, loaded_state: Dict[str, Any]) -> None:\n \"\"\"Load the checkpoint from the given path\n\n Args:\n loaded_state: dictionary of pre-trained model states\n \"\"\"\n state = {key.replace(\"module.\", \"\"): value for key, value in loaded_state[\"model\"].items()}\n self.load_state_dict(state) # type: ignore\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/models/base_surface_model.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nImplementation of Base surface model.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom abc import abstractmethod\nfrom dataclasses import dataclass, field\nfrom functools import partial\nfrom typing import Any, Dict, List, Optional, Tuple, Type\n\nimport numpy as np\nimport torch\nimport torch.nn.functional as F\nfrom einops import repeat\nfrom rich.console import Console\nfrom torch.nn import ModuleList, Parameter\nfrom torchmetrics import PeakSignalNoiseRatio\nfrom torchmetrics.functional import structural_similarity_index_measure\nfrom torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity\nfrom torchtyping import TensorType\nfrom typing_extensions import Literal\n\nfrom nerfstudio.cameras.rays import RayBundle\nfrom nerfstudio.engine.callbacks import (\n TrainingCallback,\n TrainingCallbackAttributes,\n TrainingCallbackLocation,\n)\nfrom nerfstudio.field_components.encodings import NeRFEncoding\nfrom nerfstudio.field_components.field_heads import FieldHeadNames\nfrom nerfstudio.field_components.spatial_distortions import (\n ForegroundAwareSceneContraction,\n SceneContraction,\n)\nfrom nerfstudio.fields.density_fields import HashMLPDensityField\nfrom nerfstudio.fields.nerfacto_field import TCNNNerfactoField\nfrom nerfstudio.fields.sdf_field import SDFFieldConfig\nfrom nerfstudio.fields.vanilla_nerf_field import NeRFField\nfrom nerfstudio.model_components.losses import (\n L1Loss,\n MSELoss,\n MultiViewLoss,\n ScaleAndShiftInvariantLoss,\n SensorDepthLoss,\n compute_scale_and_shift,\n interlevel_loss,\n monosdf_normal_loss,\n)\nfrom nerfstudio.model_components.patch_warping import PatchWarping\nfrom nerfstudio.model_components.ray_samplers import (\n LinearDisparitySampler,\n ProposalNetworkSampler,\n UniformLinDispPiecewiseSampler,\n)\nfrom nerfstudio.model_components.renderers import (\n AccumulationRenderer,\n DepthRenderer,\n RGBRenderer,\n SemanticRenderer,\n)\nfrom nerfstudio.model_components.scene_colliders import (\n AABBBoxCollider,\n AABBBoxNearFarCollider,\n NearFarCollider,\n SphereCollider,\n)\nfrom nerfstudio.models.base_model import Model, ModelConfig\nfrom nerfstudio.utils import colormaps, profiler, scheduler, writer\nfrom nerfstudio.utils.colors import get_color\nfrom nerfstudio.utils.pointclouds import BasicPointClouds\n\nCONSOLE = Console(width=120)\n\n\ndef map_range_val(input_val, input_start, input_end, output_start, output_end, type='linear', factor=1.):\n input_clamped=max(input_start, min(input_end, input_val))\n input_relative = (input_clamped - input_start) / (input_end - input_start)\n\n if type == 'linear':\n return output_start + (output_end - output_start) * input_relative\n elif type == 'exp':\n return output_start + (output_end - output_start) * (np.exp(factor*input_relative) - 1) / (np.exp(factor) - 1)\n elif type == 'log':\n return output_start + (output_end - output_start) * np.log1p(factor*input_relative) / np.log(1+factor) # log base 2\n else:\n raise NotImplementedError\n\n@dataclass\nclass SurfaceModelConfig(ModelConfig):\n \"\"\"Nerfacto Model Config\"\"\"\n\n _target: Type = field(default_factory=lambda: SurfaceModel)\n near_plane: float = 0.05\n \"\"\"How far along the ray to start sampling.\"\"\"\n far_plane: float = 4.0\n \"\"\"How far along the ray to stop sampling.\"\"\"\n far_plane_bg: float = 1000.0\n \"\"\"How far along the ray to stop sampling of the background model.\"\"\"\n background_color: Literal[\"random\", \"last_sample\", \"white\", \"black\"] = \"black\"\n \"\"\"Whether to randomize the background color.\"\"\"\n # fg_use_appearance_embedding: bool = False # NOTE: directly set \"sdf_field.use_appearance_embedding\" instead\n # \"\"\"Whether to use per-image appearance embedding in the fg model.\"\"\"\n bg_use_appearance_embedding: bool = True\n \"\"\"Whether to use per-image appearance embedding in the bg model.\"\"\"\n bg_sampler_type: Literal[\"uniform_lin_disp\", \"lin_disp\", \"proposal_network\", \"none\"] = \"lin_disp\"\n use_proposal_weight_anneal_bg: bool = True\n use_average_appearance_embedding: bool = (\n True # FIXME: False use zero appearance embedding during inference, which leads to bad bg rendering.\n )\n \"\"\"Whether to use average appearance embedding or zeros for inference.\"\"\"\n eikonal_loss_mult: float = 0.1\n \"\"\"Monocular normal consistency loss multiplier.\"\"\"\n fg_mask_loss_mult: float = 0.01\n \"\"\"Foreground mask loss multiplier.\"\"\"\n mono_normal_loss_mult: float = 0.002\n \"\"\"Monocular normal consistency loss multiplier.\"\"\"\n mono_depth_loss_mult: float = 0.0\n \"\"\"Monocular depth consistency loss multiplier.\"\"\"\n patch_warp_loss_mult: float = 0.0\n \"\"\"Multi-view consistency warping loss multiplier.\"\"\"\n patch_size: int = 11\n \"\"\"Multi-view consistency warping loss patch size.\"\"\"\n patch_warp_angle_thres: float = 0.3\n \"\"\"Threshold for valid homograph of multi-view consistency warping loss\"\"\"\n min_patch_variance: float = 0.01\n \"\"\"Threshold for minimal patch variance\"\"\"\n topk: int = 4\n \"\"\"Number of minimal patch consistency selected for training\"\"\"\n sensor_depth_truncation: float = 0.015\n \"\"\"Sensor depth trunction, default value is 0.015 which means 5cm with a rough scale value 0.3 (0.015 = 0.05 * 0.3)\"\"\"\n sensor_depth_l1_loss_mult: float = 0.0\n \"\"\"Sensor depth L1 loss multiplier.\"\"\"\n sensor_depth_freespace_loss_mult: float = 0.0\n \"\"\"Sensor depth free space loss multiplier.\"\"\"\n sensor_depth_sdf_loss_mult: float = 0.0\n \"\"\"Sensor depth sdf loss multiplier.\"\"\"\n sparse_points_sdf_loss_mult: float = 0.0\n \"\"\"sparse point sdf loss multiplier\"\"\"\n sdf_field: SDFFieldConfig = SDFFieldConfig()\n \"\"\"Config for SDF Field\"\"\"\n background_model: Literal[\"grid\", \"mlp\", \"none\"] = \"mlp\"\n \"\"\"background models\"\"\"\n render_background: bool = True\n \"\"\"optionally disable bg rendering during testing\"\"\"\n bg_hash_grid_num_levels: int = 16\n \"\"\"num_levels for background hash grid\"\"\"\n num_samples_outside: int = 32\n \"\"\"Number of samples outside the bounding sphere for backgound\"\"\"\n num_bg_proposal_samples_per_ray: Tuple[int, int] = (128, 64)\n handle_no_intersection_rays: bool = True\n \"\"\"whether to separately handle rays not intersecting the scene, usefull only when using AABBBoxNearFarCollider\"\"\"\n handle_bottom_intersection_rays: bool = False\n \"\"\"if a ray intersects the bottom plane of the fg aabb, ignore the querying of the bg model\"\"\"\n periodic_tvl_mult: float = 0.0\n \"\"\"Total variational loss mutliplier\"\"\"\n overwrite_near_far_plane: bool = False\n \"\"\"whether to use near and far collider from command line\"\"\"\n scene_contraction_order: Literal[\"none\", \"inf\"] = \"inf\"\n \"\"\"norm of the scene contraction, use l2-norm if none\"\"\"\n scene_contraction_scale_factor: float = 1.0\n \"\"\"Scale positions by this factor before applying contraction.\n This is useful when the [-1, 1] region (no contraction) is too small and only covers a small portion of the scene.\n \"\"\"\n use_fg_aware_scene_contraction: bool = False\n \"\"\"Use ForegroundAwareSceneContraction, which makes better use of the bg feature grid\"\"\"\n fg_aware_scene_contraction_alpha: float = 5.0\n \"\"\"map 1 < ||x|| <= alpha region to [[-1, -1/alpha], [1/alpha, 1]], bigger alpha leaves more space linearly uncontracted\"\"\"\n curvature_loss_mult: float = 0.0\n \"\"\"curvature loss in HashSDF\"\"\"\n curvature_loss_n_iters: Optional[int] = None\n \"\"\"number of iterations to apply curvature loss\"\"\"\n ptcd_reg_fg_mult: float = 0.0\n \"\"\"regularization with fg pointclouds\"\"\"\n ptcd_reg_fg_n_pts_per_iter: int = 5000\n ptcd_reg_plane_mult: float = 0.0\n \"\"\"regularization with plane pointclouds\"\"\"\n ptcd_reg_plane_n_pts_per_iter: int = 5000\n ptcd_reg_n_iters: Optional[int] = None # if None, apply ptcd_reg for the entire training\n ptcd_reg_weight_anneal_fn: Literal[\"cosine\", \"exponential\", \"constant\"] = \"cosine\"\n mask_beta_prior_mult: float = 0.0\n \"\"\"beta prior regularization for the fg accumulation\"\"\"\n mask_beta_prior_n_iters: Optional[int] = None\n mask_beta_prior_anneal_fn: Literal[\"sine\", \"cosine\", \"exponential\", \"constant\"] = \"sine\"\n depth_renderer_mode: Literal[\"expected\", \"median\"] = \"expected\"\n mask_depth_with_acc: bool = False\n \n mono_normal_loss_warmup_steps: int = 5000\n mono_normal_loss_steps: int = 20000\n \"\"\" normal warmup steps \"\"\"\n\n # init_normal_cos_thresh: float = 0.1 \n normal_loss_final_scale: float = 5\n # normal_decay_weight: float = 0.99995\n \"\"\" adaptive normal loss \"\"\"\n\nclass SurfaceModel(Model):\n \"\"\"Base surface model\n\n Args:\n config: Base surface model configuration to instantiate model\n \"\"\"\n\n config: SurfaceModelConfig\n\n def __init__(self, *args, **kwargs):\n super().__init__(*args, **kwargs)\n self.handle_no_intersect_rays = self.config.handle_no_intersection_rays and isinstance(\n self.collider, (AABBBoxNearFarCollider, SphereCollider)\n )\n self.handle_bottom_intersect_rays = self.config.handle_bottom_intersection_rays and isinstance(\n self.collider, AABBBoxNearFarCollider\n )\n self._step = 0\n self.ptcd_reg_anneal_fn = self._build_anneal_fn(\n self.config.ptcd_reg_weight_anneal_fn, self.config.ptcd_reg_n_iters\n )\n if self.config.ptcd_reg_plane_mult > 0 and self.config.ptcd_reg_fg_mult == 0:\n CONSOLE.print(\n \"[bold red]Using plane-ptcd-reg w/o fg-ptcd-reg, this could degenerate fg reconstruction!\"\n ) # due to imbalanced pt sampling\n # TODO: build separate loss classes intergrating loss computation & weight scheduling\n\n def populate_modules(self):\n \"\"\"Set the fields and modules.\"\"\"\n super().populate_modules()\n\n self._build_scene_contraction()\n\n self.field = self.config.sdf_field.setup(\n aabb=self.scene_box.aabb,\n spatial_distortion=self.scene_contraction,\n num_images=self.num_train_data,\n use_average_appearance_embedding=self.config.use_average_appearance_embedding,\n compute_curvature_loss=self.config.curvature_loss_mult > 0,\n )\n\n self._build_collider()\n self._build_bg_model_and_sampler()\n self._build_renderers()\n\n # patch warping\n self.patch_warping = PatchWarping(\n patch_size=self.config.patch_size, valid_angle_thres=self.config.patch_warp_angle_thres\n )\n\n # ptcd regularization\n if self.config.ptcd_reg_fg_mult > 0 or self.config.ptcd_reg_plane_mult > 0:\n self.ptcd_data: BasicPointClouds = self.kwargs[\"metadata\"][\"ptcd_data\"]\n\n # losses\n self.rgb_loss = L1Loss()\n self.eikonal_loss = MSELoss()\n self.depth_loss = ScaleAndShiftInvariantLoss(alpha=0.5, scales=1)\n self.patch_loss = MultiViewLoss(\n patch_size=self.config.patch_size, topk=self.config.topk, min_patch_variance=self.config.min_patch_variance\n )\n self.sensor_depth_loss = SensorDepthLoss(truncation=self.config.sensor_depth_truncation)\n\n # metrics\n self.psnr = PeakSignalNoiseRatio(data_range=1.0)\n self.ssim = structural_similarity_index_measure\n self.lpips = LearnedPerceptualImagePatchSimilarity()\n\n def step_cb(self, step):\n self._step = step\n\n def _build_scene_contraction(self):\n if not self.config.use_fg_aware_scene_contraction:\n self.scene_contraction = SceneContraction(\n order=float(\"inf\") if self.config.scene_contraction_order == \"inf\" else None,\n scale_factor=self.config.scene_contraction_scale_factor,\n )\n else:\n alpha = self.config.fg_aware_scene_contraction_alpha\n CONSOLE.print(f\"[bold yellow]Using ForegroundAwareSceneContraction({alpha=})\") # pyright: ignore\n self.scene_contraction = ForegroundAwareSceneContraction(\n order=float(\"inf\") if self.config.scene_contraction_order == \"inf\" else None,\n alpha=alpha,\n aabb=self.scene_box.aabb,\n scale_factor=self.config.scene_contraction_scale_factor,\n )\n\n def get_param_groups(self) -> Dict[str, List[Parameter]]:\n param_groups = {}\n param_groups[\"fields\"] = list(self.field.parameters())\n if self.config.background_model != \"none\":\n param_groups[\"field_background\"] = list(self.field_background.parameters())\n else:\n param_groups[\"field_background\"] = list(self.field_background)\n if self.config.bg_sampler_type == \"proposal_network\":\n param_groups[\"proposal_networks_bg\"] = list(self.proposal_networks_bg.parameters())\n return param_groups\n\n def get_training_callbacks(\n self, training_callback_attributes: TrainingCallbackAttributes\n ) -> List[TrainingCallback]:\n callbacks = super().get_training_callbacks(training_callback_attributes)\n\n callbacks.append(\n TrainingCallback(\n where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],\n update_every_num_iters=1,\n func=self.step_cb,\n )\n )\n\n if self.config.bg_sampler_type == \"proposal_network\" and self.config.use_proposal_weight_anneal_bg:\n # anneal the weights of the proposal network before doing PDF sampling\n N = 1000 # TODO：make configurable and tune\n proposal_weights_anneal_slope = 10.0\n\n def set_anneal(step):\n # https://arxiv.org/pdf/2111.12077.pdf eq. 18\n train_frac = np.clip(step / N, 0, 1)\n bias = lambda x, b: (b * x) / ((b - 1) * x + 1)\n anneal = bias(train_frac, proposal_weights_anneal_slope)\n self.proposal_sampler_bg.set_anneal(anneal)\n\n callbacks.append(\n TrainingCallback(\n where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],\n update_every_num_iters=1,\n func=set_anneal,\n )\n )\n callbacks.append(\n TrainingCallback(\n where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],\n update_every_num_iters=1,\n func=self.proposal_sampler_bg.step_cb,\n )\n )\n\n return callbacks\n\n def _build_collider(self) -> None:\n # Collider\n if self.scene_box.collider_type == \"near_far\":\n self.collider = NearFarCollider(near_plane=self.scene_box.near, far_plane=self.scene_box.far)\n elif self.scene_box.collider_type == \"box\":\n self.collider = AABBBoxCollider(self.scene_box, near_plane=self.scene_box.near)\n elif self.scene_box.collider_type == \"box_near_far\":\n self.collider = AABBBoxNearFarCollider(\n self.scene_box, near_plane=self.scene_box.near, far_plane=self.scene_box.far\n )\n elif self.scene_box.collider_type == \"sphere\":\n self.collider = SphereCollider(radius=1.0, near_plane=self.scene_box.near, far_plane=self.scene_box.far)\n else:\n raise NotImplementedError\n\n # command line near and far has highest priority\n if self.config.overwrite_near_far_plane:\n self.collider = NearFarCollider(near_plane=self.config.near_plane, far_plane=self.config.far_plane)\n\n def _build_bg_model_and_sampler(self) -> None:\n # background model\n if self.config.background_model == \"grid\":\n self.field_background = TCNNNerfactoField(\n self.scene_box.aabb,\n spatial_distortion=self.scene_contraction,\n num_images=self.num_train_data,\n use_appearance_embedding=self.config.bg_use_appearance_embedding,\n use_average_appearance_embedding=self.config.use_average_appearance_embedding,\n num_levels=self.config.bg_hash_grid_num_levels,\n )\n elif self.config.background_model == \"mlp\":\n position_encoding = NeRFEncoding(\n in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=9.0, include_input=True\n )\n direction_encoding = NeRFEncoding(\n in_dim=3, num_frequencies=4, min_freq_exp=0.0, max_freq_exp=3.0, include_input=True\n )\n\n self.field_background = NeRFField(\n position_encoding=position_encoding,\n direction_encoding=direction_encoding,\n spatial_distortion=self.scene_contraction,\n spatial_normalization_region=\"full\",\n )\n else:\n # dummy background model\n self.field_background = Parameter(torch.ones(1), requires_grad=False)\n\n # background sampler\n if self.config.bg_sampler_type == \"lin_disp\":\n self.sampler_bg = LinearDisparitySampler(num_samples=self.config.num_samples_outside)\n elif self.config.bg_sampler_type == \"uniform_lin_disp\":\n self.sampler_bg = UniformLinDispPiecewiseSampler(num_samples=self.config.num_samples_outside)\n elif self.config.bg_sampler_type == \"proposal_network\":\n assert isinstance(\n self.scene_contraction, ForegroundAwareSceneContraction\n ) # otherwise, setup a separate contraction for bg proposal nets\n self.density_fns_bg = []\n self.proposal_networks_bg = ModuleList()\n # assume using two proposal network with fixed hparams\n prop_net_args = [\n {\"hidden_dim\": 16, \"log2_hashmap_size\": 17, \"num_levels\": 5, \"max_res\": 64},\n {\"hidden_dim\": 16, \"log2_hashmap_size\": 17, \"num_levels\": 5, \"max_res\": 256},\n ]\n proposal_warmup, proposal_update_every = 5000, 5\n num_proposal_samples_per_ray = self.config.num_bg_proposal_samples_per_ray\n for _args in prop_net_args:\n network = HashMLPDensityField(\n self.scene_box.aabb, # incorrect aabb as placeholder\n spatial_distortion=self.scene_contraction,\n spatial_normalization_region=\"full\",\n **_args,\n )\n self.proposal_networks_bg.append(network)\n self.density_fns_bg.append(network.density_fn)\n\n update_schedule = lambda step: np.clip(\n np.interp(step, [0, proposal_warmup], [0, proposal_update_every]),\n 1,\n self.config.proposal_update_every,\n )\n self.proposal_sampler_bg = ProposalNetworkSampler(\n num_nerf_samples_per_ray=self.config.num_samples_outside,\n num_proposal_samples_per_ray=num_proposal_samples_per_ray,\n num_proposal_network_iterations=2,\n use_uniform_sampler=False, # use UniformLinDispPiecewiseSampler\n single_jitter=True,\n update_sched=update_schedule,\n )\n elif self.config.bg_sampler_type == \"none\": # share sampler with fg\n self.sampler_bg = None\n else:\n raise ValueError(f\"Unknown background sampler: {self.config.bg_sampler_type}\")\n\n def _build_renderers(self) -> None:\n background_color = (\n get_color(self.config.background_color)\n if self.config.background_color in set([\"white\", \"black\"])\n else self.config.background_color\n )\n self.renderer_rgb = RGBRenderer(background_color=background_color)\n self.renderer_accumulation = AccumulationRenderer()\n self.renderer_depth = DepthRenderer(method=self.config.depth_renderer_mode)\n self.renderer_normal = SemanticRenderer()\n\n @abstractmethod\n def sample_and_forward_field(self, ray_bundle: RayBundle) -> Dict:\n \"\"\"_summary_\n ray points sampling and field inference of the foreground field.\n\n Args:\n ray_bundle (RayBundle): _description_\n return_samples (bool, optional): _description_. Defaults to False.\n \"\"\"\n\n # @profiler.time_function\n def get_outputs(self, ray_bundle: RayBundle) -> Dict:\n outputs, samples_and_field_outputs = self.fg_fields_query_and_render(ray_bundle)\n outputs, samples_and_field_outputs_bg = self.bg_fields_query_and_render(\n ray_bundle, samples_and_field_outputs, outputs\n )\n self._update_outptus_for_vis(outputs, samples_and_field_outputs, samples_and_field_outputs_bg)\n\n # for auxilary field rendering (used by .get_output() methods of child classes)\n outputs.update(\n {\n \"samples_and_field_outputs\": samples_and_field_outputs,\n \"samples_and_field_outputs_bg\": samples_and_field_outputs_bg,\n }\n )\n\n return outputs\n\n def forward(self, ray_bundle: RayBundle) -> Dict[str, torch.Tensor]:\n \"\"\"Run forward starting with a ray bundle.\n This outputs different things depending on the configuration of the model\n and whether or not the batch is provided (whether or not we are training basically)\n\n Args:\n ray_bundle: containing all the information needed to render that ray latents included\n \"\"\"\n if self.collider is not None:\n ray_bundle = self.collider(ray_bundle)\n\n outputs = self.get_outputs(ray_bundle)\n\n # remove unnecessary items to avoid memory leak\n for del_k in [\"samples_and_field_outputs\", \"samples_and_field_outputs_bg\"]:\n if del_k in outputs:\n del outputs[del_k]\n return outputs\n\n def fg_fields_query_and_render(self, ray_bundle: RayBundle) -> Dict[str, Any]:\n \"\"\"query foreground fields and rendering\"\"\"\n samples_and_field_outputs = self.sample_and_forward_field(ray_bundle=ray_bundle)\n\n if self.handle_no_intersect_rays:\n samples_and_field_outputs = self._handle_no_intersect_rays_fg(ray_bundle, samples_and_field_outputs)\n\n field_outputs = samples_and_field_outputs[\"field_outputs\"]\n ray_samples = samples_and_field_outputs[\"ray_samples\"]\n weights = samples_and_field_outputs[\"weights\"]\n bg_transmittance = samples_and_field_outputs[\"bg_transmittance\"]\n\n rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)\n ray_dist = self.renderer_depth(weights=weights, ray_samples=ray_samples)\n depth = ray_dist / ray_bundle.directions_norm\n\n accumulation = self.renderer_accumulation(weights=weights)\n\n if self.config.mask_depth_with_acc:\n depth[accumulation < 0.5] = 0.0 # for tsdf fusion to work properly\n\n outputs = {\n \"rgb\": rgb,\n \"accumulation\": accumulation,\n \"depth\": depth,\n \"weights\": weights,\n \"directions_norm\": ray_bundle.directions_norm, # used to scale z_vals for free space and sdf loss\n \"bg_transmittance\": bg_transmittance,\n }\n\n sample_normals = field_outputs.get(FieldHeadNames.NORMAL, None)\n if sample_normals is not None:\n normal = self.renderer_normal(semantics=sample_normals, weights=weights)\n outputs.update({\"normal\": normal})\n\n self._update_outputs_for_loss(ray_bundle, outputs, samples_and_field_outputs)\n\n return outputs, samples_and_field_outputs\n\n def bg_fields_query_and_render(\n self, ray_bundle: RayBundle, samples_and_field_outputs_fg: Dict[str, Any], outputs: Dict[str, Any]\n ) -> Dict[str, Any]:\n if self.config.background_model == \"none\" or not self.config.render_background:\n return outputs, {}\n\n ray_bundle = self._reset_near_far_for_background(ray_bundle)\n\n # bg field inference\n samples_and_field_outputs_bg = self.sample_and_forward_field_bg(\n ray_bundle, samples_and_field_outputs_fg, outputs\n )\n if self.handle_bottom_intersect_rays: # TODO: ignore bg inference of rays intersecting the bottom plane\n self._handle_bottom_intersect_rays_bg(ray_bundle, samples_and_field_outputs_bg)\n weights_bg = samples_and_field_outputs_bg[\"weights\"]\n ray_samples_bg = samples_and_field_outputs_bg[\"ray_samples\"]\n\n # render background\n rgb_bg = self.renderer_rgb(rgb=samples_and_field_outputs_bg[FieldHeadNames.RGB], weights=weights_bg)\n\n depth_bg = self.renderer_depth(weights=weights_bg, ray_samples=ray_samples_bg)\n accumulation_bg = self.renderer_accumulation(weights=weights_bg)\n\n # merge background color to foregound color\n rgb = outputs[\"rgb\"] + outputs[\"bg_transmittance\"] * rgb_bg\n # TODO: update depth with background depth\n\n outputs.update(\n {\n \"rgb\": rgb,\n \"bg_rgb\": rgb_bg,\n \"bg_accumulation\": accumulation_bg,\n \"bg_depth\": depth_bg,\n \"bg_weights\": weights_bg,\n }\n )\n return outputs, samples_and_field_outputs_bg\n\n def sample_and_forward_field_bg(\n self, ray_bundle: RayBundle, samples_and_field_outputs_fg: Dict[str, Any], outputs: Dict[str, Any]\n ):\n if self.config.bg_sampler_type == \"proposal_network\":\n return self._sample_and_forward_field_bg_proposal(ray_bundle, samples_and_field_outputs_fg, outputs)\n\n ray_samples_bg = self.sampler_bg(ray_bundle)\n field_outputs_bg = self.field_background(ray_samples_bg)\n weights_bg = ray_samples_bg.get_weights(field_outputs_bg[FieldHeadNames.DENSITY])\n\n samples_and_field_outputs_bg = {\"ray_samples\": ray_samples_bg, \"weights\": weights_bg, **field_outputs_bg}\n return samples_and_field_outputs_bg\n\n def _sample_and_forward_field_bg_proposal(\n self, ray_bundle: RayBundle, samples_and_field_outputs_fg: Dict[str, Any], outputs: Dict[str, Any]\n ):\n ray_samples_bg, weights_list_bg, ray_samples_list_bg = self.proposal_sampler_bg(\n ray_bundle, density_fns=self.density_fns_bg\n )\n field_outputs_bg = self.field_background(ray_samples_bg)\n weights_bg = ray_samples_bg.get_weights(field_outputs_bg[FieldHeadNames.DENSITY])\n # NOTE: weights_bg doesn't take bg_transmittance into account (i.e., weights_bg *= bg_transmittance)\n # since the proposal_network_bg is unaware of the bg_transmittance\n weights_list_bg.append(weights_bg.clone())\n ray_samples_list_bg.append(ray_samples_bg)\n\n samples_and_field_outputs_bg = {\n \"ray_samples\": ray_samples_bg,\n \"weights\": weights_bg,\n \"weights_list\": weights_list_bg,\n \"ray_samples_list\": ray_samples_list_bg,\n **field_outputs_bg,\n }\n outputs.update(\n {\n \"weights_list_bg\": weights_list_bg,\n \"ray_samples_list_bg\": ray_samples_list_bg,\n }\n )\n return samples_and_field_outputs_bg\n\n def _update_outputs_for_loss(self, ray_bundle: RayBundle, outputs: Dict, samples_and_field_outputs: Dict):\n \"\"\"update outputs for loss computation\"\"\"\n if self.training: # compute gradients for eikonal loss\n field_outputs = samples_and_field_outputs[\"field_outputs\"]\n if FieldHeadNames.GRADIENT in field_outputs:\n grad_points = field_outputs[FieldHeadNames.GRADIENT]\n outputs.update({\"eik_grad\": grad_points})\n\n # TODO volsdf use different point set for eikonal loss\n # grad_points = self.field.gradient(eik_points)\n # outputs.update({\"eik_grad\": grad_points})\n\n outputs.update({\"curvature_losses\": field_outputs.get(\"curvature_losses\", None)})\n\n outputs.update(samples_and_field_outputs)\n\n if self.handle_no_intersect_rays: # ignore loss computation of no-intersection rays\n no_intersect_mask = ray_bundle.metadata[\"no_intersect_mask\"] # [..., 1]\n outputs.update({\"no_intersect_mask\": no_intersect_mask})\n if self.handle_bottom_intersect_rays: # ignore loss computation of rays intersecting the bottom plane\n outputs.update({\"bottom_intersect_mask\": ray_bundle.metadata[\"bottom_intersect_mask\"]})\n\n def _update_outptus_for_vis(\n self, outputs: Dict, samples_and_field_outputs: Dict, samples_and_field_outputs_bg: Dict\n ) -> None:\n # proposal network depths\n self._compute_prop_net_depths(outputs, samples_and_field_outputs)\n self._compute_prop_net_depths(outputs, samples_and_field_outputs_bg, ouptut_key_prefix=\"prop_depth_bg\")\n\n # accumulation -> mask\n outputs[\"mask\"] = outputs[\"accumulation\"] > 0.5\n\n def _compute_prop_net_depths(\n self, outputs: Dict, samples_and_field_outputs: Dict, ouptut_key_prefix: str = \"prop_depth\"\n ) -> None:\n \"\"\"compute depths of proposal networks for visualization\"\"\"\n # TODO how can we move it to neus_facto without out of memory\n if \"weights_list\" not in samples_and_field_outputs:\n return\n\n weights_list = samples_and_field_outputs[\"weights_list\"]\n ray_samples_list = samples_and_field_outputs[\"ray_samples_list\"]\n for i in range(len(weights_list) - 1):\n outputs[f\"{ouptut_key_prefix}_{i}\"] = self.renderer_depth(\n weights=weights_list[i], ray_samples=ray_samples_list[i]\n )\n\n def _handle_no_intersect_rays_fg(self, ray_bundle: RayBundle, samples_and_field_outputs: Dict) -> Dict:\n \"\"\"Handle rays not intersecting the foreground AABB.\n Assign 0 weights for all fg queries and 1 for bg_transmittance to these rays.\n\n Returns:\n samples_and_field_outputs: with zero weights for the rays not intersecting the foreground AABB.\n \"\"\"\n weights = samples_and_field_outputs[\"weights\"] # [..., num_samples, 1]\n bg_transmittance = samples_and_field_outputs[\"bg_transmittance\"] # [..., 1]\n # ray_samples_list = samples_and_field_outputs[\"ray_samples_list\"]\n num_samples = weights.shape[-2]\n no_intersect_mask = ray_bundle.metadata[\"no_intersect_mask\"] # [..., 1]\n\n bg_transmittance = torch.where(no_intersect_mask, torch.ones_like(bg_transmittance), bg_transmittance)\n no_intersect_mask = repeat(no_intersect_mask, \"... m -> ... n m\", n=num_samples)\n weights = torch.where(no_intersect_mask, torch.zeros_like(weights), weights)\n new_outputs = {\"weights\": weights, \"bg_transmittance\": bg_transmittance}\n\n if \"weights_list\" in samples_and_field_outputs:\n # TODO: delete this block since this might be meaningless, bacause rays without intersection are ignored\n # when computing the proposal loss\n weights_list = samples_and_field_outputs[\"weights_list\"]\n weights_list.pop()\n weights_list.append(weights)\n new_outputs[\"weights_list\"] = weights_list\n\n samples_and_field_outputs.update(new_outputs)\n return samples_and_field_outputs\n\n def _handle_bottom_intersect_rays_bg(self, ray_bundle: RayBundle, samples_and_field_outputs: Dict) -> Dict:\n \"\"\"ignore ray samples in bg regions if the ray intersects the bottom plane.\"\"\"\n weights = samples_and_field_outputs[\"weights\"]\n num_samples = weights.shape[-2]\n bottom_intersect_mask = ray_bundle.metadata[\"bottom_intersect_mask\"]\n bottom_intersect_mask = repeat(bottom_intersect_mask, \"... m -> ... n m\", n=num_samples)\n weights = torch.where(bottom_intersect_mask, torch.zeros_like(weights), weights)\n samples_and_field_outputs.update({\"weights\": weights})\n\n def _reset_near_far_for_background(self, ray_bundle: RayBundle) -> RayBundle:\n if self.handle_no_intersect_rays:\n # NOTE: self.config.far_plane_bg is not used, in favor of self.collider.far_plane\n _nears, _fars = ray_bundle.nears, ray_bundle.fars\n no_intersect_mask = ray_bundle.metadata[\"no_intersect_mask\"]\n nears = torch.where(no_intersect_mask, _nears, _fars)\n fars = torch.where(no_intersect_mask, _fars, torch.full_like(_fars, self.collider.far_plane))\n ray_bundle.nears = nears\n ray_bundle.fars = fars\n else:\n ray_bundle.nears = ray_bundle.fars\n ray_bundle.fars = torch.ones_like(ray_bundle.fars) * self.config.far_plane_bg\n\n return ray_bundle\n\n def get_outputs_flexible(self, ray_bundle: RayBundle, additional_inputs: Dict[str, TensorType]) -> Dict:\n \"\"\"run the model with additional inputs such as warping or rendering from unseen rays\n Args:\n ray_bundle: containing all the information needed to render that ray latents included\n additional_inputs: addtional inputs such as images, src_idx, src_cameras\n\n Returns:\n dict: information needed for compute gradients\n \"\"\"\n if self.collider is not None:\n ray_bundle = self.collider(ray_bundle)\n\n outputs = self.get_outputs(ray_bundle)\n\n ray_samples = outputs[\"ray_samples\"]\n field_outputs = outputs[\"field_outputs\"]\n\n if self.config.patch_warp_loss_mult > 0:\n # patch warping\n warped_patches, valid_mask = self.patch_warping(\n ray_samples,\n field_outputs[FieldHeadNames.SDF],\n field_outputs[FieldHeadNames.NORMAL],\n additional_inputs[\"src_cameras\"],\n additional_inputs[\"src_imgs\"],\n pix_indices=additional_inputs[\"uv\"],\n )\n\n outputs.update({\"patches\": warped_patches, \"patches_valid_mask\": valid_mask})\n\n return outputs\n\n # @profiler.time_function\n def get_loss_dict(self, outputs, batch, metrics_dict=None, step=None) -> Dict:\n loss_dict = {}\n image = batch[\"image\"].to(self.device)\n loss_dict[\"rgb_loss\"] = self.rgb_loss(image, outputs[\"rgb\"])\n if self.training:\n # eikonal loss\n # TODO: additionally sample random points in the fg aabb and apply eikonal regularization\n # especially near to the border of the fg aabb.\n if \"eik_grad\" in outputs:\n grad_theta = outputs[\"eik_grad\"]\n # handle out-of-box points\n if self.handle_no_intersect_rays:\n no_intersect_mask = outputs[\"no_intersect_mask\"].to(grad_theta.dtype)\n n_intersect_pts = (no_intersect_mask.nelement() - no_intersect_mask.sum()) * np.prod(\n grad_theta.shape[1:-1]\n )\n eikonal_loss = (\n ((grad_theta.norm(2, dim=-1) - 1) ** 2) * (1 - no_intersect_mask)\n ).sum() / n_intersect_pts\n loss_dict[\"eikonal_loss\"] = eikonal_loss * self.config.eikonal_loss_mult\n else:\n loss_dict[\"eikonal_loss\"] = (\n (grad_theta.norm(2, dim=-1) - 1) ** 2\n ).mean() * self.config.eikonal_loss_mult\n\n # foreground mask loss\n if \"fg_mask\" in batch and self.config.fg_mask_loss_mult > 0.0:\n fg_label = batch[\"fg_mask\"].float().to(self.device)\n weights_sum = outputs[\"weights\"].sum(dim=1).clip(1e-3, 1.0 - 1e-3)\n loss_dict[\"fg_mask_loss\"] = (\n F.binary_cross_entropy(weights_sum, fg_label) * self.config.fg_mask_loss_mult\n )\n\n # monocular normal loss\n if \"normal\" in batch and self.config.mono_normal_loss_mult > 0.0:\n if self._step > self.config.mono_normal_loss_warmup_steps:\n normal_gt = batch[\"normal\"].to(self.device)\n normal_pred = outputs[\"normal\"]\n normal_mask = (outputs[\"accumulation\"].detach().squeeze() > 0.95).float()\n normal_loss_scale = map_range_val(self._step, \n self.config.mono_normal_loss_warmup_steps, \n self.config.mono_normal_loss_steps, \n 1, \n self.config.normal_loss_final_scale, \n type='log', factor=1.)\n\n normal_loss = monosdf_normal_loss(normal_pred * normal_mask[...,None], normal_gt * normal_mask[...,None]) \n loss_dict[\"normal_loss\"] = normal_loss * self.config.mono_normal_loss_mult * normal_loss_scale\n\n # monocular depth loss\n if \"depth\" in batch and self.config.mono_depth_loss_mult > 0.0:\n # TODO check it's true that's we sample from only a single image\n # TODO only supervised pixel that hit the surface and remove hard-coded scaling for depth\n depth_gt = batch[\"depth\"].to(self.device)[..., None]\n depth_pred = outputs[\"depth\"]\n\n mask = torch.ones_like(depth_gt).reshape(1, 32, -1).bool()\n loss_dict[\"depth_loss\"] = (\n self.depth_loss(depth_pred.reshape(1, 32, -1), (depth_gt * 50 + 0.5).reshape(1, 32, -1), mask)\n * self.config.mono_depth_loss_mult\n )\n\n # sensor depth loss\n if \"sensor_depth\" in batch and (\n self.config.sensor_depth_l1_loss_mult > 0.0\n or self.config.sensor_depth_freespace_loss_mult > 0.0\n or self.config.sensor_depth_sdf_loss_mult > 0.0\n ):\n l1_loss, free_space_loss, sdf_loss = self.sensor_depth_loss(batch, outputs)\n\n loss_dict[\"sensor_l1_loss\"] = l1_loss * self.config.sensor_depth_l1_loss_mult\n loss_dict[\"sensor_freespace_loss\"] = free_space_loss * self.config.sensor_depth_freespace_loss_mult\n loss_dict[\"sensor_sdf_loss\"] = sdf_loss * self.config.sensor_depth_sdf_loss_mult\n\n # multi-view photoconsistency loss as Geo-NeuS\n if \"patches\" in outputs and self.config.patch_warp_loss_mult > 0.0:\n patches = outputs[\"patches\"]\n patches_valid_mask = outputs[\"patches_valid_mask\"]\n\n loss_dict[\"patch_loss\"] = (\n self.patch_loss(patches, patches_valid_mask) * self.config.patch_warp_loss_mult\n )\n\n # sparse points sdf loss\n if \"sparse_sfm_points\" in batch and self.config.sparse_points_sdf_loss_mult > 0.0:\n sparse_sfm_points = batch[\"sparse_sfm_points\"].to(self.device)\n sparse_sfm_points_sdf = self.field.forward_geonetwork(sparse_sfm_points)[:, 0].contiguous()\n loss_dict[\"sparse_sfm_points_sdf_loss\"] = (\n torch.mean(torch.abs(sparse_sfm_points_sdf)) * self.config.sparse_points_sdf_loss_mult\n )\n\n # total variational loss for multi-resolution periodic feature volume\n if self.config.periodic_tvl_mult > 0.0:\n assert self.field.config.encoding_type == \"periodic\"\n loss_dict[\"tvl_loss\"] = self.field.encoding.get_total_variation_loss() * self.config.periodic_tvl_mult\n\n self._compute_curvature_loss(outputs, loss_dict)\n self._compute_interlevel_loss_bg(outputs, loss_dict)\n self._compute_ptcd_reg_fg(outputs, loss_dict)\n self._compute_ptcd_reg_plane(outputs, loss_dict)\n\n return loss_dict\n\n # @profiler.time_function\n def get_metrics_dict(self, outputs, batch) -> Dict:\n metrics_dict = {}\n image = batch[\"image\"].to(self.device)\n metrics_dict[\"psnr\"] = self.psnr(outputs[\"rgb\"], image)\n return metrics_dict\n\n def get_image_metrics_and_images(\n self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]\n ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:\n images_dict = {}\n image = batch[\"image\"].to(self.device)\n rgb = outputs[\"rgb\"]\n acc = colormaps.apply_colormap(outputs[\"accumulation\"])\n\n combined_rgb = torch.cat([image, rgb], dim=1)\n combined_acc = torch.cat([acc], dim=1)\n if \"depth\" in batch:\n depth_gt = batch[\"depth\"].to(self.device)\n depth_pred = outputs[\"depth\"]\n\n # align to predicted depth and normalize\n scale, shift = compute_scale_and_shift(\n depth_pred[None, ..., 0], depth_gt[None, ...], depth_gt[None, ...] > 0.0\n )\n depth_pred = depth_pred * scale + shift\n\n combined_depth = torch.cat([depth_gt[..., None], depth_pred], dim=1)\n combined_depth = colormaps.apply_depth_colormap(combined_depth)\n else:\n # TODO: set depth[accumulation < 0.1] = depth[accumulation >= 0.1].min()\n depth = colormaps.apply_depth_colormap(\n outputs[\"depth\"],\n accumulation=outputs[\"accumulation\"],\n )\n combined_depth = torch.cat([depth], dim=1)\n images_dict = {\"img\": combined_rgb, \"accumulation\": combined_acc, \"depth\": combined_depth}\n\n if \"normal\" in outputs:\n normal = outputs[\"normal\"]\n # don't need to normalize here\n # normal = torch.nn.functional.normalize(normal, p=2, dim=-1)\n normal = (normal + 1.0) / 2.0\n if \"normal\" in batch:\n normal_gt = (batch[\"normal\"].to(self.device) + 1.0) / 2.0\n combined_normal = torch.cat([normal_gt, normal], dim=1)\n else:\n combined_normal = torch.cat([normal], dim=1)\n images_dict.update({\"normal\": combined_normal})\n\n if \"sensor_depth\" in batch:\n sensor_depth = batch[\"sensor_depth\"]\n depth_pred = outputs[\"depth\"]\n\n combined_sensor_depth = torch.cat([sensor_depth[..., None], depth_pred], dim=1)\n combined_sensor_depth = colormaps.apply_depth_colormap(combined_sensor_depth)\n images_dict[\"sensor_depth\"] = combined_sensor_depth\n\n if self.handle_no_intersect_rays:\n self._get_no_intersection_ray_mask(outputs, images_dict)\n\n if self.config.bg_sampler_type == \"proposal_network\":\n for i in range(self.config.num_proposal_iterations):\n key = f\"prop_depth_bg_{i}\"\n # TODO: set depth[accumulation < 0.1] = depth[accumulation >= 0.1].min()\n prop_depth_i = colormaps.apply_depth_colormap(outputs[key])\n images_dict[key] = prop_depth_i\n\n # Switch images from [H, W, C] to [1, C, H, W] for metrics computations\n image = torch.moveaxis(image, -1, 0)[None, ...]\n rgb = torch.moveaxis(rgb, -1, 0)[None, ...]\n\n psnr = self.psnr(image, rgb)\n ssim = self.ssim(image, rgb)\n lpips = self.lpips(image, rgb)\n\n # all of these metrics will be logged as scalars\n metrics_dict = {\"psnr\": float(psnr.item()), \"ssim\": float(ssim)} # type: ignore\n metrics_dict[\"lpips\"] = float(lpips)\n\n return metrics_dict, images_dict\n\n def _get_no_intersection_ray_mask(self, outputs: Dict[str, torch.Tensor], images_dict: Dict[str, torch.Tensor]):\n \"\"\"visualizing ray-box intersection results\"\"\"\n no_intersect_mask = outputs[\"no_intersect_mask\"].to(torch.float32) # (h, w, 1)\n images_dict[\"no_intersect_mask\"] = no_intersect_mask\n\n def _compute_curvature_loss(self, outputs: Dict[str, torch.Tensor], loss_dict: Dict[str, torch.Tensor]):\n \"\"\"compute curvature loss\"\"\"\n if (\n self.config.curvature_loss_mult == 0.0\n or self.config.curvature_loss_n_iters is not None\n and self._step > self.config.curvature_loss_n_iters\n ):\n return\n\n curvature_losses = outputs[\"curvature_losses\"] # (n_rays, n_pts, 1)\n if self.handle_no_intersect_rays:\n no_intersect_mask = outputs[\"no_intersect_mask\"].to(curvature_losses.dtype)\n n_intersect_pts = (no_intersect_mask.nelement() - no_intersect_mask.sum()) * np.prod(\n curvature_losses.shape[1:-1]\n )\n curvature_loss = (curvature_losses * (1 - no_intersect_mask[:, None])).sum() / n_intersect_pts\n else:\n curvature_loss = torch.mean(curvature_losses)\n loss_dict[\"curvature_loss\"] = curvature_loss * self.config.curvature_loss_mult\n return\n\n def _compute_interlevel_loss_bg(self, outputs: Dict[str, torch.Tensor], loss_dict: Dict[str, torch.Tensor]):\n \"\"\"compute proposal network loss for bg sampler\"\"\"\n if not self.training or self.config.bg_sampler_type != \"proposal_network\":\n return\n\n bg_interlevel_loss_mult = 1.0\n weights_list, ray_samples_list = self._get_interlevel_loss_inputs_bg(outputs)\n interlevel_loss_val = interlevel_loss(weights_list, ray_samples_list)\n if torch.isnan(interlevel_loss_val):\n __import__(\"ipdb\").set_trace()\n loss_dict[\"bg_interlevel_loss\"] = bg_interlevel_loss_mult * interlevel_loss_val\n\n def _get_interlevel_loss_inputs_bg(self, outputs: Dict[str, torch.Tensor]):\n weights_list = outputs[\"weights_list_bg\"] # (n_rays, n_samples, 1)\n ray_samples_list = outputs[\"ray_samples_list_bg\"] # (n_rays, n_samples)\n\n if not self.handle_bottom_intersect_rays:\n return weights_list, ray_samples_list\n\n # ignore supervision of rays intersecting the bottom plane of the fg aabb\n # we can optionally supervise the underneath region with zero weights\n bottom_intersect_mask = outputs[\"bottom_intersect_mask\"][..., 0] # (n_rays, )\n weights_list = [weights[~bottom_intersect_mask] for weights in weights_list]\n ray_samples_list = [ray_samples[~bottom_intersect_mask] for ray_samples in ray_samples_list]\n\n return weights_list, ray_samples_list\n\n def _build_anneal_fn(self, anneal_fn_name, max_step):\n if max_step is None:\n ptcd_reg_disabled = self.config.ptcd_reg_plane_mult == 0.0 and self.config.ptcd_reg_fg_mult == 0.0\n assert ptcd_reg_disabled or anneal_fn_name == \"constant\"\n return lambda x: 1.0\n\n if anneal_fn_name == \"cosine\":\n anneal_fn = scheduler.cosine_annealing\n elif anneal_fn_name == \"exponential\":\n anneal_fn = scheduler.exp_annealing\n elif anneal_fn_name == \"constant\":\n anneal_fn = scheduler.constant\n else:\n raise ValueError(f\"Unknown anneal fn: {anneal_fn_name}\")\n return partial(anneal_fn, max_step=max_step, min_val=0.0, max_val=1.0)\n\n def _compute_ptcd_reg_fg(self, outputs: Dict[str, torch.Tensor], loss_dict: Dict[str, torch.Tensor]):\n \"\"\"compute foreground point cloud regularization loss\n use point-wise threshold by default.\n \"\"\"\n if self.config.ptcd_reg_fg_mult == 0.0 or self._step > self.config.ptcd_reg_n_iters:\n return\n\n anneal_weight = self.ptcd_reg_anneal_fn(self._step)\n writer.put_scalar(name=\"ptcd_reg_fg_anneal_weight\", scalar=anneal_weight, step=self._step)\n fg_pts, fg_sdf_thrs = self.ptcd_data.sample_fg_points(self.config.ptcd_reg_fg_n_pts_per_iter, replace=True)\n fg_pts_sdf = self.field.forward_geonetwork(fg_pts)[:, 0].contiguous()\n # only support point-wise threshold for now\n fg_sdf_loss = -torch.min(fg_sdf_thrs - torch.abs(fg_pts_sdf), torch.zeros_like(fg_pts[:, 0])).mean()\n loss_dict[\"ptcd_reg_fg\"] = fg_sdf_loss * self.config.ptcd_reg_fg_mult * anneal_weight\n\n def _compute_ptcd_reg_plane(self, outputs: Dict[str, torch.Tensor], loss_dict: Dict[str, torch.Tensor]):\n \"\"\"compute plane point cloud regularization loss\"\"\"\n if (\n self.config.ptcd_reg_plane_mult == 0.0\n or self._step > self.config.ptcd_reg_n_iters\n or len(self.ptcd_data.plane_pts) == 0 # some sequence might not have plane pts\n ):\n return\n\n anneal_weight = self.ptcd_reg_anneal_fn(self._step)\n writer.put_scalar(name=\"ptcd_reg_plane_anneal_weight\", scalar=anneal_weight, step=self._step)\n # sample plane pts (optionally sample half by random sampling from parametric plane)\n plane_pts, udf_plane_to_fg, _ = self.ptcd_data.sample_plane_points(\n self.config.ptcd_reg_plane_n_pts_per_iter, replace=True\n )\n plane_pts_sdf = self.field.forward_geonetwork(plane_pts)[:, 0].contiguous()\n plane_sdf_loss = torch.max(udf_plane_to_fg - torch.abs(plane_pts_sdf), torch.zeros_like(udf_plane_to_fg)).mean()\n loss_dict[\"ptcd_reg_plane\"] = plane_sdf_loss * self.config.ptcd_reg_plane_mult * anneal_weight\n\n def _compute_mask_beta_prior(self, outputs: Dict[str, torch.Tensor], loss_dict: Dict[str, torch.Tensor]):\n \"\"\"compute mask beta prior regularization\"\"\"\n raise NotImplementedError\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/models/distilled_neus_facto.py",
"content": "\"\"\"\nTwo-stages training of NeuSFacto: 1. use NeRF in both fg and bg; 2. replace NeRF with NeuSFacto in fg.\n\"\"\"\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass, field\nfrom typing import List, Type, Tuple, Dict, Any\nfrom typing_extensions import Literal\nimport numpy as np\n\nimport torch\nfrom torch import nn\n\nfrom nerfstudio.cameras.rays import RayBundle, RaySamples\nfrom nerfstudio.engine.callbacks import (\n TrainingCallback,\n TrainingCallbackAttributes,\n TrainingCallbackLocation,\n)\nfrom nerfstudio.fields.density_fields import HashMLPDensityField\nfrom nerfstudio.fields.nerfacto_field import TCNNNerfactoField\nfrom nerfstudio.field_components.field_heads import FieldHeadNames\nfrom nerfstudio.field_components.spatial_distortions import SceneContraction\nfrom nerfstudio.models.neus_facto import NeuSFactoModelConfig, NeuSFactoModel\nfrom nerfstudio.model_components.losses import interlevel_loss\nfrom nerfstudio.model_components.ray_samplers import ProposalNetworkSampler\nfrom nerfstudio.model_components.scene_colliders import NearFarCollider\nfrom nerfstudio.utils import colormaps\n\n\n@dataclass\nclass DistilledNeuSFactoModelConfig(NeuSFactoModelConfig):\n \"\"\"Distilled NeuSFactoModel config\"\"\"\n _target: Type = field(default_factory=lambda: DistilledNeuSFactoModel)\n nerfacto_field_enabled: bool = True\n \"\"\"if True, use nerfacto field in fg; otherwise, use sdf_field in fg\"\"\"\n nerfacto_hash_grid_num_levels: int = 16\n nerfacto_hash_grid_max_res: int = 2048\n nerfacto_hash_grid_log2_hashmap_size: int = 19\n nerfacto_predict_normals: bool = False\n nerfacto_use_appearance_embedding: bool = False\n nerfacto_spatial_normalization_region: Literal[\"full\", \"fg\"] = \"fg\",\n single_field_fg_bg: bool = False\n \"\"\"if True, use a single nerfactor field for fg and bg\"\"\"\n ignore_pre_intersection_points: bool = True\n \"\"\"ignore samples between the camera origin and the first ray-aabb intersection point in background modeling.\"\"\"\n \n \nclass DistilledNeuSFactoModel(NeuSFactoModel):\n \n config: DistilledNeuSFactoModelConfig\n \n def populate_modules(self):\n super().populate_modules()\n self.aabb = nn.Parameter(self.scene_box.aabb, requires_grad=False)\n \n # init nerfacto field (override the original sdf_field)\n if self.config.nerfacto_field_enabled:\n self.field = TCNNNerfactoField(\n self.scene_box.aabb,\n num_levels=self.config.nerfacto_hash_grid_num_levels,\n max_res=self.config.nerfacto_hash_grid_max_res,\n log2_hashmap_size=self.config.nerfacto_hash_grid_log2_hashmap_size,\n spatial_distortion=self.scene_contraction,\n num_images=self.num_train_data,\n use_pred_normals=self.config.nerfacto_predict_normals,\n use_appearance_embedding=self.config.nerfacto_use_appearance_embedding,\n use_average_appearance_embedding=self.config.use_average_appearance_embedding,\n spatial_normalization_region=self.config.nerfacto_spatial_normalization_region,\n )\n \n if self.config.single_field_fg_bg and self.config.background_model != \"none\": # override the original bg field\n self.field_background = self.field\n \n assert isinstance(self.collider, NearFarCollider), \"DistilledNeuSFactoModel only supports NearFarCollider\"\n \n def sample_and_forward_field(self, ray_bundle: RayBundle):\n ray_samples, weights_list, ray_samples_list = self.proposal_sampler(\n ray_bundle, density_fns=self.density_fns\n )\n # split ray_samples into fg and bg samples and update ray_samples.metadata\n self._mask_ray_samples_with_aabb(ray_bundle, ray_samples)\n # TODO: masked inference in foreground and background fields\n weights, transmittance, field_outputs = self.forward_field(ray_samples)\n bg_transmittance = transmittance[:, -1, :]\n \n weights_list.append(weights)\n ray_samples_list.append(ray_samples)\n \n samples_and_field_outputs = {\n \"ray_samples\": ray_samples,\n \"field_outputs\": field_outputs,\n \"weights\": weights,\n \"bg_transmittance\": bg_transmittance,\n \"weights_list\": weights_list,\n \"ray_samples_list\": ray_samples_list,\n }\n return samples_and_field_outputs\n \n def sample_and_forward_field_bg(\n self, ray_bundle: RayBundle,\n samples_and_field_outputs_fg: Dict[str, Any],\n outputs: Dict[str, Any],\n ):\n ray_samples = samples_and_field_outputs_fg[\"ray_samples\"] # fg & bg share the same set of ray samples\n bg_transmittance = samples_and_field_outputs_fg[\"bg_transmittance\"]\n ray_samples_bg_mask = ray_samples.metadata[\"bg_mask\"]\n \n # forward the bg field and mask out weights of invalid positions\n field_outputs_bg = self.field_background(ray_samples)\n densities = field_outputs_bg[FieldHeadNames.DENSITY]\n densities = densities * ray_samples_bg_mask.to(densities.dtype)\n # densities = torch.where(ray_samples_bg_mask, densities, torch.zeros_like(densities))\n weights_bg = ray_samples.get_weights(densities) * bg_transmittance[:, None]\n weights_fg = samples_and_field_outputs_fg[\"weights_list\"][-1]\n weights_fg_bg = torch.where(ray_samples_bg_mask, weights_bg, weights_fg)\n samples_and_field_outputs_fg[\"weights_list\"][-1] = weights_fg_bg\n field_outputs_bg.update({FieldHeadNames.DENSITY: densities})\n \n samples_and_field_outputs_bg = {\n \"ray_samples\": ray_samples,\n \"weights\": weights_bg,\n **field_outputs_bg\n }\n return samples_and_field_outputs_bg\n \n def forward_field(self, ray_samples):\n \"\"\"Forward the fg field and mask out weights of points outside the fg aabb.\n \"\"\"\n if self.config.nerfacto_field_enabled:\n return self._forward_field_nerfacto(ray_samples)\n else:\n return self._forward_field_sdf(ray_samples)\n \n def _forward_field_sdf(self, ray_samples):\n field_outputs = self.field(ray_samples, return_alphas=True)\n alphas = field_outputs[FieldHeadNames.ALPHA]\n if ray_samples.metadata is not None and \"fg_mask\" in ray_samples.metadata:\n alphas = alphas * ray_samples.metadata[\"fg_mask\"].to(alphas.dtype)\n weights, transmittance = ray_samples.get_weights_and_transmittance_from_alphas(alphas)\n field_outputs.update({FieldHeadNames.ALPHA: alphas})\n return weights, transmittance, field_outputs\n \n def _forward_field_nerfacto(self, ray_samples):\n field_outputs = self.field(ray_samples, compute_normals=self.config.nerfacto_predict_normals)\n densities = field_outputs[FieldHeadNames.DENSITY]\n if ray_samples.metadata is not None and \"fg_mask\" in ray_samples.metadata:\n densities = densities * ray_samples.metadata[\"fg_mask\"].to(densities.dtype)\n weights, transmittance = ray_samples.get_weights_and_transmittance(densities)\n field_outputs.update({FieldHeadNames.DENSITY: densities})\n return weights, transmittance, field_outputs\n \n def _mask_ray_samples_with_aabb(self, ray_bundle: RayBundle, ray_samples: RaySamples) -> None:\n \"\"\"Compute per ray_sample fg masks and bg masks\"\"\"\n if self.config.background_model == \"none\":\n return\n \n aabb = self.aabb # (2, 3)\n \n positions = ray_samples.frustums.get_positions() # (n_rays, n_pts, 3)\n in_aabb_mask = ((positions >= aabb[[0], None]) & (positions <= aabb[[1], None])).all(-1, keepdim=True)\n out_aabb_mask = ~in_aabb_mask\n \n if self.config.ignore_pre_intersection_points:\n # TODO: support the modeling of regions b/w camera near & the 1st aabb intersection (currently ignored)\n nears = self._compute_aabb_intersections(ray_bundle)[:, None, None] # (n_rays, 1, 1)\n ray_dists = (ray_samples.frustums.starts + ray_samples.frustums.ends) / 2\n out_ignore_mask = ray_dists <= nears # (n_rays, n_pts, 1)\n out_aabb_mask[out_ignore_mask] = False\n else:\n raise NotImplementedError # three part composition\n # # additionally consider pre-intersection points\n # nears = self._compute_aabb_intersections(ray_bundle)[:, None, None] # (n_rays, 1, 1)\n # ray_dists = (ray_samples.frustums.starts + ray_samples.frustums.ends) / 2\n # in_additional_mask = ray_dists <= nears # (n_rays, n_pts, 1)\n # in_aabb_mask[in_additional_mask] = True\n # out_aabb_mask = ~in_aabb_mask\n \n if ray_samples.metadata is None:\n ray_samples.metadata = {}\n ray_samples.metadata[\"fg_mask\"] = in_aabb_mask\n ray_samples.metadata[\"bg_mask\"] = out_aabb_mask\n \n def _compute_aabb_intersections(self, ray_bundle: RayBundle):\n aabb = self.aabb\n rays_o, rays_d = ray_bundle.origins, ray_bundle.directions\n dir_fraction = 1.0 / (rays_d + 1e-6)\n \n # x\n t1 = (aabb[0, 0] - rays_o[:, 0:1]) * dir_fraction[:, 0:1]\n t2 = (aabb[1, 0] - rays_o[:, 0:1]) * dir_fraction[:, 0:1]\n # y\n t3 = (aabb[0, 1] - rays_o[:, 1:2]) * dir_fraction[:, 1:2]\n t4 = (aabb[1, 1] - rays_o[:, 1:2]) * dir_fraction[:, 1:2]\n # z\n t5 = (aabb[0, 2] - rays_o[:, 2:3]) * dir_fraction[:, 2:3]\n t6 = (aabb[1, 2] - rays_o[:, 2:3]) * dir_fraction[:, 2:3]\n \n tx_min, tx_max = torch.minimum(t1, t2), torch.maximum(t1, t2)\n ty_min, ty_max = torch.minimum(t3, t4), torch.maximum(t3, t4)\n tz_min, tz_max = torch.minimum(t5, t6), torch.maximum(t5, t6)\n \n nears = torch.max(torch.cat([tx_min, ty_min, tz_min], dim=1), dim=1).values\n fars = torch.min(torch.cat([tx_max, ty_max, tz_max], dim=1), dim=1).values\n nears[nears > fars] = -1.0 # no intersection with the first plane -> nears < 0\n return nears\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/models/dto.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nImplementation of VolSDF.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass, field\nfrom typing import Dict, List, Tuple, Type\n\nimport nerfacc\nimport torch\nimport torch.nn.functional as F\nfrom torch.nn import Parameter\nfrom torchtyping import TensorType\n\nfrom nerfstudio.cameras.rays import RayBundle\nfrom nerfstudio.engine.callbacks import (\n TrainingCallback,\n TrainingCallbackAttributes,\n TrainingCallbackLocation,\n)\nfrom nerfstudio.field_components.field_heads import FieldHeadNames\nfrom nerfstudio.field_components.spatial_distortions import SceneContraction\nfrom nerfstudio.fields.sdf_field import SDFFieldConfig\nfrom nerfstudio.model_components.losses import L1Loss\nfrom nerfstudio.model_components.ray_samplers import (\n ErrorBoundedSampler,\n LinearDisparitySampler,\n NeuSSampler,\n PDFSampler,\n UniformSampler,\n UniSurfSampler,\n save_points,\n)\nfrom nerfstudio.model_components.renderers import DepthRenderer, SemanticRenderer\nfrom nerfstudio.model_components.scene_colliders import SphereCollider\nfrom nerfstudio.models.nerfacto import NerfactoModel, NerfactoModelConfig\nfrom nerfstudio.utils import colormaps\nfrom nerfstudio.utils.colors import get_color\nfrom nerfstudio.utils.marching_cubes import get_surface_occupancy\n\n\n@dataclass\nclass DtoOModelConfig(NerfactoModelConfig):\n \"\"\"UniSurf Model Config\"\"\"\n\n _target: Type = field(default_factory=lambda: DtoOModel)\n smooth_loss_multi: float = 0.005\n \"\"\"smoothness loss on surface points in unisurf\"\"\"\n sdf_field: SDFFieldConfig = SDFFieldConfig()\n \"\"\"Config for SDF Field\"\"\"\n\n\nclass DtoOModel(NerfactoModel):\n \"\"\"VolSDF model\n\n Args:\n config: MonoSDF configuration to instantiate model\n \"\"\"\n\n config: DtoOModelConfig\n\n def populate_modules(self):\n \"\"\"Set the fields and modules.\"\"\"\n super().populate_modules()\n\n scene_contraction = SceneContraction(order=float(\"inf\"))\n\n # create occupancy grid from scene_bbox\n aabb_scale = 1.0\n aabb = [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]]\n aabb = torch.tensor(aabb, dtype=torch.float32)\n\n self.grid = nerfacc.OccupancyGrid(aabb.reshape(-1), resolution=32)\n self._binary = self.scene_box.coarse_binary_gird.reshape(32, 32, 32).contiguous()\n self._binary_fine = None\n # self.rank = torch.distributed.get_rank()\n print(\"self\", self.local_rank)\n\n # Occupancy\n self.occupancy_field = self.config.sdf_field.setup(\n aabb=aabb,\n spatial_distortion=scene_contraction,\n num_images=self.num_train_data,\n use_average_appearance_embedding=self.config.use_average_appearance_embedding,\n )\n\n self.pdf_sampler = PDFSampler(\n include_original=True,\n single_jitter=self.config.use_single_jitter,\n histogram_padding=1e-5,\n )\n self.surface_sampler = PDFSampler(include_original=False, single_jitter=False, histogram_padding=1e-5)\n\n self.uniform_sampler = UniformSampler(single_jitter=False)\n\n self.neus_sampler = NeuSSampler(\n num_samples=8, num_samples_importance=16, num_samples_outside=0, num_upsample_steps=2, base_variance=512\n )\n\n self.renderer_normal = SemanticRenderer()\n self.renderer_depth = DepthRenderer(\"expected\")\n\n # for merge samples\n self.unisurf_sampler = UniSurfSampler()\n\n # self.error_bounded_sampler = ErrorBoundedSampler()\n self.error_bounded_sampler = ErrorBoundedSampler(\n num_samples=64,\n num_samples_eval=128,\n num_samples_extra=32,\n )\n\n # sphere collider\n self.sphere_collider = SphereCollider(radius=1.0)\n\n # background model\n self.bg_sampler = LinearDisparitySampler(num_samples=4)\n self.step_counter = 0\n self.anneal_end = 20000\n\n self.use_nerfacto = False\n self.method = \"neus\"\n\n self.rgb_loss = L1Loss()\n\n def get_training_callbacks(\n self, training_callback_attributes: TrainingCallbackAttributes\n ) -> List[TrainingCallback]:\n if self.use_nerfacto:\n callbacks = super().get_training_callbacks(training_callback_attributes)\n else:\n callbacks = []\n\n # anneal for cos in NeuS\n if self.anneal_end > 0:\n\n def set_anneal(step):\n anneal = min([1.0, step / self.anneal_end])\n self.occupancy_field.set_cos_anneal_ratio(anneal)\n\n callbacks.append(\n TrainingCallback(\n where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],\n update_every_num_iters=1,\n func=set_anneal,\n )\n )\n\n return callbacks\n\n def get_param_groups(self) -> Dict[str, List[Parameter]]:\n param_groups = super().get_param_groups()\n param_groups[\"occupancy_field\"] = list(self.occupancy_field.parameters())\n return param_groups\n\n def get_outputs(self, ray_bundle: RayBundle):\n # compute near and far from from sphere with radius 1.0\n ray_bundle = self.sphere_collider(ray_bundle)\n\n # near and far from occupancy grids\n packed_info, ray_indices, t_starts, t_ends = nerfacc.cuda.ray_marching(\n ray_bundle.origins.contiguous(),\n ray_bundle.directions.contiguous(),\n ray_bundle.nears[:, 0].contiguous(),\n ray_bundle.fars[:, 0].contiguous(),\n self.grid.roi_aabb.contiguous(),\n self._binary.to(ray_bundle.origins.device),\n self.grid.contraction_type.to_cpp_version(),\n 1e-2, # large value for coarse voxels\n 0.0,\n )\n\n tt_starts = nerfacc.unpack_data(packed_info, t_starts)\n # tt_ends = nerfacc.unpack_data(packed_info, t_ends)\n\n hit_grid = (tt_starts > 0).any(dim=1)[:, 0]\n if hit_grid.float().sum() > 0:\n ray_bundle.nears[hit_grid] = tt_starts[hit_grid][:, 0]\n ray_bundle.fars[hit_grid] = tt_starts[hit_grid].max(dim=1)[0]\n\n # sample uniformly with currently nears and far\n voxel_samples = self.uniform_sampler(ray_bundle, num_samples=10)\n\n nears = ray_bundle.nears.clone()\n fars = ray_bundle.fars.clone()\n\n # bootstrap should needs longer if using only one gpus\n if self.training and self.step_counter > 5000 and self.step_counter % 5000 == 1:\n grid_size = 32\n voxel_size = 2.0 / 32\n fine_grid_size = 16\n offset = torch.linspace(-1.0, 1.0, fine_grid_size * 2 + 1, device=self.device)[1::2]\n x, y, z = torch.meshgrid(offset, offset, offset, indexing=\"ij\")\n fine_offset_cube = torch.stack([x, y, z], dim=-1).reshape(-1, 3) * voxel_size * 0.5\n\n # coarse grid coordinates\n offset = torch.linspace(-1.0 + voxel_size / 2.0, 1.0 - voxel_size / 2.0, grid_size, device=self.device)\n x, y, z = torch.meshgrid(offset, offset, offset, indexing=\"ij\")\n coarse_offset_cube = torch.stack([x, y, z], dim=-1).reshape(-1, 3)\n\n # xyz\n mask = torch.zeros((grid_size**3, fine_grid_size**3), dtype=torch.bool, device=self.device)\n\n occupied_voxel = coarse_offset_cube[self._binary.reshape(-1)]\n fine_voxel = occupied_voxel[:, None] + fine_offset_cube[None, :]\n print(fine_voxel.shape)\n fine_voxel = fine_voxel.reshape(-1, 3)\n # save_points(\"fine_voxel.ply\", fine_voxel.cpu().numpy())\n\n @torch.no_grad()\n def evaluate(points):\n z = []\n for _, pnts in enumerate(torch.split(points, 100000, dim=0)):\n z.append(self.occupancy_field.forward_geonetwork(pnts)[:, 0].contiguous())\n z = torch.cat(z, axis=0)\n return z\n\n sdf = evaluate(fine_voxel)\n sdf = sdf.reshape(occupied_voxel.shape[0], fine_grid_size**3)\n sdf_mask = sdf <= 0.0\n mask[self._binary.reshape(-1)] = sdf_mask\n\n self._binary_fine = (\n mask.reshape(grid_size, grid_size, grid_size, fine_grid_size, fine_grid_size, fine_grid_size)\n .permute(0, 3, 1, 4, 2, 5)\n .reshape(grid_size * fine_grid_size, grid_size * fine_grid_size, grid_size * fine_grid_size)\n .contiguous()\n )\n\n offset = torch.linspace(-1.0, 1.0, fine_grid_size * grid_size * 2 + 1, device=self.device)[1::2]\n x, y, z = torch.meshgrid(offset, offset, offset, indexing=\"ij\")\n grid_coord = torch.stack([x, y, z], dim=-1).reshape(-1, 3)\n if self.local_rank == 0:\n save_points(\"fine_voxel_valid.ply\", grid_coord[self._binary_fine.reshape(-1)].cpu().numpy())\n # breakpoint()\n\n if self._binary_fine is not None:\n # near and far from occupancy grids\n packed_info, ray_indices, t_starts, t_ends = nerfacc.cuda.ray_marching(\n ray_bundle.origins.contiguous(),\n ray_bundle.directions.contiguous(),\n ray_bundle.nears[:, 0].contiguous(),\n ray_bundle.fars[:, 0].contiguous(),\n self.grid.roi_aabb.contiguous(),\n self._binary_fine,\n self.grid.contraction_type.to_cpp_version(),\n 1e-3, # small value for fine voxels\n 0.0,\n )\n\n tt_starts = nerfacc.unpack_data(packed_info, t_starts)\n # tt_ends = nerfacc.unpack_data(packed_info, t_ends, n_samples=1024)\n\n # update with near and far\n hit_grid = (tt_starts > 0).any(dim=1)[:, 0]\n if hit_grid.float().sum() > 0:\n ray_bundle.nears[hit_grid] = tt_starts[hit_grid][:, 0] - 0.03\n ray_bundle.fars[hit_grid] = tt_starts[hit_grid][:, 0] + 0.03\n else:\n print(\"waring not intersection\")\n print(\"sampling around surfaces\")\n\n if self.use_nerfacto:\n outputs = super().get_outputs(ray_bundle)\n\n # weights and ray samples from nerfacto\n base_weights = outputs[\"weights_list\"][-1].detach()\n base_ray_samples = outputs[\"ray_samples_list\"][-1]\n\n # TODO maybe interative sampling to sample more points on the surface?\n # importance samples and merge\n # occupancy_samples = self.pdf_sampler(ray_bundle, base_ray_samples, base_weights, num_samples=64)\n\n occupancy_samples = self.neus_sampler(\n ray_bundle, sdf_fn=self.occupancy_field.get_sdf, ray_samples=base_ray_samples\n )\n else:\n outputs = {}\n if self.method == \"neus\":\n occupancy_samples = self.neus_sampler(ray_bundle, sdf_fn=self.occupancy_field.get_sdf)\n elif self.method == \"volsdf\":\n # VolSDF\n occupancy_samples, _ = self.error_bounded_sampler(\n ray_bundle, density_fn=self.occupancy_field.laplace_density, sdf_fn=self.occupancy_field.get_sdf\n )\n\n # save_points(\"p2.ply\", occupancy_samples.frustums.get_positions().detach().cpu().numpy().reshape(-1, 3))\n\n # merge samples\n occupancy_samples = self.unisurf_sampler.merge_ray_samples_in_eculidean(\n ray_bundle, occupancy_samples, voxel_samples\n )\n\n # save_points(\"p1.ply\", voxel_samples.frustums.get_positions().detach().cpu().numpy().reshape(-1, 3))\n\n # save_points(\"p3.ply\", occupancy_samples.frustums.get_positions().detach().cpu().numpy().reshape(-1, 3))\n # breakpoint()\n\n # save_points(\"p4.ply\", importance_samples.frustums.get_positions().detach().cpu().numpy().reshape(-1, 3))\n # if self._step > 0:\n # exit(-1)\n\n # save_points(\"merged.ply\", merged_samples.frustums.get_start_positions().reshape(-1, 3).detach().cpu().numpy())\n\n # breakpoint()\n\n # occupancy unisurf\n # field_outputs = self.occupancy_field(occupancy_samples, return_occupancy=True)\n # weights, transmittance = occupancy_samples.get_weights_and_transmittance_from_alphas(\n # field_outputs[FieldHeadNames.OCCUPANCY]\n # )\n # save_points(\"o.ply\", occupancy_samples.frustums.get_positions().reshape(-1, 3).detach().cpu().numpy())\n if self.method == \"neus\":\n field_outputs = self.occupancy_field(occupancy_samples, return_alphas=True)\n weights, transmittance = occupancy_samples.get_weights_and_transmittance_from_alphas(\n field_outputs[FieldHeadNames.ALPHA]\n )\n elif self.method == \"volsdf\":\n field_outputs = self.occupancy_field(occupancy_samples)\n weights, transmittance = occupancy_samples.get_weights_and_transmitance(\n field_outputs[FieldHeadNames.DENSITY]\n )\n\n # save_points(\"a.ply\", occupancy_samples.frustums.get_positions().reshape(-1, 3).detach().cpu().numpy())\n\n if self.training:\n # we should sample here before we change the near-far plane for ray_bundle\n # sample surface points according to distribution\n # surface_samples = self.surface_sampler(ray_bundle, occupancy_samples, weights, num_samples=8)\n\n self.step_counter += 1\n if self.step_counter % 5000 == 0 and self.local_rank == 0:\n\n save_points(\"a.ply\", occupancy_samples.frustums.get_positions().reshape(-1, 3).detach().cpu().numpy())\n get_surface_occupancy(\n occupancy_fn=lambda x: self.occupancy_field.forward_geonetwork(x)[:, 0], device=self.device\n )\n # breakpoint()\n\n bg_transmittance = transmittance[:, -1, :]\n\n rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)\n depth = self.renderer_depth(weights=weights, ray_samples=occupancy_samples)\n # the rendered depth is point-to-point distance and we should convert to depth\n depth = depth / ray_bundle.directions_norm\n\n normal = self.renderer_normal(semantics=field_outputs[FieldHeadNames.NORMAL], weights=weights)\n accumulation = self.renderer_accumulation(weights=weights)\n\n # save rendered points for visualization\n pts = self.renderer_normal(semantics=occupancy_samples.frustums.get_positions(), weights=weights)\n hit_mask = (field_outputs[FieldHeadNames.SDF] > 0.0).any(dim=1) & (field_outputs[FieldHeadNames.SDF] < 0.0).any(\n dim=1\n )\n pts = pts[hit_mask[:, 0]]\n # save_points(\"a.ply\", pts.reshape(-1, 3).detach().cpu().numpy())\n # breakpoint()\n if pts.shape[0] > 0:\n surface_sdf = self.occupancy_field.forward_geonetwork(pts)[:, 0]\n else:\n surface_sdf = None\n\n # sample inversely from far to 100 and points and forwart the bg model\n ray_bundle.nears = ray_bundle.fars\n ray_bundle.fars = torch.ones_like(ray_bundle.fars) * 1000.0\n\n bg_ray_samples = self.bg_sampler(ray_bundle)\n\n # use the same background model for both density field and occupancy field\n field_outputs_bg = self.field(bg_ray_samples)\n bg_weights = bg_ray_samples.get_weights(field_outputs_bg[FieldHeadNames.DENSITY])\n\n bg_rgb = self.renderer_rgb(rgb=field_outputs_bg[FieldHeadNames.RGB], weights=bg_weights)\n bg_depth = self.renderer_depth(weights=bg_weights, ray_samples=bg_ray_samples)\n bg_accumulation = self.renderer_accumulation(weights=bg_weights)\n\n bg_outputs = {\n \"bg_rgb\": bg_rgb,\n \"bg_accumulation\": bg_accumulation,\n \"bg_depth\": bg_depth,\n \"bg_weights\": bg_weights,\n }\n\n # merge background color to forgound color\n rgb = rgb + bg_transmittance * bg_rgb\n outputs_occupancy = {\n \"orgb\": rgb,\n \"oaccumulation\": accumulation,\n \"odepth\": depth,\n \"onormal\": normal,\n \"oweights\": weights,\n }\n\n # hitting information\n start_points = occupancy_samples.frustums.get_positions()[:, 0]\n end_points = occupancy_samples.frustums.get_positions()[:, -1]\n outputs_occupancy.update({\"start_points\": start_points, \"end_points\": end_points})\n\n if self.training:\n outputs_occupancy.update({\"surface_sdf\": surface_sdf})\n\n if self.use_nerfacto:\n # merge background color to forgound color of density field\n outputs[\"rgb\"] = outputs[\"rgb\"] + outputs[\"transmittance\"][:, -1, :] * bg_rgb\n\n if self.training:\n # surface_points = surface_samples.frustums.get_positions().reshape(-1, 3)\n\n # surface_points_neig = surface_points + (torch.rand_like(surface_points) - 0.5) * 0.001\n # pp = torch.cat([surface_points, surface_points_neig], dim=0)\n # surface_grad = self.occupancy_field.gradient(pp)\n # surface_points_normal = torch.nn.functional.normalize(surface_grad, p=2, dim=-1)\n\n # outputs_occupancy.update({\"surface_points_normal\": surface_points_normal, \"surface_grad\": surface_grad})\n outputs_occupancy.update({\"surface_grad\": field_outputs[FieldHeadNames.GRADIENT]})\n\n outputs.update(outputs_occupancy)\n outputs.update(bg_outputs)\n\n return outputs\n\n def get_metrics_dict(self, outputs, batch):\n if self.use_nerfacto:\n metrics_dict = super().get_metrics_dict(outputs, batch)\n else:\n metrics_dict = {}\n\n image = batch[\"image\"].to(self.device)\n metrics_dict[\"opsnr\"] = self.psnr(outputs[\"orgb\"], image)\n\n if self.training:\n # training statics\n metrics_dict[\"s_val\"] = self.occupancy_field.deviation_network.get_variance().item()\n metrics_dict[\"inv_s\"] = 1.0 / self.occupancy_field.deviation_network.get_variance().item()\n\n # training statics for volsdf\n metrics_dict[\"beta\"] = self.occupancy_field.laplace_density.get_beta().item()\n metrics_dict[\"alpha\"] = 1.0 / self.occupancy_field.laplace_density.get_beta().item()\n\n return metrics_dict\n\n def get_loss_dict(self, outputs, batch, metrics_dict=None):\n if self.use_nerfacto:\n loss_dict = super().get_loss_dict(outputs, batch, metrics_dict)\n else:\n loss_dict = {}\n\n image = batch[\"image\"].to(self.device)\n loss_dict[\"orgb_loss\"] = self.rgb_loss(image, outputs[\"orgb\"])\n\n if self.training and \"sky\" in batch:\n sky_label = 1.0 - batch[\"sky\"].float().to(self.device)\n if self.use_nerfacto:\n density_field_weights = outputs[\"weights_list\"][-1]\n loss_dict[\"sky_loss\"] = (\n F.binary_cross_entropy(density_field_weights.sum(dim=1).clip(1e-3, 1.0 - 1e-3), sky_label) * 0.01\n )\n\n occupancy_field_weights = outputs[\"oweights\"]\n loss_dict[\"osky_loss\"] = (\n F.binary_cross_entropy(occupancy_field_weights.sum(dim=1).clip(1e-3, 1.0 - 1e-3), sky_label) * 0.01\n )\n\n if self.training:\n # surface_points_normal = outputs[\"surface_points_normal\"]\n # N = surface_points_normal.shape[0] // 2\n\n # diff_norm = torch.norm(surface_points_normal[:N] - surface_points_normal[N:], dim=-1)\n # loss_dict[\"normal_smoothness_loss\"] = torch.mean(diff_norm) * 0.0001\n\n # eikonal loss\n surface_points_grad = outputs[\"surface_grad\"]\n loss_dict[\"eikonal_loss\"] = ((surface_points_grad.norm(2, dim=-1) - 1) ** 2).mean() * 0.0001\n\n # surface points loss\n surface_points_sdf = outputs[\"surface_sdf\"]\n if surface_points_sdf is not None:\n loss_dict[\"surface_sdf_loss\"] = torch.abs(surface_points_sdf).mean() * 0.0\n\n sparse_pts = batch[\"sparse_pts\"].to(self.device)\n sparse_pts, pts_weights = sparse_pts[:, :3], sparse_pts[:, 3:]\n\n # filter norm\n in_sphere = torch.norm(sparse_pts, dim=-1) < 1.0\n sparse_pts = sparse_pts[in_sphere]\n pts_weights = pts_weights[in_sphere]\n pts_weights = torch.ones_like(pts_weights)\n # print(in_sphere.float().mean())\n # save_points(\"sa.ply\", sparse_pts.cpu().numpy())\n # breakpoint()\n if sparse_pts.shape[0] > 0:\n sparse_pts_sdf = self.occupancy_field.forward_geonetwork(sparse_pts)[:, 0]\n loss_dict[\"sparse_pts_loss\"] = (torch.abs(sparse_pts_sdf) * pts_weights).mean() * 0.0\n\n return loss_dict\n\n def get_image_metrics_and_images(\n self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]\n ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:\n if self.use_nerfacto:\n metrics_dict, images_dict = super().get_image_metrics_and_images(outputs, batch)\n else:\n metrics_dict, images_dict = {}, {}\n\n image = batch[\"image\"].to(self.device)\n rgb = outputs[\"orgb\"]\n acc = colormaps.apply_colormap(outputs[\"oaccumulation\"])\n depth = colormaps.apply_depth_colormap(\n outputs[\"odepth\"],\n accumulation=outputs[\"oaccumulation\"],\n )\n\n normal = outputs[\"onormal\"]\n # don't need to normalize here\n # normal = torch.nn.functional.normalize(normal, p=2, dim=-1)\n normal = (normal + 1.0) / 2.0\n combined_normal = torch.cat([normal], dim=1)\n\n combined_rgb = torch.cat([image, rgb], dim=1)\n combined_acc = torch.cat([acc], dim=1)\n combined_depth = torch.cat([depth], dim=1)\n\n # Switch images from [H, W, C] to [1, C, H, W] for metrics computations\n image = torch.moveaxis(image, -1, 0)[None, ...]\n rgb = torch.moveaxis(rgb, -1, 0)[None, ...]\n\n psnr = self.psnr(image, rgb)\n ssim = self.ssim(image, rgb)\n lpips = self.lpips(image, rgb)\n\n # all of these metrics will be logged as scalars\n metrics_dict.update({\"opsnr\": float(psnr.item()), \"ossim\": float(ssim)}) # type: ignore\n metrics_dict[\"olpips\"] = float(lpips)\n\n images_dict.update(\n {\"oimg\": combined_rgb, \"oaccumulation\": combined_acc, \"odepth\": combined_depth, \"onormal\": combined_normal}\n )\n\n # save hitting\n if \"start_points\" in outputs:\n start_points = outputs[\"start_points\"]\n end_points = outputs[\"end_points\"]\n save_points(\"starts.ply\", start_points.cpu().numpy().reshape(-1, 3))\n save_points(\"ends.ply\", end_points.cpu().numpy().reshape(-1, 3))\n\n return metrics_dict, images_dict\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/models/instant_ngp.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nImplementation of Instant NGP.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass, field\nfrom typing import Dict, List, Optional, Tuple, Type\n\nimport nerfacc\nimport torch\nfrom nerfacc import ContractionType\nfrom torch.nn import Parameter\nfrom torchmetrics import PeakSignalNoiseRatio\nfrom torchmetrics.functional import structural_similarity_index_measure\nfrom torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity\n\nfrom nerfstudio.cameras.rays import RayBundle\nfrom nerfstudio.engine.callbacks import (\n TrainingCallback,\n TrainingCallbackAttributes,\n TrainingCallbackLocation,\n)\nfrom nerfstudio.field_components.field_heads import FieldHeadNames\nfrom nerfstudio.fields.instant_ngp_field import TCNNInstantNGPField\nfrom nerfstudio.model_components.losses import MSELoss\nfrom nerfstudio.model_components.ray_samplers import VolumetricSampler\nfrom nerfstudio.model_components.renderers import (\n AccumulationRenderer,\n DepthRenderer,\n RGBRenderer,\n)\nfrom nerfstudio.models.base_model import Model, ModelConfig\nfrom nerfstudio.utils import colormaps, colors\n\n\n@dataclass\nclass InstantNGPModelConfig(ModelConfig):\n \"\"\"Instant NGP Model Config\"\"\"\n\n _target: Type = field(\n default_factory=lambda: NGPModel\n ) # We can't write `NGPModel` directly, because `NGPModel` doesn't exist yet\n \"\"\"target class to instantiate\"\"\"\n enable_collider: bool = False\n \"\"\"Whether to create a scene collider to filter rays.\"\"\"\n collider_params: Optional[Dict[str, float]] = None\n \"\"\"Instant NGP doesn't use a collider.\"\"\"\n max_num_samples_per_ray: int = 24\n \"\"\"Number of samples in field evaluation.\"\"\"\n grid_resolution: int = 128\n \"\"\"Resolution of the grid used for the field.\"\"\"\n contraction_type: ContractionType = ContractionType.UN_BOUNDED_SPHERE\n \"\"\"Resolution of the grid used for the field.\"\"\"\n cone_angle: float = 0.004\n \"\"\"Should be set to 0.0 for blender scenes but 1./256 for real scenes.\"\"\"\n render_step_size: float = 0.01\n \"\"\"Minimum step size for rendering.\"\"\"\n near_plane: float = 0.05\n \"\"\"How far along ray to start sampling.\"\"\"\n far_plane: float = 1e3\n \"\"\"How far along ray to stop sampling.\"\"\"\n use_appearance_embedding: bool = False\n \"\"\"Whether to use an appearance embedding.\"\"\"\n randomize_background: bool = True\n \"\"\"Whether to randomize the background color.\"\"\"\n alpha_thre: float = 1e-2\n \"\"\"alpha thres for visibility pruning in nerfacc, should set to 0 for the nerf-synthetic dataset\"\"\"\n\n\nclass NGPModel(Model):\n \"\"\"Instant NGP model\n\n Args:\n config: instant NGP configuration to instantiate model\n \"\"\"\n\n config: InstantNGPModelConfig\n field: TCNNInstantNGPField\n\n def __init__(self, config: InstantNGPModelConfig, **kwargs) -> None:\n super().__init__(config=config, **kwargs)\n\n def populate_modules(self):\n \"\"\"Set the fields and modules.\"\"\"\n super().populate_modules()\n\n self.field = TCNNInstantNGPField(\n aabb=self.scene_box.aabb,\n contraction_type=self.config.contraction_type,\n use_appearance_embedding=self.config.use_appearance_embedding,\n num_images=self.num_train_data,\n )\n\n self.scene_aabb = Parameter(self.scene_box.aabb.flatten(), requires_grad=False)\n\n # Occupancy Grid\n self.occupancy_grid = nerfacc.OccupancyGrid(\n roi_aabb=self.scene_aabb,\n resolution=self.config.grid_resolution,\n contraction_type=self.config.contraction_type,\n )\n\n # Sampler\n vol_sampler_aabb = self.scene_box.aabb if self.config.contraction_type == ContractionType.AABB else None\n self.sampler = VolumetricSampler(\n scene_aabb=vol_sampler_aabb,\n occupancy_grid=self.occupancy_grid,\n density_fn=self.field.density_fn,\n )\n\n # renderers\n background_color = \"random\" if self.config.randomize_background else colors.WHITE\n self.renderer_rgb = RGBRenderer(background_color=background_color)\n self.renderer_accumulation = AccumulationRenderer()\n self.renderer_depth = DepthRenderer(method=\"expected\")\n\n # losses\n self.rgb_loss = MSELoss()\n\n # metrics\n self.psnr = PeakSignalNoiseRatio(data_range=1.0)\n self.ssim = structural_similarity_index_measure\n self.lpips = LearnedPerceptualImagePatchSimilarity()\n\n def get_training_callbacks(\n self, training_callback_attributes: TrainingCallbackAttributes\n ) -> List[TrainingCallback]:\n def update_occupancy_grid(step: int):\n # TODO: needs to get access to the sampler, on how the step size is determinated at each x. See\n # https://github.com/KAIR-BAIR/nerfacc/blob/127223b11401125a9fce5ce269bb0546ee4de6e8/examples/train_ngp_nerf.py#L190-L213\n self.occupancy_grid.every_n_step(\n step=step,\n occ_eval_fn=lambda x: self.field.get_opacity(x, self.config.render_step_size),\n )\n\n return [\n TrainingCallback(\n where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],\n update_every_num_iters=1,\n func=update_occupancy_grid,\n ),\n ]\n\n def get_param_groups(self) -> Dict[str, List[Parameter]]:\n param_groups = {}\n if self.field is None:\n raise ValueError(\"populate_fields() must be called before get_param_groups\")\n param_groups[\"fields\"] = list(self.field.parameters())\n return param_groups\n\n def get_outputs(self, ray_bundle: RayBundle):\n assert self.field is not None\n num_rays = len(ray_bundle)\n\n with torch.no_grad():\n ray_samples, packed_info, ray_indices = self.sampler(\n ray_bundle=ray_bundle,\n near_plane=self.config.near_plane,\n far_plane=self.config.far_plane,\n render_step_size=self.config.render_step_size,\n cone_angle=self.config.cone_angle,\n alpha_thre=self.config.alpha_thre,\n )\n\n field_outputs = self.field(ray_samples)\n\n # accumulation\n weights = nerfacc.render_weight_from_density(\n packed_info=packed_info,\n sigmas=field_outputs[FieldHeadNames.DENSITY],\n t_starts=ray_samples.frustums.starts,\n t_ends=ray_samples.frustums.ends,\n )\n\n rgb = self.renderer_rgb(\n rgb=field_outputs[FieldHeadNames.RGB],\n weights=weights,\n ray_indices=ray_indices,\n num_rays=num_rays,\n )\n depth = self.renderer_depth(\n weights=weights, ray_samples=ray_samples, ray_indices=ray_indices, num_rays=num_rays\n )\n accumulation = self.renderer_accumulation(weights=weights, ray_indices=ray_indices, num_rays=num_rays)\n alive_ray_mask = accumulation.squeeze(-1) > 0\n\n outputs = {\n \"rgb\": rgb,\n \"accumulation\": accumulation,\n \"depth\": depth,\n \"alive_ray_mask\": alive_ray_mask, # the rays we kept from sampler\n \"num_samples_per_ray\": packed_info[:, 1],\n }\n return outputs\n\n def get_metrics_dict(self, outputs, batch):\n image = batch[\"image\"].to(self.device)\n metrics_dict = {}\n metrics_dict[\"psnr\"] = self.psnr(outputs[\"rgb\"], image)\n metrics_dict[\"num_samples_per_batch\"] = outputs[\"num_samples_per_ray\"].sum()\n return metrics_dict\n\n def get_loss_dict(self, outputs, batch, metrics_dict=None):\n image = batch[\"image\"].to(self.device)\n mask = outputs[\"alive_ray_mask\"]\n rgb_loss = self.rgb_loss(image[mask], outputs[\"rgb\"][mask])\n loss_dict = {\"rgb_loss\": rgb_loss}\n return loss_dict\n\n def get_image_metrics_and_images(\n self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]\n ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:\n\n image = batch[\"image\"].to(self.device)\n rgb = outputs[\"rgb\"]\n acc = colormaps.apply_colormap(outputs[\"accumulation\"])\n depth = colormaps.apply_depth_colormap(\n outputs[\"depth\"],\n accumulation=outputs[\"accumulation\"],\n )\n alive_ray_mask = colormaps.apply_colormap(outputs[\"alive_ray_mask\"])\n\n combined_rgb = torch.cat([image, rgb], dim=1)\n combined_acc = torch.cat([acc], dim=1)\n combined_depth = torch.cat([depth], dim=1)\n combined_alive_ray_mask = torch.cat([alive_ray_mask], dim=1)\n\n # Switch images from [H, W, C] to [1, C, H, W] for metrics computations\n image = torch.moveaxis(image, -1, 0)[None, ...]\n rgb = torch.moveaxis(rgb, -1, 0)[None, ...]\n\n psnr = self.psnr(image, rgb)\n ssim = self.ssim(image, rgb)\n lpips = self.lpips(image, rgb)\n\n # all of these metrics will be logged as scalars\n metrics_dict = {\"psnr\": float(psnr.item()), \"ssim\": float(ssim), \"lpips\": float(lpips)} # type: ignore\n # TODO(ethan): return an image dictionary\n\n images_dict = {\n \"img\": combined_rgb,\n \"accumulation\": combined_acc,\n \"depth\": combined_depth,\n \"alive_ray_mask\": combined_alive_ray_mask,\n }\n\n return metrics_dict, images_dict\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/models/mipnerf.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nImplementation of mip-NeRF.\n\"\"\"\nfrom __future__ import annotations\n\nfrom typing import Dict, List, Tuple\n\nimport torch\nfrom torch.nn import Parameter\nfrom torchmetrics import PeakSignalNoiseRatio\nfrom torchmetrics.functional import structural_similarity_index_measure\nfrom torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity\n\nfrom nerfstudio.cameras.rays import RayBundle\nfrom nerfstudio.field_components.encodings import NeRFEncoding\nfrom nerfstudio.field_components.field_heads import FieldHeadNames\nfrom nerfstudio.fields.vanilla_nerf_field import NeRFField\nfrom nerfstudio.model_components.losses import MSELoss\nfrom nerfstudio.model_components.ray_samplers import PDFSampler, UniformSampler\nfrom nerfstudio.model_components.renderers import (\n AccumulationRenderer,\n DepthRenderer,\n RGBRenderer,\n)\nfrom nerfstudio.models.base_model import Model, ModelConfig\nfrom nerfstudio.utils import colormaps, colors, misc\n\n\nclass MipNerfModel(Model):\n \"\"\"mip-NeRF model\n\n Args:\n config: MipNerf configuration to instantiate model\n \"\"\"\n\n def __init__(\n self,\n config: ModelConfig,\n **kwargs,\n ) -> None:\n self.field = None\n super().__init__(config=config, **kwargs)\n\n def populate_modules(self):\n \"\"\"Set the fields and modules\"\"\"\n super().populate_modules()\n\n # setting up fields\n position_encoding = NeRFEncoding(\n in_dim=3, num_frequencies=16, min_freq_exp=0.0, max_freq_exp=16.0, include_input=True\n )\n direction_encoding = NeRFEncoding(\n in_dim=3, num_frequencies=4, min_freq_exp=0.0, max_freq_exp=4.0, include_input=True\n )\n\n self.field = NeRFField(\n position_encoding=position_encoding, direction_encoding=direction_encoding, use_integrated_encoding=True\n )\n\n # samplers\n self.sampler_uniform = UniformSampler(num_samples=self.config.num_coarse_samples)\n self.sampler_pdf = PDFSampler(num_samples=self.config.num_importance_samples, include_original=False)\n\n # renderers\n self.renderer_rgb = RGBRenderer(background_color=colors.WHITE)\n self.renderer_accumulation = AccumulationRenderer()\n self.renderer_depth = DepthRenderer()\n\n # losses\n self.rgb_loss = MSELoss()\n\n # metrics\n self.psnr = PeakSignalNoiseRatio(data_range=1.0)\n self.ssim = structural_similarity_index_measure\n self.lpips = LearnedPerceptualImagePatchSimilarity()\n\n def get_param_groups(self) -> Dict[str, List[Parameter]]:\n param_groups = {}\n if self.field is None:\n raise ValueError(\"populate_fields() must be called before get_param_groups\")\n param_groups[\"fields\"] = list(self.field.parameters())\n return param_groups\n\n def get_outputs(self, ray_bundle: RayBundle):\n\n if self.field is None:\n raise ValueError(\"populate_fields() must be called before get_outputs\")\n\n # uniform sampling\n ray_samples_uniform = self.sampler_uniform(ray_bundle)\n\n # First pass:\n field_outputs_coarse = self.field.forward(ray_samples_uniform)\n weights_coarse = ray_samples_uniform.get_weights(field_outputs_coarse[FieldHeadNames.DENSITY])\n rgb_coarse = self.renderer_rgb(\n rgb=field_outputs_coarse[FieldHeadNames.RGB],\n weights=weights_coarse,\n )\n accumulation_coarse = self.renderer_accumulation(weights_coarse)\n depth_coarse = self.renderer_depth(weights_coarse, ray_samples_uniform)\n\n # pdf sampling\n ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse)\n\n # Second pass:\n field_outputs_fine = self.field.forward(ray_samples_pdf)\n weights_fine = ray_samples_pdf.get_weights(field_outputs_fine[FieldHeadNames.DENSITY])\n rgb_fine = self.renderer_rgb(\n rgb=field_outputs_fine[FieldHeadNames.RGB],\n weights=weights_fine,\n )\n accumulation_fine = self.renderer_accumulation(weights_fine)\n depth_fine = self.renderer_depth(weights_fine, ray_samples_pdf)\n\n outputs = {\n \"rgb_coarse\": rgb_coarse,\n \"rgb_fine\": rgb_fine,\n \"accumulation_coarse\": accumulation_coarse,\n \"accumulation_fine\": accumulation_fine,\n \"depth_coarse\": depth_coarse,\n \"depth_fine\": depth_fine,\n }\n return outputs\n\n def get_loss_dict(self, outputs, batch, metrics_dict=None):\n image = batch[\"image\"].to(self.device)\n rgb_loss_coarse = self.rgb_loss(image, outputs[\"rgb_coarse\"])\n rgb_loss_fine = self.rgb_loss(image, outputs[\"rgb_fine\"])\n loss_dict = {\"rgb_loss_coarse\": rgb_loss_coarse, \"rgb_loss_fine\": rgb_loss_fine}\n loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients)\n return loss_dict\n\n def get_image_metrics_and_images(\n self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]\n ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:\n image = batch[\"image\"].to(outputs[\"rgb_coarse\"].device)\n rgb_coarse = outputs[\"rgb_coarse\"]\n rgb_fine = outputs[\"rgb_fine\"]\n acc_coarse = colormaps.apply_colormap(outputs[\"accumulation_coarse\"])\n acc_fine = colormaps.apply_colormap(outputs[\"accumulation_fine\"])\n depth_coarse = colormaps.apply_depth_colormap(\n outputs[\"depth_coarse\"],\n accumulation=outputs[\"accumulation_coarse\"],\n near_plane=self.config.collider_params[\"near_plane\"],\n far_plane=self.config.collider_params[\"far_plane\"],\n )\n depth_fine = colormaps.apply_depth_colormap(\n outputs[\"depth_fine\"],\n accumulation=outputs[\"accumulation_fine\"],\n near_plane=self.config.collider_params[\"near_plane\"],\n far_plane=self.config.collider_params[\"far_plane\"],\n )\n\n combined_rgb = torch.cat([image, rgb_coarse, rgb_fine], dim=1)\n combined_acc = torch.cat([acc_coarse, acc_fine], dim=1)\n combined_depth = torch.cat([depth_coarse, depth_fine], dim=1)\n\n # Switch images from [H, W, C] to [1, C, H, W] for metrics computations\n image = torch.moveaxis(image, -1, 0)[None, ...]\n rgb_coarse = torch.moveaxis(rgb_coarse, -1, 0)[None, ...]\n rgb_fine = torch.moveaxis(rgb_fine, -1, 0)[None, ...]\n rgb_coarse = torch.clip(rgb_coarse, min=-1, max=1)\n rgb_fine = torch.clip(rgb_fine, min=-1, max=1)\n\n coarse_psnr = self.psnr(image, rgb_coarse)\n fine_psnr = self.psnr(image, rgb_fine)\n fine_ssim = self.ssim(image, rgb_fine)\n fine_lpips = self.lpips(image, rgb_fine)\n\n metrics_dict = {\n \"psnr\": float(fine_psnr.item()),\n \"coarse_psnr\": float(coarse_psnr.item()),\n \"fine_psnr\": float(fine_psnr.item()),\n \"fine_ssim\": float(fine_ssim.item()),\n \"fine_lpips\": float(fine_lpips.item()),\n }\n images_dict = {\"img\": combined_rgb, \"accumulation\": combined_acc, \"depth\": combined_depth}\n return metrics_dict, images_dict\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/models/monosdf.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nImplementation of MonoSDF.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass, field\nfrom typing import Type\n\nfrom nerfstudio.models.volsdf import VolSDFModel, VolSDFModelConfig\n\n\n@dataclass\nclass MonoSDFModelConfig(VolSDFModelConfig):\n \"\"\"Nerfacto Model Config\"\"\"\n\n _target: Type = field(default_factory=lambda: MonoSDFModel)\n mono_normal_loss_mult: float = 0.1\n \"\"\"Monocular normal consistency loss multiplier.\"\"\"\n mono_depth_loss_mult: float = 0.05\n \"\"\"Monocular depth consistency loss multiplier.\"\"\"\n\n\nclass MonoSDFModel(VolSDFModel):\n \"\"\"MonoSDF model\n\n Args:\n config: MonoSDF configuration to instantiate model\n \"\"\"\n\n config: MonoSDFModelConfig\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/models/nerfacto.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nNeRF implementation that combines many recent advancements.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass, field\nfrom typing import Dict, List, Tuple, Type\n\nimport numpy as np\nimport torch\nfrom torch.nn import Parameter\nfrom torchmetrics import PeakSignalNoiseRatio\nfrom torchmetrics.functional import structural_similarity_index_measure\nfrom torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity\nfrom typing_extensions import Literal\n\nfrom nerfstudio.cameras.rays import RayBundle\nfrom nerfstudio.engine.callbacks import (\n TrainingCallback,\n TrainingCallbackAttributes,\n TrainingCallbackLocation,\n)\nfrom nerfstudio.field_components.field_heads import FieldHeadNames\nfrom nerfstudio.field_components.spatial_distortions import SceneContraction\nfrom nerfstudio.fields.density_fields import HashMLPDensityField\nfrom nerfstudio.fields.nerfacto_field import TCNNNerfactoField\nfrom nerfstudio.model_components.losses import (\n MSELoss,\n distortion_loss,\n interlevel_loss,\n orientation_loss,\n pred_normal_loss,\n)\nfrom nerfstudio.model_components.ray_samplers import ProposalNetworkSampler\nfrom nerfstudio.model_components.renderers import (\n AccumulationRenderer,\n DepthRenderer,\n NormalsRenderer,\n RGBRenderer,\n)\nfrom nerfstudio.model_components.scene_colliders import NearFarCollider\nfrom nerfstudio.models.base_model import Model, ModelConfig\nfrom nerfstudio.utils import colormaps\nfrom nerfstudio.utils.colors import get_color\n\n\n@dataclass\nclass NerfactoModelConfig(ModelConfig):\n \"\"\"Nerfacto Model Config\"\"\"\n\n _target: Type = field(default_factory=lambda: NerfactoModel)\n near_plane: float = 0.05\n \"\"\"How far along the ray to start sampling.\"\"\"\n far_plane: float = 1000.0\n \"\"\"How far along the ray to stop sampling.\"\"\"\n background_color: Literal[\"random\", \"last_sample\", \"white\", \"black\"] = \"last_sample\"\n \"\"\"Whether to randomize the background color.\"\"\"\n num_levels: int = 16\n \"\"\"Number of levels of the hashmap for the base mlp.\"\"\"\n max_res: int = 1024\n \"\"\"Maximum resolution of the hashmap for the base mlp.\"\"\"\n log2_hashmap_size: int = 19\n \"\"\"Size of the hashmap for the base mlp\"\"\"\n num_proposal_samples_per_ray: Tuple[int, ...] = (256, 96)\n \"\"\"Number of samples per ray for the proposal network.\"\"\"\n num_nerf_samples_per_ray: int = 48\n \"\"\"Number of samples per ray for the nerf network.\"\"\"\n proposal_update_every: int = 5\n \"\"\"Sample every n steps after the warmup\"\"\"\n proposal_warmup: int = 5000\n \"\"\"Scales n from 1 to proposal_update_every over this many steps\"\"\"\n num_proposal_iterations: int = 2\n \"\"\"Number of proposal network iterations.\"\"\"\n use_same_proposal_network: bool = False\n \"\"\"Use the same proposal network. Otherwise use different ones.\"\"\"\n proposal_net_args_list: List[Dict] = field(\n default_factory=lambda: [\n {\"hidden_dim\": 16, \"log2_hashmap_size\": 17, \"num_levels\": 5, \"max_res\": 64},\n {\"hidden_dim\": 16, \"log2_hashmap_size\": 17, \"num_levels\": 5, \"max_res\": 256},\n ]\n )\n \"\"\"Arguments for the proposal density fields.\"\"\"\n interlevel_loss_mult: float = 1.0\n \"\"\"Proposal loss multiplier.\"\"\"\n distortion_loss_mult: float = 0.002\n \"\"\"Distortion loss multiplier.\"\"\"\n orientation_loss_mult: float = 0.0001\n \"\"\"Orientation loss multipier on computed noramls.\"\"\"\n pred_normal_loss_mult: float = 0.001\n \"\"\"Predicted normal loss multiplier.\"\"\"\n use_proposal_weight_anneal: bool = True\n \"\"\"Whether to use proposal weight annealing.\"\"\"\n use_appearance_embedding: bool = True\n \"\"\"Whether to use per-frame appearance embedding.\"\"\"\n use_average_appearance_embedding: bool = True\n \"\"\"Whether to use average appearance embedding or zeros for inference.\"\"\"\n proposal_weights_anneal_slope: float = 10.0\n \"\"\"Slope of the annealing function for the proposal weights.\"\"\"\n proposal_weights_anneal_max_num_iters: int = 1000\n \"\"\"Max num iterations for the annealing function.\"\"\"\n use_single_jitter: bool = True\n \"\"\"Whether use single jitter or not for the proposal networks.\"\"\"\n predict_normals: bool = False\n \"\"\"Whether to predict normals or not.\"\"\"\n\n\nclass NerfactoModel(Model):\n \"\"\"Nerfacto model\n\n Args:\n config: Nerfacto configuration to instantiate model\n \"\"\"\n\n config: NerfactoModelConfig\n\n def populate_modules(self):\n \"\"\"Set the fields and modules.\"\"\"\n super().populate_modules()\n\n scene_contraction = SceneContraction(order=float(\"inf\"))\n\n # Fields\n self.field = TCNNNerfactoField(\n self.scene_box.aabb,\n num_levels=self.config.num_levels,\n max_res=self.config.max_res,\n log2_hashmap_size=self.config.log2_hashmap_size,\n spatial_distortion=scene_contraction,\n num_images=self.num_train_data,\n use_pred_normals=self.config.predict_normals,\n use_appearance_embedding=self.config.use_appearance_embedding,\n use_average_appearance_embedding=self.config.use_average_appearance_embedding,\n )\n\n self.density_fns = []\n num_prop_nets = self.config.num_proposal_iterations\n # Build the proposal network(s)\n self.proposal_networks = torch.nn.ModuleList()\n if self.config.use_same_proposal_network:\n assert len(self.config.proposal_net_args_list) == 1, \"Only one proposal network is allowed.\"\n prop_net_args = self.config.proposal_net_args_list[0]\n network = HashMLPDensityField(self.scene_box.aabb, spatial_distortion=scene_contraction, **prop_net_args)\n self.proposal_networks.append(network)\n self.density_fns.extend([network.density_fn for _ in range(num_prop_nets)])\n else:\n for i in range(num_prop_nets):\n prop_net_args = self.config.proposal_net_args_list[min(i, len(self.config.proposal_net_args_list) - 1)]\n network = HashMLPDensityField(\n self.scene_box.aabb,\n spatial_distortion=scene_contraction,\n **prop_net_args,\n )\n self.proposal_networks.append(network)\n self.density_fns.extend([network.density_fn for network in self.proposal_networks])\n\n # Samplers\n update_schedule = lambda step: np.clip(\n np.interp(step, [0, self.config.proposal_warmup], [0, self.config.proposal_update_every]),\n 1,\n self.config.proposal_update_every,\n )\n self.proposal_sampler = ProposalNetworkSampler(\n num_nerf_samples_per_ray=self.config.num_nerf_samples_per_ray,\n num_proposal_samples_per_ray=self.config.num_proposal_samples_per_ray,\n num_proposal_network_iterations=self.config.num_proposal_iterations,\n single_jitter=self.config.use_single_jitter,\n update_sched=update_schedule,\n )\n\n # Collider\n self.collider = NearFarCollider(near_plane=self.config.near_plane, far_plane=self.config.far_plane)\n\n # renderers\n background_color = (\n get_color(self.config.background_color)\n if self.config.background_color in set([\"white\", \"black\"])\n else self.config.background_color\n )\n self.renderer_rgb = RGBRenderer(background_color=background_color)\n self.renderer_accumulation = AccumulationRenderer()\n self.renderer_depth = DepthRenderer()\n self.renderer_normals = NormalsRenderer()\n\n # losses\n self.rgb_loss = MSELoss()\n\n # metrics\n self.psnr = PeakSignalNoiseRatio(data_range=1.0)\n self.ssim = structural_similarity_index_measure\n self.lpips = LearnedPerceptualImagePatchSimilarity()\n\n def get_param_groups(self) -> Dict[str, List[Parameter]]:\n param_groups = {}\n param_groups[\"proposal_networks\"] = list(self.proposal_networks.parameters())\n param_groups[\"fields\"] = list(self.field.parameters())\n return param_groups\n\n def get_training_callbacks(\n self, training_callback_attributes: TrainingCallbackAttributes\n ) -> List[TrainingCallback]:\n callbacks = []\n if self.config.use_proposal_weight_anneal:\n # anneal the weights of the proposal network before doing PDF sampling\n N = self.config.proposal_weights_anneal_max_num_iters\n\n def set_anneal(step):\n # https://arxiv.org/pdf/2111.12077.pdf eq. 18\n train_frac = np.clip(step / N, 0, 1)\n bias = lambda x, b: (b * x) / ((b - 1) * x + 1)\n anneal = bias(train_frac, self.config.proposal_weights_anneal_slope)\n self.proposal_sampler.set_anneal(anneal)\n\n callbacks.append(\n TrainingCallback(\n where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],\n update_every_num_iters=1,\n func=set_anneal,\n )\n )\n callbacks.append(\n TrainingCallback(\n where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],\n update_every_num_iters=1,\n func=self.proposal_sampler.step_cb,\n )\n )\n return callbacks\n\n def get_outputs(self, ray_bundle: RayBundle):\n ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=self.density_fns)\n field_outputs = self.field(ray_samples, compute_normals=self.config.predict_normals)\n weights = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY])\n weights_list.append(weights)\n ray_samples_list.append(ray_samples)\n\n rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)\n depth = self.renderer_depth(weights=weights, ray_samples=ray_samples)\n accumulation = self.renderer_accumulation(weights=weights)\n\n outputs = {\n \"rgb\": rgb,\n \"accumulation\": accumulation,\n \"depth\": depth,\n }\n\n if self.config.predict_normals:\n outputs[\"normals\"] = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights)\n outputs[\"pred_normals\"] = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights)\n\n # These use a lot of GPU memory, so we avoid storing them for eval.\n if True or self.training:\n outputs[\"weights_list\"] = weights_list\n outputs[\"ray_samples_list\"] = ray_samples_list\n\n if self.training and self.config.predict_normals:\n outputs[\"rendered_orientation_loss\"] = orientation_loss(\n weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions\n )\n\n outputs[\"rendered_pred_normal_loss\"] = pred_normal_loss(\n weights.detach(),\n field_outputs[FieldHeadNames.NORMALS].detach(),\n field_outputs[FieldHeadNames.PRED_NORMALS],\n )\n\n for i in range(self.config.num_proposal_iterations):\n outputs[f\"prop_depth_{i}\"] = self.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i])\n\n return outputs\n\n def get_metrics_dict(self, outputs, batch):\n metrics_dict = {}\n image = batch[\"image\"].to(self.device)\n metrics_dict[\"psnr\"] = self.psnr(outputs[\"rgb\"], image)\n if self.training:\n metrics_dict[\"distortion\"] = distortion_loss(outputs[\"weights_list\"], outputs[\"ray_samples_list\"])\n return metrics_dict\n\n def get_loss_dict(self, outputs, batch, metrics_dict=None):\n loss_dict = {}\n image = batch[\"image\"].to(self.device)\n loss_dict[\"rgb_loss\"] = self.rgb_loss(image, outputs[\"rgb\"])\n if self.training:\n loss_dict[\"interlevel_loss\"] = self.config.interlevel_loss_mult * interlevel_loss(\n outputs[\"weights_list\"], outputs[\"ray_samples_list\"]\n )\n assert metrics_dict is not None and \"distortion\" in metrics_dict\n loss_dict[\"distortion_loss\"] = self.config.distortion_loss_mult * metrics_dict[\"distortion\"]\n if self.config.predict_normals:\n # orientation loss for computed normals\n loss_dict[\"orientation_loss\"] = self.config.orientation_loss_mult * torch.mean(\n outputs[\"rendered_orientation_loss\"]\n )\n\n # ground truth supervision for normals\n loss_dict[\"pred_normal_loss\"] = self.config.pred_normal_loss_mult * torch.mean(\n outputs[\"rendered_pred_normal_loss\"]\n )\n return loss_dict\n\n def get_image_metrics_and_images(\n self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]\n ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:\n\n image = batch[\"image\"].to(self.device)\n rgb = outputs[\"rgb\"]\n acc = colormaps.apply_colormap(outputs[\"accumulation\"])\n depth = colormaps.apply_depth_colormap(\n outputs[\"depth\"],\n accumulation=outputs[\"accumulation\"],\n )\n\n combined_rgb = torch.cat([image, rgb], dim=1)\n combined_acc = torch.cat([acc], dim=1)\n combined_depth = torch.cat([depth], dim=1)\n\n # Switch images from [H, W, C] to [1, C, H, W] for metrics computations\n image = torch.moveaxis(image, -1, 0)[None, ...]\n rgb = torch.moveaxis(rgb, -1, 0)[None, ...]\n\n psnr = self.psnr(image, rgb)\n ssim = self.ssim(image, rgb)\n lpips = self.lpips(image, rgb)\n\n # all of these metrics will be logged as scalars\n metrics_dict = {\"psnr\": float(psnr.item()), \"ssim\": float(ssim)} # type: ignore\n metrics_dict[\"lpips\"] = float(lpips)\n\n images_dict = {\"img\": combined_rgb, \"accumulation\": combined_acc, \"depth\": combined_depth}\n\n # normals to RGB for visualization. TODO: use a colormap\n if \"normals\" in outputs:\n images_dict[\"normals\"] = (outputs[\"normals\"] + 1.0) / 2.0\n if \"pred_normals\" in outputs:\n images_dict[\"pred_normals\"] = (outputs[\"pred_normals\"] + 1.0) / 2.0\n\n for i in range(self.config.num_proposal_iterations):\n key = f\"prop_depth_{i}\"\n prop_depth_i = colormaps.apply_depth_colormap(\n outputs[key],\n accumulation=outputs[\"accumulation\"],\n )\n images_dict[key] = prop_depth_i\n\n return metrics_dict, images_dict\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/models/neuralreconW.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nImplementation of VolSDF.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass, field\nfrom typing import List, Type\n\nfrom nerfstudio.engine.callbacks import (\n TrainingCallback,\n TrainingCallbackAttributes,\n TrainingCallbackLocation,\n)\nfrom nerfstudio.model_components.ray_samplers import NeuralReconWSampler\nfrom nerfstudio.model_components.scene_colliders import SphereCollider\nfrom nerfstudio.models.neus import NeuSModel, NeuSModelConfig\n\n\n@dataclass\nclass NeuralReconWModelConfig(NeuSModelConfig):\n \"\"\"UniSurf Model Config\"\"\"\n\n _target: Type = field(default_factory=lambda: NeuralReconWModel)\n\n\nclass NeuralReconWModel(NeuSModel):\n \"\"\"VolSDF model\n\n Args:\n config: MonoSDF configuration to instantiate model\n \"\"\"\n\n config: NeuralReconWModelConfig\n\n def populate_modules(self):\n \"\"\"Set the fields and modules.\"\"\"\n super().populate_modules()\n\n # voxel surface bybrid sampler from NeuralReconW\n self.sampler = NeuralReconWSampler(\n aabb=self.scene_box.aabb, coarse_binary_grid=self.scene_box.coarse_binary_gird\n )\n # Neural Reconstruction in the wild use sphere collider so we overwrite it here\n self.collider = SphereCollider(radius=1.0, soft_intersection=False)\n\n def get_training_callbacks(\n self, training_callback_attributes: TrainingCallbackAttributes\n ) -> List[TrainingCallback]:\n callbacks = super().get_training_callbacks(training_callback_attributes)\n\n # add sampler call backs\n sdf_fn = lambda x: self.field.forward_geonetwork(x)[:, 0].contiguous()\n callbacks.append(\n TrainingCallback(\n where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],\n update_every_num_iters=1,\n func=self.sampler.update_binary_grid,\n kwargs={\"sdf_fn\": sdf_fn},\n )\n )\n\n return callbacks\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/models/neus.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nImplementation of NeuS.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass, field\nfrom typing import Dict, List, Type\n\nfrom nerfstudio.cameras.rays import RayBundle\nfrom nerfstudio.engine.callbacks import (\n TrainingCallback,\n TrainingCallbackAttributes,\n TrainingCallbackLocation,\n)\nfrom nerfstudio.field_components.field_heads import FieldHeadNames\nfrom nerfstudio.fields.sdf_field import SDFField\nfrom nerfstudio.model_components.ray_samplers import NeuSSampler\nfrom nerfstudio.models.base_surface_model import SurfaceModel, SurfaceModelConfig\n\n\n@dataclass\nclass NeuSModelConfig(SurfaceModelConfig):\n \"\"\"NeuS Model Config\"\"\"\n\n _target: Type = field(default_factory=lambda: NeuSModel)\n num_samples: int = 64\n \"\"\"Number of uniform samples\"\"\"\n num_samples_importance: int = 64\n \"\"\"Number of importance samples\"\"\"\n num_up_sample_steps: int = 4\n \"\"\"number of up sample step, 1 for simple coarse-to-fine sampling\"\"\"\n base_variance: float = 64\n \"\"\"fixed base variance in NeuS sampler, the inv_s will be base * 2 ** iter during upsample\"\"\"\n perturb: bool = True\n \"\"\"use to use perturb for the sampled points\"\"\"\n cos_anneal_end: int = 50000\n \"\"\"iterations to apply cosine annealing when computing alpha values (avoid vanishing gradient and get a coarse surface in the beginning.) \"\"\"\n\n\nclass NeuSModel(SurfaceModel):\n \"\"\"NeuS model\n\n Args:\n config: NeuS configuration to instantiate model\n \"\"\"\n\n config: NeuSModelConfig\n\n def populate_modules(self):\n \"\"\"Set the fields and modules.\"\"\"\n super().populate_modules()\n\n self.sampler = NeuSSampler(\n num_samples=self.config.num_samples,\n num_samples_importance=self.config.num_samples_importance,\n num_samples_outside=self.config.num_samples_outside,\n num_upsample_steps=self.config.num_up_sample_steps,\n base_variance=self.config.base_variance,\n )\n\n self.anneal_end = self.config.cos_anneal_end # 50000\n\n def get_training_callbacks(\n self, training_callback_attributes: TrainingCallbackAttributes\n ) -> List[TrainingCallback]:\n callbacks = super().get_training_callbacks(training_callback_attributes)\n \n # anneal for cos in NeuS\n if self.anneal_end > 0:\n\n def set_anneal(step):\n anneal = min([1.0, step / self.anneal_end])\n self.field.set_cos_anneal_ratio(anneal)\n\n callbacks.append(\n TrainingCallback(\n where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],\n update_every_num_iters=1,\n func=set_anneal,\n )\n )\n\n return callbacks\n\n def sample_and_forward_field(self, ray_bundle: RayBundle) -> Dict:\n ray_samples = self.sampler(ray_bundle, sdf_fn=self.field.get_sdf)\n # save_points(\"a.ply\", ray_samples.frustums.get_start_positions().reshape(-1, 3).detach().cpu().numpy())\n field_outputs = self.field(ray_samples, return_alphas=True)\n weights, transmittance = ray_samples.get_weights_and_transmittance_from_alphas(\n field_outputs[FieldHeadNames.ALPHA]\n )\n bg_transmittance = transmittance[:, -1, :] # the last fg transmittance -> the first bg transmittance\n\n samples_and_field_outputs = {\n \"ray_samples\": ray_samples,\n \"field_outputs\": field_outputs,\n \"weights\": weights,\n \"bg_transmittance\": bg_transmittance,\n }\n return samples_and_field_outputs\n\n def get_metrics_dict(self, outputs, batch) -> Dict:\n metrics_dict = super().get_metrics_dict(outputs, batch)\n if self.training and isinstance(self.field, SDFField):\n # training statics\n metrics_dict[\"s_val\"] = self.field.deviation_network.get_variance().item()\n metrics_dict[\"inv_s\"] = 1.0 / self.field.deviation_network.get_variance().item()\n\n return metrics_dict\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/models/neus_acc.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nImplementation of VolSDF.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass, field\nfrom typing import List, Type\n\nimport nerfacc\nimport torch\n\nfrom nerfstudio.cameras.rays import RayBundle\nfrom nerfstudio.engine.callbacks import (\n TrainingCallback,\n TrainingCallbackAttributes,\n TrainingCallbackLocation,\n)\nfrom nerfstudio.field_components.field_heads import FieldHeadNames\nfrom nerfstudio.model_components.ray_samplers import NeuSAccSampler\nfrom nerfstudio.models.neus import NeuSModel, NeuSModelConfig\n\n\n@dataclass\nclass NeuSAccModelConfig(NeuSModelConfig):\n \"\"\"UniSurf Model Config\"\"\"\n\n _target: Type = field(default_factory=lambda: NeuSAccModel)\n sky_loss_mult: float = 0.01\n \"\"\"Sky segmentation normal consistency loss multiplier.\"\"\"\n\n\nclass NeuSAccModel(NeuSModel):\n \"\"\"VolSDF model\n\n Args:\n config: MonoSDF configuration to instantiate model\n \"\"\"\n\n config: NeuSAccModelConfig\n\n def populate_modules(self):\n \"\"\"Set the fields and modules.\"\"\"\n super().populate_modules()\n\n # voxel surface bybrid sampler from NeuralReconW\n self.sampler = NeuSAccSampler(aabb=self.scene_box.aabb, neus_sampler=self.sampler)\n\n def get_training_callbacks(\n self, training_callback_attributes: TrainingCallbackAttributes\n ) -> List[TrainingCallback]:\n callbacks = super().get_training_callbacks(training_callback_attributes)\n\n # add sampler call backs\n sdf_fn = lambda x: self.field.forward_geonetwork(x)[:, 0].contiguous()\n inv_s = self.field.deviation_network.get_variance\n callbacks.append(\n TrainingCallback(\n where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],\n update_every_num_iters=1,\n func=self.sampler.update_binary_grid,\n kwargs={\"sdf_fn\": sdf_fn, \"inv_s\": inv_s},\n )\n )\n\n callbacks.append(\n TrainingCallback(\n where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],\n update_every_num_iters=1,\n func=self.sampler.update_step_size,\n kwargs={\"inv_s\": inv_s},\n )\n )\n\n return callbacks\n\n def get_outputs(self, ray_bundle: RayBundle):\n # bootstrap with original Neus\n if self.sampler._update_counter.item() <= 0:\n return super().get_outputs(ray_bundle)\n\n ray_samples, ray_indices = self.sampler(ray_bundle, sdf_fn=self.field.get_sdf, alpha_fn=self.field.get_alpha)\n\n if ray_samples.shape[0] > 0:\n field_outputs = self.field(ray_samples, return_alphas=True)\n\n n_rays = ray_bundle.shape[0]\n weights = nerfacc.render_weight_from_alpha(\n field_outputs[FieldHeadNames.ALPHA],\n ray_indices=ray_indices,\n n_rays=n_rays,\n )\n rgb = nerfacc.accumulate_along_rays(\n weights, ray_indices, values=field_outputs[FieldHeadNames.RGB], n_rays=n_rays\n )\n normal = nerfacc.accumulate_along_rays(\n weights, ray_indices, values=field_outputs[FieldHeadNames.NORMAL], n_rays=n_rays\n )\n\n accumulation = nerfacc.accumulate_along_rays(weights, ray_indices, values=None, n_rays=n_rays)\n depth = nerfacc.accumulate_along_rays(\n weights,\n ray_indices,\n values=(ray_samples.frustums.starts + ray_samples.frustums.ends) / 2,\n n_rays=n_rays,\n )\n\n outputs = {\n \"rgb\": rgb,\n \"accumulation\": accumulation,\n \"depth\": depth,\n \"normal\": normal,\n }\n\n if self.training:\n grad_points = field_outputs[FieldHeadNames.GRADIENT]\n outputs.update({\"eik_grad\": grad_points})\n else:\n zeros = torch.zeros((ray_bundle.shape[0], 3), dtype=torch.float32, device=self.device)\n outputs = {\"rgb\": zeros, \"accumulation\": zeros[:, :1], \"depth\": zeros[:, :1], \"normal\": zeros}\n if self.training:\n outputs.update({\"eik_grad\": zeros})\n\n return outputs\n\n def get_metrics_dict(self, outputs, batch):\n metrics = super().get_metrics_dict(outputs, batch)\n metrics[\"acc_step_size\"] = self.sampler.step_size\n return metrics\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/models/neus_facto.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nImplementation of VolSDF.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass, field\nfrom typing import List, Type, Tuple, Dict\nfrom typing_extensions import Literal\nimport numpy as np\n\nimport torch\nfrom torch.nn import Parameter\n\nfrom nerfstudio.cameras.rays import RayBundle\nfrom nerfstudio.engine.callbacks import (\n TrainingCallback,\n TrainingCallbackAttributes,\n TrainingCallbackLocation,\n)\nfrom nerfstudio.field_components.field_heads import FieldHeadNames\nfrom nerfstudio.field_components.spatial_distortions import SceneContraction\nfrom nerfstudio.models.neus import NeuSModel, NeuSModelConfig\nfrom nerfstudio.fields.density_fields import HashMLPDensityField\nfrom nerfstudio.model_components.losses import interlevel_loss\nfrom nerfstudio.model_components.ray_samplers import ProposalNetworkSampler\nfrom nerfstudio.utils import colormaps\n\n\n@dataclass\nclass NeuSFactoModelConfig(NeuSModelConfig):\n \"\"\"UniSurf Model Config\"\"\"\n\n _target: Type = field(default_factory=lambda: NeuSFactoModel)\n num_proposal_samples_per_ray: Tuple[int, ...] = (256, 96)\n \"\"\"Number of samples per ray for the proposal network.\"\"\"\n num_neus_samples_per_ray: int = 48\n \"\"\"Number of samples per ray for the nerf network.\"\"\"\n proposal_update_every: int = 5\n \"\"\"Sample every n steps after the warmup\"\"\"\n proposal_warmup: int = 5000 # FIXME: not used\n \"\"\"Scales n from 1 to proposal_update_every over this many steps\"\"\"\n num_proposal_iterations: int = 2\n \"\"\"Number of proposal network iterations.\"\"\"\n use_same_proposal_network: bool = False\n \"\"\"Use the same proposal network. Otherwise use different ones.\"\"\"\n proposal_net_args_list: List[Dict] = field(\n default_factory=lambda: [\n {\"hidden_dim\": 16, \"log2_hashmap_size\": 17, \"num_levels\": 5, \"max_res\": 64},\n {\"hidden_dim\": 16, \"log2_hashmap_size\": 17, \"num_levels\": 5, \"max_res\": 256},\n ]\n )\n \"\"\"Arguments for the proposal density fields.\"\"\"\n interlevel_loss_mult: float = 1.0\n \"\"\"Proposal loss multiplier.\"\"\"\n use_proposal_weight_anneal: bool = True\n \"\"\"Whether to use proposal weight annealing.\"\"\"\n proposal_weights_anneal_slope: float = 10.0\n \"\"\"Slope of the annealing function for the proposal weights.\"\"\"\n proposal_weights_anneal_max_num_iters: int = 1000\n \"\"\"Max num iterations for the annealing function.\"\"\"\n use_single_jitter: bool = True\n \"\"\"Whether use single jitter or not for the proposal networks.\"\"\"\n proposal_net_spatial_normalization_region: Literal[\"full\", \"fg\"] = \"full\"\n \"\"\"if \"full\", normalize [-2, 2] into [0, 1], otherwise, normalize [-1, 1] into [0, 1].\"\"\"\n proposal_use_uniform_sampler: bool = True\n \"\"\"whether to use UniformSampler as the initial sampler for the proposal network.\n It's better to use UniformSampler if the region for the proposal network to model\n is bounded and small (e.g., inside a object bbox). UniformLinDispPiecewiseSampler is more\n appropriate, if the proposal network needs to model a big or unbounded region.\n \"\"\"\n proposal_net_use_separate_contraction: bool = False\n \"\"\"use a separate contraction for the proposal network (modeling both fg & bg regions).\"\"\"\n proposal_net_contraction_scale_factor: float = 1.0\n \"\"\"scale_factor for the separate SceneContraction of the proposal network.\"\"\"\n num_proposal_samples_per_ray_for_eikonal_loss: int = 0\n\n\nclass NeuSFactoModel(NeuSModel):\n \"\"\"NeuS facto model\n\n Args:\n config: NeuS configuration to instantiate model\n \"\"\"\n\n config: NeuSFactoModelConfig\n\n def populate_modules(self):\n \"\"\"Set the fields and modules.\"\"\"\n super().populate_modules()\n\n self.density_fns = []\n num_prop_nets = self.config.num_proposal_iterations\n \n # Build the proposal network(s)\n if self.config.proposal_net_use_separate_contraction:\n proposal_contraction = SceneContraction(\n order=float(\"inf\"), scale_factor=self.config.proposal_net_contraction_scale_factor\n )\n else:\n proposal_contraction = self.scene_contraction\n \n self.proposal_networks = torch.nn.ModuleList()\n if self.config.use_same_proposal_network:\n assert len(self.config.proposal_net_args_list) == 1, \"Only one proposal network is allowed.\"\n prop_net_args = self.config.proposal_net_args_list[0]\n network = HashMLPDensityField(\n self.scene_box.aabb, spatial_distortion=proposal_contraction, **prop_net_args,\n spatial_normalization_region=self.config.proposal_net_spatial_normalization_region,\n )\n self.proposal_networks.append(network)\n self.density_fns.extend([network.density_fn for _ in range(num_prop_nets)])\n else:\n for i in range(num_prop_nets):\n prop_net_args = self.config.proposal_net_args_list[\n min(i, len(self.config.proposal_net_args_list) - 1)]\n network = HashMLPDensityField(\n self.scene_box.aabb, spatial_distortion=proposal_contraction, **prop_net_args,\n spatial_normalization_region=self.config.proposal_net_spatial_normalization_region,\n )\n self.proposal_networks.append(network)\n self.density_fns.extend([network.density_fn for network in self.proposal_networks])\n\n # proposal network scheduler\n # TODO: previously trained NeuSFacto models use update_scheduler=-1 -> might lead to performance change\n update_schedule = lambda step: np.clip(\n np.interp(step, [0, self.config.proposal_warmup], [0, self.config.proposal_update_every]),\n 1,\n self.config.proposal_update_every,\n )\n \n self.proposal_sampler = ProposalNetworkSampler(\n num_nerf_samples_per_ray=self.config.num_neus_samples_per_ray,\n num_proposal_samples_per_ray=self.config.num_proposal_samples_per_ray,\n num_proposal_network_iterations=self.config.num_proposal_iterations,\n use_uniform_sampler=self.config.proposal_use_uniform_sampler,\n single_jitter=self.config.use_single_jitter,\n update_sched=update_schedule,\n )\n\n def get_param_groups(self) -> Dict[str, List[Parameter]]:\n param_groups = super().get_param_groups()\n param_groups[\"proposal_networks\"] = list(self.proposal_networks.parameters())\n return param_groups\n\n def get_training_callbacks(\n self, training_callback_attributes: TrainingCallbackAttributes\n ) -> List[TrainingCallback]:\n callbacks = super().get_training_callbacks(training_callback_attributes)\n\n if self.config.use_proposal_weight_anneal:\n # anneal the weights of the proposal network before doing PDF sampling\n N = self.config.proposal_weights_anneal_max_num_iters\n\n def set_anneal(step):\n # https://arxiv.org/pdf/2111.12077.pdf eq. 18\n train_frac = np.clip(step / N, 0, 1)\n bias = lambda x, b: (b * x) / ((b - 1) * x + 1)\n anneal = bias(train_frac, self.config.proposal_weights_anneal_slope)\n self.proposal_sampler.set_anneal(anneal)\n\n callbacks.append(\n TrainingCallback(\n where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],\n update_every_num_iters=1,\n func=set_anneal,\n )\n )\n callbacks.append(\n TrainingCallback(\n where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],\n update_every_num_iters=1,\n func=self.proposal_sampler.step_cb,\n )\n )\n\n return callbacks\n\n def sample_and_forward_field(self, ray_bundle: RayBundle):\n ray_samples, weights_list, ray_samples_list = self.proposal_sampler(\n ray_bundle, density_fns=self.density_fns\n )\n\n field_outputs = self.field(ray_samples, return_alphas=True)\n weights, transmittance = ray_samples.get_weights_and_transmittance_from_alphas(\n field_outputs[FieldHeadNames.ALPHA]\n )\n bg_transmittance = transmittance[:, -1, :]\n\n weights_list.append(weights)\n ray_samples_list.append(ray_samples)\n\n samples_and_field_outputs = {\n \"ray_samples\": ray_samples,\n \"field_outputs\": field_outputs,\n \"weights\": weights,\n \"bg_transmittance\": bg_transmittance,\n \"weights_list\": weights_list,\n \"ray_samples_list\": ray_samples_list,\n }\n return samples_and_field_outputs\n\n def get_loss_dict(self, outputs, batch, metrics_dict=None):\n loss_dict = super().get_loss_dict(outputs, batch, metrics_dict)\n\n if self.training:\n self._compute_interlevel_loss(outputs, loss_dict)\n self._compute_proposal_eikonal_loss(outputs, loss_dict)\n\n return loss_dict\n\n def get_image_metrics_and_images(\n self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]\n ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:\n metrics_dict, images_dict = super().get_image_metrics_and_images(outputs, batch)\n for i in range(self.config.num_proposal_iterations):\n key = f\"prop_depth_{i}\"\n # TODO: set depth[accumulation < 0.1] = depth[accumulation >= 0.1].min()\n prop_depth_i = colormaps.apply_depth_colormap(\n outputs[key],\n accumulation=outputs[\"accumulation\"], # NOTE: this would remove bg depths\n )\n images_dict[key] = prop_depth_i\n\n return metrics_dict, images_dict\n\n def _compute_interlevel_loss(self, outputs: Dict[str, torch.Tensor], loss_dict: Dict[str, torch.Tensor]):\n weights_list, ray_samples_list = self._get_interlevel_loss_inputs(outputs)\n interlevel_loss_val = interlevel_loss(weights_list, ray_samples_list)\n loss_dict[\"interlevel_loss\"] = self.config.interlevel_loss_mult * interlevel_loss_val\n \n def _get_interlevel_loss_inputs(self, outputs):\n weights_list = outputs[\"weights_list\"] # (n_rays, n_samples, 1)\n ray_samples_list = outputs[\"ray_samples_list\"] # (n_rays, n_samples)\n \n if not self.handle_no_intersect_rays:\n return weights_list, ray_samples_list\n \n # ignore rays not intersecting the fg aabb\n no_intersect_mask = outputs[\"no_intersect_mask\"][..., 0] # (n_rays, )\n weights_list = [weights[~no_intersect_mask] for weights in weights_list]\n ray_samples_list = [ray_samples[~no_intersect_mask] for ray_samples in ray_samples_list]\n \n return weights_list, ray_samples_list\n\n def _compute_proposal_eikonal_loss(\n self,\n outputs: Dict[str, torch.Tensor],\n loss_dict: Dict[str, torch.Tensor]\n ):\n \"\"\"sample additional ray samples from the proposal samples and compute the eikonal loss.\n this aims to regularize the internal geometry of the reconstruction.\n \"\"\"\n n_eikonal_samples = self.config.num_proposal_samples_per_ray_for_eikonal_loss\n if n_eikonal_samples <= 0:\n return\n \n _, ray_samples_list = self._get_interlevel_loss_inputs(outputs) # TODO: called twice per iter, should cache results\n ray_samples_list = [s.frustums.get_start_positions() for s in ray_samples_list[:-1]]\n # TODO: maybe it's better to just sample from the uniform samples?\n ray_samples = torch.cat(ray_samples_list, dim=1) # (n_rays, n_samples, 3)\n n_rays, n_samples = ray_samples.shape[:2]\n if n_eikonal_samples < n_samples: # sample with replacement for speed\n _inds = torch.randint(n_samples, (n_rays, n_eikonal_samples), device=ray_samples.device)\n ray_samples = ray_samples.gather(1, _inds[..., None].expand(-1, -1, 3))\n eik_grads = self.field.gradient(ray_samples.view(-1, 3)) # (*, 3)\n loss_dict[\"proposal_eikonal_loss\"] = ((eik_grads.norm(2, dim=-1) - 1) ** 2).mean() * self.config.eikonal_loss_mult\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/models/neus_facto_dff.py",
"content": "\"\"\"\nNeuSFacto + DistilledFeatureField\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass, field\nfrom typing import List, Type, Tuple, Dict, Optional\nfrom torchtyping import TensorType\nfrom collections import defaultdict\nfrom tqdm import tqdm\n\nimport torch\nfrom rich.console import Console\nfrom torch.nn import Parameter, MSELoss\nfrom torch.nn import functional as F\nfrom sklearn.decomposition import PCA\nfrom einops import rearrange, repeat\n\nfrom nerfstudio.cameras.rays import RayBundle, RaySamples\nfrom nerfstudio.field_components.field_heads import FieldHeadNames\nfrom nerfstudio.field_components.spatial_distortions import SceneContraction\nfrom nerfstudio.fields.feature_field import FeatureFieldConfig, FeatureSegFieldConfig\nfrom nerfstudio.model_components.renderers import FeatureRenderer, SemanticRenderer\nfrom nerfstudio.model_components.ray_samplers import RayTracingSampler\nfrom nerfstudio.models.neus_facto import NeuSFactoModel, NeuSFactoModelConfig\nfrom nerfstudio.model_components.ray_samplers import Sampler, SpacedSampler, DummySampler\nfrom nerfstudio.model_components.renderers import DepthRenderer\nfrom nerfstudio.utils.pointclouds import BasicPointClouds\nfrom nerfstudio.utils.colormaps import apply_colormap, apply_depth_colormap\nfrom nerfstudio.utils.mask_utils import overlay_mask_on_image\n\nCONSOLE = Console(width=120)\n\n\n@dataclass\nclass NeuSFactoDFFModelConfig(NeuSFactoModelConfig):\n \"\"\"NeusFactoDFF Model Config\"\"\"\n\n _target: Type = field(default_factory=lambda: NeuSFactoDFFModel)\n\n feature_field: FeatureFieldConfig = FeatureFieldConfig()\n feature_field_scene_contraction_scale_factor: float = 1.0\n \"\"\"use a separate scene contraction for feature field, since the feature field use a simple field\n to represent the whole scene, including both fg and bg.\n \"\"\"\n feature_loss_mult: float = 0.1 # NOTE: this multiplier is meaningless if all other parts are frozen\n \n render_fg_seg: bool = True\n \"\"\"rendering foreground segmentation masks during evaluation\"\"\"\n feature_seg_field: FeatureSegFieldConfig = FeatureSegFieldConfig()\n \"\"\"segmentation field config\"\"\"\n fg_seg_binary_thr: float = 0.5\n \"\"\"binarization threshold for 2d segmentation masks\"\"\"\n fg_3d_seg_binary_thr: float = 0.5\n \"\"\"binarization threshold for 3d segmentation masks\"\"\"\n sdf_rectification_value: float = 10.0\n \"\"\"rectify absolute SDF values of positions where the segmentation field equals 0 to this value\"\"\"\n \n use_surface_rendering: bool = False\n \"\"\"learn the feature field with surface rendering\"\"\"\n # TODO: RayTracingSampler configs\n ignore_false_positive_sphere_tracing_with_accumulation: bool = True\n \"\"\"sphere tracing might produce some incorrect intersections near to the border of\n fg bboxes because sdf here are invalid, optionally filter out theses incorrect intersections\n whose accumulation is less than 0.5\n \"\"\"\n complement_false_negative_sphere_tracing_with_volrend: bool = True\n \"\"\"sphere tracing might fail to find some intersections, use depth of volume rendering\n as pseudo intersections.\n \"\"\"\n \n\nclass NeuSFactoDFFModel(NeuSFactoModel):\n \"\"\"NeuS facto model\n\n Args:\n config: NeuSFactoDFFModel configuration to instantiate model\n \"\"\"\n\n config: NeuSFactoDFFModelConfig\n\n def populate_modules(self):\n \"\"\"Set feature fields for foreground and background\"\"\"\n super().populate_modules()\n self.ptcd_data: BasicPointClouds = self.kwargs[\"metadata\"][\"ptcd_data\"]\n self.ptcd_data.build_nn_search_index(self.config.feature_seg_field.separate_bg_plane)\n \n self.feature_field_scene_contraction = SceneContraction(\n order=float(\"inf\"),\n scale_factor=self.config.feature_field_scene_contraction_scale_factor\n )\n self.feature_field = self.config.feature_field.setup(\n spatial_distortion=self.feature_field_scene_contraction\n )\n self.feature_seg_field = self.config.feature_seg_field.setup(ptcd_data=self.ptcd_data,\n aabb=self.scene_box.aabb)\n self.renderer_depth_median = DepthRenderer(\"median\")\n self.renderer_feature = FeatureRenderer()\n self.renderer_fg_seg = SemanticRenderer()\n self.feature_loss = MSELoss()\n \n if self.config.use_surface_rendering:\n self.ray_tracing_sampler = RayTracingSampler(aabb=self.scene_box.aabb)\n\n def get_param_groups(self) -> Dict[str, List[Parameter]]:\n param_groups = super().get_param_groups()\n for param_group in param_groups.values():\n for param in param_group:\n param.requires_grad_(False)\n CONSOLE.print(\"All parameters except for feature field have been frozen.\")\n\n param_groups[\"feature_field\"] = list(self.feature_field.parameters())\n\n return param_groups\n\n def sample_and_forward_field(self, ray_bundle: RayBundle):\n # if not (self.training and self.config.use_surface_rendering):\n # samples_and_field_outputs = super().sample_and_forward_field(ray_bundle)\n # ray_samples = samples_and_field_outputs[\"ray_samples\"]\n # else:\n # # TODO: not working yet (need to udpate SurfaceModel to ignore fg rendering)\n # # during training, we can avoid the ``super().sample_and_forward_field(ray_bundle)`` call\n # # to speed up inference, since we only need to render the feature field\n # samples_and_field_outputs = {}\n # ray_samples = None\n \n samples_and_field_outputs = super().sample_and_forward_field(ray_bundle)\n volrend_ray_samples = samples_and_field_outputs[\"ray_samples\"]\n volrend_weights = samples_and_field_outputs[\"weights\"]\n \n if self.config.use_surface_rendering:\n samples_and_field_outputs_rt = self._sample_and_forward_field_ray_tracing(\n ray_bundle, volrend_ray_samples, volrend_weights\n )\n samples_and_field_outputs.update(samples_and_field_outputs_rt) # \"{xxx}_st\"\n ray_samples = samples_and_field_outputs_rt[\"ray_samples_rt\"]\n else:\n ray_samples = volrend_ray_samples\n \n feature_field_outputs = self.feature_field(ray_samples)\n samples_and_field_outputs.update(feature_field_outputs)\n \n if not self.training and self.config.render_fg_seg:\n feature_field_seg_outptus = self.feature_seg_field(\n ray_samples, feature_field_outputs[FieldHeadNames.FEATURE]\n )\n samples_and_field_outputs.update(feature_field_seg_outptus)\n\n return samples_and_field_outputs\n\n def _sample_and_forward_field_ray_tracing(\n self,\n ray_bundle: RayBundle,\n volrend_ray_samples: RaySamples,\n volrend_weights: TensorType[\"bs\": ..., 1],\n ):\n sdf_fn = lambda x: self.field.forward_geonetwork(x)[..., [0]] # (..., 3) -> (..., 1)\n ray_samples, non_convergent_masks = self.ray_tracing_sampler(ray_bundle, sdf_fn) # (..., 1, 3), (..., 1, 1)\n non_convergent_masks_raw = non_convergent_masks.clone()\n\n if (\n self.config.ignore_false_positive_sphere_tracing_with_accumulation or\n self.config.complement_false_negative_sphere_tracing_with_volrend\n ):\n if ray_bundle.metadata.get(\"no_intersect_mask\", None) is not None:\n no_intersect_mask = ray_bundle.metadata[\"no_intersect_mask\"]\n no_intersect_mask = repeat(no_intersect_mask, \"... m -> ... n m\", n=volrend_weights.shape[-2])\n volrend_weights = torch.where(no_intersect_mask, torch.zeros_like(volrend_weights), volrend_weights)\n volrend_acc = self.renderer_accumulation(weights=volrend_weights)\n acc_bg_mask = volrend_acc < 0.5\n if self.config.ignore_false_positive_sphere_tracing_with_accumulation:\n non_convergent_masks |= acc_bg_mask[..., None]\n \n if self.config.complement_false_negative_sphere_tracing_with_volrend:\n volrend_ray_dist = self.renderer_depth_median(\n weights=volrend_weights, ray_samples=volrend_ray_samples\n )[..., None]\n false_negative_mask = non_convergent_masks & ~acc_bg_mask[..., None]\n non_convergent_masks[false_negative_mask] = False\n ray_samples.frustums.starts[false_negative_mask] = volrend_ray_dist[false_negative_mask]\n ray_samples.frustums.ends[false_negative_mask] = volrend_ray_dist[false_negative_mask]\n \n # assign weight = 1.0 for the single ray sample (intersection pt)\n # equivalent w/ self._handle_no_intersect_rays()\n weights = torch.where(non_convergent_masks,\n torch.zeros_like(ray_samples.frustums.starts),\n torch.ones_like(ray_samples.frustums.starts)) # (..., 1, 1)\n bg_transmittance = torch.where(non_convergent_masks[..., 0, [0]],\n torch.ones_like(weights[..., 0, [0]]),\n torch.zeros_like(weights[..., 0, [0]])) # (..., 1)\n \n # ray tracing depth\n depth = ray_samples.frustums.starts[..., 0, :] / ray_bundle.directions_norm # starts = ends\n \n samples_and_field_outputs = {\n \"non_convergent_mask_rt\": non_convergent_masks,\n \"non_convergent_mask_rt_raw\": non_convergent_masks_raw, # w/o postprocessing\n \"ray_samples_rt\": ray_samples,\n \"weights_rt\": weights,\n \"bg_transmittance_rt\": bg_transmittance,\n \"depth_rt\": depth\n }\n return samples_and_field_outputs\n \n def get_outputs(self, ray_bundle: RayBundle) -> Dict:\n outputs = super().get_outputs(ray_bundle)\n \n samples_and_field_outputs = outputs[\"samples_and_field_outputs\"]\n _weights_key = \"weights\" if not self.config.use_surface_rendering else \"weights_rt\"\n _trans_key = \"bg_transmittance\" if not self.config.use_surface_rendering else \"bg_transmittance_rt\"\n weights = samples_and_field_outputs[_weights_key]\n bg_transmittance = samples_and_field_outputs[_trans_key]\n\n # fg rendering (field query is done in sample_and_forward_field)\n feature = self.renderer_feature(\n features=samples_and_field_outputs[FieldHeadNames.FEATURE], weights=weights\n )\n\n if not self.training and self.config.render_fg_seg:\n with torch.no_grad():\n fg_seg_prob = self.renderer_fg_seg(\n samples_and_field_outputs[FieldHeadNames.FG_SEG].detach(), weights=weights.detach()\n )\n outputs[\"fg_seg_prob\"] = fg_seg_prob\n\n # bg field qeury and rendering\n if self.config.background_model != \"none\":\n samples_and_field_outputs_bg = outputs[\"samples_and_field_outputs_bg\"]\n ray_samples_bg = samples_and_field_outputs_bg[\"ray_samples\"]\n weights_bg = samples_and_field_outputs_bg[\"weights\"]\n feature_field_outputs_bg = self.feature_field(ray_samples_bg)\n feature_bg = self.renderer_feature(\n features=feature_field_outputs_bg[FieldHeadNames.FEATURE], weights=weights_bg\n )\n # composition\n feature = feature + bg_transmittance * feature_bg\n\n outputs[\"feature\"] = feature\n if self.config.use_surface_rendering:\n outputs[\"non_convergent_mask_rt\"] = samples_and_field_outputs[\"non_convergent_mask_rt\"]\n outputs[\"non_convergent_mask_rt_raw\"] = samples_and_field_outputs[\"non_convergent_mask_rt_raw\"]\n outputs[\"depth_rt\"] = samples_and_field_outputs[\"depth_rt\"]\n return outputs\n\n def get_loss_dict(self, outputs, batch, metrics_dict=None):\n loss_dict = super().get_loss_dict(outputs, batch, metrics_dict)\n\n if self.training and self.config.feature_loss_mult > 0.:\n loss_dict[\"feature_loss\"] = (\n self.feature_loss(batch[\"feature\"], outputs[\"feature\"]) *\n self.config.feature_loss_mult\n )\n\n return loss_dict\n\n def get_image_metrics_and_images(\n self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]\n ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:\n metrics_dict, images_dict = super().get_image_metrics_and_images(outputs, batch)\n\n self._get_pca_feature_rendering(outputs, batch, images_dict)\n self._get_fg_seg_rendering(outputs, batch, images_dict)\n self._get_ray_tracing_mask(outputs, batch, images_dict)\n self._get_ray_tracing_depth(outputs, batch, images_dict)\n\n return metrics_dict, images_dict\n\n def query_seg_field(self, inputs: TensorType[\"bs\": ..., 3]) -> TensorType[\"bs\": ..., 1]:\n \"\"\"query segmentation field\"\"\"\n feat = self.feature_field.query_features(inputs, None)\n seg = self.feature_seg_field.compute_seg_prob(inputs, feat)\n return seg\n \n def seg_aware_sdf(self, inputs: TensorType[\"bs\": ..., 3]) -> TensorType[\"bs\": ..., 1]:\n \"\"\"extract segmentation-awared sdf values\"\"\"\n sdf = self.field.forward_geonetwork(inputs)[..., [0]]\n seg = self.query_seg_field(inputs)\n \n # rectify sdf with seg\n rectified_sdf = self._rectify_sdf(sdf, seg)\n return rectified_sdf\n \n def _rectify_sdf(self, sdf: TensorType[\"bs\": ..., 1], seg: TensorType[\"bs\": ..., 1]):\n binary_seg = (seg > self.config.fg_3d_seg_binary_thr)[..., 0]\n bg_sdf = sdf[~binary_seg]\n min_bg_udf = self.config.sdf_rectification_value\n # sdf[~binary_seg] = torch.sign(bg_sdf) * torch.clamp(bg_sdf.abs(), min=min_bg_udf) # Bad results, the ground plane is kept\n sdf[~binary_seg] = torch.clamp(bg_sdf.abs(), min=min_bg_udf) # This might cause noises inside the fg region\n # TODO: visualize raw sdf (sdf viewer)\n return sdf\n \n def _get_pca_feature_rendering(\n self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor],\n images_dict: Dict[str, torch.Tensor]\n ):\n feature = outputs[\"feature\"] # (h, w, c)\n feature_gt = batch[\"feature\"] # (c, h_gt, w_gt) - low res\n feature_gt_pca = batch[\"feature_pca\"] # (h_gt, w_gt, 3)\n all_gt_pca: PCA = batch[\"gt_features_pca\"]\n h, w = feature.shape[:2]\n h_gt, w_gt = feature.shape[1:]\n min_val, max_val = batch[\"feature_pca_min\"], batch[\"feature_pca_max\"]\n\n # joint visualization of pred & gt (use PCA params fitted on gt features from all frames)\n feature_pca = all_gt_pca.transform(rearrange(feature.cpu().numpy(), 'h w c -> (h w) c'))\n feature_pca = rearrange(feature.new_tensor(feature_pca), '(h w) c -> h w c', h=h, w=w)\n feature_gt_pca = F.interpolate(\n feature_gt_pca.permute(2, 0, 1)[None], size=(h, w), mode=\"bilinear\"\n )[0].permute(1, 2, 0) # (h, w, 3)\n\n feature_pca_gt_pred = torch.cat([feature_gt_pca, feature_pca], dim=1)\n feature_pca_gt_pred = (feature_pca_gt_pred - min_val) / (max_val - min_val)\n\n images_dict[\"feature_pca_gt_pred\"] = feature_pca_gt_pred\n # TODO: separate visualization of pred (use PCA params fitted on pred only)\n \n # pre-exp: compute a segmentation mask with feature_gt and self.ptcd_data\n # self._get_pseudo_gt_seg_mask(batch[\"image\"].to(self.device), feature_gt, images_dict)\n\n def _get_fg_seg_rendering(\n self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor],\n images_dict: Dict[str, torch.Tensor]\n ):\n # plot soft fg seg\n fg_seg_prob = outputs['fg_seg_prob'] # (h, w, 1)\n soft_fg_seg = apply_colormap(fg_seg_prob, cmap=\"viridis\")\n \n # plot binary fg seg\n fg_seg_binary = (fg_seg_prob >= self.config.fg_seg_binary_thr).float()\n binary_fg_seg = fg_seg_binary.repeat(1, 1, 3)\n \n # overlay fg seg rendering on gt image\n image = batch[\"image\"].to(self.device)\n # fg_seg_overlay = overlay_mask_on_image(soft_fg_seg, image, fg_seg_prob)\n fg_seg_overlay = overlay_mask_on_image(binary_fg_seg, image, fg_seg_prob)\n\n images_dict.update({\n \"fg_seg_prob\": soft_fg_seg,\n \"fg_seg\": binary_fg_seg,\n \"fg_seg_overlay\": fg_seg_overlay\n })\n\n def _get_pseudo_gt_seg_mask(self,\n image_gt: TensorType[\"h\", \"w\", 3],\n feature_gt: TensorType[\"c\", \"h_gt\", \"w_gt\"],\n images_dict: Dict[str, torch.Tensor]):\n nc_feat, h_feat, w_feat = feature_gt.shape\n h, w = image_gt.shape[:2]\n \n feature_gt = rearrange(feature_gt, 'c h w -> (h w) c')\n assert not self.feature_seg_field.config.distance_weighted_average\n soft_psuedo_gt_mask = rearrange(\n self.feature_seg_field.compute_seg_prob(None, feature_gt),\n '(h w) 1 -> h w 1', h=h_feat, w=w_feat\n )\n soft_psuedo_gt_mask_hr = F.interpolate(\n soft_psuedo_gt_mask.permute(2, 0, 1)[None], size=(h, w), mode=\"bilinear\"\n )[0].permute(1, 2, 0) # (h, w, 1)\n binary_pseudo_gt_mask_hr = (soft_psuedo_gt_mask_hr > self.config.fg_seg_binary_thr).float()\n seg_overlay = overlay_mask_on_image(binary_pseudo_gt_mask_hr, image_gt, soft_psuedo_gt_mask_hr)\n images_dict[\"pseudo_gt_seg_overlay\"] = seg_overlay\n\n def _get_ray_tracing_mask(\n self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor],\n images_dict: Dict[str, torch.Tensor]\n ) -> None:\n if not self.config.use_surface_rendering:\n return\n \n image = batch[\"image\"].to(self.device)\n \n # postprocessed non_convergent_mask\n convergent_masks = (~outputs[\"non_convergent_mask_rt\"]).float() # (h, w, 1)\n binary_rt_mask = convergent_masks.repeat(1, 1, 3)\n rt_mask_overlay = overlay_mask_on_image(binary_rt_mask, image, convergent_masks)\n \n # raw non_convergent_mask\n non_convergent_masks_raw = (~outputs[\"non_convergent_mask_rt_raw\"]).float() # (h, w, 1)\n binary_rt_mask = non_convergent_masks_raw.repeat(1, 1, 3)\n rt_mask_raw_overlay = overlay_mask_on_image(binary_rt_mask, image, non_convergent_masks_raw)\n \n # update images_dict\n images_dict.update({\n \"rt_convergence_overlay\": rt_mask_overlay,\n \"rt_convergence_raw_overlay\": rt_mask_raw_overlay,\n })\n \n def _get_ray_tracing_depth(\n self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor],\n images_dict: Dict[str, torch.Tensor]\n ) -> None:\n if not self.config.use_surface_rendering:\n return\n \n depth = outputs[\"depth_rt\"]\n convergent_mask = (~outputs[\"non_convergent_mask_rt\"])\n min_depth = depth[convergent_mask].min()\n depth[~convergent_mask] = min_depth\n \n depth = apply_depth_colormap(depth, convergent_mask.float())\n images_dict[\"depth_rt\"] = depth\n \n @torch.no_grad()\n def get_outputs_for_mesh_culling(\n self,\n ray_bundle: RayBundle,\n ray_sampler: Optional[Sampler] = None\n ) -> Dict[str, torch.Tensor]:\n \"\"\"Render segmentation with foreground model for mesh culling.\"\"\"\n self.handle_no_intersect_rays = False\n \n num_rays_per_chunk = self.config.eval_num_rays_per_chunk\n num_rays = len(ray_bundle)\n outputs_lists = defaultdict(list)\n \n # chunked inference\n for i in tqdm(range(0, num_rays, num_rays_per_chunk), desc=\"Chunked inference for mesh culling\"):\n start_idx = i\n end_idx = i + num_rays_per_chunk\n ray_bundle_chunk = ray_bundle.get_row_major_sliced_ray_bundle(start_idx, end_idx)\n \n outputs = self.render_fg_seg_for_mesh_culling(ray_bundle_chunk, ray_sampler=ray_sampler)\n \n for output_name, output in outputs.items(): # type: ignore\n outputs_lists[output_name].append(output)\n \n # merge chunks\n outputs = {}\n for output_name, outputs_list in outputs_lists.items():\n if not torch.is_tensor(outputs_list[0]):\n # TODO: handle lists of non-tensors as well\n continue\n outputs[output_name] = torch.cat(outputs_list)\n return outputs\n\n @torch.no_grad()\n def render_fg_seg_for_mesh_culling(\n self,\n ray_bundle: RayBundle,\n ray_sampler: Optional[Sampler] = None\n ) -> Dict[str, torch.Tensor]:\n # query seg_probs & weights\n if ray_sampler is None: # use the same sampler as for training\n samples_and_field_outputs = self.sample_and_forward_field(ray_bundle=ray_bundle)\n _weights_key = \"weights\" if not self.config.use_surface_rendering else \"weights_rt\"\n weights = samples_and_field_outputs[_weights_key].detach()\n fg_seg_probs = samples_and_field_outputs[FieldHeadNames.FG_SEG].detach()\n else:\n if self.config.use_surface_rendering:\n assert isinstance(ray_sampler, DummySampler)\n ray_samples = ray_sampler(ray_bundle)\n weights = torch.ones_like(ray_samples.frustums.starts)\n else:\n assert isinstance(ray_sampler, SpacedSampler)\n ray_samples = ray_sampler(ray_bundle, num_samples=ray_sampler.num_samples)\n field_outputs = self.field(ray_samples, return_alphas=True)\n weights = ray_samples.get_weights_from_alphas(field_outputs[FieldHeadNames.ALPHA])\n feature_field_outputs = self.feature_field(ray_samples)\n feature_field_seg_outptus = self.feature_seg_field(\n ray_samples, feature_field_outputs[FieldHeadNames.FEATURE]\n )\n fg_seg_probs = feature_field_seg_outptus[FieldHeadNames.FG_SEG].detach()\n \n # rendering\n fg_seg_prob_rendering = self.renderer_fg_seg(fg_seg_probs, weights=weights.detach())\n outputs = {\n \"fg_seg_prob\": fg_seg_prob_rendering,\n \"fg_seg_probs\": fg_seg_probs,\n \"weights\": weights\n }\n return outputs\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/models/neus_facto_reg.py",
"content": "\"\"\"\nNeuSFacto w/ explicit point cloud regularization for scene decomposition.\n\"\"\"\nfrom __future__ import annotations\n\nfrom copy import deepcopy\nfrom dataclasses import dataclass, field\nfrom typing import Any, Dict, List, Literal, Optional, Tuple, Type\n\nimport numpy as np\nimport torch\nfrom rich.console import Console\nfrom torch.nn import Parameter\nfrom torch.nn import functional as F\nfrom torchtyping import TensorType\n\nfrom nerfstudio.cameras.rays import RayBundle, RaySamples\nfrom nerfstudio.field_components.field_heads import FieldHeadNames\nfrom nerfstudio.fields.nerfacto_field import TCNNNerfactoField\nfrom nerfstudio.fields.sdf_field import SDFField, SDFFieldConfig\nfrom nerfstudio.model_components.ray_samplers import NeuSSampler, UniformSampler\nfrom nerfstudio.models.neus_facto import NeuSFactoModel, NeuSFactoModelConfig\nfrom nerfstudio.utils import colormaps, writer\n\nCONSOLE = Console(width=120)\n\n\n@dataclass\nclass NeuSFactoRegModelConfig(NeuSFactoModelConfig):\n \"\"\"NeusFactoReg Model Config\"\"\"\n\n _target: Type = field(default_factory=lambda: NeuSFactoRegModel)\n\n plane_field_type: Literal[\"nerf\", \"sdf\"] = \"nerf\"\n plane_sdf_field: SDFFieldConfig = SDFFieldConfig(\n spatial_normalization_region=\"aabb\",\n use_normalized_raw_coords=False, # False might cause bad geo-init, all initial sdf are large and leads to low weights\n use_position_encoding=False,\n use_grid_feature=True,\n hash_grid_num_levels=14, # 5\n hash_grid_log2_hashmap_size=17,\n hash_grid_max_res=2048,\n num_layers=1,\n num_layers_color=2,\n hidden_dim=64, # 32\n geo_feat_dim=15,\n direction_encoding_type=\"sh\",\n hidden_dim_color=64, # 32\n color_network_include_sdf=False,\n bias=0.5, # 0.5 -> 0.1 TODO: initialize sdf_field as a plane\n beta_init=0.3,\n use_appearance_embedding=False,\n inside_outside=False,\n weight_norm=False,\n )\n plane_height_ratio: float = 0.1\n \"\"\"ratio of the scene box height to use as the height of plane aabb\"\"\"\n num_samples_plane: int = 16\n plane_hierarchical_sampling: bool = False\n \"\"\"number of samples to use for plane sdf field\"\"\"\n ptcd_reg_plane_field_mult: float = 0.0\n ptcd_reg_plane_field_n_pts_per_iter: int = 5000\n\n\nclass NeuSFactoRegModel(NeuSFactoModel):\n \"\"\"NeusFacto w/ explicit point cloud regularization and separate modeling of the ground plane for scene decomposition.\"\"\"\n\n config: NeuSFactoRegModelConfig\n\n def populate_modules(self):\n super().populate_modules()\n\n # build aabb for plane region\n scene_box_height = (self.scene_box.aabb[1, 1] - self.scene_box.aabb[0, 1]).abs()\n self.plane_box_height = scene_box_height * self.config.plane_height_ratio\n plane_aabb = self.scene_box.aabb.clone()\n plane_aabb[0, 1] = plane_aabb[1, 1] - self.plane_box_height # assume +y points downward\n # TODO: use the exact plane center and smaller plane bbox extent\n\n # build plane field\n if self.config.plane_field_type == \"sdf\":\n self.field_plane = self.config.plane_sdf_field.setup(\n aabb=plane_aabb,\n spatial_distortion=self.scene_contraction,\n num_images=self.num_train_data,\n use_average_appearance_embedding=self.config.use_average_appearance_embedding,\n compute_curvature_loss=self.config.curvature_loss_mult > 0,\n )\n elif self.config.plane_field_type == \"nerf\":\n self.field_plane = TCNNNerfactoField(\n plane_aabb,\n spatial_distortion=self.scene_contraction,\n num_images=self.num_train_data,\n use_appearance_embedding=self.config.bg_use_appearance_embedding,\n use_average_appearance_embedding=self.config.use_average_appearance_embedding,\n num_levels=16,\n max_res=256,\n log2_hashmap_size=17,\n hidden_dim=32,\n num_layers_color=3,\n hidden_dim_color=32,\n )\n else:\n raise ValueError(f\"Unknown plane field type: {self.config.plane_field_type}\")\n\n # build plane ray sampler\n if not self.config.plane_hierarchical_sampling:\n self.sampler_plane = UniformSampler(num_samples=self.config.num_samples_plane)\n else:\n assert self.config.plane_field_type == \"sdf\"\n self.sampler_plane = NeuSSampler(\n num_samples=self.config.num_samples_plane * 2,\n num_samples_importance=self.config.num_samples_plane,\n num_upsample_steps=4,\n )\n\n def get_param_groups(self) -> Dict[str, List[Parameter]]:\n param_groups = super().get_param_groups()\n param_groups[\"field_plane\"] = list(self.field_plane.parameters())\n return param_groups\n\n def sample_and_forward_field(self, ray_bundle: RayBundle):\n samples_and_field_outputs = super().sample_and_forward_field(ray_bundle)\n\n # sample additional points on top of the bottom plane\n mask = ray_bundle.metadata[\"bottom_intersect_mask\"][..., 0] # (n,)\n if not mask.any():\n return samples_and_field_outputs\n\n ray_bundle_plane = deepcopy(ray_bundle)[mask]\n ray_d, far, near = ray_bundle_plane.directions, ray_bundle_plane.fars, ray_bundle_plane.nears # (n, *)\n y_dir = ray_d.new_tensor([[0.0, 1.0, 0.0]])\n _cos = (ray_d * y_dir).sum(-1) / ray_d.norm(dim=-1)\n ray_dist_offset = self.plane_box_height / _cos[..., None] # (n, 1)\n new_near = far - ray_dist_offset\n ray_bundle_plane.nears = torch.max(new_near, near)\n\n plane_field_outputs = {}\n if self.config.plane_hierarchical_sampling:\n ray_samples_plane = self.sampler_plane(ray_bundle_plane, self.field_plane.get_sdf)\n else:\n ray_samples_plane = self.sampler_plane(ray_bundle_plane)\n\n if self.config.plane_field_type == \"sdf\":\n field_outputs = self.field_plane(ray_samples_plane, return_alphas=True)\n weights, transmittance = ray_samples_plane.get_weights_and_transmittance_from_alphas(\n field_outputs[FieldHeadNames.ALPHA]\n ) # (n_bottom, n_samples_plane, 1)\n plane_field_outputs.update(\n {\n \"plane_normals\": field_outputs[FieldHeadNames.NORMAL],\n \"plane_eik_grad\": field_outputs[FieldHeadNames.GRADIENT],\n }\n )\n else: # nerf\n field_outputs = self.field_plane(ray_samples_plane)\n densities = field_outputs[FieldHeadNames.DENSITY]\n weights, transmittance = ray_samples_plane.get_weights_and_transmittance(densities)\n\n samples_and_field_outputs.update(\n {\n **plane_field_outputs,\n \"plane_ray_samples\": ray_samples_plane,\n \"plane_weights\": weights,\n \"plane_transmittance\": transmittance[:, -1, :],\n \"plane_rgbs\": field_outputs[FieldHeadNames.RGB],\n }\n )\n\n return samples_and_field_outputs\n\n def fg_fields_query_and_render(self, ray_bundle: RayBundle) -> Dict[str, Any]:\n outputs, samples_and_field_outputs = super().fg_fields_query_and_render(ray_bundle)\n if not self.config.render_background:\n return outputs, samples_and_field_outputs\n\n # render plane field\n mask = ray_bundle.metadata[\"bottom_intersect_mask\"][..., 0] # (n,)\n normal = torch.zeros_like(outputs[\"normal\"]) # vis\n depth = torch.zeros_like(outputs[\"depth\"]) # vis\n acc = torch.zeros_like(outputs[\"accumulation\"])\n plane_rgb = torch.ones_like(outputs[\"rgb\"])\n if not mask.any():\n outputs.update(\n {\"plane_normal\": normal, \"plane_depth\": depth, \"plane_accumulation\": acc, \"plane_rgb\": plane_rgb}\n )\n return outputs, samples_and_field_outputs\n\n # rendering the plane field\n plane_weights, plane_trans, plane_rgbs = map(\n lambda x: samples_and_field_outputs.get(x), [\"plane_weights\", \"plane_transmittance\", \"plane_rgbs\"]\n )\n _plane_rgb = self.renderer_rgb(rgb=plane_rgbs, weights=plane_weights)\n\n # composite with the fg rendering\n fg_rgb, fg_trans = outputs[\"rgb\"][mask], outputs[\"bg_transmittance\"][mask]\n _fg_rgb = fg_rgb + fg_trans * _plane_rgb\n _bg_trans = fg_trans * plane_trans\n bg_trans = torch.ones_like(outputs[\"bg_transmittance\"])\n fg_rgb = torch.zeros_like(outputs[\"rgb\"])\n bg_trans[mask] = _bg_trans\n fg_rgb[mask] = _fg_rgb\n\n # plane normals (& depths) rendering\n if self.config.plane_field_type == \"sdf\":\n plane_normals = samples_and_field_outputs[\"plane_normals\"]\n _normal = self.renderer_normal(semantics=plane_normals, weights=plane_weights)\n normal[mask] = _normal\n\n plane_ray_samples = samples_and_field_outputs[\"plane_ray_samples\"]\n _depth = self.renderer_depth(weights=plane_weights, ray_samples=plane_ray_samples)\n _accumulation = self.renderer_accumulation(weights=plane_weights)\n depth[mask], acc[mask], plane_rgb[mask] = _depth, _accumulation, _plane_rgb\n\n outputs.update(\n {\"plane_normal\": normal, \"plane_depth\": depth, \"plane_accumulation\": acc, \"plane_rgb\": plane_rgb}\n )\n\n # update\n if self.training and self.config.plane_field_type == \"sdf\":\n _plane_eik_grad = samples_and_field_outputs[\"plane_eik_grad\"]\n plane_eik_grad = _plane_eik_grad.new_zeros((normal.shape[0], _plane_eik_grad.shape[-2], 3))\n plane_eik_grad[mask] = _plane_eik_grad\n outputs[\"plane_eik_grad\"] = plane_eik_grad\n outputs[\"plane_mask\"] = mask # for loss computation\n outputs[\"bg_transmittance\"] = torch.where(mask[..., None], bg_trans, outputs[\"bg_transmittance\"])\n outputs[\"rgb\"] = torch.where(mask[..., None].repeat(1, 3), fg_rgb, outputs[\"rgb\"])\n return outputs, samples_and_field_outputs\n\n def get_loss_dict(self, outputs, batch, metrics_dict=None, step=None):\n loss_dict = super().get_loss_dict(outputs, batch, metrics_dict)\n\n # eikonal loss for plane sdf field\n if \"plane_eik_grad\" in outputs:\n grad_theta = outputs[\"plane_eik_grad\"]\n plane_mask = outputs[\"plane_mask\"][..., None].to(grad_theta.dtype)\n n_pts = plane_mask.sum() * np.prod(grad_theta.shape[1:-1])\n eikonal_loss = (((grad_theta.norm(2, dim=-1) - 1) ** 2) * plane_mask).sum() / n_pts\n loss_dict[\"plane_eikonal_loss\"] = (\n eikonal_loss * self.config.eikonal_loss_mult\n ) # donnot use a separate weight for plane eikonal loss\n\n self._compute_ptcd_reg_plane_field(outputs, loss_dict)\n\n return loss_dict\n\n def get_image_metrics_and_images(\n self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]\n ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:\n metrics_dict, images_dict = super().get_image_metrics_and_images(outputs, batch)\n\n # plane normal\n if \"plane_normal\" in outputs:\n images_dict.update({\"plane_normal\": (outputs[\"plane_normal\"] + 1.0) / 2.0})\n if \"plane_depth\" in outputs:\n depth = colormaps.apply_depth_colormap(\n outputs[\"plane_depth\"],\n accumulation=outputs[\"plane_accumulation\"],\n )\n acc = colormaps.apply_colormap(outputs[\"plane_accumulation\"])\n images_dict.update({\"plane_depth\": depth, \"plane_accumulation\": acc})\n return metrics_dict, images_dict\n\n def _compute_ptcd_reg_plane_field(self, outputs: Dict[str, torch.Tensor], loss_dict: Dict[str, torch.Tensor]):\n \"\"\"compute point cloud regularization loss on the plane sdf field\"\"\"\n if self.config.ptcd_reg_plane_field_mult == 0.0 or self._step > self.config.ptcd_reg_n_iters:\n return\n\n anneal_weight = self.ptcd_reg_anneal_fn(self._step)\n writer.put_scalar(name=\"ptcd_reg_plane_field_anneal_weight\", scalar=anneal_weight, step=self._step)\n plane_pts, _, plane_sdf_thrs = self.ptcd_data.sample_plane_points(\n self.config.ptcd_reg_plane_field_n_pts_per_iter, replace=True\n )\n plane_pts_sdf = self.field_plane.forward_geonetwork(plane_pts)[:, 0].contiguous()\n # only support point-wise threshold for now\n plane_sdf_loss = -torch.min(plane_sdf_thrs - torch.abs(plane_pts_sdf), torch.zeros_like(plane_pts[:, 0])).mean()\n loss_dict[\"ptcd_reg_plane_field\"] = plane_sdf_loss * self.config.ptcd_reg_plane_field_mult * anneal_weight\n\n def get_metrics_dict(self, outputs, batch) -> Dict:\n metrics_dict = super().get_metrics_dict(outputs, batch)\n if self.training and isinstance(self.field_plane, SDFField):\n # training statics\n metrics_dict[\"s_val_plane\"] = self.field_plane.deviation_network.get_variance().item()\n metrics_dict[\"inv_s_plane\"] = 1.0 / self.field_plane.deviation_network.get_variance().item()\n\n return metrics_dict\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/models/semantic_nerfw.py",
"content": "# Copyright 2022 The nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nSemantic NeRF-W implementation which should be fast enough to view in the viewer.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass, field\nfrom typing import Dict, List, Tuple, Type\n\nimport numpy as np\nimport torch\nfrom torch.nn import Parameter\nfrom torchmetrics import PeakSignalNoiseRatio\nfrom torchmetrics.functional import structural_similarity_index_measure\nfrom torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity\n\nfrom nerfstudio.cameras.rays import RayBundle\nfrom nerfstudio.data.dataparsers.base_dataparser import Semantics\nfrom nerfstudio.engine.callbacks import (\n TrainingCallback,\n TrainingCallbackAttributes,\n TrainingCallbackLocation,\n)\nfrom nerfstudio.field_components.field_heads import FieldHeadNames\nfrom nerfstudio.field_components.spatial_distortions import SceneContraction\nfrom nerfstudio.fields.density_fields import HashMLPDensityField\nfrom nerfstudio.fields.nerfacto_field import TCNNNerfactoField\nfrom nerfstudio.model_components.losses import MSELoss, distortion_loss, interlevel_loss\nfrom nerfstudio.model_components.ray_samplers import ProposalNetworkSampler\nfrom nerfstudio.model_components.renderers import (\n AccumulationRenderer,\n DepthRenderer,\n RGBRenderer,\n SemanticRenderer,\n UncertaintyRenderer,\n)\nfrom nerfstudio.model_components.scene_colliders import NearFarCollider\nfrom nerfstudio.models.base_model import Model\nfrom nerfstudio.models.nerfacto import NerfactoModelConfig\nfrom nerfstudio.utils import colormaps\n\n\n@dataclass\nclass SemanticNerfWModelConfig(NerfactoModelConfig):\n \"\"\"Nerfacto Model Config\"\"\"\n\n _target: Type = field(default_factory=lambda: SemanticNerfWModel)\n semantic_mult: float = 1.0\n use_transient_embedding: bool = False\n \"\"\"Whether to use transient embedding.\"\"\"\n scene_contraction_scale_factor: float = 1.0\n \"\"\"Scale positions by this factor before applying contraction.\n This is useful when the [-1, 1] region (no contraction) is too small and only covers a small portion of the scene.\n \"\"\"\n\n\nclass SemanticNerfWModel(Model):\n \"\"\"Nerfacto model\n\n Args:\n config: Nerfacto configuration to instantiate model\n \"\"\"\n\n config: SemanticNerfWModelConfig\n\n def __init__(self, config: SemanticNerfWModelConfig, metadata: Dict, **kwargs) -> None:\n assert \"semantics\" in metadata.keys() and isinstance(metadata[\"semantics\"], Semantics)\n self.semantics = metadata[\"semantics\"]\n super().__init__(config=config, **kwargs)\n\n def populate_modules(self):\n \"\"\"Set the fields and modules.\"\"\"\n super().populate_modules()\n\n scene_contraction = SceneContraction(order=float(\"inf\"),\n scale_factor=self.config.scene_contraction_scale_factor)\n\n if self.config.use_transient_embedding:\n raise ValueError(\"Transient embedding is not fully working for semantic nerf-w.\")\n\n # Fields\n self.field = TCNNNerfactoField(\n self.scene_box.aabb,\n num_levels=self.config.num_levels,\n max_res=self.config.max_res,\n log2_hashmap_size=self.config.log2_hashmap_size,\n spatial_distortion=scene_contraction,\n num_images=self.num_train_data,\n use_average_appearance_embedding=self.config.use_average_appearance_embedding,\n use_transient_embedding=self.config.use_transient_embedding,\n use_semantics=True,\n num_semantic_classes=len(self.semantics.classes),\n )\n\n # Build the proposal network(s)\n self.proposal_networks = torch.nn.ModuleList()\n if self.config.use_same_proposal_network:\n network = HashMLPDensityField(self.scene_box.aabb, spatial_distortion=scene_contraction)\n self.proposal_networks.append(network)\n self.density_fns = [network.density_fn for _ in range(self.config.num_proposal_iterations)]\n else:\n for _ in range(self.config.num_proposal_iterations):\n network = HashMLPDensityField(self.scene_box.aabb, spatial_distortion=scene_contraction)\n self.proposal_networks.append(network)\n self.density_fns = [network.density_fn for network in self.proposal_networks]\n\n # Collider\n self.collider = NearFarCollider(near_plane=self.config.near_plane, far_plane=self.config.far_plane)\n\n # Samplers\n self.proposal_sampler = ProposalNetworkSampler(\n num_nerf_samples_per_ray=self.config.num_nerf_samples_per_ray,\n num_proposal_samples_per_ray=self.config.num_proposal_samples_per_ray,\n num_proposal_network_iterations=self.config.num_proposal_iterations,\n single_jitter=self.config.use_single_jitter,\n )\n\n # renderers\n self.renderer_rgb = RGBRenderer(background_color=self.config.background_color)\n self.renderer_accumulation = AccumulationRenderer()\n self.renderer_depth = DepthRenderer()\n self.renderer_uncertainty = UncertaintyRenderer()\n self.renderer_semantics = SemanticRenderer()\n\n # losses\n self.rgb_loss = MSELoss()\n self.cross_entropy_loss = torch.nn.CrossEntropyLoss(reduction=\"mean\")\n\n # metrics\n self.psnr = PeakSignalNoiseRatio(data_range=1.0)\n self.ssim = structural_similarity_index_measure\n self.lpips = LearnedPerceptualImagePatchSimilarity()\n\n def get_param_groups(self) -> Dict[str, List[Parameter]]:\n param_groups = {}\n param_groups[\"proposal_networks\"] = list(self.proposal_networks.parameters())\n param_groups[\"fields\"] = list(self.field.parameters())\n return param_groups\n\n def get_training_callbacks(\n self, training_callback_attributes: TrainingCallbackAttributes\n ) -> List[TrainingCallback]:\n callbacks = []\n if self.config.use_proposal_weight_anneal:\n # anneal the weights of the proposal network before doing PDF sampling\n N = self.config.proposal_weights_anneal_max_num_iters\n\n def set_anneal(step):\n # https://arxiv.org/pdf/2111.12077.pdf eq. 18\n train_frac = np.clip(step / N, 0, 1)\n bias = lambda x, b: (b * x) / ((b - 1) * x + 1)\n anneal = bias(train_frac, self.config.proposal_weights_anneal_slope)\n self.proposal_sampler.set_anneal(anneal)\n\n callbacks.append(\n TrainingCallback(\n where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],\n update_every_num_iters=1,\n func=set_anneal,\n )\n )\n return callbacks\n\n def get_outputs(self, ray_bundle: RayBundle):\n ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=self.density_fns)\n field_outputs = self.field(ray_samples)\n\n if self.training and self.config.use_transient_embedding:\n density = field_outputs[FieldHeadNames.DENSITY] + field_outputs[FieldHeadNames.TRANSIENT_DENSITY]\n weights = ray_samples.get_weights(density)\n weights_static = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY])\n rgb_static_component = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)\n rgb_transient_component = self.renderer_rgb(\n rgb=field_outputs[FieldHeadNames.TRANSIENT_RGB], weights=weights\n )\n rgb = rgb_static_component + rgb_transient_component\n else:\n weights_static = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY])\n weights = weights_static\n rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)\n weights_list.append(weights_static)\n ray_samples_list.append(ray_samples)\n\n depth = self.renderer_depth(weights=weights_static, ray_samples=ray_samples)\n accumulation = self.renderer_accumulation(weights=weights_static)\n\n outputs = {\"rgb\": rgb, \"accumulation\": accumulation, \"depth\": depth}\n \n if self.training:\n outputs[\"weights_list\"] = weights_list\n outputs[\"ray_samples_list\"] = ray_samples_list\n\n for i in range(self.config.num_proposal_iterations):\n outputs[f\"prop_depth_{i}\"] = self.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i])\n\n # transients\n if self.training and self.config.use_transient_embedding:\n weights_transient = ray_samples.get_weights(field_outputs[FieldHeadNames.TRANSIENT_DENSITY])\n uncertainty = self.renderer_uncertainty(field_outputs[FieldHeadNames.UNCERTAINTY], weights_transient)\n outputs[\"uncertainty\"] = uncertainty + 0.03 # NOTE(ethan): this is the uncertainty min\n outputs[\"density_transient\"] = field_outputs[FieldHeadNames.TRANSIENT_DENSITY]\n\n # semantics\n outputs[\"semantics\"] = self.renderer_semantics(\n field_outputs[FieldHeadNames.SEMANTICS], weights=weights_static.detach()\n )\n\n # semantics colormaps\n semantic_labels = torch.argmax(torch.nn.functional.softmax(outputs[\"semantics\"], dim=-1), dim=-1)\n # outputs[\"semantics_colormap\"] = self.semantics.colors[semantic_labels]\n\n return outputs\n\n def get_metrics_dict(self, outputs, batch):\n metrics_dict = {}\n image = batch[\"image\"].to(self.device)\n metrics_dict[\"psnr\"] = self.psnr(outputs[\"rgb\"], image)\n if self.training:\n metrics_dict[\"distortion\"] = distortion_loss(outputs[\"weights_list\"], outputs[\"ray_samples_list\"])\n return metrics_dict\n\n def get_loss_dict(self, outputs, batch, metrics_dict=None):\n loss_dict = {}\n image = batch[\"image\"].to(self.device)\n loss_dict[\"rgb_loss\"] = self.rgb_loss(image, outputs[\"rgb\"])\n if self.training:\n loss_dict[\"interlevel_loss\"] = self.config.interlevel_loss_mult * interlevel_loss(\n outputs[\"weights_list\"], outputs[\"ray_samples_list\"]\n )\n assert metrics_dict is not None and \"distortion\" in metrics_dict\n loss_dict[\"distortion_loss\"] = self.config.distortion_loss_mult * metrics_dict[\"distortion\"]\n if self.config.use_transient_embedding:\n betas = outputs[\"uncertainty\"]\n loss_dict[\"uncertainty_loss\"] = 3 + torch.log(betas).mean()\n loss_dict[\"density_loss\"] = 0.01 * outputs[\"density_transient\"].mean()\n loss_dict[\"rgb_loss\"] = (((image - outputs[\"rgb\"]) ** 2).sum(-1) / (betas[..., 0] ** 2)).mean() \n\n # semantic loss\n loss_dict[\"semantics_loss\"] = self.config.semantic_mult * self.cross_entropy_loss(outputs[\"semantics\"], batch[\"semantics\"][..., 0].long())\n return loss_dict\n\n def get_image_metrics_and_images(\n self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]\n ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:\n\n image = batch[\"image\"].to(self.device)\n rgb = outputs[\"rgb\"]\n rgb = torch.clamp(rgb, min=0, max=1)\n acc = colormaps.apply_colormap(outputs[\"accumulation\"])\n depth = colormaps.apply_depth_colormap(\n outputs[\"depth\"],\n accumulation=outputs[\"accumulation\"],\n )\n\n combined_rgb = torch.cat([image, rgb], dim=1)\n combined_acc = torch.cat([acc], dim=1)\n combined_depth = torch.cat([depth], dim=1)\n\n # Switch images from [H, W, C] to [1, C, H, W] for metrics computations\n image = torch.moveaxis(image, -1, 0)[None, ...]\n rgb = torch.moveaxis(rgb, -1, 0)[None, ...]\n\n psnr = self.psnr(image, rgb)\n ssim = self.ssim(image, rgb)\n lpips = self.lpips(image, rgb)\n\n # all of these metrics will be logged as scalars\n metrics_dict = {\"psnr\": float(psnr.item()), \"ssim\": float(ssim)} # type: ignore\n metrics_dict[\"lpips\"] = float(lpips)\n\n images_dict = {\"img\": combined_rgb, \"accumulation\": combined_acc, \"depth\": combined_depth}\n\n for i in range(self.config.num_proposal_iterations):\n key = f\"prop_depth_{i}\"\n prop_depth_i = colormaps.apply_depth_colormap(\n outputs[key],\n accumulation=outputs[\"accumulation\"],\n )\n images_dict[key] = prop_depth_i\n\n # semantics\n semantic_labels = torch.argmax(torch.nn.functional.softmax(outputs[\"semantics\"], dim=-1), dim=-1)\n images_dict[\"semantics_colormap\"] = self.semantics.colors.to(self.device)[semantic_labels]\n\n # valid mask\n if \"mask\" in batch:\n images_dict[\"mask\"] = batch[\"mask\"].repeat(1, 1, 3)\n\n return metrics_dict, images_dict\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/models/tensorf.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nTensorRF implementation.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass, field\nfrom typing import Dict, List, Tuple, Type\n\nimport numpy as np\nimport torch\nfrom torch.nn import Parameter\nfrom torchmetrics import PeakSignalNoiseRatio\nfrom torchmetrics.functional import structural_similarity_index_measure\nfrom torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity\n\nfrom nerfstudio.cameras.rays import RayBundle\nfrom nerfstudio.configs.config_utils import to_immutable_dict\nfrom nerfstudio.engine.callbacks import (\n TrainingCallback,\n TrainingCallbackAttributes,\n TrainingCallbackLocation,\n)\nfrom nerfstudio.field_components.encodings import NeRFEncoding, TensorVMEncoding\nfrom nerfstudio.field_components.field_heads import FieldHeadNames\nfrom nerfstudio.fields.tensorf_field import TensoRFField\nfrom nerfstudio.model_components.losses import MSELoss\nfrom nerfstudio.model_components.ray_samplers import PDFSampler, UniformSampler\nfrom nerfstudio.model_components.renderers import (\n AccumulationRenderer,\n DepthRenderer,\n RGBRenderer,\n)\nfrom nerfstudio.model_components.scene_colliders import AABBBoxCollider\nfrom nerfstudio.models.base_model import Model, ModelConfig\nfrom nerfstudio.utils import colormaps, colors, misc\n\n\n@dataclass\nclass TensoRFModelConfig(ModelConfig):\n \"\"\"TensoRF model config\"\"\"\n\n _target: Type = field(default_factory=lambda: TensoRFModel)\n \"\"\"target class to instantiate\"\"\"\n init_resolution: int = 128\n \"\"\"initial render resolution\"\"\"\n final_resolution: int = 300\n \"\"\"final render resolution\"\"\"\n upsampling_iters: Tuple[int, ...] = (2000, 3000, 4000, 5500, 7000)\n \"\"\"specifies a list of iteration step numbers to perform upsampling\"\"\"\n loss_coefficients: Dict[str, float] = to_immutable_dict({\"rgb_loss\": 1.0})\n \"\"\"Loss specific weights.\"\"\"\n num_samples: int = 256\n \"\"\"Number of samples in field evaluation\"\"\"\n num_den_components: int = 16\n \"\"\"Number of components in density encoding\"\"\"\n num_color_components: int = 48\n \"\"\"Number of components in color encoding\"\"\"\n appearance_dim: int = 27\n \"\"\"Number of channels for color encoding\"\"\"\n\n\nclass TensoRFModel(Model):\n \"\"\"TensoRF Model\n\n Args:\n config: TensoRF configuration to instantiate model\n \"\"\"\n\n def __init__(\n self,\n config: TensoRFModelConfig,\n **kwargs,\n ) -> None:\n self.init_resolution = config.init_resolution\n self.upsampling_iters = config.upsampling_iters\n self.num_den_components = config.num_den_components\n self.num_color_components = config.num_color_components\n self.appearance_dim = config.appearance_dim\n self.upsampling_steps = (\n np.round(\n np.exp(\n np.linspace(\n np.log(config.init_resolution),\n np.log(config.final_resolution),\n len(config.upsampling_iters) + 1,\n )\n )\n )\n .astype(\"int\")\n .tolist()[1:]\n )\n super().__init__(config=config, **kwargs)\n\n def get_training_callbacks(\n self, training_callback_attributes: TrainingCallbackAttributes\n ) -> List[TrainingCallback]:\n\n # the callback that we want to run every X iterations after the training iteration\n def reinitialize_optimizer(\n self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument\n ):\n resolution = self.upsampling_steps.pop(0)\n\n # upsample the position and direction grids\n self.field.density_encoding.upsample_grid(resolution)\n self.field.color_encoding.upsample_grid(resolution)\n\n # reinitialize the encodings optimizer\n optimizers_config = training_callback_attributes.optimizers.config\n enc = training_callback_attributes.pipeline.get_param_groups()[\"encodings\"]\n lr_init = optimizers_config[\"encodings\"][\"optimizer\"].lr\n\n training_callback_attributes.optimizers.optimizers[\"encodings\"] = optimizers_config[\"encodings\"][\n \"optimizer\"\n ].setup(params=enc)\n if optimizers_config[\"encodings\"][\"scheduler\"]:\n training_callback_attributes.optimizers.schedulers[\"encodings\"] = optimizers_config[\"encodings\"][\n \"scheduler\"\n ].setup(optimizer=training_callback_attributes.optimizers.optimizers[\"encodings\"], lr_init=lr_init)\n\n callbacks = [\n TrainingCallback(\n where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],\n iters=self.upsampling_iters,\n func=reinitialize_optimizer,\n args=[self, training_callback_attributes],\n )\n ]\n return callbacks\n\n def populate_modules(self):\n \"\"\"Set the fields and modules\"\"\"\n super().populate_modules()\n\n # setting up fields\n density_encoding = TensorVMEncoding(\n resolution=self.init_resolution,\n num_components=self.num_den_components,\n )\n color_encoding = TensorVMEncoding(\n resolution=self.init_resolution,\n num_components=self.num_color_components,\n )\n\n feature_encoding = NeRFEncoding(in_dim=self.appearance_dim, num_frequencies=2, min_freq_exp=0, max_freq_exp=2)\n direction_encoding = NeRFEncoding(in_dim=3, num_frequencies=2, min_freq_exp=0, max_freq_exp=2)\n\n self.field = TensoRFField(\n self.scene_box.aabb,\n feature_encoding=feature_encoding,\n direction_encoding=direction_encoding,\n density_encoding=density_encoding,\n color_encoding=color_encoding,\n appearance_dim=self.appearance_dim,\n head_mlp_num_layers=2,\n head_mlp_layer_width=128,\n use_sh=False,\n )\n\n # samplers\n self.sampler_uniform = UniformSampler(num_samples=self.config.num_samples, single_jitter=True)\n self.sampler_pdf = PDFSampler(num_samples=self.config.num_samples // 2, single_jitter=True)\n\n # renderers\n self.renderer_rgb = RGBRenderer(background_color=colors.WHITE)\n self.renderer_accumulation = AccumulationRenderer()\n self.renderer_depth = DepthRenderer()\n\n # losses\n self.rgb_loss = MSELoss()\n\n # metrics\n self.psnr = PeakSignalNoiseRatio(data_range=1.0)\n self.ssim = structural_similarity_index_measure\n self.lpips = LearnedPerceptualImagePatchSimilarity()\n\n # colliders\n if self.config.enable_collider:\n self.collider = AABBBoxCollider(scene_box=self.scene_box)\n\n def get_param_groups(self) -> Dict[str, List[Parameter]]:\n param_groups = {}\n\n param_groups[\"fields\"] = (\n list(self.field.mlp_head.parameters())\n + list(self.field.B.parameters())\n + list(self.field.field_output_rgb.parameters())\n )\n param_groups[\"encodings\"] = list(self.field.color_encoding.parameters()) + list(\n self.field.density_encoding.parameters()\n )\n\n return param_groups\n\n def get_outputs(self, ray_bundle: RayBundle):\n # uniform sampling\n ray_samples_uniform = self.sampler_uniform(ray_bundle)\n dens = self.field.get_density(ray_samples_uniform)\n weights = ray_samples_uniform.get_weights(dens)\n coarse_accumulation = self.renderer_accumulation(weights)\n acc_mask = torch.where(coarse_accumulation < 0.0001, False, True).reshape(-1)\n\n # pdf sampling\n ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights)\n\n # fine field:\n field_outputs_fine = self.field.forward(\n ray_samples_pdf, mask=acc_mask, bg_color=colors.WHITE.to(weights.device)\n )\n\n weights_fine = ray_samples_pdf.get_weights(field_outputs_fine[FieldHeadNames.DENSITY])\n\n accumulation = self.renderer_accumulation(weights_fine)\n depth = self.renderer_depth(weights_fine, ray_samples_pdf)\n\n rgb = self.renderer_rgb(\n rgb=field_outputs_fine[FieldHeadNames.RGB],\n weights=weights_fine,\n )\n\n rgb = torch.where(accumulation < 0, colors.WHITE.to(rgb.device), rgb)\n accumulation = torch.clamp(accumulation, min=0)\n\n outputs = {\"rgb\": rgb, \"accumulation\": accumulation, \"depth\": depth}\n return outputs\n\n def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]:\n # Scaling metrics by coefficients to create the losses.\n device = outputs[\"rgb\"].device\n image = batch[\"image\"].to(device)\n\n rgb_loss = self.rgb_loss(image, outputs[\"rgb\"])\n\n loss_dict = {\"rgb_loss\": rgb_loss}\n loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients)\n return loss_dict\n\n def get_image_metrics_and_images(\n self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]\n ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:\n image = batch[\"image\"].to(outputs[\"rgb\"].device)\n rgb = outputs[\"rgb\"]\n acc = colormaps.apply_colormap(outputs[\"accumulation\"])\n depth = colormaps.apply_depth_colormap(\n outputs[\"depth\"],\n accumulation=outputs[\"accumulation\"],\n near_plane=self.config.collider_params[\"near_plane\"],\n far_plane=self.config.collider_params[\"far_plane\"],\n )\n\n combined_rgb = torch.cat([image, rgb], dim=1)\n\n # Switch images from [H, W, C] to [1, C, H, W] for metrics computations\n image = torch.moveaxis(image, -1, 0)[None, ...]\n rgb = torch.moveaxis(rgb, -1, 0)[None, ...]\n\n psnr = self.psnr(image, rgb)\n ssim = self.ssim(image, rgb)\n lpips = self.lpips(image, rgb)\n\n metrics_dict = {\n \"psnr\": float(psnr.item()),\n \"ssim\": float(ssim.item()),\n \"lpips\": float(lpips.item()),\n }\n images_dict = {\"img\": combined_rgb, \"accumulation\": acc, \"depth\": depth}\n return metrics_dict, images_dict\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/models/unisurf.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nImplementation of VolSDF.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass, field\nfrom typing import Dict, List, Type\n\nimport torch\n\nfrom nerfstudio.engine.callbacks import (\n TrainingCallback,\n TrainingCallbackAttributes,\n TrainingCallbackLocation,\n)\nfrom nerfstudio.cameras.rays import RayBundle\nfrom nerfstudio.field_components.field_heads import FieldHeadNames\nfrom nerfstudio.model_components.ray_samplers import UniSurfSampler\nfrom nerfstudio.models.base_surface_model import SurfaceModel, SurfaceModelConfig\n\n\n@dataclass\nclass UniSurfModelConfig(SurfaceModelConfig):\n \"\"\"UniSurf Model Config\"\"\"\n\n _target: Type = field(default_factory=lambda: UniSurfModel)\n eikonal_loss_mult: float = 0.0\n \"\"\"overwirte eikonal loss because it's not need for unisurf\"\"\"\n smooth_loss_multi: float = 0.005\n \"\"\"smoothness loss on surface points in unisurf\"\"\"\n num_samples_interval: int = 64\n \"\"\"Number of uniform samples\"\"\"\n num_samples_importance: int = 32\n \"\"\"Number of important samples\"\"\"\n num_marching_steps: int = 256\n \"\"\"number of up sample step, 1 for simple coarse-to-fine sampling\"\"\"\n perturb: bool = True\n \"\"\"use to use perturb for the sampled points\"\"\"\n\n\nclass UniSurfModel(SurfaceModel):\n \"\"\"VolSDF model\n\n Args:\n config: MonoSDF configuration to instantiate model\n \"\"\"\n\n config: UniSurfModelConfig\n\n def populate_modules(self):\n \"\"\"Set the fields and modules.\"\"\"\n super().populate_modules()\n\n # can't use eikonal loss in Unisurf? or we could use learnable paramter to transform sdf to occupancy\n assert self.config.eikonal_loss_mult == 0.0\n\n self.sampler = UniSurfSampler(\n num_samples_interval=self.config.num_samples_interval,\n num_samples_outside=self.config.num_samples_outside,\n num_samples_importance=self.config.num_samples_importance,\n num_marching_steps=self.config.num_marching_steps,\n )\n\n def get_training_callbacks(\n self, training_callback_attributes: TrainingCallbackAttributes\n ) -> List[TrainingCallback]:\n callbacks = []\n callbacks.append(\n TrainingCallback(\n where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],\n update_every_num_iters=1,\n func=self.sampler.step_cb,\n )\n )\n return callbacks\n\n def sample_and_forward_field(self, ray_bundle: RayBundle) -> Dict:\n ray_samples, surface_points = self.sampler(\n ray_bundle, occupancy_fn=self.field.get_occupancy, sdf_fn=self.field.get_sdf, return_surface_points=True\n )\n field_outputs = self.field(ray_samples, return_occupancy=True)\n weights, transmittance = ray_samples.get_weights_and_transmittance_from_alphas(\n field_outputs[FieldHeadNames.OCCUPANCY]\n )\n bg_transmittance = transmittance[:, -1, :]\n\n samples_and_field_outputs = {\n \"ray_samples\": ray_samples,\n \"surface_points\": surface_points,\n \"field_outputs\": field_outputs,\n \"weights\": weights,\n \"bg_transmittance\": bg_transmittance,\n }\n return samples_and_field_outputs\n\n def get_metrics_dict(self, outputs, batch) -> Dict:\n metrics_dict = super().get_metrics_dict(outputs, batch)\n if self.training:\n # training statics\n metrics_dict[\"delta\"] = self.sampler.delta\n\n return metrics_dict\n\n def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict:\n loss_dict = super().get_loss_dict(outputs, batch, metrics_dict)\n\n # TODO move to base model as other model could also use it?\n if self.training and self.config.smooth_loss_multi > 0.0:\n surface_points = outputs[\"surface_points\"]\n\n surface_points_neig = surface_points + (torch.rand_like(surface_points) - 0.5) * 0.01\n pp = torch.cat([surface_points, surface_points_neig], dim=0)\n surface_grad = self.field.gradient(pp)\n surface_points_normal = torch.nn.functional.normalize(surface_grad, p=2, dim=-1)\n\n N = surface_points_normal.shape[0] // 2\n\n diff_norm = torch.norm(surface_points_normal[:N] - surface_points_normal[N:], dim=-1)\n loss_dict[\"normal_smoothness_loss\"] = torch.mean(diff_norm) * self.config.smooth_loss_multi\n\n return loss_dict\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/models/vanilla_nerf.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nImplementation of vanilla nerf.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass, field\nfrom typing import Any, Dict, List, Tuple, Type\n\nimport torch\nfrom torch.nn import Parameter\nfrom torchmetrics import PeakSignalNoiseRatio\nfrom torchmetrics.functional import structural_similarity_index_measure\nfrom torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity\n\nfrom nerfstudio.cameras.rays import RayBundle\nfrom nerfstudio.configs.config_utils import to_immutable_dict\nfrom nerfstudio.field_components.encodings import NeRFEncoding\nfrom nerfstudio.field_components.field_heads import FieldHeadNames\nfrom nerfstudio.field_components.temporal_distortions import TemporalDistortionKind\nfrom nerfstudio.fields.vanilla_nerf_field import NeRFField\nfrom nerfstudio.model_components.losses import MSELoss\nfrom nerfstudio.model_components.ray_samplers import PDFSampler, UniformSampler\nfrom nerfstudio.model_components.renderers import (\n AccumulationRenderer,\n DepthRenderer,\n RGBRenderer,\n)\nfrom nerfstudio.models.base_model import Model, ModelConfig\nfrom nerfstudio.utils import colormaps, colors, misc\n\n\n@dataclass\nclass VanillaModelConfig(ModelConfig):\n \"\"\"Vanilla Model Config\"\"\"\n\n _target: Type = field(default_factory=lambda: NeRFModel)\n num_coarse_samples: int = 64\n \"\"\"Number of samples in coarse field evaluation\"\"\"\n num_importance_samples: int = 128\n \"\"\"Number of samples in fine field evaluation\"\"\"\n\n enable_temporal_distortion: bool = False\n \"\"\"Specifies whether or not to include ray warping based on time.\"\"\"\n temporal_distortion_params: Dict[str, Any] = to_immutable_dict({\"kind\": TemporalDistortionKind.DNERF})\n \"\"\"Parameters to instantiate temporal distortion with\"\"\"\n\n\nclass NeRFModel(Model):\n \"\"\"Vanilla NeRF model\n\n Args:\n config: Basic NeRF configuration to instantiate model\n \"\"\"\n\n def __init__(\n self,\n config: VanillaModelConfig,\n **kwargs,\n ) -> None:\n self.field_coarse = None\n self.field_fine = None\n self.temporal_distortion = None\n\n super().__init__(\n config=config,\n **kwargs,\n )\n\n def populate_modules(self):\n \"\"\"Set the fields and modules\"\"\"\n super().populate_modules()\n\n # fields\n position_encoding = NeRFEncoding(\n in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0, include_input=True\n )\n direction_encoding = NeRFEncoding(\n in_dim=3, num_frequencies=4, min_freq_exp=0.0, max_freq_exp=4.0, include_input=True\n )\n\n self.field_coarse = NeRFField(\n position_encoding=position_encoding,\n direction_encoding=direction_encoding,\n )\n\n self.field_fine = NeRFField(\n position_encoding=position_encoding,\n direction_encoding=direction_encoding,\n )\n\n # samplers\n self.sampler_uniform = UniformSampler(num_samples=self.config.num_coarse_samples)\n self.sampler_pdf = PDFSampler(num_samples=self.config.num_importance_samples)\n\n # renderers\n self.renderer_rgb = RGBRenderer(background_color=colors.WHITE)\n self.renderer_accumulation = AccumulationRenderer()\n self.renderer_depth = DepthRenderer()\n\n # losses\n self.rgb_loss = MSELoss()\n\n # metrics\n self.psnr = PeakSignalNoiseRatio(data_range=1.0)\n self.ssim = structural_similarity_index_measure\n self.lpips = LearnedPerceptualImagePatchSimilarity()\n\n if getattr(self.config, \"enable_temporal_distortion\", False):\n params = self.config.temporal_distortion_params\n kind = params.pop(\"kind\")\n self.temporal_distortion = kind.to_temporal_distortion(params)\n\n def get_param_groups(self) -> Dict[str, List[Parameter]]:\n param_groups = {}\n if self.field_coarse is None or self.field_fine is None:\n raise ValueError(\"populate_fields() must be called before get_param_groups\")\n param_groups[\"fields\"] = list(self.field_coarse.parameters()) + list(self.field_fine.parameters())\n if self.temporal_distortion is not None:\n param_groups[\"temporal_distortion\"] = list(self.temporal_distortion.parameters())\n return param_groups\n\n def get_outputs(self, ray_bundle: RayBundle):\n\n if self.field_coarse is None or self.field_fine is None:\n raise ValueError(\"populate_fields() must be called before get_outputs\")\n\n # uniform sampling\n ray_samples_uniform = self.sampler_uniform(ray_bundle)\n if self.temporal_distortion is not None:\n offsets = self.temporal_distortion(ray_samples_uniform.frustums.get_positions(), ray_samples_uniform.times)\n ray_samples_uniform.frustums.set_offsets(offsets)\n\n # coarse field:\n field_outputs_coarse = self.field_coarse.forward(ray_samples_uniform)\n weights_coarse = ray_samples_uniform.get_weights(field_outputs_coarse[FieldHeadNames.DENSITY])\n rgb_coarse = self.renderer_rgb(\n rgb=field_outputs_coarse[FieldHeadNames.RGB],\n weights=weights_coarse,\n )\n accumulation_coarse = self.renderer_accumulation(weights_coarse)\n depth_coarse = self.renderer_depth(weights_coarse, ray_samples_uniform)\n\n # pdf sampling\n ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse)\n if self.temporal_distortion is not None:\n offsets = self.temporal_distortion(ray_samples_pdf.frustums.get_positions(), ray_samples_pdf.times)\n ray_samples_pdf.frustums.set_offsets(offsets)\n\n # fine field:\n field_outputs_fine = self.field_fine.forward(ray_samples_pdf)\n weights_fine = ray_samples_pdf.get_weights(field_outputs_fine[FieldHeadNames.DENSITY])\n rgb_fine = self.renderer_rgb(\n rgb=field_outputs_fine[FieldHeadNames.RGB],\n weights=weights_fine,\n )\n accumulation_fine = self.renderer_accumulation(weights_fine)\n depth_fine = self.renderer_depth(weights_fine, ray_samples_pdf)\n\n outputs = {\n \"rgb_coarse\": rgb_coarse,\n \"rgb_fine\": rgb_fine,\n \"accumulation_coarse\": accumulation_coarse,\n \"accumulation_fine\": accumulation_fine,\n \"depth_coarse\": depth_coarse,\n \"depth_fine\": depth_fine,\n }\n return outputs\n\n def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]:\n # Scaling metrics by coefficients to create the losses.\n device = outputs[\"rgb_coarse\"].device\n image = batch[\"image\"].to(device)\n\n rgb_loss_coarse = self.rgb_loss(image, outputs[\"rgb_coarse\"])\n rgb_loss_fine = self.rgb_loss(image, outputs[\"rgb_fine\"])\n\n loss_dict = {\"rgb_loss_coarse\": rgb_loss_coarse, \"rgb_loss_fine\": rgb_loss_fine}\n loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients)\n return loss_dict\n\n def get_image_metrics_and_images(\n self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]\n ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:\n image = batch[\"image\"].to(outputs[\"rgb_coarse\"].device)\n rgb_coarse = outputs[\"rgb_coarse\"]\n rgb_fine = outputs[\"rgb_fine\"]\n acc_coarse = colormaps.apply_colormap(outputs[\"accumulation_coarse\"])\n acc_fine = colormaps.apply_colormap(outputs[\"accumulation_fine\"])\n depth_coarse = colormaps.apply_depth_colormap(\n outputs[\"depth_coarse\"],\n accumulation=outputs[\"accumulation_coarse\"],\n near_plane=self.config.collider_params[\"near_plane\"],\n far_plane=self.config.collider_params[\"far_plane\"],\n )\n depth_fine = colormaps.apply_depth_colormap(\n outputs[\"depth_fine\"],\n accumulation=outputs[\"accumulation_fine\"],\n near_plane=self.config.collider_params[\"near_plane\"],\n far_plane=self.config.collider_params[\"far_plane\"],\n )\n\n combined_rgb = torch.cat([image, rgb_coarse, rgb_fine], dim=1)\n combined_acc = torch.cat([acc_coarse, acc_fine], dim=1)\n combined_depth = torch.cat([depth_coarse, depth_fine], dim=1)\n\n # Switch images from [H, W, C] to [1, C, H, W] for metrics computations\n image = torch.moveaxis(image, -1, 0)[None, ...]\n rgb_coarse = torch.moveaxis(rgb_coarse, -1, 0)[None, ...]\n rgb_fine = torch.moveaxis(rgb_fine, -1, 0)[None, ...]\n\n coarse_psnr = self.psnr(image, rgb_coarse)\n fine_psnr = self.psnr(image, rgb_fine)\n fine_ssim = self.ssim(image, rgb_fine)\n fine_lpips = self.lpips(image, rgb_fine)\n\n metrics_dict = {\n \"psnr\": float(fine_psnr.item()),\n \"coarse_psnr\": float(coarse_psnr),\n \"fine_psnr\": float(fine_psnr),\n \"fine_ssim\": float(fine_ssim),\n \"fine_lpips\": float(fine_lpips),\n }\n images_dict = {\"img\": combined_rgb, \"accumulation\": combined_acc, \"depth\": combined_depth}\n return metrics_dict, images_dict\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/models/volsdf.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nImplementation of VolSDF.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass, field\nfrom typing import Dict, Type\n\nfrom nerfstudio.cameras.rays import RayBundle\nfrom nerfstudio.field_components.field_heads import FieldHeadNames\nfrom nerfstudio.model_components.ray_samplers import ErrorBoundedSampler\nfrom nerfstudio.models.base_surface_model import SurfaceModel, SurfaceModelConfig\n\n\n@dataclass\nclass VolSDFModelConfig(SurfaceModelConfig):\n \"\"\"VolSDF Model Config\"\"\"\n\n _target: Type = field(default_factory=lambda: VolSDFModel)\n num_samples: int = 64\n \"\"\"Number of samples after error bounded sampling\"\"\"\n num_samples_eval: int = 128\n \"\"\"Number of samples per iteration used in error bounded sampling\"\"\"\n num_samples_extra: int = 32\n \"\"\"Number of uniformly sampled points for training\"\"\"\n\n\nclass VolSDFModel(SurfaceModel):\n \"\"\"VolSDF model\n\n Args:\n config: VolSDF configuration to instantiate model\n \"\"\"\n\n config: VolSDFModelConfig\n\n def populate_modules(self):\n \"\"\"Set the fields and modules.\"\"\"\n super().populate_modules()\n\n self.sampler = ErrorBoundedSampler(\n num_samples=self.config.num_samples,\n num_samples_eval=self.config.num_samples_eval,\n num_samples_extra=self.config.num_samples_extra,\n )\n\n def sample_and_forward_field(self, ray_bundle: RayBundle) -> Dict:\n ray_samples, eik_points = self.sampler(\n ray_bundle, density_fn=self.field.laplace_density, sdf_fn=self.field.get_sdf\n )\n field_outputs = self.field(ray_samples)\n weights, transmittance = ray_samples.get_weights_and_transmittance(field_outputs[FieldHeadNames.DENSITY])\n bg_transmittance = transmittance[:, -1, :]\n\n samples_and_field_outputs = {\n \"ray_samples\": ray_samples,\n \"eik_points\": eik_points,\n \"field_outputs\": field_outputs,\n \"weights\": weights,\n \"bg_transmittance\": bg_transmittance,\n }\n return samples_and_field_outputs\n\n def get_metrics_dict(self, outputs, batch) -> Dict:\n metrics_dict = super().get_metrics_dict(outputs, batch)\n if self.training:\n # training statics\n metrics_dict[\"beta\"] = self.field.laplace_density.get_beta().item()\n metrics_dict[\"alpha\"] = 1.0 / self.field.laplace_density.get_beta().item()\n\n return metrics_dict\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/pipelines/__init__.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/pipelines/base_pipeline.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nAbstracts for the Pipeline class.\n\"\"\"\nfrom __future__ import annotations\n\nimport typing\nfrom abc import abstractmethod\nfrom dataclasses import dataclass, field\nfrom time import time\nfrom typing import Any, Dict, List, Optional, Type, Union, cast\n\nimport torch\nimport torch.distributed as dist\nfrom rich.progress import (\n BarColumn,\n MofNCompleteColumn,\n Progress,\n TextColumn,\n TimeElapsedColumn,\n)\nfrom torch import nn\nfrom torch.nn import Parameter\nfrom torch.nn.parallel import DistributedDataParallel as DDP\nfrom torchinfo import summary\nfrom typing_extensions import Literal\n\nfrom nerfstudio.configs import base_config as cfg\nfrom nerfstudio.data.datamanagers.base_datamanager import (\n DataManager,\n FlexibleDataManager,\n FlexibleDataManagerConfig,\n VanillaDataManager,\n VanillaDataManagerConfig,\n)\nfrom nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes\nfrom nerfstudio.models.base_model import Model, ModelConfig\nfrom nerfstudio.utils import profiler\nfrom nerfstudio.utils.images import BasicImages\n\n\ndef module_wrapper(ddp_or_model: Union[DDP, Model]) -> Model:\n \"\"\"\n If DDP, then return the .module. Otherwise, return the model.\n \"\"\"\n if isinstance(ddp_or_model, DDP):\n return cast(Model, ddp_or_model.module)\n return ddp_or_model\n\n\nclass Pipeline(nn.Module):\n \"\"\"The intent of this class is to provide a higher level interface for the Model\n that will be easy to use for our Trainer class.\n\n This class will contain high level functions for the model like getting the loss\n dictionaries and visualization code. It should have ways to get the next iterations\n training loss, evaluation loss, and generate whole images for visualization. Each model\n class should be 1:1 with a pipeline that can act as a standardized interface and hide\n differences in how each model takes in and outputs data.\n\n This class's function is to hide the data manager and model classes from the trainer,\n worrying about:\n 1) Fetching data with the data manager\n 2) Feeding the model the data and fetching the loss\n Hopefully this provides a higher level interface for the trainer to use, and\n simplifying the model classes, which each may have different forward() methods\n and so on.\n\n Args:\n config: configuration to instantiate pipeline\n device: location to place model and data\n test_mode:\n 'train': loads train/eval datasets into memory\n 'test': loads train/test datset into memory\n 'inference': does not load any dataset into memory\n world_size: total number of machines available\n local_rank: rank of current machine\n\n Attributes:\n datamanager: The data manager that will be used\n model: The model that will be used\n \"\"\"\n\n # pylint: disable=abstract-method\n\n datamanager: DataManager\n _model: Model\n\n @property\n def model(self):\n \"\"\"Returns the unwrapped model if in ddp\"\"\"\n return module_wrapper(self._model)\n\n @property\n def device(self):\n \"\"\"Returns the device that the model is on.\"\"\"\n return self.model.device\n\n @profiler.time_function\n def get_train_loss_dict(self, step: int):\n \"\"\"This function gets your training loss dict. This will be responsible for\n getting the next batch of data from the DataManager and interfacing with the\n Model class, feeding the data to the model's forward function.\n\n Args:\n step: current iteration step to update sampler if using DDP (distributed)\n \"\"\"\n if self.world_size > 1 and step:\n assert self.datamanager.train_sampler is not None\n self.datamanager.train_sampler.set_epoch(step)\n ray_bundle, batch = self.datamanager.next_train(step)\n model_outputs = self.model(ray_bundle, batch)\n metrics_dict = self.model.get_metrics_dict(model_outputs, batch)\n loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict)\n\n return model_outputs, loss_dict, metrics_dict\n\n @profiler.time_function\n def get_eval_loss_dict(self, step: int):\n \"\"\"This function gets your evaluation loss dict. It needs to get the data\n from the DataManager and feed it to the model's forward function\n\n Args:\n step: current iteration step\n \"\"\"\n self.eval()\n if self.world_size > 1:\n assert self.datamanager.eval_sampler is not None\n self.datamanager.eval_sampler.set_epoch(step)\n ray_bundle, batch = self.datamanager.next_eval(step)\n model_outputs = self.model(ray_bundle, batch)\n metrics_dict = self.model.get_metrics_dict(model_outputs, batch)\n loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict)\n self.train()\n return model_outputs, loss_dict, metrics_dict\n\n @abstractmethod\n @profiler.time_function\n def get_eval_image_metrics_and_images(self, step: int):\n \"\"\"This function gets your evaluation loss dict. It needs to get the data\n from the DataManager and feed it to the model's forward function\n\n Args:\n step: current iteration step\n \"\"\"\n\n @abstractmethod\n @profiler.time_function\n def get_average_eval_image_metrics(self, step: Optional[int] = None):\n \"\"\"Iterate over all the images in the eval dataset and get the average.\"\"\"\n\n def load_pipeline(self, loaded_state: Dict[str, Any]) -> None:\n \"\"\"Load the checkpoint from the given path\n\n Args:\n loaded_state: pre-trained model state dict\n \"\"\"\n\n def get_training_callbacks(\n self, training_callback_attributes: TrainingCallbackAttributes\n ) -> List[TrainingCallback]:\n \"\"\"Returns the training callbacks from both the Dataloader and the Model.\"\"\"\n\n def get_param_groups(self) -> Dict[str, List[Parameter]]:\n \"\"\"Get the param groups for the pipeline.\n\n Returns:\n A list of dictionaries containing the pipeline's param groups.\n \"\"\"\n\n\n@dataclass\nclass VanillaPipelineConfig(cfg.InstantiateConfig):\n \"\"\"Configuration for pipeline instantiation\"\"\"\n\n _target: Type = field(default_factory=lambda: VanillaPipeline)\n \"\"\"target class to instantiate\"\"\"\n datamanager: VanillaDataManagerConfig = VanillaDataManagerConfig()\n \"\"\"specifies the datamanager config\"\"\"\n model: ModelConfig = ModelConfig()\n \"\"\"specifies the model config\"\"\"\n\n\nclass VanillaPipeline(Pipeline):\n \"\"\"The pipeline class for the vanilla nerf setup of multiple cameras for one or a few scenes.\n\n config: configuration to instantiate pipeline\n device: location to place model and data\n test_mode:\n 'val': loads train/val datasets into memory\n 'test': loads train/test datset into memory\n 'inference': does not load any dataset into memory\n world_size: total number of machines available\n local_rank: rank of current machine\n\n Attributes:\n datamanager: The data manager that will be used\n model: The model that will be used\n \"\"\"\n\n def __init__(\n self,\n config: VanillaPipelineConfig,\n device: str,\n test_mode: Literal[\"test\", \"val\", \"inference\"] = \"val\",\n world_size: int = 1,\n local_rank: int = 0,\n ):\n super().__init__()\n self.config = config\n self.test_mode = test_mode\n self.datamanager: VanillaDataManager = config.datamanager.setup(\n device=device, test_mode=test_mode, world_size=world_size, local_rank=local_rank\n )\n self.datamanager.to(device)\n # TODO(ethan): get rid of scene_bounds from the model\n assert self.datamanager.train_dataset is not None, \"Missing input dataset\"\n\n self._model = config.model.setup(\n scene_box=self.datamanager.train_dataset.scene_box,\n num_train_data=len(self.datamanager.train_dataset),\n metadata=self.datamanager.train_dataset.metadata,\n world_size=world_size,\n local_rank=local_rank,\n )\n # print(self._model)\n # model_stats = summary(self._model, col_names=[\"kernel_size\", \"num_params\"])\n # TODO: write model_stats to file (pass in the exp dir from trainer)\n self.model.to(device)\n\n self.world_size = world_size\n if world_size > 1:\n self._model = typing.cast(Model, DDP(self._model, device_ids=[local_rank], find_unused_parameters=True))\n dist.barrier(device_ids=[local_rank])\n\n @property\n def device(self):\n \"\"\"Returns the device that the model is on.\"\"\"\n return self.model.device\n\n @profiler.time_function\n def get_train_loss_dict(self, step: int):\n \"\"\"This function gets your training loss dict. This will be responsible for\n getting the next batch of data from the DataManager and interfacing with the\n Model class, feeding the data to the model's forward function.\n\n Args:\n step: current iteration step to update sampler if using DDP (distributed)\n \"\"\"\n ray_bundle, batch = self.datamanager.next_train(step)\n model_outputs = self._model(ray_bundle)\n metrics_dict = self.model.get_metrics_dict(model_outputs, batch)\n\n camera_opt_param_group = self.config.datamanager.camera_optimizer.param_group\n if camera_opt_param_group in self.datamanager.get_param_groups():\n # Report the camera optimization metrics\n metrics_dict[\"camera_opt_translation\"] = (\n self.datamanager.get_param_groups()[camera_opt_param_group][0].data[:, :3].norm()\n )\n metrics_dict[\"camera_opt_rotation\"] = (\n self.datamanager.get_param_groups()[camera_opt_param_group][0].data[:, 3:].norm()\n )\n\n loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict, step=step) # FIXME: donnot use step here\n\n return model_outputs, loss_dict, metrics_dict\n\n def forward(self):\n \"\"\"Blank forward method\n\n This is an nn.Module, and so requires a forward() method normally, although in our case\n we do not need a forward() method\"\"\"\n raise NotImplementedError\n\n @profiler.time_function\n def get_eval_loss_dict(self, step: int):\n \"\"\"This function gets your evaluation loss dict. It needs to get the data\n from the DataManager and feed it to the model's forward function\n\n Args:\n step: current iteration step\n \"\"\"\n self.eval()\n ray_bundle, batch = self.datamanager.next_eval(step)\n model_outputs = self.model(ray_bundle)\n metrics_dict = self.model.get_metrics_dict(model_outputs, batch)\n loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict)\n self.train()\n return model_outputs, loss_dict, metrics_dict\n\n @profiler.time_function\n def get_eval_image_metrics_and_images(self, step: int):\n \"\"\"This function gets your evaluation loss dict. It needs to get the data\n from the DataManager and feed it to the model's forward function\n\n Args:\n step: current iteration step\n \"\"\"\n self.eval()\n image_idx, camera_ray_bundle, batch = self.datamanager.next_eval_image(step)\n outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) # (h, w, c)\n metrics_dict, images_dict = self.model.get_image_metrics_and_images(outputs, batch)\n assert \"image_idx\" not in metrics_dict\n metrics_dict[\"image_idx\"] = image_idx\n assert \"num_rays\" not in metrics_dict\n metrics_dict[\"num_rays\"] = len(camera_ray_bundle)\n self.train()\n return metrics_dict, images_dict\n\n @profiler.time_function\n def get_average_eval_image_metrics(self, step: Optional[int] = None):\n \"\"\"Iterate over all the images in the eval dataset and get the average.\n\n Returns:\n metrics_dict: dictionary of metrics\n \"\"\"\n self.eval()\n metrics_dict_list = []\n num_images = len(self.datamanager.fixed_indices_eval_dataloader)\n with Progress(\n TextColumn(\"[progress.description]{task.description}\"),\n BarColumn(),\n TimeElapsedColumn(),\n MofNCompleteColumn(),\n transient=True,\n ) as progress:\n task = progress.add_task(\"[green]Evaluating all eval images...\", total=num_images)\n for camera_ray_bundle, batch in self.datamanager.fixed_indices_eval_dataloader:\n isbasicimages = False\n if isinstance(\n batch[\"image\"], BasicImages\n ): # If this is a generalized dataset, we need to get image tensor\n isbasicimages = True\n batch[\"image\"] = batch[\"image\"].images[0]\n camera_ray_bundle = camera_ray_bundle.reshape((*batch[\"image\"].shape[:-1],))\n # time this the following line\n inner_start = time()\n height, width = camera_ray_bundle.shape\n num_rays = height * width\n outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle)\n metrics_dict, _ = self.model.get_image_metrics_and_images(outputs, batch)\n assert \"num_rays_per_sec\" not in metrics_dict\n metrics_dict[\"num_rays_per_sec\"] = num_rays / (time() - inner_start)\n fps_str = \"fps\"\n assert fps_str not in metrics_dict\n metrics_dict[fps_str] = metrics_dict[\"num_rays_per_sec\"] / (height * width)\n metrics_dict_list.append(metrics_dict)\n progress.advance(task)\n # average the metrics list\n metrics_dict = {}\n for key in metrics_dict_list[0].keys():\n metrics_dict[key] = float(\n torch.mean(torch.tensor([metrics_dict[key] for metrics_dict in metrics_dict_list]))\n )\n self.train()\n return metrics_dict\n\n def load_pipeline(self, loaded_state: Dict[str, Any], strict=True) -> None:\n \"\"\"Load the checkpoint from the given path\n\n Args:\n loaded_state: pre-trained model state dict\n \"\"\"\n state = {key.replace(\"module.\", \"\"): value for key, value in loaded_state.items()}\n if self.test_mode == \"inference\":\n state.pop(\"datamanager.train_camera_optimizer.pose_adjustment\", None)\n state.pop(\"datamanager.train_ray_generator.image_coords\", None)\n state.pop(\"datamanager.train_ray_generator.pose_optimizer.pose_adjustment\", None)\n state.pop(\"datamanager.eval_ray_generator.image_coords\", None)\n state.pop(\"datamanager.eval_ray_generator.pose_optimizer.pose_adjustment\", None)\n elif not strict:\n # TODO: is it ok? why do we save image_coords anyway?\n state.pop(\"datamanager.train_ray_generator.image_coords\", None)\n state.pop(\"datamanager.eval_ray_generator.image_coords\", None)\n \n # TODO: print unloaded state_dict keys\n self.load_state_dict(state, strict=strict) # type: ignore\n\n def get_training_callbacks(\n self, training_callback_attributes: TrainingCallbackAttributes\n ) -> List[TrainingCallback]:\n \"\"\"Returns the training callbacks from both the Dataloader and the Model.\"\"\"\n datamanager_callbacks = self.datamanager.get_training_callbacks(training_callback_attributes)\n model_callbacks = self.model.get_training_callbacks(training_callback_attributes)\n callbacks = datamanager_callbacks + model_callbacks\n return callbacks\n\n def get_param_groups(self) -> Dict[str, List[Parameter]]:\n \"\"\"Get the param groups for the pipeline.\n\n Returns:\n A list of dictionaries containing the pipeline's param groups.\n \"\"\"\n datamanager_params = self.datamanager.get_param_groups()\n model_params = self.model.get_param_groups()\n # TODO(ethan): assert that key names don't overlap\n return {**datamanager_params, **model_params}\n\n\n@dataclass\nclass FlexibleInputPipelineConfig(VanillaPipelineConfig):\n \"\"\"Configuration for pipeline instantiation\"\"\"\n\n _target: Type = field(default_factory=lambda: FlexibleInputPipeline)\n \"\"\"target class to instantiate\"\"\"\n datamanager: FlexibleDataManagerConfig = FlexibleDataManagerConfig()\n \"\"\"specifies the datamanager config\"\"\"\n model: ModelConfig = ModelConfig()\n \"\"\"specifies the model config\"\"\"\n\n\nclass FlexibleInputPipeline(VanillaPipeline):\n @profiler.time_function\n def get_train_loss_dict(self, step: int):\n \"\"\"This function gets your training loss dict. This will be responsible for\n getting the next batch of data from the DataManager and interfacing with the\n Model class, feeding the data to the model's forward function.\n\n Args:\n step: current iteration step to update sampler if using DDP (distributed)\n \"\"\"\n ray_bundle, batch, additional_input = self.datamanager.next_train(step)\n model_outputs = self.model.get_outputs_flexible(ray_bundle, additional_input)\n # TODO: add additional_input to model get_metrics_dict or get_loss_dict?\n metrics_dict = self.model.get_metrics_dict(model_outputs, batch)\n\n camera_opt_param_group = self.config.datamanager.camera_optimizer.param_group\n if camera_opt_param_group in self.datamanager.get_param_groups():\n # Report the camera optimization metrics\n metrics_dict[\"camera_opt_translation\"] = (\n self.datamanager.get_param_groups()[camera_opt_param_group][0].data[:, :3].norm()\n )\n metrics_dict[\"camera_opt_rotation\"] = (\n self.datamanager.get_param_groups()[camera_opt_param_group][0].data[:, 3:].norm()\n )\n\n loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict)\n\n return model_outputs, loss_dict, metrics_dict\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/pipelines/dynamic_batch.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nA pipeline that dynamically chooses the number of rays to sample.\n\"\"\"\n\nfrom dataclasses import dataclass, field\nfrom typing import Type\n\nimport torch\nfrom typing_extensions import Literal\n\nfrom nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager\nfrom nerfstudio.pipelines.base_pipeline import VanillaPipeline, VanillaPipelineConfig\n\n\n@dataclass\nclass DynamicBatchPipelineConfig(VanillaPipelineConfig):\n \"\"\"Dynamic Batch Pipeline Config\"\"\"\n\n _target: Type = field(default_factory=lambda: DynamicBatchPipeline)\n target_num_samples: int = 262144 # 1 << 18\n \"\"\"The target number of samples to use for an entire batch of rays.\"\"\"\n max_num_samples_per_ray: int = 1024 # 1 << 10\n \"\"\"The maximum number of samples to be placed along a ray.\"\"\"\n\n\nclass DynamicBatchPipeline(VanillaPipeline):\n \"\"\"Pipeline with logic for changing the number of rays per batch.\"\"\"\n\n # pylint: disable=abstract-method\n\n config: DynamicBatchPipelineConfig\n datamanager: VanillaDataManager\n dynamic_num_rays_per_batch: int\n\n def __init__(\n self,\n config: DynamicBatchPipelineConfig,\n device: str,\n test_mode: Literal[\"test\", \"val\", \"inference\"] = \"val\",\n world_size: int = 1,\n local_rank: int = 0,\n ):\n super().__init__(config, device, test_mode, world_size, local_rank)\n assert isinstance(\n self.datamanager, VanillaDataManager\n ), \"DynamicBatchPipeline only works with VanillaDataManager.\"\n\n self.dynamic_num_rays_per_batch = self.config.target_num_samples // self.config.max_num_samples_per_ray\n self._update_pixel_samplers()\n\n def _update_pixel_samplers(self):\n \"\"\"Update the pixel samplers for train and eval with the dynamic number of rays per batch.\"\"\"\n self.datamanager.train_pixel_sampler.set_num_rays_per_batch(self.dynamic_num_rays_per_batch)\n self.datamanager.eval_pixel_sampler.set_num_rays_per_batch(self.dynamic_num_rays_per_batch)\n\n def _update_dynamic_num_rays_per_batch(self, num_samples_per_batch: int):\n \"\"\"Updates the dynamic number of rays per batch variable,\n based on the total number of samples in the last batch of rays.\"\"\"\n self.dynamic_num_rays_per_batch = int(\n self.dynamic_num_rays_per_batch * (self.config.target_num_samples / num_samples_per_batch)\n )\n\n def get_train_loss_dict(self, step: int):\n model_outputs, loss_dict, metrics_dict = super().get_train_loss_dict(step)\n\n # update the number of rays for the next step\n if \"num_samples_per_batch\" not in metrics_dict:\n raise ValueError(\n \"'num_samples_per_batch' is not in metrics_dict.\"\n \"Please return 'num_samples_per_batch' in the models get_metrics_dict function to use this method.\"\n )\n self._update_dynamic_num_rays_per_batch(metrics_dict[\"num_samples_per_batch\"])\n self._update_pixel_samplers()\n\n # add the number of rays\n assert \"num_rays_per_batch\" not in metrics_dict\n metrics_dict[\"num_rays_per_batch\"] = torch.tensor(self.datamanager.train_pixel_sampler.num_rays_per_batch)\n\n return model_outputs, loss_dict, metrics_dict\n\n def get_eval_loss_dict(self, step: int):\n model_outputs, loss_dict, metrics_dict = super().get_eval_loss_dict(step)\n\n # add the number of rays\n assert \"num_rays_per_batch\" not in metrics_dict\n metrics_dict[\"num_rays_per_batch\"] = torch.tensor(self.datamanager.eval_pixel_sampler.num_rays_per_batch)\n\n return model_outputs, loss_dict, metrics_dict\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/process_data/__init__.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/process_data/colmap_utils.py",
"content": "\"\"\"\nHere we have modified code taken from COLMAP for parsing data in the COLMAP format.\nOriginal file at:\nhttps://github.com/colmap/colmap/blob/1a4d0bad2e90aa65ce997c9d1779518eaed998d5/scripts/python/read_write_model.py.\n\"\"\"\n\n# Copyright (c) 2022, ETH Zurich and UNC Chapel Hill.\n# All rights reserved.\n#\n# Redistribution and use in source and binary forms, with or without\n# modification, are permitted provided that the following conditions are met:\n#\n# * Redistributions of source code must retain the above copyright\n# notice, this list of conditions and the following disclaimer.\n#\n# * Redistributions in binary form must reproduce the above copyright\n# notice, this list of conditions and the following disclaimer in the\n# documentation and/or other materials provided with the distribution.\n#\n# * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of\n# its contributors may be used to endorse or promote products derived\n# from this software without specific prior written permission.\n#\n# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\n# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE\n# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE\n# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR\n# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF\n# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS\n# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN\n# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)\n# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE\n# POSSIBILITY OF SUCH DAMAGE.\n#\n# Author: Johannes L. Schoenberger (jsch-at-demuc-dot-de)\n\nimport json\nimport os\nimport struct\nfrom dataclasses import dataclass\nfrom io import BufferedReader\nfrom pathlib import Path\nfrom typing import Dict, Optional, Tuple\n\nimport appdirs\nimport numpy as np\nimport requests\nfrom rich.console import Console\nfrom rich.progress import track\nfrom typing_extensions import Literal\n\nfrom nerfstudio.process_data.process_data_utils import CameraModel\nfrom nerfstudio.utils.rich_utils import status\nfrom nerfstudio.utils.scripts import run_command\n\nCONSOLE = Console(width=120)\n\n\n@dataclass\nclass ColmapCameraModel:\n \"\"\"Camera model\"\"\"\n\n model_id: int\n \"\"\"Model identifier\"\"\"\n model_name: str\n \"\"\"Model name\"\"\"\n num_params: int\n \"\"\"Number of parameters\"\"\"\n\n\n@dataclass\nclass Camera:\n \"\"\"Camera\"\"\"\n\n id: int\n \"\"\"Camera identifier\"\"\"\n model: str\n \"\"\"Camera model\"\"\"\n width: int\n \"\"\"Image width\"\"\"\n height: int\n \"\"\"Image height\"\"\"\n params: np.ndarray\n \"\"\"Camera parameters\"\"\"\n\n\n@dataclass\nclass Image:\n \"\"\"Data the corresponds to a single image\"\"\"\n\n id: int\n \"\"\"Image identifier\"\"\"\n qvec: np.ndarray\n \"\"\"Quaternion vector\"\"\"\n tvec: np.ndarray\n \"\"\"Translation vector\"\"\"\n camera_id: int\n \"\"\"Camera identifier\"\"\"\n name: str\n \"\"\"Image name\"\"\"\n xys: np.ndarray\n \"\"\"2D points\"\"\"\n point3d_ids: np.ndarray\n \"\"\"Point3D identifiers\"\"\"\n\n\n@dataclass\nclass Point3D:\n \"\"\"Data that corresponds to a single 3D point\"\"\"\n\n id: int\n \"\"\"Point3D identifier\"\"\"\n xyz: np.ndarray\n \"\"\"3D point\"\"\"\n rgb: np.ndarray\n \"\"\"Color\"\"\"\n error: float\n \"\"\"Reconstruction error\"\"\"\n image_ids: np.ndarray\n \"\"\"Image identifiers\"\"\"\n point2d_idxs: np.ndarray\n \"\"\"Point2D indices\"\"\"\n\n\nCOLMAP_CAMERA_MODELS = [\n ColmapCameraModel(model_id=0, model_name=\"SIMPLE_PINHOLE\", num_params=3),\n ColmapCameraModel(model_id=1, model_name=\"PINHOLE\", num_params=4),\n ColmapCameraModel(model_id=2, model_name=\"SIMPLE_RADIAL\", num_params=4),\n ColmapCameraModel(model_id=3, model_name=\"RADIAL\", num_params=5),\n ColmapCameraModel(model_id=4, model_name=\"OPENCV\", num_params=8),\n ColmapCameraModel(model_id=5, model_name=\"OPENCV_FISHEYE\", num_params=8),\n ColmapCameraModel(model_id=6, model_name=\"FULL_OPENCV\", num_params=12),\n ColmapCameraModel(model_id=7, model_name=\"FOV\", num_params=5),\n ColmapCameraModel(model_id=8, model_name=\"SIMPLE_RADIAL_FISHEYE\", num_params=4),\n ColmapCameraModel(model_id=9, model_name=\"RADIAL_FISHEYE\", num_params=5),\n ColmapCameraModel(model_id=10, model_name=\"THIN_PRISM_FISHEYE\", num_params=12),\n]\nCOLMAP_CAMERA_MODEL_IDS = {camera_model.model_id: camera_model for camera_model in COLMAP_CAMERA_MODELS}\nCOLMAP_CAMERA_MODEL_NAMES = {camera_model.model_name: camera_model for camera_model in COLMAP_CAMERA_MODELS}\n\n\ndef get_colmap_version(colmap_cmd: str, default_version=3.8) -> float:\n \"\"\"Returns the version of COLMAP.\n This code assumes that colmap returns a version string of the form\n \"COLMAP 3.8 ...\" which may not be true for all versions of COLMAP.\n\n Args:\n default_version: Default version to return if COLMAP version can't be determined.\n Returns:\n The version of COLMAP.\n \"\"\"\n output = run_command(colmap_cmd, verbose=False)\n assert output is not None\n for line in output.split(\"\\n\"):\n if line.startswith(\"COLMAP\"):\n return float(line.split(\" \")[1])\n CONSOLE.print(f\"[bold red]Could not find COLMAP version. Using default {default_version}\")\n return default_version\n\n\ndef read_next_bytes(fid: BufferedReader, num_bytes: int, format_char_sequence, endian_character: str = \"<\"):\n \"\"\"Read and unpack the next bytes from a binary file.\n\n Args:\n fid: Open file\n num_bytes: Sum of combination of {2, 4, 8}, e.g. 2, 6, 16, 30, etc.\n format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}.\n endian_character: Any of {@, =, <, >, !}\n Tuple of read and unpacked values.\n \"\"\"\n data = fid.read(num_bytes)\n return struct.unpack(endian_character + format_char_sequence, data)\n\n\ndef read_cameras_text(path: Path) -> Dict[int, Camera]:\n \"\"\"Parse COLMAP cameras.txt file into a dictionary of Camera objects.\n\n Args:\n path: Path to cameras.txt file.\n Returns:\n Dictionary of Camera objects.\n \"\"\"\n cameras = {}\n with open(path, encoding=\"utf-8\") as fid:\n while True:\n line = fid.readline()\n if not line:\n break\n line = line.strip()\n if len(line) > 0 and line[0] != \"#\":\n elems = line.split()\n camera_id = int(elems[0])\n model = elems[1]\n width = int(elems[2])\n height = int(elems[3])\n params = np.array(tuple(map(float, elems[4:])))\n cameras[camera_id] = Camera(id=camera_id, model=model, width=width, height=height, params=params)\n return cameras\n\n\ndef read_cameras_binary(path_to_model_file: Path) -> Dict[int, Camera]:\n \"\"\"Parse COLMAP cameras.bin file into a dictionary of Camera objects.\n\n Args:\n path_to_model_file: Path to cameras.bin file.\n Returns:\n Dictionary of Camera objects.\n \"\"\"\n cameras = {}\n with open(path_to_model_file, \"rb\") as fid:\n num_cameras = read_next_bytes(fid, 8, \"Q\")[0]\n for _ in range(num_cameras):\n camera_properties = read_next_bytes(fid, num_bytes=24, format_char_sequence=\"iiQQ\")\n camera_id = camera_properties[0]\n model_id = camera_properties[1]\n model_name = COLMAP_CAMERA_MODEL_IDS[camera_properties[1]].model_name\n width = camera_properties[2]\n height = camera_properties[3]\n num_params = COLMAP_CAMERA_MODEL_IDS[model_id].num_params\n params = read_next_bytes(fid, num_bytes=8 * num_params, format_char_sequence=\"d\" * num_params)\n cameras[camera_id] = Camera(\n id=camera_id, model=model_name, width=width, height=height, params=np.array(params)\n )\n assert len(cameras) == num_cameras\n return cameras\n\n\ndef read_images_text(path: Path) -> Dict[int, Image]:\n \"\"\"Parse COLMAP images.txt file into a dictionary of Image objects.\n\n Args:\n path: Path to images.txt file.\n Returns:\n Dictionary of Image objects.\n \"\"\"\n images = {}\n with open(path, encoding=\"utf-8\") as fid:\n while True:\n line = fid.readline()\n if not line:\n break\n line = line.strip()\n if len(line) > 0 and line[0] != \"#\":\n elems = line.split()\n image_id = int(elems[0])\n qvec = np.array(tuple(map(float, elems[1:5])))\n tvec = np.array(tuple(map(float, elems[5:8])))\n camera_id = int(elems[8])\n image_name = elems[9]\n elems = fid.readline().split()\n xys = np.column_stack([tuple(map(float, elems[0::3])), tuple(map(float, elems[1::3]))])\n point3d_ids = np.array(tuple(map(int, elems[2::3])))\n images[image_id] = Image(\n id=image_id,\n qvec=qvec,\n tvec=tvec,\n camera_id=camera_id,\n name=image_name,\n xys=xys,\n point3d_ids=point3d_ids,\n )\n return images\n\n\ndef read_images_binary(path_to_model_file: Path) -> Dict[int, Image]:\n \"\"\"Parse COLMAP images.bin file into a dictionary of Image objects.\n\n Args:\n path_to_model_file: Path to images.bin file.\n Returns:\n Dictionary of Image objects.\n \"\"\"\n images = {}\n with open(path_to_model_file, \"rb\") as fid:\n num_reg_images = read_next_bytes(fid, 8, \"Q\")[0]\n for _ in range(num_reg_images):\n binary_image_properties = read_next_bytes(fid, num_bytes=64, format_char_sequence=\"idddddddi\")\n image_id = binary_image_properties[0]\n qvec = np.array(binary_image_properties[1:5])\n tvec = np.array(binary_image_properties[5:8])\n camera_id = binary_image_properties[8]\n image_name = \"\"\n current_char = read_next_bytes(fid, 1, \"c\")[0]\n while current_char != b\"\\x00\": # look for the ASCII 0 entry\n image_name += current_char.decode(\"utf-8\")\n current_char = read_next_bytes(fid, 1, \"c\")[0]\n num_points2d = read_next_bytes(fid, num_bytes=8, format_char_sequence=\"Q\")[0]\n x_y_id_s = read_next_bytes(fid, num_bytes=24 * num_points2d, format_char_sequence=\"ddq\" * num_points2d)\n xys = np.column_stack([tuple(map(float, x_y_id_s[0::3])), tuple(map(float, x_y_id_s[1::3]))])\n point3d_ids = np.array(tuple(map(int, x_y_id_s[2::3])))\n images[image_id] = Image(\n id=image_id,\n qvec=qvec,\n tvec=tvec,\n camera_id=camera_id,\n name=image_name,\n xys=xys,\n point3d_ids=point3d_ids,\n )\n return images\n\n\ndef read_points3d_text(path) -> Dict[int, Point3D]:\n \"\"\"Parse COLMAP points3D.txt file into a dictionary of Point3D objects.\n\n Args:\n path: Path to points3D.txt file.\n Returns:\n Dictionary of Point3D objects.\n \"\"\"\n points3d = {}\n with open(path, encoding=\"utf-8\") as fid:\n while True:\n line = fid.readline()\n if not line:\n break\n line = line.strip()\n if len(line) > 0 and line[0] != \"#\":\n elems = line.split()\n point3d_id = int(elems[0])\n xyz = np.array(tuple(map(float, elems[1:4])))\n rgb = np.array(tuple(map(int, elems[4:7])))\n error = float(elems[7])\n image_ids = np.array(tuple(map(int, elems[8::2])))\n point2d_idxs = np.array(tuple(map(int, elems[9::2])))\n points3d[point3d_id] = Point3D(\n id=point3d_id, xyz=xyz, rgb=rgb, error=error, image_ids=image_ids, point2d_idxs=point2d_idxs\n )\n return points3d\n\n\ndef read_points3d_binary(path_to_model_file: Path) -> Dict[int, Point3D]:\n \"\"\"Parse COLMAP points3D.bin file into a dictionary of Point3D objects.\n\n Args:\n path_to_model_file: Path to points3D.bin file.\n Returns:\n Dictionary of Point3D objects.\n \"\"\"\n points3d = {}\n with open(path_to_model_file, \"rb\") as fid:\n num_points = read_next_bytes(fid, 8, \"Q\")[0]\n for _ in range(num_points):\n binary_point_line_properties = read_next_bytes(fid, num_bytes=43, format_char_sequence=\"QdddBBBd\")\n point3d_id = binary_point_line_properties[0]\n xyz = np.array(binary_point_line_properties[1:4])\n rgb = np.array(binary_point_line_properties[4:7])\n error = np.array(binary_point_line_properties[7])\n track_length = read_next_bytes(fid, num_bytes=8, format_char_sequence=\"Q\")[0]\n track_elems = read_next_bytes(fid, num_bytes=8 * track_length, format_char_sequence=\"ii\" * track_length)\n image_ids = np.array(tuple(map(int, track_elems[0::2])))\n point2d_idxs = np.array(tuple(map(int, track_elems[1::2])))\n points3d[point3d_id] = Point3D(\n id=point3d_id, xyz=xyz, rgb=rgb, error=float(error), image_ids=image_ids, point2d_idxs=point2d_idxs\n )\n return points3d\n\n\ndef detect_model_format(path: Path, ext: str) -> bool:\n \"\"\"Detect the format of the model file.\n\n Args:\n path: Path to the model file.\n ext: Extension to test.\n Returns:\n True if the model file is the tested extenstion, False otherwise.\n \"\"\"\n\n if (\n os.path.isfile(path / f\"cameras{ext}\")\n and os.path.isfile(path / f\"images{ext}\")\n and os.path.isfile(path / f\"points3D{ext}\")\n ):\n print(\"Detected model format: '\" + ext + \"'\")\n return True\n\n return False\n\n\ndef read_model(path: Path, ext: Optional[str] = None) -> Tuple[Dict[int, Camera], Dict[int, Image], Dict[int, Point3D]]:\n \"\"\"Read a COLMAP model from a directory.\n\n Args:\n path: Path to the model directory.\n ext: Extension of the model files. If None, the function will try to detect the format.\n Returns:\n Tuple of dictionaries of Camera, Image, and Point3D objects.\n \"\"\"\n # try to detect the extension automatically\n if ext is None:\n if detect_model_format(path, \".bin\"):\n ext = \".bin\"\n elif detect_model_format(path, \".txt\"):\n ext = \".txt\"\n else:\n raise ValueError(\"Provide model format: '.bin' or '.txt'\")\n\n if ext == \".txt\":\n cameras = read_cameras_text(path / f\"cameras{ext}\")\n images = read_images_text(path / f\"images{ext}\")\n points3d = read_points3d_text(path / f\"points3D{ext}\")\n else:\n cameras = read_cameras_binary(path / f\"cameras{ext}\")\n images = read_images_binary(path / f\"images{ext}\")\n points3d = read_points3d_binary(path / f\"points3D{ext}\")\n return cameras, images, points3d\n\n\ndef qvec2rotmat(qvec) -> np.ndarray:\n \"\"\"Convert quaternion to rotation matrix.\n\n Args:\n qvec: Quaternion vector of shape (4,).\n Returns:\n Rotation matrix of shape (3, 3).\n \"\"\"\n return np.array(\n [\n [\n 1 - 2 * qvec[2] ** 2 - 2 * qvec[3] ** 2,\n 2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3],\n 2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2],\n ],\n [\n 2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3],\n 1 - 2 * qvec[1] ** 2 - 2 * qvec[3] ** 2,\n 2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1],\n ],\n [\n 2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2],\n 2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1],\n 1 - 2 * qvec[1] ** 2 - 2 * qvec[2] ** 2,\n ],\n ]\n )\n\n\ndef rotmat2qvec(R):\n \"\"\"Convert rotation matrix to quaternion.\n\n Args:\n R: Rotation matrix of shape (3, 3).\n Returns:\n Quaternion vector of shape (4,).\n \"\"\"\n rxx, ryx, rzx, rxy, ryy, rzy, rxz, ryz, rzz = R.flat\n K = (\n np.array(\n [\n [rxx - ryy - rzz, 0, 0, 0],\n [ryx + rxy, ryy - rxx - rzz, 0, 0],\n [rzx + rxz, rzy + ryz, rzz - rxx - ryy, 0],\n [ryz - rzy, rzx - rxz, rxy - ryx, rxx + ryy + rzz],\n ]\n )\n / 3.0\n )\n eigvals, eigvecs = np.linalg.eigh(K)\n qvec = eigvecs[np.array([3, 0, 1, 2]), np.argmax(eigvals)]\n if qvec[0] < 0:\n qvec *= -1\n return qvec\n\n\ndef get_vocab_tree() -> Path:\n \"\"\"Return path to vocab tree. Downloads vocab tree if it doesn't exist.\n\n Returns:\n The path to the vocab tree.\n \"\"\"\n vocab_tree_filename = Path(appdirs.user_data_dir(\"nerfstudio\")) / \"vocab_tree.fbow\"\n\n if not vocab_tree_filename.exists():\n r = requests.get(\"https://demuc.de/colmap/vocab_tree_flickr100K_words32K.bin\", stream=True)\n vocab_tree_filename.parent.mkdir(parents=True, exist_ok=True)\n with open(vocab_tree_filename, \"wb\") as f:\n total_length = r.headers.get(\"content-length\")\n assert total_length is not None\n for chunk in track(\n r.iter_content(chunk_size=1024),\n total=int(total_length) / 1024 + 1,\n description=\"Downloading vocab tree...\",\n ):\n if chunk:\n f.write(chunk)\n f.flush()\n return vocab_tree_filename\n\n\ndef run_colmap(\n image_dir: Path,\n colmap_dir: Path,\n camera_model: CameraModel,\n gpu: bool = True,\n verbose: bool = False,\n matching_method: Literal[\"vocab_tree\", \"exhaustive\", \"sequential\"] = \"vocab_tree\",\n colmap_cmd: str = \"colmap\",\n) -> None:\n \"\"\"Runs COLMAP on the images.\n\n Args:\n image_dir: Path to the directory containing the images.\n colmap_dir: Path to the output directory.\n camera_model: Camera model to use.\n gpu: If True, use GPU.\n verbose: If True, logs the output of the command.\n \"\"\"\n\n colmap_version = get_colmap_version(colmap_cmd)\n\n colmap_database_path = colmap_dir / \"database.db\"\n if colmap_database_path.exists():\n # Can't use missing_ok argument because of Python 3.7 compatibility.\n colmap_database_path.unlink()\n\n # Feature extraction\n feature_extractor_cmd = [\n f\"{colmap_cmd} feature_extractor\",\n f\"--database_path {colmap_dir / 'database.db'}\",\n f\"--image_path {image_dir}\",\n \"--ImageReader.single_camera 1\",\n f\"--ImageReader.camera_model {camera_model.value}\",\n f\"--SiftExtraction.use_gpu {int(gpu)}\",\n ]\n feature_extractor_cmd = \" \".join(feature_extractor_cmd)\n with status(msg=\"[bold yellow]Running COLMAP feature extractor...\", spinner=\"moon\", verbose=verbose):\n run_command(feature_extractor_cmd, verbose=verbose)\n\n CONSOLE.log(\"[bold green]:tada: Done extracting COLMAP features.\")\n\n # Feature matching\n feature_matcher_cmd = [\n f\"{colmap_cmd} {matching_method}_matcher\",\n f\"--database_path {colmap_dir / 'database.db'}\",\n f\"--SiftMatching.use_gpu {int(gpu)}\",\n ]\n if matching_method == \"vocab_tree\":\n vocab_tree_filename = get_vocab_tree()\n feature_matcher_cmd.append(f\"--VocabTreeMatching.vocab_tree_path {vocab_tree_filename}\")\n feature_matcher_cmd = \" \".join(feature_matcher_cmd)\n with status(msg=\"[bold yellow]Running COLMAP feature matcher...\", spinner=\"runner\", verbose=verbose):\n run_command(feature_matcher_cmd, verbose=verbose)\n CONSOLE.log(\"[bold green]:tada: Done matching COLMAP features.\")\n\n # Bundle adjustment\n sparse_dir = colmap_dir / \"sparse\"\n sparse_dir.mkdir(parents=True, exist_ok=True)\n mapper_cmd = [\n f\"{colmap_cmd} mapper\",\n f\"--database_path {colmap_dir / 'database.db'}\",\n f\"--image_path {image_dir}\",\n f\"--output_path {sparse_dir}\",\n ]\n if colmap_version >= 3.7:\n mapper_cmd.append(\"--Mapper.ba_global_function_tolerance 1e-6\")\n\n mapper_cmd = \" \".join(mapper_cmd)\n\n with status(\n msg=\"[bold yellow]Running COLMAP bundle adjustment... (This may take a while)\",\n spinner=\"circle\",\n verbose=verbose,\n ):\n run_command(mapper_cmd, verbose=verbose)\n CONSOLE.log(\"[bold green]:tada: Done COLMAP bundle adjustment.\")\n with status(msg=\"[bold yellow]Refine intrinsics...\", spinner=\"dqpb\", verbose=verbose):\n bundle_adjuster_cmd = [\n f\"{colmap_cmd} bundle_adjuster\",\n f\"--input_path {sparse_dir}/0\",\n f\"--output_path {sparse_dir}/0\",\n \"--BundleAdjustment.refine_principal_point 1\",\n ]\n run_command(\" \".join(bundle_adjuster_cmd), verbose=verbose)\n CONSOLE.log(\"[bold green]:tada: Done refining intrinsics.\")\n\n\ndef colmap_to_json(cameras_path: Path, images_path: Path, output_dir: Path, camera_model: CameraModel) -> int:\n \"\"\"Converts COLMAP's cameras.bin and images.bin to a JSON file.\n\n Args:\n cameras_path: Path to the cameras.bin file.\n images_path: Path to the images.bin file.\n output_dir: Path to the output directory.\n camera_model: Camera model used.\n\n Returns:\n The number of registered images.\n \"\"\"\n\n cameras = read_cameras_binary(cameras_path)\n images = read_images_binary(images_path)\n\n # Only supports one camera\n camera_params = cameras[1].params\n\n frames = []\n for _, im_data in images.items():\n rotation = qvec2rotmat(im_data.qvec)\n translation = im_data.tvec.reshape(3, 1)\n w2c = np.concatenate([rotation, translation], 1)\n w2c = np.concatenate([w2c, np.array([[0, 0, 0, 1]])], 0)\n c2w = np.linalg.inv(w2c)\n # Convert from COLMAP's camera coordinate system to ours\n c2w[0:3, 1:3] *= -1\n c2w = c2w[np.array([1, 0, 2, 3]), :]\n c2w[2, :] *= -1\n\n name = Path(f\"./images/{im_data.name}\")\n\n frame = {\n \"file_path\": name.as_posix(),\n \"transform_matrix\": c2w.tolist(),\n }\n frames.append(frame)\n\n out = {\n \"fl_x\": float(camera_params[0]),\n \"fl_y\": float(camera_params[1]),\n \"cx\": float(camera_params[2]),\n \"cy\": float(camera_params[3]),\n \"w\": cameras[1].width,\n \"h\": cameras[1].height,\n \"camera_model\": camera_model.value,\n }\n\n if camera_model == CameraModel.OPENCV:\n out.update(\n {\n \"k1\": float(camera_params[4]),\n \"k2\": float(camera_params[5]),\n \"p1\": float(camera_params[6]),\n \"p2\": float(camera_params[7]),\n }\n )\n if camera_model == CameraModel.OPENCV_FISHEYE:\n out.update(\n {\n \"k1\": float(camera_params[4]),\n \"k2\": float(camera_params[5]),\n \"k3\": float(camera_params[6]),\n \"k4\": float(camera_params[7]),\n }\n )\n\n out[\"frames\"] = frames\n\n with open(output_dir / \"transforms.json\", \"w\", encoding=\"utf-8\") as f:\n json.dump(out, f, indent=4)\n\n return len(frames)\n\n\ndef get_matching_summary(num_intial_frames: int, num_matched_frames: int) -> str:\n \"\"\"Returns a summary of the matching results.\n\n Args:\n num_intial_frames: The number of initial frames.\n num_matched_frames: The number of matched frames.\n\n Returns:\n A summary of the matching results.\n \"\"\"\n match_ratio = num_matched_frames / num_intial_frames\n if match_ratio == 1:\n return \"[bold green]COLAMP found poses for all images, CONGRATS!\"\n if match_ratio < 0.4:\n result = f\"[bold red]COLMAP only found poses for {num_matched_frames / num_intial_frames * 100:.2f}%\"\n result += \" of the images. This is low.\\nThis can be caused by a variety of reasons,\"\n result += \" such poor scene coverage, blurry images, or large exposure changes.\"\n return result\n if match_ratio < 0.8:\n result = f\"[bold yellow]COLMAP only found poses for {num_matched_frames / num_intial_frames * 100:.2f}%\"\n result += \" of the images.\\nThis isn't great, but may be ok.\"\n result += \"\\nMissing poses can be caused by a variety of reasons, such poor scene coverage, blurry images,\"\n result += \" or large exposure changes.\"\n return result\n return f\"[bold green]COLMAP found poses for {num_matched_frames / num_intial_frames * 100:.2f}% of the images.\"\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/process_data/hloc_utils.py",
"content": "\"\"\"\nCode that uses the hierarchical localization toolbox (hloc)\nto extract and match image features, estimate camera poses,\nand do sparse reconstruction.\nRequires hloc module from : https://github.com/cvg/Hierarchical-Localization\n\"\"\"\n\n# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport sys\nfrom pathlib import Path\n\nfrom rich.console import Console\nfrom typing_extensions import Literal\n\nfrom nerfstudio.process_data.process_data_utils import CameraModel\n\ntry:\n import pycolmap\n from hloc import (\n extract_features,\n match_features,\n pairs_from_exhaustive,\n pairs_from_retrieval,\n reconstruction,\n )\nexcept ImportError:\n _HAS_HLOC = False\nelse:\n _HAS_HLOC = True\n\nCONSOLE = Console(width=120)\n\n\ndef run_hloc(\n image_dir: Path,\n colmap_dir: Path,\n camera_model: CameraModel,\n verbose: bool = False,\n matching_method: Literal[\"vocab_tree\", \"exhaustive\", \"sequential\"] = \"vocab_tree\",\n feature_type: Literal[\n \"sift\", \"superpoint_aachen\", \"superpoint_max\", \"superpoint_inloc\", \"r2d2\", \"d2net-ss\", \"sosnet\", \"disk\"\n ] = \"superpoint_aachen\",\n matcher_type: Literal[\n \"superglue\", \"superglue-fast\", \"NN-superpoint\", \"NN-ratio\", \"NN-mutual\", \"adalam\"\n ] = \"superglue\",\n num_matched: int = 50,\n) -> None:\n \"\"\"Runs hloc on the images.\n\n Args:\n image_dir: Path to the directory containing the images.\n colmap_dir: Path to the output directory.\n camera_model: Camera model to use.\n gpu: If True, use GPU.\n verbose: If True, logs the output of the command.\n \"\"\"\n if not _HAS_HLOC:\n CONSOLE.print(\n f\"[bold red]Error: To use this set of parameters ({feature_type}/{matcher_type}/hloc), \"\n \"you must install hloc toolbox!!\"\n )\n sys.exit(1)\n\n outputs = colmap_dir\n sfm_pairs = outputs / \"pairs-netvlad.txt\"\n sfm_dir = outputs / \"sparse\" / \"0\"\n features = outputs / \"features.h5\"\n matches = outputs / \"matches.h5\"\n\n retrieval_conf = extract_features.confs[\"netvlad\"]\n feature_conf = extract_features.confs[feature_type]\n matcher_conf = match_features.confs[matcher_type]\n\n references = [p.relative_to(image_dir).as_posix() for p in image_dir.iterdir()]\n extract_features.main(feature_conf, image_dir, image_list=references, feature_path=features)\n if matching_method == \"exhaustive\":\n pairs_from_exhaustive.main(sfm_pairs, image_list=references)\n else:\n retrieval_path = extract_features.main(retrieval_conf, image_dir, outputs)\n if num_matched >= len(references):\n num_matched = len(references)\n pairs_from_retrieval.main(retrieval_path, sfm_pairs, num_matched=num_matched)\n match_features.main(matcher_conf, sfm_pairs, features=features, matches=matches)\n\n image_options = pycolmap.ImageReaderOptions(camera_model=camera_model.value)\n reconstruction.main(\n sfm_dir,\n image_dir,\n sfm_pairs,\n features,\n matches,\n camera_mode=pycolmap.CameraMode.SINGLE,\n image_options=image_options,\n verbose=verbose,\n )\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/process_data/insta360_utils.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Helper utils for processing insta360 data.\"\"\"\n\nimport sys\nfrom pathlib import Path\nfrom typing import List, Tuple\n\nfrom rich.console import Console\n\nfrom nerfstudio.process_data.process_data_utils import get_num_frames_in_video\nfrom nerfstudio.utils.rich_utils import status\nfrom nerfstudio.utils.scripts import run_command\n\nCONSOLE = Console(width=120)\n\n\ndef get_insta360_filenames(data: Path) -> Tuple[Path, Path]:\n \"\"\"Returns the filenames of the Insta360 videos from a single video file.\n\n Example input name: VID_20220212_070353_00_003.insv\n\n Args:\n data: Path to a Insta360 file.\n\n Returns:\n The filenames of the Insta360 videios.\n \"\"\"\n if data.suffix != \".insv\":\n raise ValueError(\"The input file must be an .insv file.\")\n file_parts = data.stem.split(\"_\")\n\n stem_back = f\"VID_{file_parts[1]}_{file_parts[2]}_00_{file_parts[4]}.insv\"\n stem_front = f\"VID_{file_parts[1]}_{file_parts[2]}_10_{file_parts[4]}.insv\"\n\n filename_back = data.parent / stem_back\n filename_front = data.parent / stem_front\n\n return filename_back, filename_front\n\n\ndef convert_insta360_to_images(\n video_front: Path,\n video_back: Path,\n image_dir: Path,\n num_frames_target: int,\n crop_percentage: float = 0.7,\n verbose: bool = False,\n) -> Tuple[List[str], int]:\n \"\"\"Converts a video into a sequence of images.\n\n Args:\n video_front: Path to the front video.\n video_back: Path to the back video.\n output_dir: Path to the output directory.\n num_frames_target: Number of frames to extract.\n crop_percentage: Percentage used to calculate the cropped dimentions of extracted frames. Currently used to crop\n out the curved portions of the fish-eye lens.\n verbose: If True, logs the output of the command.\n Returns:\n A tuple containing summary of the conversion and the number of extracted frames.\n \"\"\"\n\n with status(msg=\"Converting video to images...\", spinner=\"bouncingBall\", verbose=verbose):\n # delete existing images in folder\n for img in image_dir.glob(\"*.png\"):\n if verbose:\n CONSOLE.log(f\"Deleting {img}\")\n img.unlink()\n\n num_frames_front = get_num_frames_in_video(video_front)\n num_frames_back = get_num_frames_in_video(video_back)\n if num_frames_front == 0:\n CONSOLE.print(f\"[bold red]Error: Video has no frames: {video_front}\")\n sys.exit(1)\n if num_frames_back == 0:\n CONSOLE.print(f\"[bold red]Error: Video has no frames: {video_front}\")\n sys.exit(1)\n\n spacing = num_frames_front // (num_frames_target // 2)\n vf_cmds = []\n if spacing > 1:\n vf_cmds = [f\"thumbnail={spacing}\", \"setpts=N/TB\"]\n else:\n CONSOLE.print(\"[bold red]Can't satify requested number of frames. Extracting all frames.\")\n\n vf_cmds.append(f\"crop=iw*{crop_percentage}:ih*{crop_percentage}\")\n\n front_vf_cmds = vf_cmds + [\"transpose=2\"]\n back_vf_cmds = vf_cmds + [\"transpose=1\"]\n\n front_ffmpeg_cmd = f\"ffmpeg -i {video_front} -vf {','.join(front_vf_cmds)} -r 1 {image_dir / 'frame_%05d.png'}\"\n back_ffmpeg_cmd = (\n f\"ffmpeg -i {video_back} -vf {','.join(back_vf_cmds)} -r 1 {image_dir / 'back_frame_%05d.png'}\"\n )\n\n run_command(front_ffmpeg_cmd, verbose=verbose)\n run_command(back_ffmpeg_cmd, verbose=verbose)\n\n num_extracted_front_frames = len(list(image_dir.glob(\"frame*.png\")))\n for i, img in enumerate(image_dir.glob(\"back_frame_*.png\")):\n img.rename(image_dir / f\"frame_{i+1+num_extracted_front_frames:05d}.png\")\n\n num_final_frames = len(list(image_dir.glob(\"*.png\")))\n summary_log = []\n summary_log.append(f\"Starting with {num_frames_front + num_frames_back} video frames\")\n summary_log.append(f\"We extracted {num_final_frames} images\")\n CONSOLE.log(\"[bold green]:tada: Done converting insta360 to images.\")\n\n return summary_log, num_final_frames\n\n\ndef convert_insta360_single_file_to_images(\n video: Path,\n image_dir: Path,\n num_frames_target: int,\n crop_percentage: float = 0.7,\n verbose: bool = False,\n) -> Tuple[List[str], int]:\n \"\"\"Converts a video into a sequence of images.\n\n Args:\n video: Path to the video.\n output_dir: Path to the output directory.\n num_frames_target: Number of frames to extract.\n crop_percentage: Percentage used to calculate the cropped dimentions of extracted frames. Currently used to crop\n out the curved portions of the fish-eye lens.\n verbose: If True, logs the output of the command.\n Returns:\n A tuple containing summary of the conversion and the number of extracted frames.\n \"\"\"\n\n with status(msg=\"Converting video to images...\", spinner=\"bouncingBall\", verbose=verbose):\n # delete existing images in folder\n for img in image_dir.glob(\"*.png\"):\n if verbose:\n CONSOLE.log(f\"Deleting {img}\")\n img.unlink()\n\n num_frames = get_num_frames_in_video(video)\n if num_frames == 0:\n CONSOLE.print(f\"[bold red]Error: Video has no frames: {video}\")\n sys.exit(1)\n\n spacing = num_frames // (num_frames_target // 2)\n vf_cmds = []\n if spacing > 1:\n vf_cmds = [f\"thumbnail={spacing}\", \"setpts=N/TB\"]\n else:\n CONSOLE.print(\"[bold red]Can't satify requested number of frames. Extracting all frames.\")\n\n vf_cmds_back = vf_cmds.copy()\n vf_cmds_front = vf_cmds.copy()\n\n vf_cmds_back.append(\n f\"crop=ih*{crop_percentage}:ih*{crop_percentage}:ih*({crop_percentage}/4):ih*({crop_percentage}/4)\"\n )\n vf_cmds_front.append(\n f\"crop=ih*{crop_percentage}:ih*{crop_percentage}:iw/2+ih*{crop_percentage/4}:ih*{crop_percentage/4}\"\n )\n\n front_ffmpeg_cmd = f\"ffmpeg -i {video} -vf {','.join(vf_cmds_front)} -r 1 {image_dir / 'frame_%05d.png'}\"\n back_ffmpeg_cmd = f\"ffmpeg -i {video} -vf {','.join(vf_cmds_back)} -r 1 {image_dir / 'back_frame_%05d.png'}\"\n\n run_command(back_ffmpeg_cmd, verbose=verbose)\n run_command(front_ffmpeg_cmd, verbose=verbose)\n\n num_extracted_frames = len(list(image_dir.glob(\"frame*.png\")))\n for i, img in enumerate(image_dir.glob(\"back_frame_*.png\")):\n img.rename(image_dir / f\"frame_{i+1+num_extracted_frames:05d}.png\")\n\n num_final_frames = len(list(image_dir.glob(\"*.png\")))\n summary_log = []\n summary_log.append(f\"Starting with {num_frames} video frames\")\n summary_log.append(f\"We extracted {num_final_frames} images\")\n CONSOLE.log(\"[bold green]:tada: Done converting insta360 to images.\")\n\n return summary_log, num_final_frames\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/process_data/metashape_utils.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Helper utils for processing metashape data into the nerfstudio format.\"\"\"\n\nimport json\nimport xml.etree.ElementTree as ET\nfrom pathlib import Path\nfrom typing import Dict, List\n\nimport numpy as np\nfrom rich.console import Console\n\nfrom nerfstudio.process_data.process_data_utils import CAMERA_MODELS\n\nCONSOLE = Console(width=120)\n\n\ndef _find_distortion_param(calib_xml: ET.Element, param_name: str):\n param = calib_xml.find(param_name)\n if param is not None:\n return float(param.text) # type: ignore\n return 0.0\n\n\ndef metashape_to_json( # pylint: disable=too-many-statements\n image_filename_map: Dict[str, Path],\n xml_filename: Path,\n output_dir: Path,\n verbose: bool = False,\n) -> List[str]:\n \"\"\"Convert Metashape data into a nerfstudio dataset.\n\n Args:\n image_filename_map: Mapping of original image filenames to their saved locations.\n xml_filename: Path to the metashap cameras xml file.\n output_dir: Path to the output directory.\n verbose: Whether to print verbose output.\n\n Returns:\n Summary of the conversion.\n \"\"\"\n\n xml_tree = ET.parse(xml_filename)\n root = xml_tree.getroot()\n chunk = root[0]\n sensors = chunk.find(\"sensors\")\n\n # TODO Add support for per-frame intrinsics\n if sensors is None or len(sensors) != 1:\n raise ValueError(\"Only one sensor is supported for now\")\n\n sensor = sensors.find(\"sensor\")\n\n data = {}\n\n assert sensor is not None, \"Sensor not found in Metashape XML\"\n resolution = sensor.find(\"resolution\")\n assert resolution is not None, \"Resolution not found in Metashape xml\"\n data[\"w\"] = int(resolution.get(\"width\")) # type: ignore\n data[\"h\"] = int(resolution.get(\"height\")) # type: ignore\n\n calib = sensor.find(\"calibration\")\n assert calib is not None, \"Calibration not found in Metashape xml\"\n data[\"fl_x\"] = float(calib.find(\"f\").text) # type: ignore\n data[\"fl_y\"] = float(calib.find(\"f\").text) # type: ignore\n data[\"cx\"] = float(calib.find(\"cx\").text) + data[\"w\"] / 2.0 # type: ignore\n data[\"cy\"] = float(calib.find(\"cy\").text) + data[\"h\"] / 2.0 # type: ignore\n\n data[\"k1\"] = _find_distortion_param(calib, \"k1\")\n data[\"k2\"] = _find_distortion_param(calib, \"k2\")\n data[\"k3\"] = _find_distortion_param(calib, \"k3\")\n data[\"k4\"] = _find_distortion_param(calib, \"k4\")\n data[\"p1\"] = _find_distortion_param(calib, \"p1\")\n data[\"p2\"] = _find_distortion_param(calib, \"p2\")\n\n data[\"camera_model\"] = CAMERA_MODELS[\"perspective\"].value\n\n frames = []\n cameras = chunk.find(\"cameras\")\n assert cameras is not None, \"Cameras not found in Metashape xml\"\n num_skipped = 0\n for camera in cameras:\n frame = {}\n # Labels sometimes have a file extension. We remove it for consistency.\n camera_label = camera.get(\"label\").split(\".\")[0] # type: ignore\n if camera_label not in image_filename_map:\n continue\n frame[\"file_path\"] = image_filename_map[camera_label].as_posix()\n if camera.find(\"transform\") is None:\n if verbose:\n CONSOLE.print(f\"Missing transforms data for {camera.get('label')}, Skipping\")\n num_skipped += 1\n continue\n t = [float(x) for x in camera.find(\"transform\").text.split()] # type: ignore\n transform = np.array(\n [\n [t[8], -t[9], -t[10], t[11]],\n [t[0], -t[1], -t[2], t[3]],\n [t[4], -t[5], -t[6], t[7]],\n [t[12], -t[13], -t[14], t[15]],\n ]\n )\n frame[\"transform_matrix\"] = transform.tolist()\n frames.append(frame)\n\n data[\"frames\"] = frames\n\n with open(output_dir / \"transforms.json\", \"w\", encoding=\"utf-8\") as f:\n json.dump(data, f, indent=4)\n\n summary = []\n if num_skipped == 1:\n summary.append(f\"{num_skipped} image skipped because it was missing its camera pose.\")\n if num_skipped > 1:\n summary.append(f\"{num_skipped} images were skipped because they were missing camera poses.\")\n\n summary.append(f\"Final dataset is {len(data['frames'])} frames.\")\n\n return summary\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/process_data/polycam_utils.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Helper utils for processing polycam data into the nerfstudio format.\"\"\"\n\nimport json\nimport sys\nfrom pathlib import Path\nfrom typing import List\n\nfrom rich.console import Console\n\nfrom nerfstudio.process_data.process_data_utils import CAMERA_MODELS\nfrom nerfstudio.utils import io\n\nCONSOLE = Console(width=120)\n\n\ndef polycam_to_json(\n image_filenames: List[Path],\n cameras_dir: Path,\n output_dir: Path,\n min_blur_score: float = 0.0,\n crop_border_pixels: int = 0,\n) -> List[str]:\n \"\"\"Convert Polycam data into a nerfstudio dataset.\n\n Args:\n image_filenames: List of paths to the original images.\n cameras_dir: Path to the polycam cameras directory.\n output_dir: Path to the output directory.\n min_blur_score: Minimum blur score to use an image. Images below this value will be skipped.\n crop_border_pixels: Number of pixels to crop from each border of the image.\n\n Returns:\n Summary of the conversion.\n \"\"\"\n data = {}\n data[\"camera_model\"] = CAMERA_MODELS[\"perspective\"].value\n # Needs to be a string for camera_utils.auto_orient_and_center_poses\n data[\"orientation_override\"] = \"none\"\n\n frames = []\n skipped_frames = 0\n for i, image_filename in enumerate(image_filenames):\n json_filename = cameras_dir / f\"{image_filename.stem}.json\"\n frame_json = io.load_from_json(json_filename)\n if \"blur_score\" in frame_json and frame_json[\"blur_score\"] < min_blur_score:\n skipped_frames += 1\n continue\n frame = {}\n frame[\"fl_x\"] = frame_json[\"fx\"]\n frame[\"fl_y\"] = frame_json[\"fy\"]\n frame[\"cx\"] = frame_json[\"cx\"] - crop_border_pixels\n frame[\"cy\"] = frame_json[\"cy\"] - crop_border_pixels\n frame[\"w\"] = frame_json[\"width\"] - crop_border_pixels * 2\n frame[\"h\"] = frame_json[\"height\"] - crop_border_pixels * 2\n frame[\"file_path\"] = f\"./images/frame_{i+1:05d}{image_filename.suffix}\"\n # Transform matrix to nerfstudio format. Please refer to the documentation for coordinate system conventions.\n frame[\"transform_matrix\"] = [\n [frame_json[\"t_20\"], frame_json[\"t_21\"], frame_json[\"t_22\"], frame_json[\"t_23\"]],\n [frame_json[\"t_00\"], frame_json[\"t_01\"], frame_json[\"t_02\"], frame_json[\"t_03\"]],\n [frame_json[\"t_10\"], frame_json[\"t_11\"], frame_json[\"t_12\"], frame_json[\"t_13\"]],\n [0.0, 0.0, 0.0, 1.0],\n ]\n frames.append(frame)\n data[\"frames\"] = frames\n\n with open(output_dir / \"transforms.json\", \"w\", encoding=\"utf-8\") as f:\n json.dump(data, f, indent=4)\n\n summary = []\n if skipped_frames > 0:\n summary.append(f\"Skipped {skipped_frames} frames due to low blur score.\")\n summary.append(f\"Final dataset is {len(image_filenames) - skipped_frames} frames.\")\n\n if len(image_filenames) - skipped_frames == 0:\n CONSOLE.print(\"[bold red]No images remain after filtering, exiting\")\n sys.exit(1)\n\n return summary\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/process_data/process_data_utils.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Helper utils for processing data into the nerfstudio format.\"\"\"\n\nimport shutil\nimport sys\nfrom enum import Enum\nfrom pathlib import Path\nfrom typing import List, Optional, Tuple\n\nfrom rich.console import Console\nfrom typing_extensions import Literal\n\nfrom nerfstudio.utils.rich_utils import status\nfrom nerfstudio.utils.scripts import run_command\n\nCONSOLE = Console(width=120)\n\n\nclass CameraModel(Enum):\n \"\"\"Enum for camera types.\"\"\"\n\n OPENCV = \"OPENCV\"\n OPENCV_FISHEYE = \"OPENCV_FISHEYE\"\n\n\nCAMERA_MODELS = {\n \"perspective\": CameraModel.OPENCV,\n \"fisheye\": CameraModel.OPENCV_FISHEYE,\n}\n\n\ndef get_num_frames_in_video(video: Path) -> int:\n \"\"\"Returns the number of frames in a video.\n\n Args:\n video: Path to a video.\n\n Returns:\n The number of frames in a video.\n \"\"\"\n cmd = f\"ffprobe -v error -select_streams v:0 -count_packets \\\n -show_entries stream=nb_read_packets -of csv=p=0 {video}\"\n output = run_command(cmd)\n assert output is not None\n output = output.strip(\" ,\\t\\n\\r\")\n return int(output)\n\n\ndef convert_video_to_images(\n video_path: Path, image_dir: Path, num_frames_target: int, verbose: bool = False\n) -> Tuple[List[str], int]:\n \"\"\"Converts a video into a sequence of images.\n\n Args:\n video_path: Path to the video.\n output_dir: Path to the output directory.\n num_frames_target: Number of frames to extract.\n verbose: If True, logs the output of the command.\n Returns:\n A tuple containing summary of the conversion and the number of extracted frames.\n \"\"\"\n\n with status(msg=\"Converting video to images...\", spinner=\"bouncingBall\", verbose=verbose):\n # delete existing images in folder\n for img in image_dir.glob(\"*.png\"):\n if verbose:\n CONSOLE.log(f\"Deleting {img}\")\n img.unlink()\n\n num_frames = get_num_frames_in_video(video_path)\n if num_frames == 0:\n CONSOLE.print(f\"[bold red]Error: Video has no frames: {video_path}\")\n sys.exit(1)\n print(\"Number of frames in video:\", num_frames)\n\n out_filename = image_dir / \"frame_%05d.png\"\n ffmpeg_cmd = f\"ffmpeg -i {video_path}\"\n spacing = num_frames // num_frames_target\n\n if spacing > 1:\n ffmpeg_cmd += f\" -vf thumbnail={spacing},setpts=N/TB -r 1\"\n else:\n CONSOLE.print(\"[bold red]Can't satify requested number of frames. Extracting all frames.\")\n\n ffmpeg_cmd += f\" {out_filename}\"\n\n run_command(ffmpeg_cmd, verbose=verbose)\n\n num_final_frames = len(list(image_dir.glob(\"*.png\")))\n summary_log = []\n summary_log.append(f\"Starting with {num_frames} video frames\")\n summary_log.append(f\"We extracted {num_final_frames} images\")\n CONSOLE.log(\"[bold green]:tada: Done converting video to images.\")\n\n return summary_log, num_final_frames\n\n\ndef copy_images_list(\n image_paths: List[Path], image_dir: Path, crop_border_pixels: Optional[int] = None, verbose: bool = False\n) -> List[Path]:\n \"\"\"Copy all images in a list of Paths. Useful for filtering from a directory.\n Args:\n image_paths: List of Paths of images to copy to a new directory.\n image_dir: Path to the output directory.\n crop_border_pixels: If not None, crops each edge by the specified number of pixels.\n verbose: If True, print extra logging.\n Returns:\n A list of the copied image Paths.\n \"\"\"\n\n # Remove original directory only if we provide a proper image folder path\n if image_dir.is_dir() and len(image_paths):\n shutil.rmtree(image_dir, ignore_errors=True)\n image_dir.mkdir(exist_ok=True, parents=True)\n\n copied_image_paths = []\n\n # Images should be 1-indexed for the rest of the pipeline.\n for idx, image_path in enumerate(image_paths):\n if verbose:\n CONSOLE.log(f\"Copying image {idx + 1} of {len(image_paths)}...\")\n copied_image_path = image_dir / f\"frame_{idx + 1:05d}{image_path.suffix}\"\n shutil.copy(image_path, copied_image_path)\n copied_image_paths.append(copied_image_path)\n\n if crop_border_pixels is not None:\n file_type = image_paths[0].suffix\n filename = f\"frame_%05d{file_type}\"\n crop = f\"crop=iw-{crop_border_pixels*2}:ih-{crop_border_pixels*2}\"\n ffmpeg_cmd = f\"ffmpeg -y -i {image_dir / filename} -q:v 2 -vf {crop} {image_dir / filename}\"\n run_command(ffmpeg_cmd, verbose=verbose)\n\n num_frames = len(image_paths)\n\n if num_frames == 0:\n CONSOLE.log(\"[bold red]:skull: No usable images in the data folder.\")\n else:\n CONSOLE.log(\"[bold green]:tada: Done copying images.\")\n\n return copied_image_paths\n\n\ndef copy_images(data: Path, image_dir: Path, verbose) -> int:\n \"\"\"Copy images from a directory to a new directory.\n\n Args:\n data: Path to the directory of images.\n image_dir: Path to the output directory.\n verbose: If True, print extra logging.\n Returns:\n The number of images copied.\n \"\"\"\n with status(msg=\"[bold yellow]Copying images...\", spinner=\"bouncingBall\", verbose=verbose):\n allowed_exts = [\".jpg\", \".jpeg\", \".png\", \".tif\", \".tiff\"]\n image_paths = sorted([p for p in data.glob(\"[!.]*\") if p.suffix.lower() in allowed_exts])\n\n num_frames = len(copy_images_list(image_paths, image_dir, verbose))\n\n return num_frames\n\n\ndef downscale_images(image_dir: Path, num_downscales: int, verbose: bool = False) -> str:\n \"\"\"Downscales the images in the directory. Uses FFMPEG.\n\n Assumes images are named frame_00001.png, frame_00002.png, etc.\n\n Args:\n image_dir: Path to the directory containing the images.\n num_downscales: Number of times to downscale the images. Downscales by 2 each time.\n verbose: If True, logs the output of the command.\n\n Returns:\n Summary of downscaling.\n \"\"\"\n\n if num_downscales == 0:\n return \"No downscaling performed.\"\n\n with status(msg=\"[bold yellow]Downscaling images...\", spinner=\"growVertical\", verbose=verbose):\n downscale_factors = [2**i for i in range(num_downscales + 1)[1:]]\n for downscale_factor in downscale_factors:\n assert downscale_factor > 1\n assert isinstance(downscale_factor, int)\n downscale_dir = image_dir.parent / f\"images_{downscale_factor}\"\n downscale_dir.mkdir(parents=True, exist_ok=True)\n file_type = image_dir.glob(\"frame_*\").__next__().suffix\n filename = f\"frame_%05d{file_type}\"\n ffmpeg_cmd = [\n f\"ffmpeg -i {image_dir / filename} \",\n f\"-q:v 2 -vf scale=iw/{downscale_factor}:ih/{downscale_factor} \",\n f\"{downscale_dir / filename}\",\n ]\n ffmpeg_cmd = \" \".join(ffmpeg_cmd)\n run_command(ffmpeg_cmd, verbose=verbose)\n\n CONSOLE.log(\"[bold green]:tada: Done downscaling images.\")\n downscale_text = [f\"[bold blue]{2**(i+1)}x[/bold blue]\" for i in range(num_downscales)]\n downscale_text = \", \".join(downscale_text[:-1]) + \" and \" + downscale_text[-1]\n return f\"We downsampled the images by {downscale_text}\"\n\n\ndef find_tool_feature_matcher_combination(\n sfm_tool: Literal[\"any\", \"colmap\", \"hloc\"],\n feature_type: Literal[\n \"any\",\n \"sift\",\n \"superpoint\",\n \"superpoint_aachen\",\n \"superpoint_max\",\n \"superpoint_inloc\",\n \"r2d2\",\n \"d2net-ss\",\n \"sosnet\",\n \"disk\",\n ],\n matcher_type: Literal[\n \"any\", \"NN\", \"superglue\", \"superglue-fast\", \"NN-superpoint\", \"NN-ratio\", \"NN-mutual\", \"adalam\"\n ],\n):\n \"\"\"Find a valid combination of sfm tool, feature type, and matcher type.\n Basically, replace the default parameters 'any' by usable value\n\n Args:\n sfm_tool: Sfm tool name (any, colmap, hloc)\n feature_type: Type of image features (any, sift, superpoint, ...)\n matcher_type: Type of matching algorithm (any, NN, superglue,...)\n\n Returns:\n Tuple of sfm tool, feature type, and matcher type.\n Returns (None,None,None) if no valid combination can be found\n \"\"\"\n if sfm_tool == \"any\":\n if (feature_type in (\"any\", \"sift\")) and (matcher_type in (\"any\", \"NN\")):\n sfm_tool = \"colmap\"\n else:\n sfm_tool = \"hloc\"\n\n if sfm_tool == \"colmap\":\n if (feature_type not in (\"any\", \"sift\")) or (matcher_type not in (\"any\", \"NN\")):\n return (None, None, None)\n return (\"colmap\", \"sift\", \"NN\")\n if sfm_tool == \"hloc\":\n if feature_type in (\"any\", \"superpoint\"):\n feature_type = \"superpoint_aachen\"\n\n if matcher_type == \"any\":\n matcher_type = \"superglue\"\n elif matcher_type == \"NN\":\n matcher_type = \"NN-mutual\"\n\n return (sfm_tool, feature_type, matcher_type)\n return (None, None, None)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/process_data/record3d_utils.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Helper functions for processing record3d data.\"\"\"\n\nimport json\nfrom pathlib import Path\nfrom typing import List\n\nimport numpy as np\nfrom rich.console import Console\nfrom scipy.spatial.transform import Rotation\n\nfrom nerfstudio.process_data.process_data_utils import CAMERA_MODELS\nfrom nerfstudio.utils import io\n\nCONSOLE = Console(width=120)\n\n\ndef record3d_to_json(images_paths: List[Path], metadata_path: Path, output_dir: Path, indices: np.ndarray) -> int:\n \"\"\"Converts Record3D's metadata and image paths to a JSON file.\n\n Args:\n images_paths: list if image paths.\n metadata_path: Path to the Record3D metadata JSON file.\n output_dir: Path to the output directory.\n indices: Indices to sample the metadata_path. Should be the same length as images_paths.\n\n Returns:\n The number of registered images.\n \"\"\"\n\n assert len(images_paths) == len(indices)\n\n metadata_dict = io.load_from_json(metadata_path)\n\n poses_data = np.array(metadata_dict[\"poses\"]) # (N, 3, 4)\n camera_to_worlds = np.concatenate(\n [Rotation.from_quat(poses_data[:, :4]).as_matrix(), poses_data[:, 4:, None]],\n axis=-1,\n ).astype(np.float32)\n camera_to_worlds = camera_to_worlds[indices]\n\n homogeneous_coord = np.zeros_like(camera_to_worlds[..., :1, :])\n homogeneous_coord[..., :, 3] = 1\n camera_to_worlds = np.concatenate([camera_to_worlds, homogeneous_coord], -2)\n\n frames = []\n for i, im_path in enumerate(images_paths):\n c2w = camera_to_worlds[i]\n frame = {\n \"file_path\": im_path.as_posix(),\n \"transform_matrix\": c2w.tolist(),\n }\n frames.append(frame)\n\n # Camera intrinsics\n K = np.array(metadata_dict[\"K\"]).reshape((3, 3)).T\n focal_length = K[0, 0]\n\n H = metadata_dict[\"h\"]\n W = metadata_dict[\"w\"]\n\n # TODO(akristoffersen): The metadata dict comes with principle points,\n # but caused errors in image coord indexing. Should update once that is fixed.\n cx, cy = W / 2, H / 2\n\n out = {\n \"fl_x\": focal_length,\n \"fl_y\": focal_length,\n \"cx\": cx,\n \"cy\": cy,\n \"w\": W,\n \"h\": H,\n \"camera_model\": CAMERA_MODELS[\"perspective\"].name,\n }\n\n out[\"frames\"] = frames\n\n with open(output_dir / \"transforms.json\", \"w\", encoding=\"utf-8\") as f:\n json.dump(out, f, indent=4)\n\n return len(frames)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/py.typed",
"content": ""
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/__init__.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/bilateral_solver.py",
"content": "import numpy as np\nfrom scipy import ndimage\nfrom scipy.sparse import diags\nfrom scipy.sparse.linalg import cg\nfrom scipy.sparse import csr_matrix\n\n\nRGB_TO_YUV = np.array(\n [[0.299, 0.587, 0.114], [-0.168736, -0.331264, 0.5], [0.5, -0.418688, -0.081312]]\n)\nYUV_TO_RGB = np.array([[1.0, 0.0, 1.402], [1.0, -0.34414, -0.71414], [1.0, 1.772, 0.0]])\nYUV_OFFSET = np.array([0, 128.0, 128.0]).reshape(1, 1, -1)\nMAX_VAL = 255.0\n\n\ndef rgb2yuv(im):\n return np.tensordot(im, RGB_TO_YUV, ([2], [1])) + YUV_OFFSET\n\n\ndef yuv2rgb(im):\n return np.tensordot(im.astype(float) - YUV_OFFSET, YUV_TO_RGB, ([2], [1]))\n\n\ndef get_valid_idx(valid, candidates):\n \"\"\"Find which values are present in a list and where they are located\"\"\"\n locs = np.searchsorted(valid, candidates)\n # Handle edge case where the candidate is larger than all valid values\n locs = np.clip(locs, 0, len(valid) - 1)\n # Identify which values are actually present\n valid_idx = np.flatnonzero(valid[locs] == candidates)\n locs = locs[valid_idx]\n return valid_idx, locs\n\n\nclass BilateralGrid(object):\n def __init__(self, im, sigma_spatial=32, sigma_luma=8, sigma_chroma=8):\n im_yuv = rgb2yuv(im)\n # Compute 5-dimensional XYLUV bilateral-space coordinates\n Iy, Ix = np.mgrid[:im.shape[0], :im.shape[1]]\n x_coords = (Ix / sigma_spatial).astype(int)\n y_coords = (Iy / sigma_spatial).astype(int)\n luma_coords = (im_yuv[..., 0] / sigma_luma).astype(int)\n chroma_coords = (im_yuv[..., 1:] / sigma_chroma).astype(int)\n coords = np.dstack((x_coords, y_coords, luma_coords, chroma_coords))\n coords_flat = coords.reshape(-1, coords.shape[-1])\n self.npixels, self.dim = coords_flat.shape\n # Hacky \"hash vector\" for coordinates,\n # Requires all scaled coordinates be < MAX_VAL\n self.hash_vec = (MAX_VAL**np.arange(self.dim))\n # Construct S and B matrix\n self._compute_factorization(coords_flat)\n\n def _compute_factorization(self, coords_flat):\n # Hash each coordinate in grid to a unique value\n hashed_coords = self._hash_coords(coords_flat)\n unique_hashes, unique_idx, idx = \\\n np.unique(hashed_coords, return_index=True, return_inverse=True)\n # Identify unique set of vertices\n unique_coords = coords_flat[unique_idx]\n self.nvertices = len(unique_coords)\n # Construct sparse splat matrix that maps from pixels to vertices\n self.S = csr_matrix((np.ones(self.npixels), (idx, np.arange(self.npixels))))\n # Construct sparse blur matrices.\n # Note that these represent [1 0 1] blurs, excluding the central element\n self.blurs = []\n for d in range(self.dim):\n blur = 0.0\n for offset in (-1, 1):\n offset_vec = np.zeros((1, self.dim))\n offset_vec[:, d] = offset\n neighbor_hash = self._hash_coords(unique_coords + offset_vec)\n valid_coord, idx = get_valid_idx(unique_hashes, neighbor_hash)\n blur = blur + csr_matrix(\n (np.ones((len(valid_coord), )), (valid_coord, idx)),\n shape=(self.nvertices, self.nvertices)\n )\n self.blurs.append(blur)\n\n def _hash_coords(self, coord):\n \"\"\"Hacky function to turn a coordinate into a unique value\"\"\"\n return np.dot(coord.reshape(-1, self.dim), self.hash_vec)\n\n def splat(self, x):\n return self.S.dot(x)\n\n def slice(self, y):\n return self.S.T.dot(y)\n\n def blur(self, x):\n \"\"\"Blur a bilateral-space vector with a 1 2 1 kernel in each dimension\"\"\"\n assert x.shape[0] == self.nvertices\n out = 2 * self.dim * x\n for blur in self.blurs:\n out = out + blur.dot(x)\n return out\n\n def filter(self, x):\n \"\"\"Apply bilateral filter to an input x\"\"\"\n return self.slice(self.blur(self.splat(x))) / \\\n self.slice(self.blur(self.splat(np.ones_like(x))))\n\n\ndef bistochastize(grid, maxiter=10):\n \"\"\"Compute diagonal matrices to bistochastize a bilateral grid\"\"\"\n m = grid.splat(np.ones(grid.npixels))\n n = np.ones(grid.nvertices)\n for i in range(maxiter):\n n = np.sqrt(n * m / grid.blur(n))\n # Correct m to satisfy the assumption of bistochastization regardless\n # of how many iterations have been run.\n m = n * grid.blur(n)\n Dm = diags(m, 0)\n Dn = diags(n, 0)\n return Dn, Dm\n\n\nclass BilateralSolver(object):\n def __init__(self, grid, params):\n self.grid = grid\n self.params = params\n self.Dn, self.Dm = bistochastize(grid)\n\n def solve(self, x, w):\n # Check that w is a vector or a nx1 matrix\n if w.ndim == 2:\n assert (w.shape[1] == 1)\n elif w.dim == 1:\n w = w.reshape(w.shape[0], 1)\n A_smooth = (self.Dm - self.Dn.dot(self.grid.blur(self.Dn)))\n w_splat = self.grid.splat(w)\n A_data = diags(w_splat[:, 0], 0)\n A = self.params[\"lam\"] * A_smooth + A_data\n xw = x * w\n b = self.grid.splat(xw)\n # Use simple Jacobi preconditioner\n A_diag = np.maximum(A.diagonal(), self.params[\"A_diag_min\"])\n M = diags(1 / A_diag, 0)\n # Flat initialization\n y0 = self.grid.splat(xw) / w_splat\n yhat = np.empty_like(y0)\n for d in range(x.shape[-1]):\n yhat[..., d], info = cg(\n A,\n b[..., d],\n x0=y0[..., d],\n M=M,\n maxiter=self.params[\"cg_maxiter\"],\n tol=self.params[\"cg_tol\"]\n )\n xhat = self.grid.slice(yhat)\n return xhat\n\n\ndef bilateral_refine(image, mask, sigma_spatial=24, sigma_luma=4, sigma_chroma=4,\n soft_mask=False, binary_thresh=0.5, confidence=0.999):\n reference = image\n h, w = mask.shape\n \n if not soft_mask:\n confidence = np.ones((h, w)) * confidence\n else:\n assert mask.max() <= 1.0\n confidence = mask.clone()\n mask = (mask > binary_thresh).astype(np.float64)\n\n grid_params = {\n 'sigma_luma': sigma_luma, # Brightness bandwidth\n 'sigma_chroma': sigma_chroma, # Color bandwidth\n 'sigma_spatial': sigma_spatial # Spatial bandwidth\n }\n\n bs_params = {\n 'lam': 256, # The strength of the smoothness parameter\n 'A_diag_min': 1e-5, # Clamp the diagonal of the A diagonal in the Jacobi preconditioner.\n 'cg_tol': 1e-5, # The tolerance on the convergence in PCG\n 'cg_maxiter': 25 # The number of PCG iterations\n }\n\n grid = BilateralGrid(reference, **grid_params)\n\n t = mask.reshape(-1, 1).astype(np.double)\n c = confidence.reshape(-1, 1).astype(np.double)\n\n # output solver, which is a soft value\n output_solver = BilateralSolver(grid, bs_params).solve(t, c).reshape((h, w))\n\n binary_solver = ndimage.binary_fill_holes(output_solver > 0.5)\n labeled, nr_objects = ndimage.label(binary_solver)\n\n nb_pixel = [np.sum(labeled == i) for i in range(nr_objects + 1)]\n pixel_order = np.argsort(nb_pixel)\n try: # take the second-largest components (assume the largest component is bg)\n binary_solver = labeled == pixel_order[-2]\n except:\n binary_solver = np.ones((h, w), dtype=bool)\n\n return output_solver, binary_solver\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/colormaps.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\" Helper functions for visualizing outputs \"\"\"\n\nfrom typing import Optional\n\nimport torch\nfrom matplotlib import cm\nfrom torchtyping import TensorType\n\nfrom nerfstudio.utils import colors\n\n\ndef apply_colormap(image: TensorType[\"bs\":..., 1], cmap=\"viridis\") -> TensorType[\"bs\":..., \"rgb\":3]:\n \"\"\"Convert single channel to a color image.\n\n Args:\n image: Single channel image.\n cmap: Colormap for image.\n\n Returns:\n TensorType: Colored image\n \"\"\"\n\n # TODO: support more colormaps\n colormap = cm.get_cmap(cmap)\n colormap = torch.tensor(colormap.colors).to(image.device) # type: ignore\n image_long = (image * 255).long()\n image_long_min = torch.min(image_long)\n image_long_max = torch.max(image_long)\n assert image_long_min >= 0, f\"the min value is {image_long_min}\"\n assert image_long_max <= 255, f\"the max value is {image_long_max}\"\n return colormap[image_long[..., 0]]\n\n\ndef apply_depth_colormap(\n depth: TensorType[\"bs\":..., 1],\n accumulation: Optional[TensorType[\"bs\":..., 1]] = None,\n near_plane: Optional[float] = None,\n far_plane: Optional[float] = None,\n cmap=\"turbo\",\n) -> TensorType[\"bs\":..., \"rgb\":3]:\n \"\"\"Converts a depth image to color for easier analysis.\n\n Args:\n depth: Depth image.\n accumulation: Ray accumulation used for masking vis.\n near_plane: Closest depth to consider. If None, use min image value.\n far_plane: Furthest depth to consider. If None, use max image value.\n cmap: Colormap to apply.\n\n Returns:\n Colored depth image\n \"\"\"\n\n near_plane = near_plane or float(torch.min(depth))\n far_plane = far_plane or float(torch.max(depth))\n\n depth = (depth - near_plane) / (far_plane - near_plane + 1e-10)\n depth = torch.clip(depth, 0, 1)\n # depth = torch.nan_to_num(depth, nan=0.0) # TODO(ethan): remove this\n\n colored_image = apply_colormap(depth, cmap=cmap)\n\n if accumulation is not None:\n colored_image = colored_image * accumulation + (1 - accumulation)\n\n return colored_image\n\n\ndef apply_boolean_colormap(\n image: TensorType[\"bs\":..., 1, bool],\n true_color: TensorType[\"bs\":..., \"rgb\":3] = colors.WHITE,\n false_color: TensorType[\"bs\":..., \"rgb\":3] = colors.BLACK,\n) -> TensorType[\"bs\":..., \"rgb\":3]:\n \"\"\"Converts a depth image to color for easier analysis.\n\n Args:\n image: Boolean image.\n true_color: Color to use for True.\n false_color: Color to use for False.\n\n Returns:\n Colored boolean image\n \"\"\"\n\n colored_image = torch.ones(image.shape[:-1] + (3,))\n colored_image[image[..., 0], :] = true_color\n colored_image[~image[..., 0], :] = false_color\n return colored_image\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/colors.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Common Colors\"\"\"\nfrom typing import Union\n\nimport torch\nfrom torchtyping import TensorType\n\nWHITE = torch.tensor([1.0, 1.0, 1.0])\nBLACK = torch.tensor([0.0, 0.0, 0.0])\nRED = torch.tensor([1.0, 0.0, 0.0])\nGREEN = torch.tensor([0.0, 1.0, 0.0])\nBLUE = torch.tensor([0.0, 0.0, 1.0])\n\nCOLORS_DICT = {\n \"white\": WHITE,\n \"black\": BLACK,\n \"red\": RED,\n \"green\": GREEN,\n \"blue\": BLUE,\n}\n\n\ndef get_color(color: Union[str, list]) -> TensorType[3]:\n \"\"\"\n Args:\n color (Union[str, list]): Color as a string or a rgb list\n\n Returns:\n TensorType[3]: Parsed color\n \"\"\"\n if isinstance(color, str):\n color = color.lower()\n if color not in COLORS_DICT:\n raise ValueError(f\"{color} is not a valid preset color\")\n return COLORS_DICT[color]\n if isinstance(color, list):\n if len(color) != 3:\n raise ValueError(f\"Color should be 3 values (RGB) instead got {color}\")\n return torch.tensor(color)\n\n raise ValueError(f\"Color should be an RGB list or string, instead got {type(color)}\")\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/comms.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"functionality to handle multiprocessing syncing and communicating\"\"\"\nimport torch.distributed as dist\n\nLOCAL_PROCESS_GROUP = None\n\n\ndef is_dist_avail_and_initialized() -> bool:\n \"\"\"Returns True if distributed is available and initialized.\"\"\"\n return dist.is_available() and dist.is_initialized()\n\n\ndef get_world_size() -> int:\n \"\"\"Get total number of available gpus\"\"\"\n if not is_dist_avail_and_initialized():\n return 1\n return dist.get_world_size()\n\n\ndef get_rank() -> int:\n \"\"\"Get global rank of current thread\"\"\"\n if not is_dist_avail_and_initialized():\n return 0\n return dist.get_rank()\n\n\ndef get_local_rank() -> int:\n \"\"\"The rank of the current process within the local (per-machine) process group.\"\"\"\n if not is_dist_avail_and_initialized():\n return 0\n assert (\n LOCAL_PROCESS_GROUP is not None\n ), \"Local process group is not created! Please use launch() to spawn processes!\"\n return dist.get_rank(group=LOCAL_PROCESS_GROUP)\n\n\ndef get_local_size() -> int:\n \"\"\"\n The size of the per-machine process group,\n i.e. the number of processes per machine.\n \"\"\"\n if not is_dist_avail_and_initialized():\n return 1\n return dist.get_world_size(group=LOCAL_PROCESS_GROUP)\n\n\ndef is_main_process() -> bool:\n \"\"\"check to see if you are currently on the main process\"\"\"\n return get_rank() == 0\n\n\ndef synchronize():\n \"\"\"\n Helper function to synchronize (barrier) among all processes when\n using distributed training\n \"\"\"\n if dist.get_world_size() == 1:\n return\n if dist.get_backend() == dist.Backend.NCCL:\n # This argument is needed to avoid warnings.\n # It's valid only for NCCL backend.\n dist.barrier(device_ids=[get_local_rank()])\n else:\n dist.barrier()\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/decorators.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nDecorator definitions\n\"\"\"\nfrom typing import Callable, List\n\nfrom nerfstudio.utils import comms\n\n\ndef decorate_all(decorators: List[Callable]) -> Callable:\n \"\"\"A decorator to decorate all member functions of a class\n\n Args:\n decorators: list of decorators to add to all functions in the class\n \"\"\"\n\n def decorate(cls):\n for attr in cls.__dict__:\n if callable(getattr(cls, attr)) and attr != \"__init__\":\n for decorator in decorators:\n setattr(cls, attr, decorator(getattr(cls, attr)))\n return cls\n\n return decorate\n\n\ndef check_profiler_enabled(func: Callable) -> Callable:\n \"\"\"Decorator: check if profiler is enabled\"\"\"\n\n def wrapper(self, *args, **kwargs):\n ret = None\n if self.config.enable_profiler:\n ret = func(self, *args, **kwargs)\n return ret\n\n return wrapper\n\n\ndef check_viewer_enabled(func: Callable) -> Callable:\n \"\"\"Decorator: check if viewer is enabled and only run on main process\"\"\"\n\n def wrapper(self, *args, **kwargs):\n ret = None\n if self.config.is_viewer_enabled() and comms.is_main_process():\n ret = func(self, *args, **kwargs)\n return ret\n\n return wrapper\n\n\ndef check_eval_enabled(func: Callable) -> Callable:\n \"\"\"Decorator: check if evaluation step is enabled\"\"\"\n\n def wrapper(self, *args, **kwargs):\n ret = None\n if self.config.is_wandb_enabled() or self.config.is_tensorboard_enabled():\n ret = func(self, *args, **kwargs)\n return ret\n\n return wrapper\n\n\ndef check_main_thread(func: Callable) -> Callable:\n \"\"\"Decorator: check if you are on main thread\"\"\"\n\n def wrapper(*args, **kwargs):\n ret = None\n if comms.is_main_process():\n ret = func(*args, **kwargs)\n return ret\n\n return wrapper\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/eval_utils.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nEvaluation utils\n\"\"\"\nfrom __future__ import annotations\n\nimport os\nimport sys\nfrom pathlib import Path\nfrom typing import Optional, Tuple, Callable\n\nimport torch\nimport yaml\nfrom rich.console import Console\nfrom typing_extensions import Literal\n\nfrom nerfstudio.configs import base_config as cfg\nfrom nerfstudio.pipelines.base_pipeline import Pipeline\n\nCONSOLE = Console(width=120)\n\n\ndef eval_load_checkpoint(config: cfg.TrainerConfig, pipeline: Pipeline) -> Path:\n ## TODO: ideally eventually want to get this to be the same as whatever is used to load train checkpoint too\n \"\"\"Helper function to load checkpointed pipeline\n\n Args:\n config (DictConfig): Configuration of pipeline to load\n pipeline (Pipeline): Pipeline instance of which to load weights\n \"\"\"\n assert config.load_dir is not None\n if config.load_step is None:\n CONSOLE.print(\"Loading latest checkpoint from load_dir\")\n # NOTE: this is specific to the checkpoint name format\n if not os.path.exists(config.load_dir):\n CONSOLE.rule(\"Error\", style=\"red\")\n CONSOLE.print(f\"No checkpoint directory found at {config.load_dir}, \", justify=\"center\")\n CONSOLE.print(\n \"Please make sure the checkpoint exists, they should be generated periodically during training\",\n justify=\"center\",\n )\n sys.exit(1)\n load_step = sorted(int(x[x.find(\"-\") + 1 : x.find(\".\")]) for x in os.listdir(config.load_dir))[-1]\n else:\n load_step = config.load_step\n load_path = config.load_dir / f\"step-{load_step:09d}.ckpt\"\n assert load_path.exists(), f\"Checkpoint {load_path} does not exist\"\n loaded_state = torch.load(load_path, map_location=\"cpu\")\n if getattr(config, \"load_pipeline_ckpt_strict\", True):\n pipeline.load_pipeline(loaded_state[\"pipeline\"])\n else:\n pipeline.load_pipeline(loaded_state[\"pipeline\"], strict=False)\n CONSOLE.print(\"[bold yellow]Pipeline checkpoint loaded with strict=False\")\n CONSOLE.print(f\":white_check_mark: Done loading checkpoint from {load_path}\")\n return load_path\n\n\ndef eval_setup(\n config_path: Path,\n eval_num_rays_per_chunk: Optional[int] = None,\n test_mode: Literal[\"test\", \"val\", \"inference\"] = \"test\",\n cache_images: bool = True,\n config_override_fn: Optional[Callable] = None\n) -> Tuple[cfg.Config, Pipeline, Path]:\n \"\"\"Shared setup for loading a saved pipeline for evaluation.\n\n Args:\n config_path: Path to config YAML file.\n eval_num_rays_per_chunk: Number of rays per forward pass\n test_mode:\n 'val': loads train/val datasets into memory\n 'test': loads train/test datset into memory\n 'inference': does not load any dataset into memory\n\n\n Returns:\n Loaded config, pipeline module, and corresponding checkpoint.\n \"\"\"\n # load save config\n config = yaml.load(config_path.read_text(), Loader=yaml.Loader)\n assert isinstance(config, cfg.Config)\n\n if eval_num_rays_per_chunk:\n config.pipeline.model.eval_num_rays_per_chunk = eval_num_rays_per_chunk\n\n # load checkpoints from wherever they were saved\n # TODO: expose the ability to choose an arbitrary checkpoint\n config.trainer.load_dir = config.get_checkpoint_dir()\n config.pipeline.datamanager.eval_image_indices = None\n if not cache_images:\n config.pipeline.datamanager.train_num_images_to_sample_from = 0\n config.pipeline.datamanager.eval_num_images_to_sample_from = 0\n \n if config_override_fn is not None:\n CONSOLE.print(f\"Overriding config with {config_override_fn}\")\n config_override_fn(config)\n\n # setup pipeline (which includes the DataManager)\n device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n pipeline = config.pipeline.setup(device=device, test_mode=test_mode)\n assert isinstance(pipeline, Pipeline)\n pipeline.eval()\n\n # load checkpointed information\n checkpoint_path = eval_load_checkpoint(config.trainer, pipeline)\n\n return config, pipeline, checkpoint_path\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/func_utils.py",
"content": "\nfrom typing import Sequence\n\n\ndef get_first_element(elems: Sequence):\n return elems[0]\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/images.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nDefines an image that can be batched with the default nerfstudio collate fn, even if the images\naren't of the same height and width.\n\"\"\"\n\nfrom typing import List\n\nimport torch\n\n\n# pylint: disable=too-few-public-methods\nclass BasicImages:\n \"\"\"This is a very primitive struct for holding images, especially for when these images\n are of different heights / widths.\n\n The purpose of this is to have a special struct wrapping around a list so that the\n nerfstudio_collate fn and other parts of the code recognise this as a struct to leave alone\n instead of reshaping or concatenating into a single tensor (since this will likely be used\n for cases where we have images of different sizes and shapes).\n\n This only has one batch dimension and will likely be replaced down the line with some\n TensorDataclass alternative that supports arbitrary batches.\n \"\"\"\n\n def __init__(self, images: List):\n assert isinstance(images, List)\n assert not images or isinstance(\n images[0], torch.Tensor\n ), f\"Input should be a list of tensors, not {type(images[0]) if isinstance(images, List) else type(images)}\"\n self.images = images\n\n def to(self, device):\n \"\"\"Move the images to the given device.\"\"\"\n assert isinstance(device, (torch.device, str))\n return BasicImages([image.to(device) for image in self.images])\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/install_checks.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Helpers for checking if programs are installed\"\"\"\n\nimport shutil\nimport sys\n\nfrom rich.console import Console\n\nCONSOLE = Console(width=120)\n\n\ndef check_ffmpeg_installed():\n \"\"\"Checks if ffmpeg is installed.\"\"\"\n ffmpeg_path = shutil.which(\"ffmpeg\")\n if ffmpeg_path is None:\n CONSOLE.print(\"[bold red]Could not find ffmpeg. Please install ffmpeg.\")\n print(\"See https://ffmpeg.org/download.html for installation instructions.\")\n print(\"ffmpeg is only necessary if using videos as input.\")\n sys.exit(1)\n\n\ndef check_colmap_installed():\n \"\"\"Checks if colmap is installed.\"\"\"\n colmap_path = shutil.which(\"colmap\")\n if colmap_path is None:\n CONSOLE.print(\"[bold red]Could not find COLMAP. Please install COLMAP.\")\n print(\"See https://colmap.github.io/install.html for installation instructions.\")\n sys.exit(1)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/io.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nInput/output utils.\n\"\"\"\n\nimport json\nfrom pathlib import Path\n\n\ndef load_from_json(filename: Path):\n \"\"\"Load a dictionary from a JSON filename.\n\n Args:\n filename: The filename to load from.\n \"\"\"\n assert filename.suffix == \".json\"\n with open(filename, encoding=\"UTF-8\") as file:\n return json.load(file)\n\n\ndef write_to_json(filename: Path, content: dict):\n \"\"\"Write data to a JSON file.\n\n Args:\n filename: The filename to write to.\n content: The dictionary data to write.\n \"\"\"\n assert filename.suffix == \".json\"\n with open(filename, \"w\", encoding=\"UTF-8\") as file:\n json.dump(content, file)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/marching_cubes.py",
"content": "from pathlib import Path\nfrom rich import console\n\nimport numpy as np\nimport pymeshlab\nimport torch\nimport trimesh\nfrom skimage import measure\n\nCONSOLE = console.Console(width=120)\n\navg_pool_3d = torch.nn.AvgPool3d(2, stride=2)\nupsample = torch.nn.Upsample(scale_factor=2, mode=\"nearest\")\n\n\n@torch.no_grad()\ndef get_surface_sliding(\n sdf,\n resolution=512,\n bounding_box_min=(-1.0, -1.0, -1.0),\n bounding_box_max=(1.0, 1.0, 1.0),\n return_mesh=False,\n level=0,\n coarse_mask=None,\n output_path: Path = Path(\"test.ply\"),\n simplify_mesh=True,\n chunk_size=100000,\n store_float16=False, # store intermediate results with float16 to save VRAM\n):\n dtype = torch.float32 if not store_float16 else torch.float16\n if coarse_mask is not None:\n # we need to permute here as pytorch's grid_sample use (z, y, x)\n coarse_mask = coarse_mask.permute(2, 1, 0)[None, None].cuda().float()\n\n resN = resolution\n cropN = min(resN, 512)\n level = 0\n N = resN // cropN\n\n grid_min = bounding_box_min\n grid_max = bounding_box_max\n xs = np.linspace(grid_min[0], grid_max[0], N + 1)\n ys = np.linspace(grid_min[1], grid_max[1], N + 1)\n zs = np.linspace(grid_min[2], grid_max[2], N + 1)\n\n meshes = []\n # TODO: add a status bar\n for i in range(N):\n for j in range(N):\n for k in range(N):\n # print(i, j, k)\n x_min, x_max = xs[i], xs[i + 1]\n y_min, y_max = ys[j], ys[j + 1]\n z_min, z_max = zs[k], zs[k + 1]\n\n x = np.linspace(x_min, x_max, cropN)\n y = np.linspace(y_min, y_max, cropN)\n z = np.linspace(z_min, z_max, cropN)\n\n xx, yy, zz = np.meshgrid(x, y, z, indexing=\"ij\")\n points = torch.tensor(np.vstack([xx.ravel(), yy.ravel(), zz.ravel()]).T, dtype=torch.float).cuda()\n\n def evaluate(points):\n z = []\n for _, pnts in enumerate(torch.split(points, chunk_size, dim=0)):\n z.append(sdf(pnts))\n z = torch.cat(z, axis=0)\n sdf_val = z.to(dtype)\n del z\n torch.cuda.empty_cache()\n return sdf_val\n\n # construct point pyramids\n points = points.reshape(cropN, cropN, cropN, 3).permute(3, 0, 1, 2)\n if coarse_mask is not None:\n # breakpoint()\n points_tmp = points.permute(1, 2, 3, 0)[None].cuda()\n current_mask = torch.nn.functional.grid_sample(coarse_mask, points_tmp)\n current_mask = (current_mask > 0.0).cpu().numpy()[0, 0]\n else:\n current_mask = None\n\n points_pyramid = [points]\n for _ in range(3):\n points = avg_pool_3d(points[None])[0]\n points_pyramid.append(points)\n points_pyramid = points_pyramid[::-1]\n\n # evalute pyramid with mask\n mask = None\n threshold = 2 * (x_max - x_min) / cropN * 8\n for pid, pts in enumerate(points_pyramid):\n coarse_N = pts.shape[-1]\n pts = pts.reshape(3, -1).permute(1, 0).contiguous()\n\n if mask is None:\n # only evaluate\n if coarse_mask is not None:\n pts_sdf = torch.ones_like(pts[:, 1], dtpye=dtype)\n valid_mask = (\n torch.nn.functional.grid_sample(coarse_mask, pts[None, None, None])[0, 0, 0, 0] > 0\n )\n if valid_mask.any():\n pts_sdf[valid_mask] = evaluate(pts[valid_mask].contiguous())\n else:\n pts_sdf = evaluate(pts)\n else:\n mask = mask.reshape(-1)\n pts_to_eval = pts[mask]\n\n if pts_to_eval.shape[0] > 0:\n pts_sdf_eval = evaluate(pts_to_eval.contiguous())\n pts_sdf[mask] = pts_sdf_eval\n # print(\"ratio\", pts_to_eval.shape[0] / pts.shape[0])\n\n if pid < 3:\n # update mask\n mask = torch.abs(pts_sdf) < threshold\n mask = mask.reshape(coarse_N, coarse_N, coarse_N)[None, None]\n mask = upsample(mask.float()).bool()\n\n pts_sdf = pts_sdf.reshape(coarse_N, coarse_N, coarse_N)[None, None]\n pts_sdf = upsample(pts_sdf)\n pts_sdf = pts_sdf.reshape(-1)\n\n threshold /= 2.0\n\n z = pts_sdf.detach().cpu().numpy()\n\n # skip if no surface found\n if current_mask is not None:\n valid_z = z.reshape(cropN, cropN, cropN)[current_mask]\n if valid_z.shape[0] <= 0 or (np.min(valid_z) > level or np.max(valid_z) < level):\n continue\n\n if not (np.min(z) > level or np.max(z) < level):\n z = z.astype(np.float32)\n verts, faces, vertex_normals, _ = measure.marching_cubes(\n volume=z.reshape(cropN, cropN, cropN), # .transpose([1, 0, 2]),\n level=level,\n spacing=(\n (x_max - x_min) / (cropN - 1),\n (y_max - y_min) / (cropN - 1),\n (z_max - z_min) / (cropN - 1),\n ),\n mask=current_mask,\n )\n # print(nzp.array([x_min, y_min, z_min]))\n # print(verts.min(), verts.max())\n verts = verts + np.array([x_min, y_min, z_min])\n # print(verts.min(), verts.max())\n meshcrop = trimesh.Trimesh(verts, faces,\n vertex_normals=vertex_normals)\n # meshcrop.export(f\"{i}_{j}_{k}.ply\")\n meshes.append(meshcrop)\n\n combined = trimesh.util.concatenate(meshes)\n\n if return_mesh:\n return combined\n else:\n filename = str(output_path)\n filename_simplify = str(output_path).replace(\".ply\", \"-simplify.ply\")\n combined.export(filename)\n CONSOLE.print(f\"Marching cube mesh saved: {filename}\")\n if simplify_mesh:\n ms = pymeshlab.MeshSet()\n ms.load_new_mesh(filename)\n\n print(\"simply mesh\")\n ms.meshing_decimation_quadric_edge_collapse(targetfacenum=2000000)\n ms.save_current_mesh(filename_simplify, save_face_color=False)\n CONSOLE.print(f\"Simplified mesh saved: {filename_simplify}\")\n\n\n@torch.no_grad()\ndef get_surface_occupancy(\n occupancy_fn,\n resolution=512,\n bounding_box_min=(-1.0, -1.0, -1.0),\n bounding_box_max=(1.0, 1.0, 1.0),\n return_mesh=False,\n level=0.5,\n device=None,\n output_path: Path = Path(\"test.ply\"),\n chunk_size=100000,\n):\n grid_min = bounding_box_min\n grid_max = bounding_box_max\n N = resolution\n xs = np.linspace(grid_min[0], grid_max[0], N)\n ys = np.linspace(grid_min[1], grid_max[1], N)\n zs = np.linspace(grid_min[2], grid_max[2], N)\n\n xx, yy, zz = np.meshgrid(xs, ys, zs, indexing=\"ij\")\n points = torch.tensor(np.vstack([xx.ravel(), yy.ravel(), zz.ravel()]).T, dtype=torch.float).to(device=device)\n\n def evaluate(points):\n z = []\n for _, pnts in enumerate(torch.split(points, chunk_size, dim=0)):\n z.append(occupancy_fn(pnts.contiguous()).contiguous())\n z = torch.cat(z, axis=0)\n return z\n\n z = evaluate(points).detach().cpu().numpy()\n\n if not (np.min(z) > level or np.max(z) < level):\n verts, faces, normals, _ = measure.marching_cubes(\n volume=z.reshape(resolution, resolution, resolution),\n level=level,\n spacing=(\n (grid_max[0] - grid_min[0]) / (N - 1),\n (grid_max[1] - grid_min[1]) / (N - 1),\n (grid_max[2] - grid_min[2]) / (N - 1),\n ),\n )\n verts = verts + np.array(grid_min)\n output_path.parent.mkdir(parents=True, exist_ok=True)\n meshexport = trimesh.Trimesh(verts, faces, normals)\n meshexport.export(str(output_path))\n else:\n print(\"=================================================no surface skip\")\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/mask_utils.py",
"content": "\nfrom typing import Dict, Any\n\nimport torch\nfrom torchtyping import TensorType\n\n\ndef alpha_composite(rgb0: torch.Tensor, rgb1: torch.Tensor,\n alpha0: torch.Tensor, alpha1: torch.Tensor):\n \"\"\"Alpha composite two RGB images.\"\"\"\n alpha = alpha0 + alpha1 * (1 - alpha0)\n rgb = (rgb0 * alpha0 + rgb1 * alpha1 * (1 - alpha0)) / alpha\n rgba = torch.cat([rgb, alpha], -1)\n return rgba\n\n\ndef overlay_mask_on_image(mask_rgb: TensorType[\"h\", \"w\", 3],\n img_rgb: TensorType[\"h\", \"w\", 3],\n mask_prob: TensorType[\"h\", \"w\", 1]):\n mask_alpha = mask_prob\n img_alpha = torch.ones_like(mask_alpha)\n return alpha_composite(mask_rgb, img_rgb, mask_alpha, img_alpha)\n\n\ndef compute_mask_indices(collated_batch: Dict[str, Any]) -> Dict[str, Any]:\n \"\"\"precompute zero and non-zero mask indices for fast sampling from masked & non-masked regions\"\"\"\n if \"mask\" not in collated_batch:\n return collated_batch\n if not isinstance(collated_batch[\"mask\"], torch.Tensor):\n raise NotImplementedError # e.g., handling ``BasicImages``\n \n mask = collated_batch[\"mask\"][..., 0] # (n_frames, h, w)\n assert mask.ndim == 3\n nonzero_indices = torch.nonzero(mask, as_tuple=False).to(torch.int16)\n zero_indices = torch.nonzero(~mask, as_tuple=False).to(torch.int16)\n collated_batch.update({\n \"mask_nonzero_indices\": nonzero_indices,\n \"mask_zero_indices\": zero_indices\n })\n return collated_batch\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/math.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\" Math Helper Functions \"\"\"\n\nfrom dataclasses import dataclass\n\nimport torch\nfrom torchtyping import TensorType\n\n\ndef components_from_spherical_harmonics(levels: int, directions: TensorType[..., 3]) -> TensorType[..., \"components\"]:\n \"\"\"\n Returns value for each component of spherical harmonics.\n\n Args:\n levels: Number of spherical harmonic levels to compute.\n directions: Spherical hamonic coefficients\n \"\"\"\n num_components = levels**2\n components = torch.zeros((*directions.shape[:-1], num_components), device=directions.device)\n\n assert 1 <= levels <= 5, f\"SH levels must be in [1,4], got {levels}\"\n assert directions.shape[-1] == 3, f\"Direction input should have three dimensions. Got {directions.shape[-1]}\"\n\n x = directions[..., 0]\n y = directions[..., 1]\n z = directions[..., 2]\n\n xx = x**2\n yy = y**2\n zz = z**2\n\n # l0\n components[..., 0] = 0.28209479177387814\n\n # l1\n if levels > 1:\n components[..., 1] = 0.4886025119029199 * y\n components[..., 2] = 0.4886025119029199 * z\n components[..., 3] = 0.4886025119029199 * x\n\n # l2\n if levels > 2:\n components[..., 4] = 1.0925484305920792 * x * y\n components[..., 5] = 1.0925484305920792 * y * z\n components[..., 6] = 0.9461746957575601 * zz - 0.31539156525251999\n components[..., 7] = 1.0925484305920792 * x * z\n components[..., 8] = 0.5462742152960396 * (xx - yy)\n\n # l3\n if levels > 3:\n components[..., 9] = 0.5900435899266435 * y * (3 * xx - yy)\n components[..., 10] = 2.890611442640554 * x * y * z\n components[..., 11] = 0.4570457994644658 * y * (5 * zz - 1)\n components[..., 12] = 0.3731763325901154 * z * (5 * zz - 3)\n components[..., 13] = 0.4570457994644658 * x * (5 * zz - 1)\n components[..., 14] = 1.445305721320277 * z * (xx - yy)\n components[..., 15] = 0.5900435899266435 * x * (xx - 3 * yy)\n\n # l4\n if levels > 4:\n components[..., 16] = 2.5033429417967046 * x * y * (xx - yy)\n components[..., 17] = 1.7701307697799304 * y * z * (3 * xx - yy)\n components[..., 18] = 0.9461746957575601 * x * y * (7 * zz - 1)\n components[..., 19] = 0.6690465435572892 * y * (7 * zz - 3)\n components[..., 20] = 0.10578554691520431 * (35 * zz * zz - 30 * zz + 3)\n components[..., 21] = 0.6690465435572892 * x * z * (7 * zz - 3)\n components[..., 22] = 0.47308734787878004 * (xx - yy) * (7 * zz - 1)\n components[..., 23] = 1.7701307697799304 * x * z * (xx - 3 * yy)\n components[..., 24] = 0.4425326924449826 * (xx * (xx - 3 * yy) - yy * (3 * xx - yy))\n\n return components\n\n\n@dataclass\nclass Gaussians:\n \"\"\"Stores Gaussians\n\n Args:\n mean: Mean of multivariate Gaussian\n cov: Covariance of multivariate Gaussian.\n \"\"\"\n\n mean: TensorType[..., \"dim\"]\n cov: TensorType[..., \"dim\", \"dim\"]\n\n\ndef compute_3d_gaussian(\n directions: TensorType[..., 3],\n means: TensorType[..., 3],\n dir_variance: TensorType[..., 1],\n radius_variance: TensorType[..., 1],\n) -> Gaussians:\n \"\"\"Compute guassian along ray.\n\n Args:\n directions: Axis of Gaussian.\n means: Mean of Gaussian.\n dir_variance: Variance along direction axis.\n radius_variance: Variance tangent to direction axis.\n\n Returns:\n Gaussians: Oriented 3D gaussian.\n \"\"\"\n\n dir_outer_product = directions[..., :, None] * directions[..., None, :]\n eye = torch.eye(directions.shape[-1], device=directions.device)\n dir_mag_sq = torch.clamp(torch.sum(directions**2, dim=-1, keepdim=True), min=1e-10)\n null_outer_product = eye - directions[..., :, None] * (directions / dir_mag_sq)[..., None, :]\n dir_cov_diag = dir_variance[..., None] * dir_outer_product[..., :, :]\n radius_cov_diag = radius_variance[..., None] * null_outer_product[..., :, :]\n cov = dir_cov_diag + radius_cov_diag\n return Gaussians(mean=means, cov=cov)\n\n\ndef cylinder_to_gaussian(\n origins: TensorType[..., 3],\n directions: TensorType[..., 3],\n starts: TensorType[..., 1],\n ends: TensorType[..., 1],\n radius: TensorType[..., 1],\n) -> Gaussians:\n \"\"\"Approximates cylinders with a Gaussian distributions.\n\n Args:\n origins: Origins of cylinders.\n directions: Direction (axis) of cylinders.\n starts: Start of cylinders.\n ends: End of cylinders.\n radius: Radii of cylinders.\n\n Returns:\n Gaussians: Approximation of cylinders\n \"\"\"\n means = origins + directions * ((starts + ends) / 2.0)\n dir_variance = (ends - starts) ** 2 / 12\n radius_variance = radius**2 / 4.0\n return compute_3d_gaussian(directions, means, dir_variance, radius_variance)\n\n\ndef conical_frustum_to_gaussian(\n origins: TensorType[..., 3],\n directions: TensorType[..., 3],\n starts: TensorType[..., 1],\n ends: TensorType[..., 1],\n radius: TensorType[..., 1],\n) -> Gaussians:\n \"\"\"Approximates conical frustums with a Gaussian distributions.\n\n Uses stable parameterization described in mip-NeRF publication.\n\n Args:\n origins: Origins of cones.\n directions: Direction (axis) of frustums.\n starts: Start of conical frustums.\n ends: End of conical frustums.\n radius: Radii of cone a distance of 1 from the origin.\n\n Returns:\n Gaussians: Approximation of conical frustums\n \"\"\"\n mu = (starts + ends) / 2.0\n hw = (ends - starts) / 2.0\n means = origins + directions * (mu + (2.0 * mu * hw**2.0) / (3.0 * mu**2.0 + hw**2.0))\n dir_variance = (hw**2) / 3 - (4 / 15) * ((hw**4 * (12 * mu**2 - hw**2)) / (3 * mu**2 + hw**2) ** 2)\n radius_variance = radius**2 * ((mu**2) / 4 + (5 / 12) * hw**2 - 4 / 15 * (hw**4) / (3 * mu**2 + hw**2))\n return compute_3d_gaussian(directions, means, dir_variance, radius_variance)\n\n\ndef expected_sin(x_means: torch.Tensor, x_vars: torch.Tensor) -> torch.Tensor:\n \"\"\"Computes the expected value of sin(y) where y ~ N(x_means, x_vars)\n\n Args:\n x_means: Mean values.\n x_vars: Variance of values.\n\n Returns:\n torch.Tensor: The expected value of sin.\n \"\"\"\n\n return torch.exp(-0.5 * x_vars) * torch.sin(x_means)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/misc.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nMiscellaneous helper code.\n\"\"\"\n\nfrom typing import Any, Callable, Dict, List, Optional, Union\n\nimport torch\nfrom nerfstudio.utils.pointclouds import BasicPointClouds\n\n\ndef get_dict_to_torch(stuff: Any, device: Union[torch.device, str] = \"cpu\", exclude: Optional[List[str]] = None):\n \"\"\"Set everything in the dict to the specified torch device.\n\n Args:\n stuff: things to convert to torch\n device: machine to put the \"stuff\" on\n exclude: list of keys to skip over transferring to device\n \"\"\"\n if isinstance(stuff, dict):\n for k, v in stuff.items():\n if exclude and k in exclude:\n stuff[k] = v\n else:\n stuff[k] = get_dict_to_torch(v, device)\n return stuff\n if isinstance(stuff, (torch.Tensor, BasicPointClouds)): # TODO: add BasicImages?\n return stuff.to(device)\n return stuff\n\n\ndef get_dict_to_cpu(stuff: Any):\n \"\"\"Set everything in the dict to CPU.\n\n Args:\n stuff: things to place onto cpu\n \"\"\"\n if isinstance(stuff, dict):\n for k, v in stuff.items():\n stuff[k] = get_dict_to_cpu(v)\n return stuff\n if isinstance(stuff, torch.Tensor):\n return stuff.detach().cpu()\n return stuff\n\n\ndef get_masked_dict(d, mask):\n \"\"\"Return a masked dictionary.\n TODO(ethan): add more asserts/checks so this doesn't have unpredictable behavior.\n\n Args:\n d: dict to process\n mask: mask to apply to values in dictionary\n \"\"\"\n masked_dict = {}\n for key, value in d.items():\n masked_dict[key] = value[mask]\n return masked_dict\n\n\nclass IterableWrapper: # pylint: disable=too-few-public-methods\n \"\"\"A helper that will allow an instance of a class to return multiple kinds of iterables bound\n to different functions of that class.\n\n To use this, take an instance of a class. From that class, pass in the .\n and . to the IterableWrapper constructor. By passing in the instance's\n functions instead of just the class's functions, the self argument should automatically be accounted\n for.\n\n Args:\n new_iter: function that will be called instead as the __iter__() function\n new_next: function that will be called instead as the __next__() function\n length: length of the iterable. If -1, the iterable will be infinite.\n\n\n Attributes:\n new_iter: object's pointer to the function we are calling for __iter__()\n new_next: object's pointer to the function we are calling for __next__()\n length: length of the iterable. If -1, the iterable will be infinite.\n i: current index of the iterable.\n\n \"\"\"\n\n i: int\n\n def __init__(self, new_iter: Callable, new_next: Callable, length: int = -1):\n self.new_iter = new_iter\n self.new_next = new_next\n self.length = length\n\n def __next__(self):\n if self.length != -1 and self.i >= self.length:\n raise StopIteration\n self.i += 1\n return self.new_next()\n\n def __iter__(self):\n self.new_iter()\n self.i = 0\n return self\n\n\ndef scale_dict(dictionary: Dict[Any, Any], coefficients: Dict[str, float]) -> Dict[Any, Any]:\n \"\"\"Scale a dictionary in-place given a coefficients dictionary.\n\n Args:\n dictionary: input dict to be scaled.\n coefficients: scalar dict config for holding coefficients.\n\n Returns:\n Input dict scaled by coefficients.\n \"\"\"\n for key in dictionary:\n if key in coefficients:\n dictionary[key] *= coefficients[key]\n return dictionary\n\n\ndef step_check(step, step_size, run_at_zero=False) -> bool:\n \"\"\"Returns true based on current step and step interval.\"\"\"\n if step_size == 0:\n return False\n return (run_at_zero or step != 0) and step % step_size == 0\n\n\ndef update_avg(prev_avg: float, new_val: float, step: int) -> float:\n \"\"\"helper to calculate the running average\n\n Args:\n prev_avg (float): previous average value\n new_val (float): new value to update the average with\n step (int): current step number\n\n Returns:\n float: new updated average\n \"\"\"\n return (step * prev_avg + new_val) / (step + 1)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/plotly_utils.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nVisualization code for plotly.\nThe function use prefix conventions in the following way:\n - 'get_*' functions (e.g., 'get_camera_frustums')\n return data that can be ploted with plotly\n - 'vis_*' functions (e.g., 'vis_camera_rays')\n return 'go.Figure' objects which are the plots. Go Figure! :')\n\"\"\"\n\nfrom typing import Any, List, Optional, Union\n\nimport numpy as np\nimport plotly.graph_objects as go\nimport torch\nimport torch.nn.functional as F\nfrom plotly import express as ex\nfrom torchtyping import TensorType\n\nfrom nerfstudio.cameras.cameras import Cameras\nfrom nerfstudio.cameras.rays import Frustums, RayBundle\nfrom nerfstudio.utils.math import Gaussians\n\n\ndef color_str(color):\n \"\"\"Plotly color string.\n\n Args:\n color: list [r, g, b] in [0, 1] range\n\n Returns:\n str: plotly-formatted color string\n \"\"\"\n color = list((np.array(color) * 255.0).astype(\"int\"))\n return f\"\"\"rgb({color[0]}, {color[1]}, {color[2]})\"\"\"\n\n\ndef get_line_segments_from_lines(\n lines: TensorType[\"num_rays\", 2, 3],\n color: str = color_str((1, 0, 0)),\n marker_color: str = color_str((1, 0, 0)),\n colors: Optional[List[str]] = None,\n draw_marker: bool = True,\n draw_line: bool = True,\n marker_size: float = 4,\n line_width: float = 10,\n) -> List[Any]:\n \"\"\"Returns a list of Scatter3D objects for creating lines with plotly.\n # TODO(ethan): make this function more efficient instead of having a list of objects.\n\n Args:\n lines: Tensor of lines.\n color: Color of the lines. Defaults to red.\n marker_color: Color of the markers. Defaults to red.\n colors: List of colors for each line. Defaults to None.\n draw_marker: Whether to draw markers. Defaults to True.\n draw_line: Whether to draw lines. Defaults to True.\n marker_size: Size of the markers. Defaults to 4.\n line_width: Width of the lines. Defaults to 10.\n\n Returns:\n Scatter3D object on lines.\n \"\"\"\n data = []\n for idx, line in enumerate(lines):\n thiscolor = color if draw_line else \"rgba(0, 0, 0, 0)\"\n if colors is not None:\n marker_color = colors[idx]\n thiscolor = colors[idx]\n data.append(\n go.Scatter3d(\n x=line[:, 0],\n y=line[:, 1],\n z=line[:, 2],\n showlegend=False,\n marker=dict(\n size=marker_size,\n color=marker_color,\n colorscale=\"Viridis\",\n )\n if draw_marker\n else dict(color=\"rgba(0, 0, 0, 0)\"),\n line=dict(color=thiscolor, width=line_width),\n )\n )\n return data\n\n\ndef vis_dataset(camera_origins: TensorType[\"num_cameras\", 3], ray_bundle: RayBundle) -> go.FigureWidget: # type: ignore\n \"\"\"Visualize a dataset with plotly using our cameras and generated rays.\n\n Args:\n camera_origins: Tensor of camera origins.\n ray_bundle: Ray bundle.\n\n Returns:\n plotly figure.\n \"\"\"\n\n skip = 1\n size = 8\n assert len(ray_bundle) < 500, \"Let's not break plotly by plotting too many rays!\"\n\n data = []\n data += [\n go.Scatter3d(\n x=camera_origins[::skip, 0],\n y=camera_origins[::skip, 1],\n z=camera_origins[::skip, 2],\n mode=\"markers\",\n name=\"camera origins\",\n marker=dict(color=\"rgba(0, 0, 0, 1)\", size=size),\n )\n ]\n\n length = 2.0\n lines = torch.stack(\n [ray_bundle.origins, ray_bundle.origins + ray_bundle.directions * length], dim=1\n ) # (num_rays, 2, 3)\n\n data += get_line_segments_from_lines(lines)\n\n layout = go.Layout(\n autosize=False,\n width=1000,\n height=1000,\n margin=go.layout.Margin(l=50, r=50, b=100, t=100, pad=4), # type: ignore\n scene=go.layout.Scene( # type: ignore\n aspectmode=\"data\",\n camera=dict(up=dict(x=0, y=0, z=1), center=dict(x=0, y=0, z=0), eye=dict(x=1.25, y=1.25, z=1.25)),\n ),\n )\n fig = go.Figure(data=data, layout=layout)\n return fig\n\n\ndef get_random_color(colormap: Optional[List[str]] = None, idx: Optional[int] = None) -> str:\n \"\"\"Get a random color from a colormap.\n\n Args:\n colormap: List of colors. Defaults to Plotly colors.\n idx: Index of color to return. Defaults to None.\n\n Returns:\n random color string\n \"\"\"\n if colormap is None:\n colormap = ex.colors.qualitative.Plotly\n if idx is None:\n return colormap[np.random.randint(0, len(colormap))]\n return colormap[idx % len(colormap)]\n\n\ndef get_sphere(\n radius: float, center: TensorType[3] = None, color: str = \"black\", opacity: float = 1.0, resolution: int = 32\n) -> go.Mesh3d: # type: ignore\n \"\"\"Returns a sphere object for plotting with plotly.\n\n Args:\n radius: radius of sphere.\n center: center of sphere. Defaults to origin.\n color: color of sphere. Defaults to \"black\".\n opacity: opacity of sphere. Defaults to 1.0.\n resolution: resolution of sphere. Defaults to 32.\n\n Returns:\n sphere object.\n \"\"\"\n phi = torch.linspace(0, 2 * torch.pi, resolution)\n theta = torch.linspace(-torch.pi / 2, torch.pi / 2, resolution)\n phi, theta = torch.meshgrid(phi, theta, indexing=\"ij\")\n\n x = torch.cos(theta) * torch.sin(phi)\n y = torch.cos(theta) * torch.cos(phi)\n z = torch.sin(theta)\n pts = torch.stack((x, y, z), dim=-1)\n\n pts *= radius\n if center is not None:\n pts += center\n\n return go.Mesh3d(\n {\n \"x\": pts[:, :, 0].flatten(),\n \"y\": pts[:, :, 1].flatten(),\n \"z\": pts[:, :, 2].flatten(),\n \"alphahull\": 0,\n \"opacity\": opacity,\n \"color\": color,\n }\n )\n\n\ndef get_cube(\n side_length: float,\n center: TensorType[3] = None,\n color: str = \"black\",\n opacity: float = 1.0,\n) -> go.Mesh3d: # type: ignore\n \"\"\"Returns a cube object for plotting with plotly.\n\n Args:\n side_length: side_length of cube.\n center: center of cube.\n color: color of cube.\n opacity: opacity of cube.\n\n Returns:\n cube object.\n \"\"\"\n\n x = np.array([-1.0, -1.0, 1.0, 1.0, -1.0, -1.0, 1.0, 1.0])\n y = np.array([-1.0, 1.0, 1.0, -1.0, -1.0, 1.0, 1.0, -1.0])\n z = np.array([-1.0, -1.0, -1.0, -1.0, 1.0, 1.0, 1.0, 1.0])\n\n pts = np.stack((x, y, z), axis=0)\n\n pts *= side_length / 2.0\n if center is not None:\n pts += center\n\n return go.Mesh3d(\n {\n \"x\": pts[0],\n \"y\": pts[1],\n \"z\": pts[2],\n \"alphahull\": 0,\n \"opacity\": opacity,\n \"color\": color,\n }\n )\n\n\ndef get_gaussian_ellipsiod(\n mean: TensorType[3],\n cov: TensorType[3, 3],\n n_std: int = 2,\n color=\"lightblue\",\n opacity: float = 0.5,\n resolution: int = 20,\n name: str = \"ellipse\",\n) -> go.Mesh3d: # type: ignore\n \"\"\"Get a plotly ellipsoid for a Gaussian.\n\n Args:\n mean: mean of the Gaussian.\n cov: covariance of the Gaussian.\n n_std: Standard devation to visualize. Defaults to 2 (95% confidence).\n color: Color of the ellipsoid. Defaults to None.\n opacity: Opacity of the ellipsoid. Defaults to 0.5.\n resolution: Resolution of the ellipsoid. Defaults to 20.\n name: Name of the ellipsoid. Defaults to \"ellipse\".\n\n Returns:\n ellipsoid object.\n \"\"\"\n\n phi = torch.linspace(0, 2 * torch.pi, resolution)\n theta = torch.linspace(-torch.pi / 2, torch.pi / 2, resolution)\n phi, theta = torch.meshgrid(phi, theta, indexing=\"ij\")\n\n x = torch.cos(theta) * torch.sin(phi)\n y = torch.cos(theta) * torch.cos(phi)\n z = torch.sin(theta)\n pts = torch.stack((x, y, z), dim=-1)\n\n eigenvals, eigenvecs = torch.linalg.eigh(cov)\n idx = torch.sum(cov, dim=0).argsort()\n idx = eigenvals[idx].argsort()\n eigenvals = eigenvals[idx]\n eigenvecs = eigenvecs[:, idx]\n\n scaling = torch.sqrt(eigenvals) * n_std\n pts = pts * scaling\n\n pts = pts @ eigenvecs.t()\n\n pts += mean\n\n return go.Mesh3d(\n {\n \"x\": pts[:, :, 0].flatten(),\n \"y\": pts[:, :, 1].flatten(),\n \"z\": pts[:, :, 2].flatten(),\n \"alphahull\": 0,\n \"opacity\": opacity,\n \"color\": color,\n \"name\": name,\n }\n )\n\n\ndef get_gaussian_ellipsoids_list(\n gaussians: Gaussians, opacity: float = 0.5, color: str = \"random\", resolution: int = 20\n) -> List[Union[go.Mesh3d, go.Scatter3d]]: # type: ignore\n \"\"\"Get a list of plotly meshes for frustums.\n\n Args:\n gaussians (Gaussians): Gaussians to visualize.\n opacity (float, optional): Opacity of the mesh. Defaults to 0.3.\n color (str, optional): Color of the mesh. Defaults to \"random\".\n resolution: Resolution of the mesh. Defaults to 20.\n\n Returns:\n List of plotly meshes\n \"\"\"\n data = []\n\n vis_means = go.Scatter3d(\n x=gaussians.mean[:, 0],\n y=gaussians.mean[:, 1],\n z=gaussians.mean[:, 2],\n mode=\"markers\",\n marker=dict(size=2, color=\"black\"),\n name=\"Means\",\n )\n data.append(vis_means)\n\n for i in range(gaussians.mean.shape[0]):\n if color == \"random\":\n c = get_random_color()\n else:\n c = color\n ellipse = get_gaussian_ellipsiod(\n gaussians.mean[i],\n cov=gaussians.cov[i],\n color=c,\n opacity=opacity,\n resolution=resolution,\n )\n data.append(ellipse)\n\n return data\n\n\ndef get_frustum_mesh(\n frustum: Frustums, opacity: float = 0.3, color: str = \"#DC203C\", resolution: int = 20\n) -> go.Mesh3d: # type: ignore\n \"\"\"Get a plotly mesh for a single frustum.\n\n Args:\n frustum: Single frustum\n opacity: Opacity of the mesh. Defaults to 0.3.\n color: Color of the mesh. Defaults to \"#DC203C\".\n resolution: Resolution of the mesh. Defaults to 20.\n\n Returns:\n Plotly mesh\n \"\"\"\n\n if frustum.ndim > 1:\n raise ValueError(\"Frustum must be a single Frustum object.\")\n\n base_radius = torch.sqrt(frustum.pixel_area / torch.pi)\n f_radius = frustum.starts * base_radius\n b_radius = frustum.ends * base_radius\n\n x = torch.cat([torch.ones(resolution) * frustum.starts, torch.ones(resolution) * frustum.ends])\n pts = torch.linspace(0, 2 * torch.pi, resolution)\n\n y = torch.sin(pts)\n z = torch.cos(pts)\n y = torch.cat([y * f_radius, y * b_radius])\n z = torch.cat([z * f_radius, z * b_radius])\n\n pts = torch.stack([x, y, z], dim=-1)\n\n forward = frustum.directions\n up = F.normalize(torch.cross(forward, torch.tensor([0.0, 0, 1])), dim=-1) # type: ignore\n right = F.normalize(torch.cross(forward, up), dim=-1) # type: ignore\n rotation = torch.stack([forward, up, right], dim=1)\n pts = torch.einsum(\"kj,ij->ki\", pts, rotation)\n\n pts += frustum.origins\n return go.Mesh3d(\n x=pts[..., 0],\n y=pts[..., 1],\n z=pts[..., 2],\n opacity=opacity,\n alphahull=0,\n color=color,\n flatshading=True,\n name=\"Frustums\",\n )\n\n\ndef get_frustums_mesh_list(\n frustums: Frustums, opacity: float = 1.0, color: str = \"random\", resolution: int = 20\n) -> List[go.Mesh3d]: # type: ignore\n \"\"\"Get a list of plotly meshes for a list of frustums.\n\n Args:\n frustums: List of frustums\n opacity: Opacity of the mesh. Defaults to 0.3.\n color: Color of the mesh. Defaults to \"random\".\n resolution: Resolution of the mesh. Defaults to 20.\n\n Returns:\n List of plotly meshes\n \"\"\"\n data = []\n for i, frustum in enumerate(frustums.flatten()): # type: ignore\n if color == \"random\":\n c = get_random_color(idx=i)\n else:\n c = color\n data.append(get_frustum_mesh(frustum, opacity=opacity, color=c, resolution=resolution))\n return data\n\n\ndef get_frustum_points(\n frustum: Frustums, opacity: float = 1.0, color: str = \"forestgreen\", size: float = 5\n) -> go.Scatter3d: # type: ignore\n \"\"\"Get a set plotly points for frustums centers.\n\n Args:\n frustum: Frustums to visualize.\n opacity: Opacity of the points. Defaults to 0.3.\n color: Color of the poinst. Defaults to \"forestgreen\".\n size: Size of points. Defaults to 10.\n\n Returns:\n Plotly points\n \"\"\"\n\n frustum = frustum.flatten()\n pts = frustum.get_positions()\n\n return go.Scatter3d(\n x=pts[..., 0],\n y=pts[..., 1],\n z=pts[..., 2],\n mode=\"markers\",\n marker=dict(\n size=size,\n color=color,\n opacity=opacity,\n ),\n name=\"Frustums -> Positions\",\n )\n\n\ndef get_ray_bundle_lines(\n ray_bundle: RayBundle, length: float = 1.0, color: str = \"#DC203C\", width: float = 1\n) -> go.Scatter3d: # type: ignore\n \"\"\"Get a plotly line for a ray bundle.\n\n Args:\n ray_bundle: Ray bundle\n length: Length of the line. Defaults to 1.0.\n color: Color of the line.\n width: Width of the line. Defaults to 1.\n\n Returns:\n Plotly lines\n \"\"\"\n\n origins = ray_bundle.origins.view(-1, 3)\n directions = ray_bundle.directions.view(-1, 3)\n\n lines = torch.empty((origins.shape[0] * 2, 3))\n lines[0::2] = origins\n lines[1::2] = origins + directions * length\n return go.Scatter3d(\n x=lines[..., 0],\n y=lines[..., 1],\n z=lines[..., 2],\n mode=\"lines\",\n name=\"Ray Bundle\",\n line=dict(color=color, width=width),\n )\n\n\ndef vis_camera_rays(cameras: Cameras) -> go.Figure: # type: ignore\n \"\"\"Visualize camera rays.\n\n Args:\n camera: Camera to visualize.\n\n Returns:\n Plotly lines\n \"\"\"\n\n coords = cameras.get_image_coords()\n coords[..., 0] /= cameras.image_height[0] # All the cameras have the same image height for now\n coords[..., 1] /= cameras.image_width[0] # All the cameras have the same image width for now\n coords = torch.cat([coords, torch.ones((*coords.shape[:-1], 1))], dim=-1)\n\n ray_bundle = cameras.generate_rays(camera_indices=0)\n\n origins = ray_bundle.origins.view(-1, 3)\n directions = ray_bundle.directions.view(-1, 3)\n coords = coords.view(-1, 3)\n\n lines = torch.empty((origins.shape[0] * 2, 3))\n lines[0::2] = origins\n lines[1::2] = origins + directions\n\n colors = torch.empty((coords.shape[0] * 2, 3))\n colors[0::2] = coords\n colors[1::2] = coords\n\n fig = go.Figure(\n data=go.Scatter3d(\n x=lines[:, 0],\n y=lines[:, 2],\n z=lines[:, 1],\n marker=dict(\n size=4,\n color=colors,\n ),\n line=dict(color=\"lightblue\", width=1),\n )\n )\n fig.update_layout(\n scene=dict(\n xaxis=dict(title=\"x\", showspikes=False),\n yaxis=dict(title=\"z\", showspikes=False),\n zaxis=dict(title=\"y\", showspikes=False),\n ),\n margin=dict(r=0, b=10, l=0, t=10),\n hovermode=False,\n )\n\n return fig\n\n\ndef get_camera_frustums(cameras: Cameras):\n \"\"\"Returns the camera frustums for the cameras that we are using.\n\n Args:\n cameras: The cameras that we want to plot.\n\n Returns:\n A plotly scatter that can be plotted.\n \"\"\"\n for camera_idx in range(cameras.size):\n json_ = cameras.to_json(camera_idx=camera_idx)\n print(json_)\n raise NotImplementedError\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/pointclouds.py",
"content": "from copy import deepcopy\nfrom dataclasses import dataclass, fields\nfrom typing import Optional\n\nimport torch\nfrom rich.console import Console\n\ntry:\n import faiss\n import faiss.contrib.torch_utils\nexcept ImportError as err:\n pass\nfrom torch import nn\nfrom torch.nn import functional as F\n\nCONSOLE = Console(width=120)\n\n\nclass BasicPointClouds(nn.Module):\n \"\"\"A primitive struct for holding pointcloud related data (e.g., points, colors, features, etc.)\n\n The purpose of this is to have a special struct wrapping around a list so that the\n nerfstudio_collate fn and other parts of the code recognise this as a struct to leave alone\n instead of reshaping or concatenating into a single tensor (since this will likely be used\n for cases where we have images of different sizes and shapes).\n\n We assume that all images share the same pointcloud, thus only the first element of\n the list of pointclouds will be kept by the nerfstudio_collate.\n \"\"\"\n\n # TOOD: make {fg,bg,plane}_pts non-optional\n\n def __init__(\n self,\n fg_pts: Optional[torch.Tensor] = None,\n bg_pts: Optional[torch.Tensor] = None,\n plane_pts: Optional[torch.Tensor] = None,\n fg_feats: Optional[torch.Tensor] = None,\n bg_feats: Optional[torch.Tensor] = None,\n plane_feats: Optional[torch.Tensor] = None,\n # data for ptcd regularization\n udf_plane_to_fg: Optional[torch.Tensor] = None,\n plane_sdf_deviation_thr_ptwise: Optional[torch.Tensor] = None,\n fg_sdf_deviation_thr_global: Optional[torch.Tensor] = None,\n fg_sdf_deviation_thr_ptwise: Optional[torch.Tensor] = None,\n # data for feature-field based segmentation\n fg_nn_index: Optional[faiss.GpuIndexFlatL2] = None,\n bg_nn_index: Optional[faiss.GpuIndexFlatL2] = None,\n plane_nn_index: Optional[faiss.GpuIndexFlatL2] = None,\n ):\n super().__init__()\n self.fg_pts = fg_pts\n self.bg_pts = bg_pts\n self.plane_pts = plane_pts\n self.fg_feats = fg_feats\n self.bg_feats = bg_feats\n self.plane_feats = plane_feats\n\n self.udf_plane_to_fg = udf_plane_to_fg\n self.fg_sdf_deviation_thr_global = fg_sdf_deviation_thr_global\n self.fg_sdf_deviation_thr_ptwise = fg_sdf_deviation_thr_ptwise\n self.plane_sdf_deviation_thr_ptwise = plane_sdf_deviation_thr_ptwise\n\n self.fg_nn_index = fg_nn_index\n self.bg_nn_index = bg_nn_index\n self.plane_nn_index = plane_nn_index\n\n self._register_all_tensors_as_buffers()\n\n def build_nn_search_index(self, separate_bg_plane: bool = False):\n \"\"\"Build faiss gpu index for fast nn search\"\"\"\n for f in [\"fg_feats\", \"bg_feats\", \"plane_feats\"]:\n assert getattr(self, f, None) is not None, f\"BasicPointClouds.{f} is None\"\n nc_feat = self.fg_feats.shape[-1]\n\n res = faiss.StandardGpuResources()\n cfg = faiss.GpuIndexFlatConfig()\n cfg.useFloat16 = True\n\n # build fg inex\n self.fg_nn_index = faiss.GpuIndexFlatL2(res, nc_feat, cfg)\n self.fg_nn_index.add(F.normalize(self.fg_feats, p=2, dim=-1))\n CONSOLE.print(f\"Using GpuIndexFlatL2 for fg NN search. (#fg: {len(self.fg_feats)})\")\n\n # build bg index\n if not separate_bg_plane:\n self.bg_nn_index = faiss.GpuIndexFlatL2(res, nc_feat, cfg)\n self.bg_nn_index.add(F.normalize(torch.cat([self.bg_feats, self.plane_feats], 0), p=2, dim=-1))\n CONSOLE.print(\n f\"Using GpuIndexFlatL2 for (bg, plane) NN search.\"\n f\"(#bg: {len(self.bg_feats)} | #plane: {len(self.plane_pts)}))\"\n )\n # TODO: build a mapper to distinguish bg and plane pts\n # TODO: tune the number of fg, bg and plane pts\n else:\n self.bg_nn_index = faiss.GpuIndexFlatL2(res, nc_feat, cfg)\n self.bg_nn_index.add(F.normalize(self.bg_feats, p=2, dim=-1))\n CONSOLE.print(f\"Using GpuIndexFlatL2 for bg NN search. (#bg: {len(self.bg_feats)})\")\n self.plane_nn_index = faiss.GpuIndexFlatL2(res, nc_feat, cfg)\n self.plane_nn_index.add(F.normalize(self.plane_feats, p=2, dim=-1))\n CONSOLE.print(f\"Using GpuIndexFlatL2 for bg NN search. (#bg: {len(self.plane_feats)})\")\n\n def sample_fg_points(self, n_pts: int, replace: bool = True):\n if len(self.fg_pts) <= n_pts:\n return self.fg_pts, self.fg_sdf_deviation_thr_ptwise\n\n rand_inds = (\n torch.randint(0, len(self.fg_pts), (n_pts,), device=self.device)\n if replace\n else torch.randperm(len(self.fg_pts), device=self.device)[:n_pts]\n )\n fg_pts = self.fg_pts[rand_inds]\n fg_sdf_thrs = self.fg_sdf_deviation_thr_ptwise[rand_inds]\n return fg_pts, fg_sdf_thrs\n\n def sample_plane_points(self, n_pts: int, replace: bool = True, rand_sample_half: bool = False):\n if rand_sample_half:\n raise NotImplementedError\n\n if len(self.plane_pts) <= n_pts:\n return self.plane_pts, self.udf_plane_to_fg, self.plane_sdf_deviation_thr_ptwise\n\n rand_inds = (\n torch.randint(0, len(self.plane_pts), (n_pts,), device=self.device)\n if replace\n else torch.randperm(len(self.plane_pts), device=self.device)[:n_pts]\n )\n plane_pts = self.plane_pts[rand_inds]\n udf_plane_to_fg = self.udf_plane_to_fg[rand_inds]\n plane_sdf_thrs = self.plane_sdf_deviation_thr_ptwise[rand_inds]\n\n return plane_pts, udf_plane_to_fg, plane_sdf_thrs\n\n @property\n def device(self):\n return self.fg_pts.device\n\n def _register_all_tensors_as_buffers(self):\n \"\"\"register all Tensor attributes as buffers to automatically move them to the target device\"\"\"\n _attr_names = list(vars(self).keys())\n for _name in _attr_names:\n _attr = getattr(self, _name)\n if not isinstance(_attr, torch.Tensor):\n continue\n\n delattr(self, _name)\n self.register_buffer(_name, _attr, persistent=False)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/poses.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nCommon 3D pose methods\n\"\"\"\n\nimport torch\nfrom torchtyping import TensorType\n\n\ndef to4x4(pose: TensorType[..., 3, 4]) -> TensorType[..., 4, 4]:\n \"\"\"Convert 3x4 pose matrices to a 4x4 with the addition of a homogeneous coordinate.\n\n Args:\n pose: Camera pose without homogenous coordinate.\n\n Returns:\n Camera poses with additional homogenous coordinate added.\n \"\"\"\n constants = torch.zeros_like(pose[..., :1, :], device=pose.device)\n constants[..., :, 3] = 1\n return torch.cat([pose, constants], dim=-2)\n\n\ndef inverse(pose: TensorType[..., 3, 4]) -> TensorType[..., 3, 4]:\n \"\"\"Invert provided pose matrix.\n\n Args:\n pose: Camera pose without homogenous coordinate.\n\n Returns:\n Inverse of pose.\n \"\"\"\n R = pose[..., :3, :3]\n t = pose[..., :3, 3:]\n R_inverse = R.transpose(-2, -1) # pylint: disable=invalid-name\n t_inverse = -R_inverse.matmul(t)\n return torch.cat([R_inverse, t_inverse], dim=-1)\n\n\ndef multiply(pose_a: TensorType[..., 3, 4], pose_b: TensorType[..., 3, 4]) -> TensorType[..., 3, 4]:\n \"\"\"Multiply two pose matrices, A @ B.\n\n Args:\n pose_a: Left pose matrix, usually a transformation applied to the right.\n pose_b: Right pose matrix, usually a camera pose that will be tranformed by pose_a.\n\n Returns:\n Camera pose matrix where pose_a was applied to pose_b.\n \"\"\"\n R1, t1 = pose_a[..., :3, :3], pose_a[..., :3, 3:]\n R2, t2 = pose_b[..., :3, :3], pose_b[..., :3, 3:]\n R = R1.matmul(R2)\n t = t1 + R1.matmul(t2)\n return torch.cat([R, t], dim=-1)\n\n\ndef normalize(poses: TensorType[..., 3, 4]) -> TensorType[..., 3, 4]:\n \"\"\"Normalize the XYZs of poses to fit within a unit cube ([-1, 1]). Note: This operation is not in-place.\n\n Args:\n poses: A collection of poses to be normalized.\n\n Returns;\n Normalized collection of poses.\n \"\"\"\n pose_copy = torch.clone(poses)\n pose_copy[..., :3, 3] /= torch.max(torch.abs(poses[..., :3, 3]))\n\n return pose_copy\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/printing.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"A collection of common strings and print statements used throughout the codebase.\"\"\"\n\nfrom math import floor, log\n\nfrom rich.console import Console\n\nCONSOLE = Console(width=120)\n\n\ndef print_tcnn_speed_warning(method_name: str):\n \"\"\"Prints a warning about the speed of the TCNN.\"\"\"\n CONSOLE.line()\n CONSOLE.print(f\"[bold yellow]WARNING: Using a slow implementation of {method_name}. \")\n CONSOLE.print(\n \"[bold yellow]:person_running: :person_running: \"\n + \"Install tcnn for speedups :person_running: :person_running:\"\n )\n CONSOLE.print(\"[yellow]pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch\")\n CONSOLE.line()\n\n\ndef human_format(num):\n \"\"\"Format a number in a more human readable way\n\n Args:\n num: number to format\n \"\"\"\n units = [\"\", \"K\", \"M\", \"B\", \"T\", \"P\"]\n k = 1000.0\n magnitude = int(floor(log(num, k)))\n return f\"{(num / k**magnitude):.2f} {units[magnitude]}\"\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/profiler.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nProfiler base class and functionality\n\"\"\"\nfrom __future__ import annotations\n\nimport time\nfrom typing import Callable\n\nfrom rich.console import Console\n\nimport torch\nfrom nerfstudio.configs import base_config as cfg\nfrom nerfstudio.utils import comms\nfrom nerfstudio.utils.decorators import (\n check_main_thread,\n check_profiler_enabled,\n decorate_all,\n)\n\nCONSOLE = Console(width=120)\n\nPROFILER = []\n\n\ndef time_function(func: Callable) -> Callable:\n \"\"\"Decorator: time a function call\"\"\"\n def wrapper(*args, **kwargs):\n if PROFILER:\n # torch.cuda.synchronize()\n start = time.time()\n \n ret = func(*args, **kwargs)\n \n if PROFILER:\n # torch.cuda.synchronize()\n class_str = func.__qualname__\n PROFILER[0].update_time(class_str, start, time.time())\n return ret\n\n return wrapper\n\n\ndef time_function_cuda(func: Callable) -> Callable:\n \"\"\"Decorator: time a function call\"\"\"\n def wrapper(*args, **kwargs):\n if PROFILER:\n torch.cuda.synchronize()\n start = time.time()\n \n ret = func(*args, **kwargs)\n \n if PROFILER:\n torch.cuda.synchronize()\n class_str = func.__qualname__\n PROFILER[0].update_time(class_str, start, time.time())\n return ret\n\n return wrapper\n\n\ndef flush_profiler(config: cfg.LoggingConfig):\n \"\"\"Method that checks if profiler is enabled before flushing\"\"\"\n if config.enable_profiler and PROFILER:\n PROFILER[0].print_profile()\n\n\ndef setup_profiler(config: cfg.LoggingConfig):\n \"\"\"Initialization of profilers\"\"\"\n if comms.is_main_process():\n PROFILER.append(Profiler(config))\n\n\n@decorate_all([check_profiler_enabled, check_main_thread])\nclass Profiler:\n \"\"\"Profiler class\"\"\"\n\n def __init__(self, config: cfg.LoggingConfig):\n self.config = config\n self.profiler_dict = {}\n\n def update_time(self, func_name: str, start_time: float, end_time: float):\n \"\"\"update the profiler dictionary with running averages of durations\n\n Args:\n func_name: the function name that is being profiled\n start_time: the start time when function is called\n end_time: the end time when function terminated\n \"\"\"\n val = end_time - start_time\n func_dict = self.profiler_dict.get(func_name, {\"val\": 0, \"step\": 0})\n prev_val = func_dict[\"val\"]\n prev_step = func_dict[\"step\"]\n self.profiler_dict[func_name] = {\"val\": (prev_val * prev_step + val) / (prev_step + 1), \"step\": prev_step + 1}\n\n def print_profile(self):\n \"\"\"helper to print out the profiler stats\"\"\"\n CONSOLE.print(\"Printing profiling stats, from longest to shortest duration in seconds\")\n sorted_keys = sorted(\n self.profiler_dict.keys(),\n key=lambda k: self.profiler_dict[k][\"val\"],\n reverse=True,\n )\n for k in sorted_keys:\n val = f\"{self.profiler_dict[k]['val']:0.4f}\"\n CONSOLE.print(f\"{k:<20}: {val:<20}\")\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/rich_utils.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Additional rich ui components\"\"\"\n\nfrom contextlib import nullcontext\nfrom typing import Optional\n\nfrom rich.console import Console\nfrom rich.progress import (\n BarColumn,\n Progress,\n ProgressColumn,\n TaskProgressColumn,\n TextColumn,\n TimeRemainingColumn,\n)\nfrom rich.text import Text\n\nCONSOLE = Console(width=120)\n\n\nclass ItersPerSecColumn(ProgressColumn):\n \"\"\"Renders the iterations per second for a progress bar.\"\"\"\n\n def __init__(self, suffix=\"it/s\") -> None:\n super().__init__()\n self.suffix = suffix\n\n def render(self, task: \"Task\") -> Text:\n \"\"\"Show data transfer speed.\"\"\"\n speed = task.finished_speed or task.speed\n if speed is None:\n return Text(\"?\", style=\"progress.data.speed\")\n return Text(f\"{speed:.2f} {self.suffix}\", style=\"progress.data.speed\")\n\n\ndef status(msg: str, spinner: str = \"bouncingBall\", verbose: bool = False):\n \"\"\"A context manager that does nothing is verbose is True. Otherwise it hides logs under a message.\n\n Args:\n msg: The message to log.\n spinner: The spinner to use.\n verbose: If True, print all logs, else hide them.\n \"\"\"\n if verbose:\n return nullcontext()\n return CONSOLE.status(msg, spinner=spinner)\n\n\ndef get_progress(description: str, suffix: Optional[str] = None):\n \"\"\"Helper function to return a rich Progress object.\"\"\"\n progress_list = [TextColumn(description), BarColumn(), TaskProgressColumn(show_speed=True)]\n progress_list += [ItersPerSecColumn(suffix=suffix)] if suffix else []\n progress_list += [TimeRemainingColumn(elapsed_when_finished=True, compact=True)]\n progress = Progress(*progress_list)\n return progress\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/scheduler.py",
"content": "\"\"\"Scheduler for arbitrary weights\"\"\"\nimport math\nimport numpy as np\nfrom functools import lru_cache\n\n\ndef cosine_annealing(cur_step, max_step=None, min_val=0.0, max_val=1.0):\n cur_val = min_val + (max_val - min_val) \\\n * (1 + math.cos(cur_step * math.pi / max_step)) / 2\n return cur_val\n\n\n@lru_cache(maxsize=128)\ndef _compute_exp_anneal_gamma(max_step, min_val, max_val):\n return np.power(min_val / max_val, 1 / max_step)\n\n\ndef exp_annealing(cur_step, max_step=None, min_val=0.0, max_val=1.0):\n gamma = _compute_exp_anneal_gamma(max_step, min_val, max_val)\n return max_val * np.power(gamma, cur_step)\n\n\ndef constant(cur_step, max_step=None, min_val=0.0, max_val=1.0):\n return max_val\n\n\n# TODO: sine scheduler\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/scripts.py",
"content": "\"\"\"Helpers for running script commands.\"\"\"\n# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport subprocess\nimport sys\nfrom typing import Optional\n\nfrom rich.console import Console\n\nCONSOLE = Console(width=120)\n\n\ndef run_command(cmd: str, verbose=False) -> Optional[str]:\n \"\"\"Runs a command and returns the output.\n\n Args:\n cmd: Command to run.\n verbose: If True, logs the output of the command.\n Returns:\n The output of the command if return_output is True, otherwise None.\n \"\"\"\n out = subprocess.run(cmd, capture_output=not verbose, shell=True, check=False)\n if out.returncode != 0:\n CONSOLE.rule(\"[bold red] :skull: :skull: :skull: ERROR :skull: :skull: :skull: \", style=\"red\")\n CONSOLE.print(f\"[bold red]Error running command: {cmd}\")\n CONSOLE.rule(style=\"red\")\n CONSOLE.print(out.stderr.decode(\"utf-8\"))\n sys.exit(1)\n if out.stdout is not None:\n return out.stdout.decode(\"utf-8\")\n return out\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/tensor_dataclass.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tensor dataclass\"\"\"\n\nimport dataclasses\nfrom copy import deepcopy\nfrom typing import Callable, Dict, List, NoReturn, Optional, Tuple, TypeVar, Union\n\nimport numpy as np\nimport torch\n\nTensorDataclassT = TypeVar(\"TensorDataclassT\", bound=\"TensorDataclass\")\n\n\nclass TensorDataclass:\n \"\"\"@dataclass of tensors with the same size batch. Allows indexing and standard tensor ops.\n Fields that are not Tensors will not be batched unless they are also a TensorDataclass.\n Any fields that are dictionaries will have their Tensors or TensorDataclasses batched, and\n dictionaries will have their tensors or TensorDataclasses considered in the initial broadcast.\n Tensor fields must have at least 1 dimension, meaning that you must convert a field like torch.Tensor(1)\n to torch.Tensor([1])\n\n Example:\n\n .. code-block:: python\n\n @dataclass\n class TestTensorDataclass(TensorDataclass):\n a: torch.Tensor\n b: torch.Tensor\n c: torch.Tensor = None\n\n # Create a new tensor dataclass with batch size of [2,3,4]\n test = TestTensorDataclass(a=torch.ones((2, 3, 4, 2)), b=torch.ones((4, 3)))\n\n test.shape # [2, 3, 4]\n test.a.shape # [2, 3, 4, 2]\n test.b.shape # [2, 3, 4, 3]\n\n test.reshape((6,4)).shape # [6, 4]\n test.flatten().shape # [24,]\n\n test[..., 0].shape # [2, 3]\n test[:, 0, :].shape # [2, 4]\n \"\"\"\n\n _shape: tuple\n\n # A mapping from field-name (str): n (int)\n # Any field OR any key in a dictionary field with this name (field-name) and a corresponding\n # torch.Tensor will be assumed to have n dimensions after the batch dims. These n final dimensions\n # will remain the same shape when doing reshapes, broadcasting, etc on the tensordataclass\n _field_custom_dimensions: Dict[str, int] = {}\n\n def __post_init__(self) -> None:\n \"\"\"Finishes setting up the TensorDataclass\n\n This will 1) find the broadcasted shape and 2) broadcast all fields to this shape 3)\n set _shape to be the broadcasted shape.\n \"\"\"\n if self._field_custom_dimensions is not None:\n for k, v in self._field_custom_dimensions.items():\n assert (\n isinstance(v, int) and v > 1\n ), f\"Custom dimensions must be an integer greater than 1, since 1 is the default, received {k}: {v}\"\n\n if not dataclasses.is_dataclass(self):\n raise TypeError(\"TensorDataclass must be a dataclass\")\n\n batch_shapes = self._get_dict_batch_shapes(\n {f.name: self.__getattribute__(f.name) for f in dataclasses.fields(self)}\n )\n if len(batch_shapes) == 0:\n raise ValueError(\"TensorDataclass must have at least one tensor\")\n batch_shape = torch.broadcast_shapes(*batch_shapes)\n\n broadcasted_fields = self._broadcast_dict_fields(\n {f.name: self.__getattribute__(f.name) for f in dataclasses.fields(self)}, batch_shape\n )\n for f, v in broadcasted_fields.items():\n self.__setattr__(f, v)\n\n self.__setattr__(\"_shape\", batch_shape)\n\n def _get_dict_batch_shapes(self, dict_: Dict) -> List:\n \"\"\"Returns batch shapes of all tensors in a dictionary\n\n Args:\n dict_: The dictionary to get the batch shapes of.\n\n Returns:\n The batch shapes of all tensors in the dictionary.\n \"\"\"\n batch_shapes = []\n for k, v in dict_.items():\n if isinstance(v, torch.Tensor):\n if isinstance(self._field_custom_dimensions, dict) and k in self._field_custom_dimensions:\n # pylint: disable=unsubscriptable-object\n batch_shapes.append(v.shape[: -self._field_custom_dimensions[k]])\n else:\n batch_shapes.append(v.shape[:-1])\n elif isinstance(v, TensorDataclass):\n batch_shapes.append(v.shape)\n elif isinstance(v, Dict):\n batch_shapes.extend(self._get_dict_batch_shapes(v))\n return batch_shapes\n\n def _broadcast_dict_fields(self, dict_: Dict, batch_shape) -> Dict:\n \"\"\"Broadcasts all tensors in a dictionary according to batch_shape\n\n Args:\n dict_: The dictionary to broadcast.\n\n Returns:\n The broadcasted dictionary.\n \"\"\"\n new_dict = {}\n for k, v in dict_.items():\n if isinstance(v, torch.Tensor):\n # If custom dimension key, then we need to\n if isinstance(self._field_custom_dimensions, dict) and k in self._field_custom_dimensions:\n # pylint: disable=unsubscriptable-object\n new_dict[k] = v.broadcast_to(\n (\n *batch_shape,\n *v.shape[-self._field_custom_dimensions[k] :],\n )\n )\n else:\n new_dict[k] = v.broadcast_to((*batch_shape, v.shape[-1]))\n elif isinstance(v, TensorDataclass):\n new_dict[k] = v.broadcast_to(batch_shape)\n elif isinstance(v, Dict):\n new_dict[k] = self._broadcast_dict_fields(v, batch_shape)\n return new_dict\n\n def __getitem__(self: TensorDataclassT, indices) -> TensorDataclassT:\n if isinstance(indices, (torch.Tensor)):\n return self._apply_fn_to_fields(lambda x: x[indices])\n if isinstance(indices, (int, slice, type(Ellipsis))):\n indices = (indices,)\n assert isinstance(indices, tuple)\n tensor_fn = lambda x: x[indices + (slice(None),)]\n dataclass_fn = lambda x: x[indices]\n\n def custom_tensor_dims_fn(k, v):\n custom_dims = self._field_custom_dimensions[k] # pylint: disable=unsubscriptable-object\n return v[indices + ((slice(None),) * custom_dims)]\n\n return self._apply_fn_to_fields(tensor_fn, dataclass_fn, custom_tensor_dims_fn=custom_tensor_dims_fn)\n\n def __setitem__(self, indices, value) -> NoReturn:\n raise RuntimeError(\"Index assignment is not supported for TensorDataclass\")\n\n def __len__(self) -> int:\n if len(self._shape) == 0:\n raise TypeError(\"len() of a 0-d tensor\")\n return self.shape[0]\n\n def __bool__(self) -> bool:\n if len(self) == 0:\n raise ValueError(\n f\"The truth value of {self.__class__.__name__} when `len(x) == 0` \"\n \"is ambiguous. Use `len(x)` or `x is not None`.\"\n )\n return True\n\n @property\n def shape(self) -> Tuple[int, ...]:\n \"\"\"Returns the batch shape of the tensor dataclass.\"\"\"\n return self._shape\n\n @property\n def size(self) -> int:\n \"\"\"Returns the number of elements in the tensor dataclass batch dimension.\"\"\"\n if len(self._shape) == 0:\n return 1\n return int(np.prod(self._shape))\n\n @property\n def ndim(self) -> int:\n \"\"\"Returns the number of dimensions of the tensor dataclass.\"\"\"\n return len(self._shape)\n\n def reshape(self: TensorDataclassT, shape: Tuple[int, ...]) -> TensorDataclassT:\n \"\"\"Returns a new TensorDataclass with the same data but with a new shape.\n\n This should deepcopy as well.\n\n Args:\n shape: The new shape of the tensor dataclass.\n\n Returns:\n A new TensorDataclass with the same data but with a new shape.\n \"\"\"\n if isinstance(shape, int):\n shape = (shape,)\n tensor_fn = lambda x: x.reshape((*shape, x.shape[-1]))\n dataclass_fn = lambda x: x.reshape(shape)\n\n def custom_tensor_dims_fn(k, v):\n custom_dims = self._field_custom_dimensions[k] # pylint: disable=unsubscriptable-object\n return v.reshape((*shape, *v.shape[-custom_dims:]))\n\n return self._apply_fn_to_fields(tensor_fn, dataclass_fn, custom_tensor_dims_fn=custom_tensor_dims_fn)\n\n def flatten(self: TensorDataclassT) -> TensorDataclassT:\n \"\"\"Returns a new TensorDataclass with flattened batch dimensions\n\n Returns:\n TensorDataclass: A new TensorDataclass with the same data but with a new shape.\n \"\"\"\n return self.reshape((-1,))\n\n def broadcast_to(self: TensorDataclassT, shape: Union[torch.Size, Tuple[int, ...]]) -> TensorDataclassT:\n \"\"\"Returns a new TensorDataclass broadcast to new shape.\n\n Changes to the original tensor dataclass should effect the returned tensor dataclass,\n meaning it is NOT a deepcopy, and they are still linked.\n\n Args:\n shape: The new shape of the tensor dataclass.\n\n Returns:\n A new TensorDataclass with the same data but with a new shape.\n \"\"\"\n\n def custom_tensor_dims_fn(k, v):\n custom_dims = self._field_custom_dimensions[k] # pylint: disable=unsubscriptable-object\n return v.broadcast_to((*shape, *v.shape[-custom_dims:]))\n\n return self._apply_fn_to_fields(\n lambda x: x.broadcast_to((*shape, x.shape[-1])), custom_tensor_dims_fn=custom_tensor_dims_fn\n )\n\n def to(self: TensorDataclassT, device) -> TensorDataclassT:\n \"\"\"Returns a new TensorDataclass with the same data but on the specified device.\n\n Args:\n device: The device to place the tensor dataclass.\n\n Returns:\n A new TensorDataclass with the same data but on the specified device.\n \"\"\"\n return self._apply_fn_to_fields(lambda x: x.to(device))\n\n def _apply_fn_to_fields(\n self: TensorDataclassT,\n fn: Callable,\n dataclass_fn: Optional[Callable] = None,\n custom_tensor_dims_fn: Optional[Callable] = None,\n ) -> TensorDataclassT:\n \"\"\"Applies a function to all fields of the tensor dataclass.\n\n TODO: Someone needs to make a high level design choice for whether not not we want this\n to apply the function to any fields in arbitray superclasses. This is an edge case until we\n upgrade to python 3.10 and dataclasses can actually be subclassed with vanilla python and no\n janking, but if people try to jank some subclasses that are grandchildren of TensorDataclass\n (imagine if someone tries to subclass the RayBundle) this will matter even before upgrading\n to 3.10 . Currently we aren't going to be able to work properly for grandchildren, but you\n want to use self.__dict__ if you want to apply this to grandchildren instead of our dictionary\n from dataclasses.fields(self) as we do below and in other places.\n\n Args:\n fn: The function to apply to tensor fields.\n dataclass_fn: The function to apply to TensorDataclass fields.\n\n Returns:\n A new TensorDataclass with the same data but with a new shape.\n \"\"\"\n\n new_fields = self._apply_fn_to_dict(\n {f.name: self.__getattribute__(f.name) for f in dataclasses.fields(self)},\n fn,\n dataclass_fn,\n custom_tensor_dims_fn,\n )\n\n return dataclasses.replace(self, **new_fields)\n\n def _apply_fn_to_dict(\n self,\n dict_: Dict,\n fn: Callable,\n dataclass_fn: Optional[Callable] = None,\n custom_tensor_dims_fn: Optional[Callable] = None,\n ) -> Dict:\n \"\"\"A helper function for _apply_fn_to_fields, applying a function to all fields of dict_\n\n Args:\n dict_: The dictionary to apply the function to.\n fn: The function to apply to tensor fields.\n dataclass_fn: The function to apply to TensorDataclass fields.\n\n Returns:\n A new dictionary with the same data but with a new shape. Will deep copy\"\"\"\n\n field_names = dict_.keys()\n new_dict = {}\n for f in field_names:\n v = dict_[f]\n if v is not None:\n if isinstance(v, TensorDataclass) and dataclass_fn is not None:\n new_dict[f] = dataclass_fn(v)\n # This is the case when we have a custom dimensions tensor\n elif (\n isinstance(v, torch.Tensor)\n and isinstance(self._field_custom_dimensions, dict)\n and f in self._field_custom_dimensions\n and custom_tensor_dims_fn is not None\n ):\n new_dict[f] = custom_tensor_dims_fn(f, v)\n elif isinstance(v, (torch.Tensor, TensorDataclass)):\n new_dict[f] = fn(v)\n elif isinstance(v, Dict):\n new_dict[f] = self._apply_fn_to_dict(v, fn, dataclass_fn)\n else:\n new_dict[f] = deepcopy(v)\n\n return new_dict\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/vis_utils.py",
"content": "\nimport cv2\nimport numpy as np\nimport torch\nfrom rich.console import Console\nfrom scipy.interpolate import interp1d\nfrom scipy.spatial.transform import Rotation, Slerp\n\nfrom nerfstudio.cameras.cameras import Cameras\n\nCONSOLE = Console(width=120)\n\n\ndef interpolate_trajectory(cameras: Cameras, num_views: int = 300):\n \"\"\"calculate interpolate path\"\"\"\n\n c2ws = np.stack(cameras.camera_to_worlds.cpu().numpy())\n\n key_rots = Rotation.from_matrix(c2ws[:, :3, :3])\n key_times = list(range(len(c2ws)))\n slerp = Slerp(key_times, key_rots)\n interp = interp1d(key_times, c2ws[:, :3, 3], axis=0)\n \n if num_views == -1:\n render_c2ws = torch.from_numpy(c2ws)\n else:\n render_c2ws = []\n for i in range(num_views):\n time = float(i) / num_views * (len(c2ws) - 1)\n cam_location = interp(time)\n cam_rot = slerp(time).as_matrix()\n c2w = np.eye(4)\n c2w[:3, :3] = cam_rot\n c2w[:3, 3] = cam_location\n render_c2ws.append(c2w)\n render_c2ws = torch.from_numpy(np.stack(render_c2ws, axis=0))\n\n # use intrinsic of first camera\n camera_path = Cameras(\n fx=cameras[0].fx,\n fy=cameras[0].fy,\n cx=cameras[0].cx,\n cy=cameras[0].cy,\n height=cameras[0].height,\n width=cameras[0].width,\n camera_to_worlds=render_c2ws[:, :3, :4],\n camera_type=cameras[0].camera_type,\n )\n return camera_path\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/utils/writer.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nGeneric Writer class\n\"\"\"\nfrom __future__ import annotations\n\nimport enum\nfrom abc import abstractmethod\nfrom pathlib import Path\nfrom time import time\nfrom typing import Any, Dict, List, Optional, Union\n\nimport torch\nimport wandb\nfrom rich.console import Console\nfrom torch.utils.tensorboard import SummaryWriter\nfrom torchtyping import TensorType\n\nfrom nerfstudio.configs import base_config as cfg\nfrom nerfstudio.utils.decorators import check_main_thread, decorate_all\nfrom nerfstudio.utils.printing import human_format\n\nCONSOLE = Console(width=120, soft_wrap=True)\nto8b = lambda x: (255 * torch.clamp(x, min=0, max=1)).to(torch.uint8)\nEVENT_WRITERS = []\nEVENT_STORAGE = []\nGLOBAL_BUFFER = {}\n\n\nclass EventName(enum.Enum):\n \"\"\"Names of possible events that can be logged via Local Writer for convenience.\n see config/logging/default_logging.yaml\"\"\"\n\n ITER_TRAIN_TIME = \"Train Iter (time)\"\n TOTAL_TRAIN_TIME = \"Train Total (time)\"\n ITER_VIS_TIME = \"Viewer Rendering (time)\"\n ETA = \"ETA (time)\"\n TRAIN_RAYS_PER_SEC = \"Train Rays / Sec\"\n TEST_RAYS_PER_SEC = \"Test Rays / Sec\"\n VIS_RAYS_PER_SEC = \"Vis Rays / Sec\"\n CURR_TEST_PSNR = \"Test PSNR\"\n\n\nclass EventType(enum.Enum):\n \"\"\"Possible Event types and their associated write function\"\"\"\n\n IMAGE = \"write_image\"\n SCALAR = \"write_scalar\"\n DICT = \"write_scalar_dict\"\n CONFIG = \"write_config\"\n\n\n@check_main_thread\ndef put_image(name, image: TensorType[\"H\", \"W\", \"C\"], step: int):\n \"\"\"Setter function to place images into the queue to be written out\n\n Args:\n image: image to write out\n step: step associated with image\n \"\"\"\n if isinstance(name, EventName):\n name = name.value\n\n EVENT_STORAGE.append({\"name\": name, \"write_type\": EventType.IMAGE, \"event\": image.detach().cpu(), \"step\": step})\n\n\n@check_main_thread\ndef put_scalar(name: str, scalar: Any, step: int):\n \"\"\"Setter function to place scalars into the queue to be written out\n\n Args:\n name: name of scalar\n scalar: value\n step: step associated with scalar\n \"\"\"\n if isinstance(name, EventName):\n name = name.value\n\n EVENT_STORAGE.append({\"name\": name, \"write_type\": EventType.SCALAR, \"event\": scalar, \"step\": step})\n\n\n@check_main_thread\ndef put_dict(name: str, scalar_dict: Dict[str, Any], step: int):\n \"\"\"Setter function to place a dictionary of scalars into the queue to be written out\n\n Args:\n name: name of scalar dictionary\n scalar_dict: values to write out\n step: step associated with dict\n \"\"\"\n EVENT_STORAGE.append({\"name\": name, \"write_type\": EventType.DICT, \"event\": scalar_dict, \"step\": step})\n\n\n@check_main_thread\ndef put_config(name: str, config_dict: Dict[str, Any], step: int):\n \"\"\"Setter function to place a dictionary of scalars into the queue to be written out\n\n Args:\n name: name of scalar dictionary\n scalar_dict: values to write out\n step: step associated with dict\n \"\"\"\n EVENT_STORAGE.append({\"name\": name, \"write_type\": EventType.CONFIG, \"event\": config_dict, \"step\": step})\n\n\n@check_main_thread\ndef put_time(name: str, duration: float, step: int, avg_over_steps: bool = True, update_eta: bool = False):\n \"\"\"Setter function to place a time element into the queue to be written out.\n Processes the time info according to the options.\n\n Args:\n name: name of time item\n duration: value\n step: step associated with value\n avg_over_steps: if True, calculate and record a running average of the times\n update_eta: if True, update the ETA. should only be set for the training iterations/s\n \"\"\"\n if isinstance(name, EventName):\n name = name.value\n\n if avg_over_steps:\n GLOBAL_BUFFER[\"step\"] = step\n curr_event = GLOBAL_BUFFER[\"events\"].get(name, {\"buffer\": [], \"avg\": 0})\n curr_buffer = curr_event[\"buffer\"]\n if len(curr_buffer) >= GLOBAL_BUFFER[\"max_buffer_size\"]:\n curr_buffer.pop(0)\n curr_buffer.append(duration)\n curr_avg = sum(curr_buffer) / len(curr_buffer)\n put_scalar(name, curr_avg, step)\n GLOBAL_BUFFER[\"events\"][name] = {\"buffer\": curr_buffer, \"avg\": curr_avg}\n else:\n put_scalar(name, duration, step)\n\n if update_eta:\n ## NOTE: eta should be called with avg train iteration time\n remain_iter = GLOBAL_BUFFER[\"max_iter\"] - step\n remain_time = remain_iter * GLOBAL_BUFFER[\"events\"][name][\"avg\"]\n put_scalar(EventName.ETA, remain_time, step)\n GLOBAL_BUFFER[\"events\"][EventName.ETA.value] = _format_time(remain_time)\n\n\n@check_main_thread\ndef write_out_storage():\n \"\"\"Function that writes all the events in storage to all the writer locations\"\"\"\n for writer in EVENT_WRITERS:\n if isinstance(writer, LocalWriter) and len(EVENT_STORAGE) > 0:\n writer.write_stats_log(EVENT_STORAGE[0][\"step\"])\n continue\n for event in EVENT_STORAGE:\n write_func = getattr(writer, event[\"write_type\"].value)\n write_func(event[\"name\"], event[\"event\"], event[\"step\"])\n\n EVENT_STORAGE.clear()\n\n\ndef setup_local_writer(config: cfg.LoggingConfig, max_iter: int, banner_messages: Optional[List[str]] = None) -> None:\n \"\"\"Initialization of all event writers specified in config\n\n Args:\n config: configuration to instantiate loggers\n max_iter: maximum number of train iterations\n banner_messages: list of messages to always display at bottom of screen\n \"\"\"\n if config.local_writer.enable:\n curr_writer = config.local_writer.setup(banner_messages=banner_messages)\n EVENT_WRITERS.append(curr_writer)\n else:\n CONSOLE.log(\"disabled local writer\")\n\n ## configure all the global buffer basic information\n GLOBAL_BUFFER[\"max_iter\"] = max_iter\n GLOBAL_BUFFER[\"max_buffer_size\"] = config.max_buffer_size\n GLOBAL_BUFFER[\"steps_per_log\"] = config.steps_per_log\n GLOBAL_BUFFER[\"events\"] = {}\n\n\n@check_main_thread\ndef setup_event_writer(config: cfg.Config, log_dir: Path) -> None:\n \"\"\"Initialization of all event writers specified in config\n\n Args:\n config: configuration to instantiate loggers\n max_iter: maximum number of train iterations\n banner_messages: list of messages to always display at bottom of screen\n \"\"\"\n using_event_writer = False\n if config.is_wandb_enabled():\n curr_writer = WandbWriter(log_dir=log_dir, experiment_name=config.experiment_name)\n EVENT_WRITERS.append(curr_writer)\n using_event_writer = True\n if config.is_tensorboard_enabled():\n curr_writer = TensorboardWriter(log_dir=log_dir)\n EVENT_WRITERS.append(curr_writer)\n using_event_writer = True\n if using_event_writer:\n string = f\"logging events to: {log_dir}\"\n else:\n string = \"Disabled tensorboard/wandb event writers\"\n CONSOLE.print(f\"[bold yellow]{string}\")\n\n\nclass Writer:\n \"\"\"Writer class\"\"\"\n\n @abstractmethod\n def write_image(self, name: str, image: TensorType[\"H\", \"W\", \"C\"], step: int) -> None:\n \"\"\"method to write out image\n\n Args:\n name: data identifier\n image: rendered image to write\n step: the time step to log\n \"\"\"\n raise NotImplementedError\n\n @abstractmethod\n def write_scalar(self, name: str, scalar: Union[float, torch.Tensor], step: int) -> None:\n \"\"\"Required method to write a single scalar value to the logger\n\n Args:\n name: data identifier\n scalar: value to write out\n step: the time step to log\n \"\"\"\n raise NotImplementedError\n\n @check_main_thread\n def write_scalar_dict(self, name: str, scalar_dict: Dict[str, Any], step: int) -> None:\n \"\"\"Function that writes out all scalars from a given dictionary to the logger\n\n Args:\n scalar_dict: dictionary containing all scalar values with key names and quantities\n step: the time step to log\n \"\"\"\n for key, scalar in scalar_dict.items():\n self.write_scalar(name + \"/\" + key, float(scalar), step)\n\n\nclass TimeWriter:\n \"\"\"Timer context manager that calculates duration around wrapped functions\"\"\"\n\n def __init__(self, writer, name, step=None, write=True):\n self.writer = writer\n self.name = name\n self.step = step\n self.write = write\n\n self.start: float = 0.0\n self.duration: float = 0.0\n\n def __enter__(self):\n torch.cuda.synchronize()\n self.start = time()\n return self\n\n def __exit__(self, *args):\n torch.cuda.synchronize()\n self.duration = time() - self.start\n update_step = self.step is not None\n if self.write:\n self.writer.put_time(\n name=self.name,\n duration=self.duration,\n step=self.step if update_step else GLOBAL_BUFFER[\"max_iter\"],\n avg_over_steps=update_step,\n update_eta=self.name == EventName.ITER_TRAIN_TIME,\n )\n\n\n@decorate_all([check_main_thread])\nclass WandbWriter(Writer):\n \"\"\"WandDB Writer Class\"\"\"\n\n def __init__(self, log_dir: Path, experiment_name: str):\n wandb.init(project=\"sdfstudio\", name=experiment_name, dir=str(log_dir), reinit=True)\n\n def write_image(self, name: str, image: TensorType[\"H\", \"W\", \"C\"], step: int) -> None:\n image = torch.permute(image, (2, 0, 1))\n wandb.log({name: wandb.Image(image)}, step=step)\n\n def write_scalar(self, name: str, scalar: Union[float, torch.Tensor], step: int) -> None:\n wandb.log({name: scalar}, step=step)\n\n def write_config(self, name: str, config_dict: Dict[str, Any], step: int):\n # pylint: disable=unused-argument\n # pylint: disable=no-self-use\n \"\"\"Function that writes out the config to wandb\n\n Args:\n config: config dictionary to write out\n \"\"\"\n wandb.config.update(config_dict)\n\n\n@decorate_all([check_main_thread])\nclass TensorboardWriter(Writer):\n \"\"\"Tensorboard Writer Class\"\"\"\n\n def __init__(self, log_dir: Path):\n self.tb_writer = SummaryWriter(log_dir=log_dir)\n\n def write_image(self, name: str, image: TensorType[\"H\", \"W\", \"C\"], step: int) -> None:\n image = to8b(image)\n self.tb_writer.add_image(name, image, step, dataformats=\"HWC\")\n\n def write_scalar(self, name: str, scalar: Union[float, torch.Tensor], step: int) -> None:\n self.tb_writer.add_scalar(name, scalar, step)\n\n def write_config(self, name: str, config_dict: Dict[str, Any], step: int): # pylint: disable=unused-argument\n \"\"\"Function that writes out the config to tensorboard\n\n Args:\n config: config dictionary to write out\n \"\"\"\n self.tb_writer.add_text(\"config\", str(config_dict))\n\n\ndef _cursorup(x: int):\n \"\"\"utility tool to move the cursor up on the terminal\n\n Args:\n x: amount of lines to move cursor upward\n \"\"\"\n print(f\"\\r\\033[{x}A\", end=\"\\x1b[1K\\r\")\n\n\ndef _format_time(seconds):\n \"\"\"utility tool to format time in human readable form given seconds\"\"\"\n ms = seconds % 1\n ms = ms * 1e3\n seconds = int(seconds)\n days, seconds = divmod(seconds, 86400)\n hours, seconds = divmod(seconds, 3600)\n minutes, seconds = divmod(seconds, 60)\n if days > 0:\n return f\"{days} d, {hours} h, {minutes} m, {seconds} s\"\n if hours > 0:\n return f\"{hours} h, {minutes} m, {seconds} s\"\n if minutes > 0:\n return f\"{minutes} m, {seconds} s\"\n if seconds > 0:\n return f\"{seconds} s, {ms:0.3f} ms\"\n\n return f\"{ms:0.3f} ms\"\n\n\n@decorate_all([check_main_thread])\nclass LocalWriter:\n \"\"\"Local Writer Class\n TODO: migrate to prettyprint\n\n Args:\n config: configuration to instatiate class\n banner_messages: list of messages to always display at bottom of screen\n \"\"\"\n\n def __init__(self, config: cfg.LocalWriterConfig, banner_messages: Optional[List[str]] = None):\n self.config = config\n self.stats_to_track = [name.value for name in config.stats_to_track]\n self.keys = set()\n self.past_mssgs = [\"\", \"\"]\n self.banner_len = 0 if banner_messages is None else len(banner_messages) + 1\n if banner_messages:\n self.past_mssgs.extend([\"-\" * 100])\n self.past_mssgs.extend(banner_messages)\n self.has_printed = False\n\n def write_stats_log(self, step: int) -> None:\n \"\"\"Function to write out scalars to terminal\n\n Args:\n step: current train step\n \"\"\"\n valid_step = step > 0 and step % GLOBAL_BUFFER[\"steps_per_log\"] == 0\n if valid_step:\n if not self.has_printed and self.config.max_log_size:\n CONSOLE.log(\n f\"Printing max of {self.config.max_log_size} lines. \"\n \"Set flag çççç[yellow]--logging.local-writer.max-log-size=0[/yellow] \"\n \"to disable line wrapping.\"\n )\n latest_map, new_key = self._consolidate_events()\n self._update_header(latest_map, new_key)\n self._print_stats(latest_map)\n\n def write_config(self, name: str, config_dict: Dict[str, Any], step: int):\n \"\"\"Function that writes out the config to local\n\n Args:\n config: config dictionary to write out\n \"\"\"\n # TODO: implement this\n\n def _consolidate_events(self):\n latest_map = {}\n new_key = False\n for event in EVENT_STORAGE:\n name = event[\"name\"]\n if name not in self.keys:\n self.keys.add(name)\n new_key = True\n latest_map[name] = event[\"event\"]\n return latest_map, new_key\n\n def _update_header(self, latest_map, new_key):\n \"\"\"helper to handle the printing of the header labels\n\n Args:\n latest_map: the most recent dictionary of stats that have been recorded\n new_key: indicator whether or not there is a new key added to logger\n \"\"\"\n full_log_cond = not self.config.max_log_size and GLOBAL_BUFFER[\"step\"] <= GLOBAL_BUFFER[\"steps_per_log\"]\n capped_log_cond = self.config.max_log_size and (len(self.past_mssgs) - self.banner_len <= 2 or new_key)\n if full_log_cond or capped_log_cond:\n mssg = f\"{'Step (% Done)':<20}\"\n for name, _ in latest_map.items():\n if name in self.stats_to_track:\n mssg += f\"{name:<20} \"\n self.past_mssgs[0] = mssg\n self.past_mssgs[1] = \"-\" * len(mssg)\n if full_log_cond or not self.has_printed:\n print(mssg)\n print(\"-\" * len(mssg))\n\n def _print_stats(self, latest_map, padding=\" \"):\n \"\"\"helper to print out the stats in a readable format\n\n Args:\n latest_map: the most recent dictionary of stats that have been recorded\n padding: type of characters to print to pad open space\n \"\"\"\n step = GLOBAL_BUFFER[\"step\"]\n fraction_done = step / GLOBAL_BUFFER[\"max_iter\"]\n curr_mssg = f\"{step} ({fraction_done*100:.02f}%)\"\n curr_mssg = f\"{curr_mssg:<20}\"\n for name, v in latest_map.items():\n if name in self.stats_to_track:\n if \"(time)\" in name:\n v = _format_time(v)\n elif \"Rays\" in name:\n v = human_format(v)\n else:\n v = f\"{v:0.4f}\"\n curr_mssg += f\"{v:<20} \"\n\n # update the history buffer\n if self.config.max_log_size:\n if not self.has_printed:\n cursor_idx = len(self.past_mssgs) - self.banner_len\n self.has_printed = True\n else:\n cursor_idx = len(self.past_mssgs)\n if len(self.past_mssgs[2:]) - self.banner_len >= self.config.max_log_size:\n self.past_mssgs.pop(2)\n self.past_mssgs.insert(len(self.past_mssgs) - self.banner_len, curr_mssg)\n _cursorup(cursor_idx)\n\n for i, mssg in enumerate(self.past_mssgs):\n pad_len = len(max(self.past_mssgs, key=len))\n style = \"\\x1b[6;30;42m\" if self.banner_len and i >= len(self.past_mssgs) - self.banner_len + 1 else \"\"\n print(f\"{style}{mssg:{padding}<{pad_len}} \\x1b[0m\")\n else:\n print(curr_mssg)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/viewer/__init__.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/viewer/app/.eslintrc.json",
"content": "{\n \"env\": {\n \"browser\": true,\n \"es2021\": true\n },\n \"extends\": [\n \"eslint:recommended\",\n \"plugin:react/recommended\",\n \"airbnb\",\n \"prettier\"\n ],\n \"parser\": \"@typescript-eslint/parser\",\n \"parserOptions\": {\n \"ecmaFeatures\": {\n \"jsx\": true\n },\n \"ecmaVersion\": \"latest\",\n \"sourceType\": \"module\"\n },\n \"plugins\": [\"react\", \"@typescript-eslint\", \"unused-imports\"],\n \"rules\": {\n \"arrow-body-style\": \"off\",\n \"camelcase\": \"off\",\n \"import/prefer-default-export\": \"off\",\n \"no-alert\": \"off\",\n \"no-console\": \"off\",\n \"prefer-destructuring\": \"off\",\n \"react/destructuring-assignment\": \"off\",\n \"react/prop-types\": 0,\n \"unused-imports/no-unused-imports-ts\": 2\n }\n}\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/viewer/app/.gitignore",
"content": ".idea/\n.vscode/\nnode_modules/\nbuild/\n.DS_Store\n*.tgz\nmy-app*\ntemplate/src/__tests__/__snapshots__/\nlerna-debug.log\nnpm-debug.log*\nyarn-debug.log*\nyarn-error.log*\n/.changelog\n.npm/"
},
{
"path": "AutoReconForDens3R/nerfstudio/viewer/app/package.json",
"content": "{\n \"name\": \"viewer\",\n \"homepage\": \".\",\n \"version\": \"22-12-02-0\",\n \"private\": true,\n \"dependencies\": {\n \"@emotion/react\": \"^11.10.4\",\n \"@emotion/styled\": \"^11.10.4\",\n \"@mui/icons-material\": \"^5.10.3\",\n \"@mui/lab\": \"^5.0.0-alpha.98\",\n \"@mui/material\": \"^5.10.3\",\n \"@mui/system\": \"^5.10.6\",\n \"@mui/x-date-pickers\": \"^5.0.0\",\n \"@reduxjs/toolkit\": \"^1.8.3\",\n \"@testing-library/jest-dom\": \"^5.16.4\",\n \"@testing-library/react\": \"^13.3.0\",\n \"@testing-library/user-event\": \"^14.2.0\",\n \"camera-controls\": \"^1.37.2\",\n \"classnames\": \"^2.3.1\",\n \"dat.gui\": \"^0.7.9\",\n \"dayjs\": \"^1.11.5\",\n \"eslint-config-prettier\": \"^8.5.0\",\n \"eslint-plugin-unused-imports\": \"^2.0.0\",\n \"leva\": \"^0.9.29\",\n \"meshline\": \"^2.0.4\",\n \"msgpack-lite\": \"^0.1.26\",\n \"prop-types\": \"^15.8.1\",\n \"re-resizable\": \"^6.9.9\",\n \"react\": \"^18.1.0\",\n \"react-dom\": \"^18.1.0\",\n \"react-icons\": \"^4.4.0\",\n \"react-pro-sidebar\": \"^0.7.1\",\n \"react-redux\": \"^8.0.2\",\n \"redux\": \"^4.2.0\",\n \"sass\": \"^1.54.8\",\n \"socket.io-client\": \"^4.5.1\",\n \"three\": \"^0.142.0\",\n \"three-wtm\": \"^1.0\",\n \"websocket\": \"^1.0.34\",\n \"wwobjloader2\": \"^4.0\"\n },\n \"scripts\": {\n \"start\": \"react-scripts start\",\n \"build\": \"react-scripts build\",\n \"test\": \"react-scripts test\",\n \"eject\": \"react-scripts eject\",\n \"electron\": \"electron .\",\n \"lint\": \"eslint --ext .js,.jsx .\",\n \"lint:fix\": \"eslint --fix --ext .js,.jsx .\"\n },\n \"eslintConfig\": {\n \"extends\": \"react-app\"\n },\n \"browserslist\": {\n \"production\": [\n \">0.2%\",\n \"not dead\",\n \"not op_mini all\"\n ],\n \"development\": [\n \"last 1 chrome version\",\n \"last 1 firefox version\",\n \"last 1 safari version\"\n ]\n },\n \"main\": \"public/electron.js\",\n \"author\": \"\",\n \"license\": \"ISC\",\n \"description\": \"\",\n \"devDependencies\": {\n \"concurrently\": \"^7.2.1\",\n \"eslint\": \"^8.2.0\",\n \"eslint-config-airbnb\": \"19.0.4\",\n \"eslint-plugin-import\": \"^2.25.3\",\n \"eslint-plugin-jsx-a11y\": \"^6.5.1\",\n \"eslint-plugin-react\": \"^7.28.0\",\n \"eslint-plugin-react-hooks\": \"^4.3.0\",\n \"prettier\": \"2.7.1\",\n \"react-scripts\": \"^5.0.1\",\n \"typescript\": \"^4.7.3\",\n \"wait-on\": \"^6.0.1\"\n },\n \"resolutions\": {\n \"nth-check\": \"^2.0.1\"\n }\n}\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/viewer/app/public/electron.js",
"content": "const path = require('path');\n\nconst { app, BrowserWindow } = require('electron');\nconst isDev = require('electron-is-dev');\n\nfunction createWindow() {\n // Create the browser window.\n const win = new BrowserWindow({\n width: 800,\n height: 600,\n webPreferences: {\n nodeIntegration: true,\n },\n });\n\n // and load the index.html of the app.\n // win.loadFile(\"index.html\");\n win.loadURL(\n isDev\n ? 'http://localhost:3000'\n : `file://${path.join(__dirname, '../build/index.html')}`,\n );\n // Open the DevTools.\n if (isDev) {\n win.webContents.openDevTools({ mode: 'detach' });\n }\n}\n\n// This method will be called when Electron has finished\n// initialization and is ready to create browser windows.\n// Some APIs can only be used after this event occurs.\napp.whenReady().then(createWindow);\n\n// Quit when all windows are closed, except on macOS. There, it's common\n// for applications and their menu bar to stay active until the user quits\n// explicitly with Cmd + Q.\napp.on('window-all-closed', () => {\n if (process.platform !== 'darwin') {\n app.quit();\n }\n});\n\napp.on('activate', () => {\n if (BrowserWindow.getAllWindows().length === 0) {\n createWindow();\n }\n});\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/viewer/app/public/index.html",
"content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n nerfstudio\n \n \n \n \n \n \n \n \n \n\n \n \n \n \n \n\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/viewer/app/public/manifest.json",
"content": "{\n \"short_name\": \"nerfstudio viewer\",\n \"name\": \"Interactive NeRF viewer\",\n \"start_url\": \".\",\n \"display\": \"standalone\",\n \"theme_color\": \"#000000\",\n \"background_color\": \"#ffffff\"\n}\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/viewer/app/public/robots.txt",
"content": "# https://www.robotstxt.org/robotstxt.html\nUser-agent: *\nDisallow:\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/viewer/app/requirements.txt",
"content": "tyro>=0.3.22\nsshconf==0.2.5\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/viewer/app/run_deploy.py",
"content": "# Copyright 2022 The Nerfstudio Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nCode for deploying the built viewer folder to a server and handing versioning.\nWe use the library sshconf (https://github.com/sorend/sshconf) for working with the ssh config file.\n\"\"\"\nimport json\nimport os\nimport subprocess\nfrom os.path import expanduser\nfrom pathlib import Path\nfrom typing import Optional\n\nimport tyro\nfrom sshconf import empty_ssh_config_file, read_ssh_config\n\n\ndef run_cmd(cmd: str):\n \"\"\"Run a command in the terminal.\"\"\"\n print(\"cmd:\", cmd)\n print(\"output:\")\n subprocess.Popen(cmd, shell=True).wait()\n\n\ndef main(\n branch_name: str = \"\",\n ssh_key_string: Optional[str] = None,\n ssh_key_filename: str = \" ~/.ssh/github_actions_user_key_filename\",\n hostname_or_ip_address: str = \"34.102.68.79\",\n local_folder: str = \"/home/eweb0124/build\",\n remote_folder: str = \"/home/eweb0124/viewer\",\n host: str = \"viewer_deploy_host\",\n user: str = \"eweb0124\",\n package_json_filename: str = \"package.json\",\n increment_version: str = \"False\",\n):\n \"\"\"Copy a local folder to a remote machine and handle versioning.\n\n Args:\n ssh_key: The private ssh key needed to ssh.\n hostname_or_ip_address: The hostname or ip_address of the remote machine.\n \"\"\"\n\n print()\n print(\"branch_name\", branch_name)\n print(\"ssh_key_string\", ssh_key_string)\n print(\"ssh_key_filename\", ssh_key_filename)\n print(\"hostname_or_ip_address\", hostname_or_ip_address)\n print(\"local_folder\", local_folder)\n print(\"remote_folder\", remote_folder)\n print(\"host\", host)\n print(\"user\", user)\n print(\"package_json_filename\", package_json_filename)\n print(\"increment_version\", increment_version)\n print()\n\n # save the ssh key to a file\n run_cmd(\"mkdir ~/.ssh\")\n if ssh_key_string:\n run_cmd(f\"\"\"rm -f {ssh_key_filename}\"\"\")\n run_cmd(f\"\"\"echo \"{ssh_key_string}\" >> {ssh_key_filename}\"\"\")\n run_cmd(f\"chmod 400 {ssh_key_filename}\")\n\n # setup the config in ~/.ssh/config\n config_filename = expanduser(\"~/.ssh/config\")\n Path(config_filename).parent.mkdir(exist_ok=True)\n try:\n config = read_ssh_config(config_filename)\n except FileNotFoundError:\n config = empty_ssh_config_file()\n config.write(config_filename)\n config = read_ssh_config(config_filename)\n\n # add the host if it doesn't exist\n if not config.host(host):\n config.add(host)\n\n config.set(\n host,\n Hostname=hostname_or_ip_address,\n User=user,\n IdentityFile=ssh_key_filename,\n StrictHostKeyChecking=\"No\",\n )\n\n # show that the config is correct\n run_cmd(\"cat ~/.ssh/config\")\n\n # save the config file\n config.save()\n\n # get the version\n with open(package_json_filename, \"r\", encoding=\"utf-8\") as f:\n package_json = json.load(f)\n\n # TODO: add logic to increment the version number\n if increment_version == \"True\":\n raise NotImplementedError()\n else:\n version = package_json[\"version\"]\n\n print(f\"\\nusing version: {version}\")\n\n # write to the /home/eweb0124/build folder\n run_cmd(f\"\"\"ssh {host} 'rm -rf /home/eweb0124/build'\"\"\")\n run_cmd(f\"\"\"scp -r {local_folder} {host}:/home/eweb0124/build\"\"\")\n\n # update the symlink of latest\n if branch_name == \"main\":\n\n # move the build folder to the correct location\n run_cmd(f\"\"\"ssh {host} 'rm -rf {remote_folder}/{version}'\"\"\")\n run_cmd(f\"\"\"ssh {host} 'cp -R /home/eweb0124/build {remote_folder}/{version}'\"\"\")\n\n run_cmd(f\"\"\"ssh {host} 'rm {remote_folder}/latest'\"\"\")\n run_cmd(f\"\"\"ssh {host} 'ln -s {remote_folder}/{version} {remote_folder}/latest'\"\"\")\n\n # otherwise just move to some branch folder\n else:\n updated_branch_name = branch_name.replace(\"/\", \"-\")\n run_cmd(f\"\"\"ssh {host} 'rm -rf {remote_folder}/branch/{updated_branch_name}'\"\"\")\n run_cmd(f\"\"\"ssh {host} 'cp -R /home/eweb0124/build {remote_folder}/branch/{updated_branch_name}'\"\"\")\n\n\nif __name__ == \"__main__\":\n tyro.extras.set_accent_color(\"bright_yellow\")\n tyro.cli(main)\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/viewer/app/src/App.jsx",
"content": "import { CssBaseline, ThemeProvider } from '@mui/material';\nimport React from 'react';\nimport {\n SceneTreeWebSocketListener,\n get_scene_tree,\n} from './modules/Scene/Scene';\n\nimport Banner from './modules/Banner';\nimport { BasicTabs } from './modules/SidePanel/SidePanel';\nimport ViewerWindow from './modules/ViewerWindow/ViewerWindow';\nimport { appTheme } from './themes/theme.ts';\n\nexport default function App() {\n // The scene tree won't rerender but it will listen to changes\n // from the redux store and draw three.js objects.\n // In particular, it listens to changes to 'sceneState' coming over the websocket.\n const sceneTree = get_scene_tree();\n\n return (\n \n \n
\n {/* Listens for websocket 'write' messages and updates the redux store. */}\n \n {/* The banner at the top of the page. */}\n \n
\n {/* Order matters here. The viewer window must be rendered first. */}\n \n
\n \n
\n
\n
\n \n );\n}\n"
},
{
"path": "AutoReconForDens3R/nerfstudio/viewer/app/src/SceneNode.js",
"content": "/* eslint-disable no-underscore-dangle */\n/* eslint-disable no-restricted-syntax */\nimport * as THREE from 'three';\n\nfunction dispose(object) {\n if (!object) {\n return;\n }\n if (object.geometry) {\n object.geometry.dispose();\n }\n if (object.material) {\n if (Array.isArray(object.material)) {\n for (const material of object.material) {\n if (material.map) {\n material.map.dispose();\n }\n material.dispose();\n }\n } else {\n if (object.material.map) {\n object.material.map.dispose();\n }\n object.material.dispose();\n }\n }\n}\n\nexport default class SceneNode {\n constructor(object, scene_state) {\n this.metadata = {}; // relevant metadata for the object\n this.object = object; // three.js objects\n this.children = {};\n this.scene_state = scene_state;\n for (const c of this.object.children) {\n this.add_child(c);\n }\n }\n\n get_scene_state() {\n return this.scene_state;\n }\n\n add_child(object) {\n const node = new SceneNode(object);\n this.children[object.name] = node;\n return node;\n }\n\n create_child(name) {\n const obj = new THREE.Group();\n obj.name = name;\n this.object.add(obj);\n return this.add_child(obj);\n }\n\n find(path) {\n if (path.length === 0) {\n return this;\n }\n const name = path[0];\n let child = this.children[name];\n if (child === undefined) {\n child = this.create_child(name);\n }\n return child.find(path.slice(1));\n }\n\n find_object(path) {\n return this.find(path.concat(['
![\\\"Open](\\\"https://colab.research.google.com/assets/colab-badge.svg\\\"){kind=icon}
![\\\"nerfstudio\\\"](\\\"https://docs.nerf.studio/en/latest/_images/logo.png\\\"){kind=logo}
\\n\",\n \"